tilelang.contrib.dlpack module#

Wrapping functions to bridge frameworks with DLPack support to TVM

tilelang.contrib.dlpack.convert_func(tvm_func, tensor_type, to_dlpack_func)#
Convert a tvm function into one that accepts a tensor from another

framework, provided the other framework supports DLPACK

Parameters:

  • tvm_func (Function) – Built tvm function operating on arrays

  • tensor_type (Type) – Type of the tensors of the target framework

  • to_dlpack_func (Function) – Function to convert the source tensors to DLPACK

tilelang.contrib.dlpack.to_pytorch_func(tvm_func)#

Convert a tvm function into one that accepts PyTorch tensors

Parameters:

tvm_func (Function) – Built tvm function operating on arrays

Returns:

wrapped_func – Wrapped tvm function that operates on PyTorch tensors

Return type:

Function