Adds external library route to existing TOPI ops

AtuNuka · April 3, 2018, 12:38am

I would like to add new routes to the existing TOPI ops for calling functions implemented in external library (ex. cudnn), in order to compare the performance between tvm.compute and external library. For example, ReLU and Softmax are only available for tvm.compute not for cuDNN or other external libraries.

I know that there is already for realizing this route. (ex. Conv2D ops)

github.com/apache/tvm

topi/python/topi/cuda/conv2d.py

dccc1e3fc


      
          if "cudnn" in target.libs:
              assert layout != 'HWCN', "HWCN layout not supported with CUDNN."
              tensor_format = 0 # CUDNN_TENSOR_NCHW
              if layout == 'NHWC':
                  tensor_format = 1 # CUDNN_TENSOR_NHWC
              return cudnn.conv2d_forward(data,
                                          kernel,
                                          stride_h,
                                          stride_w,
                                          pad_h,
                                          pad_w,
                                          1,  # dilation_h
                                          1,  # dilation_w
                                          conv_mode=1,
                                          tensor_format=tensor_format,
                                          algo=-1) # let CUDNN choose the best algo

I recently noticed that TOPI is now implemented with C++ and Python.

Which language should I implement the external library route? (C++ or Python) And is this fix contribution welcomed?

Thanks

tqchen · April 3, 2018, 8:57pm

for conv2ds, we are still using the python part. I think it is always good to start from python route and eventually add c++ support when they are stable

kevinthesun · April 3, 2018, 11:47pm

Python is good for quick prototyping. With nnvm register function, you can quickly implement your customized compute/schedule for any existing operator and register them with higher level. Then simply importing your registration module will override the default implementation.

AtuNuka · April 4, 2018, 8:17am

@tqchen @kevinthesun

I could make sense the relationship between Python and C++ implementation.
I will start implement the operator with Python at first.

Thanks for your fast reply and great advice.