[TASK] TextPrinter Support for Unified TVM IRs

[tqchen]

So far we have a text printer for relay. which allows us to print an IRModule into text format. On the TIR side, we still relies on the ReprPrinter.

This is issue is for upgrading the text printer so that we can print an IRModule that include PrimFunc(tir::Function in the upstream) as a text format. This will help us to enhance the demo experience.

Possible Design Points

• Try to be as consistent as possible with relay (e.g. use %x for variables)
• Think about parsing, but the first attempt does not have to be parsiable
• Think about meta-data usage

Ideally we want to land a version in the mainline in about two to three weeks. @spectrometerHBH please see if it is possible for you and @Hzfengsy to coordinate a format and land a version, then we can pull back to the tensorIR

[spectrometerHBH]

So ideally we want text format & hybrid script to be two different printing methods for TIR which are both parseable?

[tqchen]

yes. but as a path for progression, we can

• Have a text printer that is not parsable
• Support hybrid bidirection
• Support text parser

[spectrometerHBH]

What is hybrid bidirection? Bidirectional translation between text format & hybrid?

[tqchen]

Sorry, i meant hybrid parsing and printing

[Hzfengsy]

Generally, it is ok to use the current ReprPrinter as long as we do some minor mutation. However, there are some details to be determined.

Function Syntax

1. As we know, relay function has its arg and return type.

   fn(%x : Tensor[(10, 10), float32], %y : Tensor[(10, 10), float32])
              -> Tensor[(10, 10), float32] {
       add(%x, %y)
   }
   

   Of course, TIR can do the same thing to show the parameter type. The problem is TIR does not have a return type even the return value. It directly changes the buffer by mutating its element. Here are two options for this:
   • Just make it without a return value. Pros: More natural to low-level code. Cons: It is not so unified in syntax, and also, the vars are no longer immutable in the "unified" syntax.
   • Change the behavior of TIR function to make it has a return value and type. Pros: Unified and make the vars immutable. Cons: It is strange that a low-level code to return a buffer. Somehow a TIR function may have more than one return buffer or even change the input buffer.

I think both are not the perfect solution but I have no other ideas. Would love to listen to your opinions

Annotation

I'd like to add a kind of annotation in function to show whether it is a Relay or TIR function. One possible syntax would be

@relay
fn(%x, %y) { add(%x, %y) }

@tir
fn(%x, %y) { X[0] = Y[0] }

Function Body

There is no doubt that the relay function and TIR function have different stmt, expr. (e.g. Let Binding in relay and BufferLoad in TIR). So I think we can ignore the difference of the function body but try to make the function itself unified.

@spectrometerHBH @tqchen

[tqchen]

• re the signature: we do have signature for TIR functions now, see the implementation of the function signature here: mhttps://github.com/apache/incubator-tvm/blob/master/include/tvm/tir/function.h#L129
  • The default return type of a PrimFunc is void, which is represented by an empty tuple type(as in swift)
  • After the function is lowered to a PackedFunc, the return type could become int, both are indicated by the ret_type
• re annotation: for now let us try to annotate the PrimFunc, and keep relay as default un-annotated(to preserve backward compact), besides the top level annotation, perhaps modifier keyword is another approach, think of inline void func, extern "c" int func, perhaps we could add a keyword(e.g. ) primfunc fn to indicate primfunc

[tqchen]

FYI, I am in the process of refactoring out tensor field(func) in Realize Provide and replace them by BufferRealize BufferStore and BufferLoad. So we don’t need to consider supporting these nodes in the text format

[tqchen]

Related change in the upstream apache/tvm#5372

[spectrometerHBH]

A comparison between repr's format and text format(which is problematic now)

[23:41:44] /home/spectrometer/tvm-upstream/src/te/schedule/bound.cc:128: not in feed graph consumer = compute(B, 0x1fb5cf0)
PrimFunc([A0, A1, C]) attrs={"tir.noalias": (bool)1, "global_symbol": "main"} {
  // attr [B.v0] storage_scope = "global"
  allocate B.v0[float32 * n]
  // attr [B.v1] storage_scope = "global"
  allocate B.v1[float32 * n]
  for (i, 0, m) {
    for (j, 0, n) {
      B.v0[j] = (A0[((i*stride) + (j*stride))] + 2f)
      B.v1[j] = (A0[((i*stride) + (j*stride))]*3f)
    }
    for (j, 0, n) {
      C[((i*stride) + (j*stride))] = (A1[((i*stride) + (j*stride))] + B.v0[j])
    }
  }
}


[23:41:44] /home/spectrometer/tvm-upstream/src/te/schedule/bound.cc:128: not in feed graph consumer = compute(B, 0x1fb5cf0)
class Module:
  PrimFunc(%A0: Buffer([%m, %n], "float32", name=%A01),
           %A1: Buffer([%m, %n], "float32", name=%A11),
           %C: Buffer([%m, %n], "float32", name=%C1)) {
    attr [%B.v0] "storage_scope" = "global" {
      allocate(%B.v0, "float32", [%n]) if "bool"(1) {
        attr [%B.v1] "storage_scope" = "global" {
          allocate(%B.v1, "float32", [%n]) if "bool"(1) {
            for (%i, 0, %m) "serial" {
              for (%j, 0, %n) "serial" {
                %B.v0[%j] = (%A0[((%i*%stride) + (%j*%stride1))] + "float32"(2)) if "bool"(1)
                %B.v1[%j] = (%A0[((%i*%stride) + (%j*%stride1))]*"float32"(3)) if "bool"(1)
              }
              for (%j1, 0, %n) "serial" {
                %C[((%i*%stride2) + (%j1*%stride3))] = (%A1[((%i*%stride4) + (%j1*%stride5))] + %B.v0[%j1]) if "bool"(1)
              }
            }
          }
        }
      }
    }
  }

Several Problems

1. The name of Vars

For readability, the name_hint of var is useful. But B.v0 may not be a legal identifier for parser since it contains .. And the user may give an arbitrary name.

2. dtype of Vars

I'm not very clear where Var is being used in the upstream AST. As long as I can see now

1. used in buffer shape declare
2. loop var
3. stride var(stride above, but I'm not clear what they are)
4. buffer var
5. Iter var(but iter var seems to only appear in attr, otherwise we use the var in iter_var)

When printing a var, the best way is to just give its name_hint, but in principle we have to determine its dtype as well. But in most cases, dtype doesn't matter. I'm not sure about this point.

[tqchen]

• 1. Given that the name_hint is only a hint, we do not have to abide by the rule, we can rewrite . to _ or other symboles
• 2. dtype should be printed during declaration(e.g. type annotation of let, for(where var is declared) ) but not at usage point

[tqchen]

Some additional thoughts after seeing the example.

• When dtype need to be printed, we can always print the Type of the(var), not dtype. Since Type always implies dtype, but not the other way around, we can use https://github.com/apache/incubator-tvm/blob/master/include/tvm/tir/op.h#L55 to obtain the type
  See discussion of dtype vs type here https://github.com/apache/incubator-tvm/blob/master/include/tvm/ir/type.h#L30
  • We can make the type printing consistent with relay
• The for_type("serial") does not need to be printed(because it is the default), we might want to discuss how to print other functions
• Constant literals "float32"(1) can be 1.0f, or float32(0)
• When predicate is const_true(in allocate) do not print the predicate.
• The % prefix seems to be quite dense and confusing when used together with operators(e.g. %), perhaps we want to allow variable names without the prefix?
• Perhaps we should print attr differently (it is a bit strange to have string literal appear on the rhs)

A strawman, I am not too happy with the syntax

attr ("storage_scope", %Bv0) = "global"

[spectrometerHBH]

For vars like %m and %n, should we print their Type/dtype? Their declaration is outside AST.

%A0: Buffer([%m, %n], "float32", name=%A0_1),
%A1: Buffer([%m, %n], "float32", name=%A1_1),
%C: Buffer([%m, %n], "float32", name=%C_1)

One way is to print Type for the first time we encounter the Var

%A0: Buffer([%m : int32, %n : int32], "float32", name=%A0_1),
%A1: Buffer([%m, %n], "float32", name=%A1_1),
%C: Buffer([%m, %n], "float32", name=%C_1)

Another way is to declare first their defs somewhere else.

The third way is to put it in meta.

I think all the ways are somewhat strange.

And for vars like stride

%A0[((%i*%stride) + (%j*%stride1))

It looks like more a placeholder for unknown buffer shape. Does their Type/dtype matter?

[tqchen]

I think the first encouter or declare def somewhere else makes sense. Formally, we should not put buffer_map in the argument list(that makes Buffer a type) but they are not types, we might want to have some seciton like c++'s initializer list (the corresponding hybrid script part is the buffer bind)