From 911152c1af57590541f257b1e5aec180a6288109 Mon Sep 17 00:00:00 2001 From: Amitesh Gupta Date: Mon, 20 Jul 2026 09:32:44 +0530 Subject: [PATCH] fix: Cold self-review found the libinfo.py fix duplicated an existing Addresses #20031. --- python/tvm/contrib/cutlass/gemm_operation.py | 2 +- python/tvm/contrib/cutlass/gen_tensor_op.py | 6 ++- python/tvm/support/libinfo.py | 26 ++++++++++- tests/python/contrib/test_cutlass_gemm.py | 46 ++++++++++++++++++++ tests/python/support/test_libinfo.py | 43 ++++++++++++++++++ tests/python/testing/test_env.py | 17 ++++++++ 6 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tests/python/support/test_libinfo.py diff --git a/python/tvm/contrib/cutlass/gemm_operation.py b/python/tvm/contrib/cutlass/gemm_operation.py index 372f26943678..164856e51da0 100644 --- a/python/tvm/contrib/cutlass/gemm_operation.py +++ b/python/tvm/contrib/cutlass/gemm_operation.py @@ -347,7 +347,7 @@ def instantiate_gemm_template(attrs): status = gemm_op.initialize(arguments, workspace.get()); TVM_FFI_ICHECK(status == cutlass::Status::kSuccess); - cudaStream_t stream = static_cast(TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id)); + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(kDLCUDA, ${lhs_arg}->device.device_id)); status = gemm_op(stream); TVM_FFI_ICHECK(status == cutlass::Status::kSuccess); diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py index 477c1ee44953..29a2d98d3649 100644 --- a/python/tvm/contrib/cutlass/gen_tensor_op.py +++ b/python/tvm/contrib/cutlass/gen_tensor_op.py @@ -490,7 +490,11 @@ def instantiate_template(func_name, annotations, func_args): if k in annotations: attrs[k] = annotations[k] - headers = ["tvm/ffi/function.h", "tvm/ffi/extra/c_env_api.h"] + headers = [ + "tvm/ffi/function.h", + "tvm/ffi/extra/c_env_api.h", + "tvm/ffi/container/tensor.h", + ] if "relu" in func_name: headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h") diff --git a/python/tvm/support/libinfo.py b/python/tvm/support/libinfo.py index 5d461236b30d..8be3408d47ce 100644 --- a/python/tvm/support/libinfo.py +++ b/python/tvm/support/libinfo.py @@ -19,12 +19,29 @@ import os +def _native_codegen_registered(global_func_name: str) -> bool: + """Return whether a native global function used only by one optional build is registered. + + Some CMake flags (e.g. ``USE_CUTLASS``) gate compilation of the C++ file + that registers a given global function, so the function's presence is a + direct signal that the corresponding flag was enabled at build time. + """ + try: + import tvm_ffi # pylint: disable=import-outside-toplevel + + return tvm_ffi.get_global_func(global_func_name, allow_missing=True) is not None + except Exception: # pylint: disable=broad-except + return False + + def libinfo(): """Returns a dictionary of compile-time info — minimal Python fallback. The native ``support.GetLibInfo`` global function is no longer registered after the upstream sync, so we synthesize the values from build-time hints - instead. + instead. Where an optional component's registration code is only + compiled in under a specific CMake flag, we probe for one of its global + functions rather than guessing. """ return { "USE_CUDA": os.environ.get("TVM_USE_CUDA", "ON"), @@ -34,7 +51,12 @@ def libinfo(): "USE_NVSHMEM": os.environ.get("TVM_USE_NVSHMEM", "OFF"), "USE_HEXAGON": "OFF", "USE_CUDNN": "OFF", - "USE_CUTLASS": "OFF", + # "relax.ext.cutlass" is the same global function tvm.contrib.cutlass + # .has_cutlass() checks; it is registered by + # src/relax/backend/contrib/cutlass/codegen.cc, which is only + # compiled when USE_CUDA AND USE_CUTLASS are both ON (see + # cmake/modules/contrib/CUTLASS.cmake). + "USE_CUTLASS": "ON" if _native_codegen_registered("relax.ext.cutlass") else "OFF", "USE_VULKAN": "OFF", "USE_OPENCL": "OFF", "USE_METAL": "OFF", diff --git a/tests/python/contrib/test_cutlass_gemm.py b/tests/python/contrib/test_cutlass_gemm.py index e2d9162c064c..f4a0716251fd 100644 --- a/tests/python/contrib/test_cutlass_gemm.py +++ b/tests/python/contrib/test_cutlass_gemm.py @@ -395,5 +395,51 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +def test_instantiate_matmul_template_regular_gemm(monkeypatch): + """Host-side codegen for a regular (non-decode) float16 2D matmul. + + This does not require a CUTLASS-enabled build: it exercises the pure + Python template-instantiation logic in ``gen_tensor_op.py`` directly, so + it always runs and catches source-generation regressions even on builds + that skip the rest of this file for lacking CUTLASS. + + Regression test for two bugs: the regular GEMM template referenced the + undefined substitution variable ``${A_arg}`` (the attribute map for this + path only provides ``${lhs_arg}``/``${rhs_arg}``, so ``${A_arg}`` was + left unsubstituted in the generated source), and the generated headers + omitted ``tvm/ffi/container/tensor.h``, which is required for + ``AnyView::cast()`` to compile in the generated wrapper. + """ + from tvm.contrib.cutlass import gen_tensor_op + + class _FakeCodegenResult: + def __init__(self, code, headers): + self.code = code + self.headers = headers + + monkeypatch.setattr(gen_tensor_op, "CodegenResult", _FakeCodegenResult) + + annotations = { + "arg0_shape": [16, 32], + "arg1_shape": [32, 64], + "arg0_dtype": "float16", + "arg1_dtype": "float16", + "ret_dtype": "float16", + "lda": "K", + "ldb": "N", + "ldc": "N", + "cutlass_op_def": "using Operation_test_op = int;", + "cutlass_op_name": "test_op", + "op_type": "cutlass.matmul", + } + result = gen_tensor_op.instantiate_template("cutlass.matmul", annotations, ["arg0", "arg1"]) + + assert "${" not in result.code, ( + f"unsubstituted template variable in generated CUTLASS source:\n{result.code}" + ) + assert "arg0->device.device_id" in result.code + assert "tvm/ffi/container/tensor.h" in result.headers + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/support/test_libinfo.py b/tests/python/support/test_libinfo.py new file mode 100644 index 000000000000..b3bbf358128c --- /dev/null +++ b/tests/python/support/test_libinfo.py @@ -0,0 +1,43 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for tvm.support.libinfo.""" + +import tvm +import tvm.testing +from tvm.contrib.cutlass import has_cutlass + + +def test_libinfo_use_cutlass_matches_codegen_registration(): + """USE_CUTLASS mirrors whether the native CUTLASS codegen is compiled in. + + Regression test: ``libinfo()`` used to hardcode USE_CUTLASS to "OFF", so + ``build_flag_enabled("USE_CUTLASS")`` always reported False -- even on + builds configured with ``USE_CUTLASS=ON`` -- silently skipping the + CUTLASS test suite. ``has_cutlass()`` probes the same "relax.ext.cutlass" + global function that ``libinfo()`` now checks; that function is + registered by ``src/relax/backend/contrib/cutlass/codegen.cc``, which + CMake only compiles under ``USE_CUTLASS=ON`` (see + ``cmake/modules/contrib/CUTLASS.cmake``), so its presence at runtime is a + direct, always-accurate signal of the flag -- unlike an environment + variable that nothing in the build sets automatically. + """ + expected = "ON" if has_cutlass() else "OFF" + assert tvm.support.libinfo()["USE_CUTLASS"] == expected + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/testing/test_env.py b/tests/python/testing/test_env.py index 6fc7697df07d..ba70fc00c26c 100644 --- a/tests/python/testing/test_env.py +++ b/tests/python/testing/test_env.py @@ -77,6 +77,23 @@ def test_has_gpu_is_raw_any_device(): assert env.has_gpu() == any_device +def test_build_flag_enabled_matches_use_cutlass_registration(): + """build_flag_enabled("USE_CUTLASS") mirrors the native codegen registration. + + Regression test: this used to always report False because + ``tvm.support.libinfo()`` hardcoded USE_CUTLASS to "OFF", silently + skipping the CUTLASS test suite even on builds configured with + USE_CUTLASS=ON. + """ + from tvm.contrib.cutlass import has_cutlass # pylint: disable=import-outside-toplevel + + env.build_flag_enabled.cache_clear() + try: + assert env.build_flag_enabled("USE_CUTLASS") == has_cutlass() + finally: + env.build_flag_enabled.cache_clear() + + def test_target_enabled_respects_tvm_test_targets(monkeypatch): """A device kind excluded from TVM_TEST_TARGETS is reported as not enabled.""" env._target_enabled.cache_clear() # pylint: disable=protected-access