[Relay] Add support for TupleGetItem in op fusion

[masahi]

A partial fix for #2890. It enables fusing through TupleGetItemNode.

The strategy is similar to the TupleNode case, but a special care is needed to handle the reference type. Since the gradient pass generates many TupleNode, TupleGetItemNode and Reference nodes under the hood, and references do not play well with fusion, I simply avoid fusing through reference nodes.

A CPU/GPU execution test is added, as requested.

Example: GRU-ish unit

Before this PR (also see #2890)

fn (%X: Tensor[(1, 10), float32],
    %y: Tensor[(30, 10), float32])
    -> Tensor[(1, 10), float32] {
  %0 = fn(%p0: Tensor[(1, 10), float32],
          %p1: Tensor[(30, 10), float32])
          -> Tensor[(1, 30), float32] {
    %1 = nn.dense(%p0, %p1, units=None) # ty=Tensor[(1, 30), float32]
    %1
  }
  %2 = %0(%X, %y) # ty=Tensor[(1, 30), float32]
  %3 = fn(%p01: Tensor[(1, 30), float32])
          -> Tuple[Tensor[(1, 10), float32], Tensor[(1, 10), float32], Tensor[(1, 10), float32]] {
    %4 = split(%p01, indices_or_sections=int64(3), axis=1) # ty=Tuple[Tensor[(1, 10), float32], Tensor[(1, 10), float32], Tensor[(1, 10), float32]]
    %4
  }
  %5 = %3(%2) # ty=Tuple[Tensor[(1, 10), float32], Tensor[(1, 10), float32], Tensor[(1, 10), float32]]
  %6 = %5.0
  %7 = %5.1
  %8 = %5.2
  %9 = fn(%p02: Tensor[(1, 10), float32],
          %p11: Tensor[(1, 10), float32],
          %p2: Tensor[(1, 10), float32])
          -> Tensor[(1, 10), float32] {
    %10 = sigmoid(%p02) # ty=Tensor[(1, 10), float32]
    %11 = tanh(%p11) # ty=Tensor[(1, 10), float32]
    %12 = exp(%p2) # ty=Tensor[(1, 10), float32]
    %13 = multiply(%11, %12) # ty=Tensor[(1, 10), float32]
    %14 = add(%10, %13) # ty=Tensor[(1, 10), float32]
    %14
  }
  %15 = %9(%6, %7, %8) # ty=Tensor[(1, 10), float32]
  %15
}

With this PR

%14 = fn (%X: Tensor[(1, 1000), float32], %y: Tensor[(3000, 1000), float32]) -> Tensor[(1, 1000), float32] {
  %1 = fn (%p0: Tensor[(1, 1000), float32], %p1: Tensor[(3000, 1000), float32], __dict__=meta[StrMap][0]) -> Tensor[(1, 3000), float32] {
    %0 = nn.dense(%p0, %p1, units=None) // ty=Tensor[(1, 3000), float32]
    %0
  }
  %2 = %1(%X, %y) // ty=Tensor[(1, 3000), float32]
  %12 = fn (%p01: Tensor[(1, 3000), float32], __dict__=meta[StrMap][1]) -> Tensor[(1, 1000), float32] {
    %3 = split(%p01, indices_or_sections=int64(3), axis=1) // ty=(Tensor[(1, 1000), float32], Tensor[(1, 1000), float32], Tensor[(1, 1000), float32])
    %4 = %3.0
    %5 = sigmoid(%4) // ty=Tensor[(1, 1000), float32]
    %6 = %3.1
    %7 = tanh(%6) // ty=Tensor[(1, 1000), float32]
    %8 = %3.2
    %9 = exp(%8) // ty=Tensor[(1, 1000), float32]
    %10 = multiply(%7, %9) // ty=Tensor[(1, 1000), float32]
    %11 = add(%5, %10) // ty=Tensor[(1, 1000), float32]
    %11
  }
  %13 = %12(%2) // ty=Tensor[(1, 1000), float32]
  %13
}

[masahi]

@tqchen @jroesch @MarisaKirisame @vinx13 can you help review?

[MarisaKirisame]

what if there is a reference inside? does it get recursively handled?

[masahi]

Tuple fields are recursively visited via call to ExprVisitor::VisitExpr_(op) below, even if it is a Reference nodes. Ref nodes are never fused with its parent TupleGetItemNode.

[tqchen]

Please add the comment block to explain this so that future developers won't be confused here

[ajtulloch]

Thanks, @masahi, this looks great and is strictly better than the current output.

I do think it would be quite interesting as a follow-up to see what it would look like in terms of perf if we enable fusing the Splits into the GRU, but it's easy to implement a injective GRUGates(..) op and schedule the fused Dense + GRUGate op.

[ajtulloch]

This looks great, thanks so much for this.

[masahi]

@MarisaKirisame this is the IR for higher order gradient of tvm.relay.log after fusion. It is identical with or without this PR.

%62 = fn () -> (Tensor[(10, 4), float32], (Tensor[(10, 4), float32],)) {
  %60 = fn (%x: Tensor[(10, 4), float32]) -> (Tensor[(10, 4), float32], (Tensor[(10, 4), float32],)) {
    %1 = fn () -> () {
      %0 = ()
      %0
    }
    %2 = ref(%1)
    %59 = {
      let %x1: ref(fn () -> ()) = %2
      %4 = fn (%p0: Tensor[(10, 4), float32], __dict__=meta[StrMap][0]) -> Tensor[(10, 4), float32] {
        %3 = zeros_like(%p0)
        %3
      }
      %5 = %4(%x)
      %6 = ref(%5)
      %7 = (%x, %6)
      %58 = {
        let %x2: (Tensor[(10, 4), float32], ref(Tensor[(10, 4), float32])) = %7
        %40 = fn (%x4: (Tensor[(10, 4), float32], ref(Tensor[(10, 4), float32]))) -> (Tensor[(10, 4), float32], ref(Tensor[(10, 4), float32])) {
          %39 = {
            let %x5: (Tensor[(10, 4), float32], ref(Tensor[(10, 4), float32])) = %x4
            %8 = %x5.0
            %10 = fn (%p01: Tensor[(10, 4), float32], __dict__=meta[StrMap][1]) -> Tensor[(10, 4), float32] {
              %9 = log(%p01)
              %9
            }
            %11 = %10(%8)
            %38 = {
              let %x6: Tensor[(10, 4), float32] = %11
              %13 = fn (%p02: Tensor[(10, 4), float32], __dict__=meta[StrMap][2]) -> Tensor[(10, 4), float32] {
                %12 = zeros_like(%p02)
                %12
              }
              %14 = %13(%x6)
              %15 = ref(%14)
              %37 = {
                let %x7: ref(Tensor[(10, 4), float32]) = %15
                %16 = %x1^
                %36 = {
                  let %x8: fn () -> () = %16
                  %32 = fn () -> () {
                    %17 = %x7^
                    %31 = {
                      let %x10: Tensor[(10, 4), float32] = %17
                      %18 = %x5.1
                      %19 = %18^
                      %30 = {
                        let %x11: Tensor[(10, 4), float32] = %19
                        %20 = %x5.1
                        %25 = fn (%p03: Tensor[(10, 4), float32], %p1: Tensor[(10, 4), float32], %p2: Tensor[(10, 4), float32], __dict__=meta[StrMap][3]) -> Tensor[(10, 4), float32] {
                          %21 = ones_like(%p2)
                          %22 = multiply(%p1, %21)
                          %23 = divide(%22, %p2)
                          %24 = add(%p03, %23)
                          %24
                        }
                        %26 = %25(%x11, %x10, %8)
                        %27 = (%20 := %26)
                        %29 = {
                          let %x12: () = %27
                          %28 = %x8() // ty=()
                          %28
                        }
                        %29
                      }
                      %30
                    }
                    %31
                  }
                  %33 = (%x1 := %32)
                  %35 = {
                    let %x9: () = %33
                    %34 = (%x6, %x7)
                    %34
                  }
                  %35
                }
                %36
              }
              %37
            }
            %38
          }
          %39
        }
        %41 = %40(%x2)
        %57 = {
          let %x3: (Tensor[(10, 4), float32], ref(Tensor[(10, 4), float32])) = %41
          %42 = %x3.1
          %43 = %x3.0
          %45 = fn (%p04: Tensor[(10, 4), float32], __dict__=meta[StrMap][4]) -> Tensor[(10, 4), float32] {
            %44 = ones_like(%p04)
            %44
          }
          %46 = %45(%43)
          %47 = (%42 := %46)
          %56 = {
            let %x13: () = %47
            %48 = %x1^
            %49 = %48() // ty=()
            %55 = {
              let %x14: () = %49
              %50 = %x3.0
              %51 = %x2.1
              %52 = %51^
              %53 = (%52,)
              %54 = (%50, %53)
              %54
            }
            %55
          }
          %56
        }
        %57
      }
      %58
    }
    %59
  }
  %61 = %60(meta[relay.Constant][0] // ty=Tensor[(10, 4), float32])
  %61
}
%62

[masahi]

@ajtulloch yes, fusing dense with split + other ops should be possible in principle, but our fusion algorithm is implemented such that dense and other "complex" ops (conv, pooling etc.) can only be fused with strictly elemwise ops. Relaxing this restriction requires major rethinking and rewrite of the algorithm.We can discuss this further in #2890

[ajtulloch]

Yeah, no problem - this is great. Thank you for this :)

[tqchen]

Please add the comment block to explain this so that future developers won't be confused here

[masahi]

@MarisaKirisame @tqchen I added a comment block on reference handling.