PyTorch: How to pass in different data types to TVM?

#1

How could we compile and run scripted functions which take non-tensor data types (such as int, bool, List[int] etc), in TVM ?

For example

graph(%a.1 : int[],
      %b.1 : Tensor):
  %4 : int = prim::Constant[value=1]()
  %3 : int = prim::Constant[value=2]() 
  %c.1 : Tensor = aten::add(%b.1, %3, %4) 
  %8 : (Tensor, int[]) = prim::TupleConstruct(%c.1, %a.1)
  return (%8)

This is a dummy graph. How do we send int[] to TVM as DLTensor… Similarly, there is a possibility that this graph could take a bool or an int or a float. How do we send those data types to TVM as DLTensors?

#2

While we do support non-float tensor inputs (at least with PyTorch-graph annotations to that effect), I’m not aware of support for non-tensor inputs. And my understanding is that it’s not just the PyTorch interface that would need to cover the prospective conversion but rather that the TVM side doesn’t have much support for non-tensor manipulation (starting with Relay).

Of course, if the evolution of the PyTorch JIT is any indication, the ongoing push for dynamic graphs will eventually lead to Relay and backends gaining support for this. (Or it might have started already and I just missed it.)

At the moment, you could try your luck with converting int[] to an int tensor. For scalars, 0d tensors might work.

Best regards

Thomas

#3

@t-vi: Thanks for the response. I am still curious to understand how these data types are handled for intermediate nodes and why would the input to a graph be different from the input to an intermediate node.

#4

I’m probably not the best person to answer, but so three things:

  • if we think this from the end, I don’t think you can set arguments other than tensors for compiled modules,
  • things like tvm.relay.op.shape_of return a tensor that can then be used as a shape, so it avoids int lists as intermediates,
  • at the compilation level (in particular the dynamic to static), TVM can reason about shapes, so in a range of applications, the shape can be fixed and eliminated (i.e. transformed into attributes to ops that are then lowered). It can not if have “shape” (or int list in general) attributes.

Now I must admit I don’t know if and how we are handling (or planning to handle) truly variable shapes (e.g. detections from maskrcnn), someone more informed will have to jump in. In the discussions I saw here on the forums, there are plans of including data types (like Array, which would be what TorchScript lists correspond to) in the runtime, but I’m not sure it’s there yet.

So to the best of my knowledge, this means

  • your program currently cannot really depend on the non-tensor inputs (this is similar to tracing in the PyTorch JIT),
  • you could try if you can achieve what you want by using integer tensors.

Best regards

Thomas

1 Like