fix: Cold self-review found the libinfo.py fix duplicated an existing

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<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${lhs_arg}->device.device_id));
 
   status = gemm_op(stream);
   TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);

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")

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",

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<DLTensor*>()`` 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()

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()

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