Closed Cookiee235 closed 1 day ago
@tqchen @Hzfengsy @junrushao
The function tir_zeros
is defined in the Module
explicitly. Why the compilation crash and throw: Cannot find PackedFunc tir_zeros
? I'm very confused. Look forward to your response. Many thanks!
I was able to reproduce the error, and it looks like it's an incompatibility between the AttachGlobalSymbols
pass and the relax.VMCodeGen
step. When the "system_lib_prefix"
is included, the "global_symbol"
attribute produced by AttachGlobalSymbols
doesn't match the name of the GlobalVar
. As a result, the compiled TIR module provides a function named "hello_tir_zeros"
, but the Relax VM attempts to find a function named "tir_zeros"
.
This would be caught by the well-formed checker, which requires the "global_symbol"
to exactly match the GlobalVar
value if present. However, this is only validated for Relax functions, not for TIR functions (implementation here).
To resolve this, and to prevent similar bugs from occurring in the future, I think we should make the following changes:
1) When AttachGlobalSymbol
changes the name of a function, it must also replace the GlobalVar
used to refer to that function. All occurrences of the old GlobalVar
across the IRModule
must then be replaced by the new GlobalVar
.
2) The well-formed checker should validate that the GlobalVar
and "global_symbol"
match across all functions, not just relax functions.
@Cookiee235 Can you test with #17202 applied? With this PR, I am able to successfully run your test case.
@Cookiee235 Can you test with #17202 applied? With this PR, I am able to successfully run your test case.
Hi @Lunderberg Thanks for your efforts. Following your advice, I pull the source code in your repository (i.e., Lunderberg:replace_global_var_on_rename), and build TVM from source code. However, it built failed. Part of the crash message is as follows
In file included from /software/tvm-lunder/src/runtime/builtin_fp16.cc:25:
/software/tvm-lunder/include/tvm/runtime/c_runtime_api.h:144:15: error: 'kDLROCMHost' was not declared in this scope; did you mean 'kDLROCM'?
144 | static_assert(kDLROCMHost == 11, TVM_HARCODED_INTEGER_CHANGED_MSG);
| ^~~~~~~~~~~
| kDLROCM
...
make[2]: *** [CMakeFiles/tvm_libinfo_objs.dir/build.make:76: CMakeFiles/tvm_libinfo_objs.dir/src/support/libinfo.cc.o] Error 1
make[1]: *** [CMakeFiles/Makefile2:905: CMakeFiles/tvm_libinfo_objs.dir/all] Error 2
Hmm. That looks like it isn't able to find any symbols from dlpack.h
. Did you checkout the submodules? (git submodule update --init --recursive
)
Hmm. That looks like it isn't able to find any symbols from
dlpack.h
. Did you checkout the submodules? (git submodule update --init --recursive
)
So sorry, I forgot. Indeed, I also successfully ran this test under the PR https://github.com/apache/tvm/pull/17202 applied! @Lunderberg Thank you again for your PR!
@Lunderberg @Cookiee235 hi, When I try to use cudnn dispatch conv2d, I encounter the same problem as you during runtime. After replacing with the code changes you submitted, I still have the same problem. Perhaps our issues are similar but not exactly the same Could you please give me some suggestions?
1、Actual behavior:
vm::runtime::Optional<tvm::runtime::Session> const&, bool&>(tvm::runtime::String&, tvm::runtime::String&, picojson::object_with_ordered_keys const&, DLDevice&, tvm::runtime::Optional<tvm::runtime::Session> const&, bool&) at mlc-llm/3rdparty/tvm/include/tvm/runtime/memory.h:196 5: tvm::runtime::ObjectPtr<mlc::llm::serve::ModelImpl> tvm::runtime::ObjAllocatorBase<tvm::runtime::SimpleObjAllocator>::make_object<mlc::llm::serve::ModelImpl, tvm::runtime::String&, tvm::runtime::String&, picojson::object_with_ordered_keys const&, DLDevice&, tvm::runtime::Optional<tvm::runtime::Session> const&, bool&>(tvm::runtime::String&, tvm::runtime::String&, picojson::object_with_ordered_keys const&, DLDevice&, tvm::runtime::Optional<tvm::runtime::Session> const&, bool&) at mlc-llm/3rdparty/tvm/include/tvm/runtime/memory.h:72 4: mlc::llm::serve::ModelImpl* tvm::runtime::SimpleObjAllocator::Handler<mlc::llm::serve::ModelImpl>::New<tvm::runtime::String&, tvm::runtime::String&, picojson::object_with_ordered_keys const&, DLDevice&, tvm::runtime::Optional<tvm::runtime::Session> const&, bool&>(tvm::runtime::SimpleObjAllocator*, tvm::runtime::String&, tvm::runtime::String&, picojson::object_with_ordered_keys const&, DLDevice&, tvm::runtime::Optional<tvm::runtime::Session> const&, bool&) at mlc-llm/3rdparty/tvm/include/tvm/runtime/memory.h:122 3: mlc::llm::serve::ModelImpl::ModelImpl(tvm::runtime::String, tvm::runtime::String, picojson::object_with_ordered_keys, DLDevice, tvm::runtime::Optional<tvm::runtime::Session> const&, bool) at /mlc-llm/cpp/serve/model.cc:66 2: mlc::llm::serve::FunctionTable::Init(tvm::runtime::String, DLDevice, picojson::object_with_ordered_keys, tvm::runtime::Optional<tvm::runtime::Session>) at /mlc-llm/cpp/serve/function_table.cc:133 1: tvm::runtime::relax_vm::VirtualMachineImpl::_Init(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) 0: tvm::runtime::relax_vm::VirtualMachineImpl::InitFuncPool() File "/mlc-llm/3rdparty/tvm/src/runtime/relax_vm/vm.cc", line 707 InternalError:
Check failed: (func.defined()) is false: Error: Cannot find PackedFunc fused_relax_nn_conv2d_cudnn in either Relax VM kernel library, or in TVM runtime PackedFunc registry, or in global Relax functions of the VM executable
2、But when I run the test case of TVM's cudnn conv2d separately, it works normally, and I can find 'fused_relax_nn_conv2d_cudnn'.
3、I printed out the mod before and after the cudnn conv2d patten processing, and the names are all the same with tvm cudnn test, such as fused_relax_nn_conv2d_cudnn."
@jpf888 This sounds like a similar bug, but would depend on how the dispatch is implemented. When the pattern-matching is applied, is the call to the fused kernel generated as module.fused_function(args)
, or as R.call_extern("fused_function", args)
? The fix applied in #17202 would only apply to calls that are known to be within the same IRModule (the module.fused_function(args)
version), and not to calls that may be defined outside of the IRModule (the R.call_extern
version).
@Lunderberg 1、When I apply dispatch in mlcllm, this issue occurs,and When the pattern-matching is applied, is the call to the fused kernel generated as “R.call_dps_packed("fused_relax_nn_conv2d_cudnn", args) of the class Module“ ,
2、When in the TVM test case, it works fine, log : before pattern :
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def main(data: R.Tensor((16, 32, 32, 16), dtype="float16"), weight: R.Tensor((32, 3, 3, 16), dtype="float16")) -> R.Tensor((16, 32, 32, 32), dtype="float16"):
with R.dataflow():
lv: R.Tensor((16, 32, 32, 32), dtype="float16") = R.nn.conv2d(data, weight, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout="NHWC", kernel_layout="OHWI", out_layout="NHWC", out_dtype="float16")
R.output(lv)
return lv
after pattern:
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def fused_relax_nn_conv2d_cudnn(data: R.Tensor((16, 32, 32, 16), dtype="float16"), weight: R.Tensor((32, 3, 3, 16), dtype="float16")) -> R.Tensor((16, 32, 32, 32), dtype="float16"):
R.func_attr({"Codegen": "cudnn"})
# from tvm.script import relax as R
@R.function
def local_func(data_1: R.Tensor((16, 32, 32, 16), dtype="float16"), weight_1: R.Tensor((32, 3, 3, 16), dtype="float16")) -> R.Tensor((16, 32, 32, 32), dtype="float16"):
R.func_attr({"Composite": "cudnn.conv2d.nhwc_ohwi"})
with R.dataflow():
gv: R.Tensor((16, 32, 32, 32), dtype="float16") = R.nn.conv2d(data_1, weight_1, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout="NHWC", kernel_layout="OHWI", out_layout="NHWC", out_dtype="float16")
R.output(gv)
return gv
output: R.Tensor((16, 32, 32, 32), dtype="float16") = local_func(data, weight)
return output
@R.function
def main(data: R.Tensor((16, 32, 32, 16), dtype="float16"), weight: R.Tensor((32, 3, 3, 16), dtype="float16")) -> R.Tensor((16, 32, 32, 32), dtype="float16"):
cls = Module
with R.dataflow():
gv: R.Tensor((16, 32, 32, 32), dtype="float16") = cls.fused_relax_nn_conv2d_cudnn(data, weight)
R.output(gv)
return gv
mlc-llm cudnn dispatch
import tvm
import tvm.relax.backend.contrib.cudnn as _cudnn
from tvm.relax.backend.contrib.cudnn import partition_for_cudnn
from tvm import IRModule, relax
from tvm.relax.backend import get_patterns_with_prefix
@tvm.transform.module_pass(opt_level=0, name="CudnnDispatch")
class CudnnDispatch: # pylint: disable=too-few-public-methods,broad-exception-raised
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
has_cudnn = tvm.get_global_func("relax.ext.cudnn", True)
if not has_cudnn:
raise Exception("CUDNN is not enabled.")
patterns = get_patterns_with_prefix("cudnn")
model_names = [
gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function)
]
# exclude single batch decode
model_names = [name for name in model_names if "image_embed" in name]
mod = tvm.transform.Sequential(
[
relax.transform.FuseOpsByPattern(
patterns,
bind_constants=False,
annotate_codegen=True,
entry_functions=model_names,
),
relax.transform.RunCodegen({}, entry_functions=model_names),
]
)(mod)
return mod
@Lunderberg I resolved the issue by adding multiple dispatch to the mlcllm compile pipeline. The problem consistently occurred with the first dispatch(cudnn), but leaving just one dispatch, any dispatch, it worked !
Actual behavior
Environment
Steps to reproduce
cc @junrushao @Hzfengsy