[Relay] Partial Evaluation

[MarisaKirisame]

During the pass two weeks, i had prototyped a partial evaluator in metaocaml for a toy DSL with similar characteristic to relay.

It works by taking a evaluator, lifting the value to partially-static domain, reifying the store, and anfing the code generated using let list to avoid code duplication and get capture avoidance substitution for free.

It is written such that it could be translated line by line to tvm.

I had wrote some test, and it can indeed remove the closure and reference #2496 generate.
However, it does generate lots of dead code (closure, reference) that need to be remove by the dead code elimination pass. We probably need an ad hoc naive escape analysis as only write to dead variable is ok to remove.

I will port this but the high level idea is ready for review: the prototype is only 336 loc (including all the DSL definition), and the core function is less then 100 line.

Thanks to @weberlo for reviewing the doc

[MarisaKirisame]

This is now feature complete. I will add some test and it will be good.
@icemelon9 @masahi @junrushao1994 @FrozenGene @nhynes @wweic can you guys give a round of review?

[nhynes]

Cool PR! Just a few minor suggestions on my part.

[tqchen]

tagging more reviewers, @wweic @junrushao1994

[junrushao]

This PR enables compile-time partial evaluation. LGTM. This could slow down compilation, especially when all computation is done on CPU. But it should be fine for most of applications.

Only some minor nits.

[junrushao]

Why not just inherit the TVM's node system? Should be similar i guess

[MarisaKirisame]

It is considerably more work, and we has no need.

[tqchen]

While it might take a bit more work to adopt TVM's node system, it makes us share the same infra and reuse later. I will recommend using node instead here.

Note that it is not too much more work, just replace shared_ptr->NodePtr and make_shared->make_node

[tqchen]

While it might take a bit more work to adopt TVM's node system, it makes us share the same infra and reuse later. I will recommend using node instead here.

Note that it is not too much more work, just replace shared_ptr->NodePtr and make_shared->make_node

[wweic]

Nice work. I'll review today. Interesting to see its effect in the future. As I understand, majority of relay function calls are ADT utilities and recursion generated by framework converter. I don't think PE can help the latter case since it's conditioned on dynamic value, not easy to handle by PE.

[MarisaKirisame]

@wweic we are planning to use PE to remove the closure/reference of higher order ad as much as possible.
additionally, it help on ADT utilities, as it can do some fusion. for example, I think this can do map f . map g = map (f . g) once it get a termination checker.

[wweic]

overall looks great. Request for comments and clarification.

[wweic]

maybe name it continuation or body to make it clear?

[MarisaKirisame]

@junrushao1994 can you approve?

[wweic]

@MarisaKirisame I modified your test case a bit, instead of creating a ref from const 1, I pass in an argument and create ref from there. I think PE did an incorrect inlining:

In [5]: print(pe_f.astext())
v0.0.1
%17 = fn (%d: bool, %i: int32) -> int32 {
  %0 = ref(%i)
  %16 = {
    let %x: ref(int32) = %0
    %4 = fn () -> () {
      %1 = add(%i, %i) // ty=int32
      %2 = (%x := %1)
      %3 = {
        let %x2: () = %2
        %x2
      }
      %3
    }
    %15 = {
      let %x1: fn () -> () = %4
      %11 = if (%d) {
        %5 = add(%i, %i) // ty=int32
        %6 = (%x := %5)
        %7 = {
          let %x4: () = %6
          %x4
        }
        %7
      } else {
        %8 = add(%i, %i) // ty=int32
        %9 = (%x := %8)
        %10 = {
          let %x5: () = %9
          %x5
        }
        %10
      }
      %14 = {
        let %x3: () = %11
        %12 = %x^
        %13 = {
          let %x6: int32 = %12
          %x6
        }
        %13
      }
      %14
    }
    %15
  }
  %16
}
%17

In [6]: print(f.astext())
v0.0.1
%13 = fn (%d: bool, %i: int32) -> int32 {
  %0 = ref(%i)
  %12 = {
    let %r: ref(int32) = %0
    %5 = fn () -> () {
      %1 = %r^
      %2 = %r^
      %3 = add(%1, %2) // ty=int32
      %4 = (%r := %3)
      %4
    }
    %11 = {
      let %u: fn () -> () = %5
      %8 = if (%d) {
        %6 = %u() // ty=()
        %6
      } else {
        %7 = %u() // ty=()
        %7
      }
      %10 = {
        let %eff: () = %8
        %9 = %r^
        %9
      }
      %10
    }
    %11
  }
  %12
}
%13

It's aggressively inlines the value into the closure, but I think closure only captures the reference, not the value, right?

[MarisaKirisame]

@wweic good catch, that is indeed incorrect. can you give me the offending test? I know how to fix it but wanna add it to the test case.

[wweic]

@MarisaKirisame

def test_if_ref():
    shape = ()
    dtype = 'bool'
    t = relay.TensorType(shape, dtype)
    d = relay.Var("d", t)
    i = relay.Var("i", relay.TensorType(shape, 'int32'))    
    r = relay.Var("r")
    update = relay.Function([], relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)))
    u = relay.Var("u")
    body = relay.If(d, u(), u())
    eff = relay.Var("eff")
    body = relay.Let(eff, body, relay.RefRead(r))
    f = relay.Function([d, i], relay.Let(r, relay.RefCreate(i), relay.Let(u, update, body)))
    f = infer_type(f)
    pe_f = infer_type(partial_eval(f))

[MarisaKirisame]

@wweic I had updated with a regression test to the error.

[MarisaKirisame]

@wweic the test is added.

[wweic]

thanks! @MarisaKirisame

[junrushao]

LGTM :-)

[MarisaKirisame]

@tqchen this is good. can you give some review?

[MarisaKirisame]

@tqchen I address your comment. can you merge?