Add TensorFlow custom op and run tvm in TensorFlow


Now tvm has supported most of TensorFlow op which can to load with tvm.relay. Users can convert TensorFlow op and run with tvm runtime but sometimes we want to run with TensorFlow session with some optimized tvm op.

It is possible to wrap tvm op as TensorFlow custom op so that tvm op can be part of TensorFlow graph. The user interfaces may be like TF-TensorRT which has TensorFlow custom op and graph editor to run optimized TensorFlow model without effort.

How to deploy tvm export model on linux gpu server

TVM runtime is designed on top of DLPack, as long as TensorFlow supports DLPack, we might be able to convert operators back and forth. PyTorch, Chainer and MXNet has been part of DLPack, but I am not sure if TF has a timeline for this…


CC @yzh119 if DGL team share the same interest


Thanks for your reply @junrushao . We have the new project which can wrap TVM Runtime API with TensorFlow custom op without native DLPack support . It is like TF-TRT and end users can run TVM op with TensorFlow graph and session easily.


The implementation looks great. Would you like to do a bit more API design proposal and contribute it back?

In particular, i can see a few places that could be improved:

  • Look into TF’s memory alignment and other status, do zero copy by creating a DLTensor instead of memcpy
  • Have better ways to specify the modules and function inside the module
  • Have a good calling convention.


This is great! Looking forward to the rfc and discussion.


I think it is a proof-of-concept which is worth more attention - the use cases are ubiquitous and we could imagine that people are absolutely interested.

  • to-tvm: TVM already has importers from many frameworks (TF, ONNX, MXNet, etc).
  • from-tvm: method to embed TVM into other systems (TF, PyTorch) is still in lack.

So I would say “from-tvm” is probably worth trying and may have great impact.


Thanks for all you interest. We have updated the code to support better way to load TVM libraries base on @tqchen 's comment. Users can load TVM op just like TensorFlow op without effort. We have scripts of graph editor to replace original TensorFlow op with TVM op as well.

import tensorflow as tf
from tvm_runtime import tvm_runtime

with tf.Session() as sess:
  a = tf.constant([10.1, 20.0, 11.2, -30.3])
  b = tvm_runtime(a, lib_path="", function_name="addone")

It would be great if someone can help us to draft the RFC together and push this to TVM community. The code is open source and any feedback is welcome.