Re: [I] [Unity] [Tracking Issue] Heterogeneous execution for Relax [tvm]

["liquanfeng (via GitHub)" <gi...@apache.org>]

liquanfeng commented on issue #15101:
URL: https://github.com/apache/tvm/issues/15101#issuecomment-1779026435

   A very useful job! And I try it on tests/python/relax/test_codegen_cudnn.py::test_conv2d_offload according to tests/python/relax/test_vm_multi_device.py::test_multi_device as shown below
   ```python
   import numpy as np
   
   import tvm
   import tvm.testing
   import tvm.topi.testing
   from tvm import relax
   from tvm.relax.backend.contrib.cudnn import partition_for_cudnn
   from tvm.script import relax as R, ir as I
   
   from tvm.script.ir_builder import IRBuilder
   from tvm.script.ir_builder import relax as relax_builder
   
   data_shape, weight_shape, dtype = (
       (16, 32, 32, 16),
       (32, 3, 3, 16),
       "float32",
   )
   
   input_np = np.random.randn(*data_shape).astype(dtype)
   weight_np = np.random.randn(*weight_shape).astype(dtype)
   
   oc = weight_shape[0]
   bias_np = np.random.randn(1, 1, 1, oc).astype(dtype)
   args = (input_np, weight_np, bias_np)
   
   with IRBuilder() as builder:
       with relax_builder.function():
           R.func_name("main")
           data = R.arg("data", R.Tensor(data_shape, dtype))
           weight = R.arg("weight", R.Tensor(weight_shape, dtype))
           bias = R.arg("bias", R.Tensor((1, 1, 1, weight_shape[0]), dtype))
   
           with R.dataflow() as frame:
               output = R.emit(
                   R.nn.conv2d(
                       data,
                       weight,
                       out_dtype=dtype,
                       padding=(1, 1),
                       data_layout="NHWC",
                       kernel_layout="OHWI",
                   )
               )
               output = R.emit(output + bias)
   
               output = R.emit(relax.op.to_vdevice(output, I.vdevice("llvm")))
               output = R.emit(R.multiply(output, R.const(2, "float32")))
               R.output(output)
   
           R.func_ret_value(frame.output_vars[0])
   
   func = builder.get()
   mod = tvm.IRModule(
       {"main": func},
       global_infos={
           "vdevice": [
               I.vdevice("cuda", 0),
               I.vdevice("llvm"),
           ]
       },
   )
   
   mod = partition_for_cudnn(mod)
   mod = relax.transform.RunCodegen()(mod)
   
   devs = [tvm.device("cuda", 0), tvm.device("llvm")]
   mod = relax.transform.RealizeVDevice()(mod)
   mod = relax.transform.LegalizeOps()(mod)
   mod = tvm.tir.transform.DefaultGPUSchedule()(mod)
   
   with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": False}):
       ex = relax.build(mod)
   vm = relax.VirtualMachine(ex, devs)
   f = vm["main"]
   inputs = [tvm.nd.array(inp, tvm.device("cuda", 0)) for inp in input_np]
   
   print(f(*inputs).numpy())
   ```
   raise following error
   ```
   Traceback (most recent call last):
     File "/workspace/yongwww/tvm/tests/byoc.py", line 77, in <module>
       ex = relax.build(mod)
     File "/workspace/yongwww/tvm/python/tvm/relax/vm_build.py", line 334, in build
       new_mod = lowering_passes(mod)
     File "/workspace/yongwww/tvm/python/tvm/ir/transform.py", line 238, in __call__
       return _ffi_transform_api.RunPass(self, mod)
     File "/workspace/yongwww/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 239, in __call__
       raise_last_ffi_error()
     File "/workspace/yongwww/tvm/python/tvm/_ffi/base.py", line 476, in raise_last_ffi_error
       raise py_err
   tvm._ffi.base.TVMError: Traceback (most recent call last):
     24: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::transform::Pass, tvm::IRModule)>::AssignTypedLambda<tvm::transform::{lambda(tvm::transform::Pass, tvm::IRModule)#7}>(tvm::transform::{lambda(tvm::transform::Pass, tvm::IRModule)#7}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tvm::runtime::TVMRetValue)
     23: tvm::transform::Pass::operator()(tvm::IRModule) const
     22: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
     21: tvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
     20: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
     19: tvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
     18: _ZN3tvm7runtime13PackedFuncObj
     17: tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::relax::transform::CallTIRRewrite()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1}>(tvm::relax::transform::CallTIRRewrite()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1})::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const, tvm::runtime::TVMRetValue) const
     16: tvm::relax::CallTIRMutator::Run()
     15: tvm::relax::ExprMutator::VisitExpr(tvm::RelayExpr const&)
     14: _ZZN3tvm5relax11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
     13: tvm::relax::ExprMutator::VisitExpr_(tvm::relax::FunctionNode const*)
     12: tvm::relax::ExprMutator::VisitWithNewScope(tvm::RelayExpr const&, tvm::runtime::Optional<tvm::runtime::Array<tvm::relax::Var, void> >)
     11: tvm::relax::ExprMutator::VisitExpr(tvm::RelayExpr const&)
     10: _ZZN3tvm5relax11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
     9: tvm::relax::ExprMutator::VisitExpr_(tvm::relax::SeqExprNode const*)
     8: tvm::relax::ExprMutator::VisitBindingBlock(tvm::relax::BindingBlock const&)
     7: tvm::relax::ExprMutator::VisitBindingBlock_(tvm::relax::BindingBlockNode const*)
     6: tvm::relax::ExprMutator::VisitBinding(tvm::relax::Binding const&)
     5: tvm::relax::ExprMutator::VisitBinding_(tvm::relax::VarBindingNode const*)
     4: tvm::relax::ExprMutator::VisitBinding_(tvm::relax::VarBindingNode const*, tvm::relax::CallNode const*)
     3: tvm::relax::ExprMutator::VisitExpr(tvm::RelayExpr const&)
     2: _ZZN3tvm5relax11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
     1: tvm::relax::CallTIRMutator::VisitExpr_(tvm::relax::CallNode const*)
     0: tvm::relax::GetDeviceIndex(tvm::IRModule const&, tvm::VDevice const&)
     File "/workspace/yongwww/tvm/src/relax/transform/utils.h", line 384
   TVMError: The vdevice is not in the ir_module.
   ```
   Is there any problem with byoc?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org