tilelang.contrib.dlpack¶

Wrapping functions to bridge frameworks with DLPack support to TVM

Functions¶

`convert_func`(tvm_func, tensor_type, to_dlpack_func)	Convert a tvm function into one that accepts a tensor from another
`to_pytorch_func`(tvm_func)	Convert a tvm function into one that accepts PyTorch tensors

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