完善 MACA 后端适配：CMake 校验、架构传递、contrib 入口与 Relax 外部库入口

cmake/modules/MACA.cmake

@@ -15,39 +15,47 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# update 3rdparty/tvm-ffi for adding kDLMACA/kDLMACAHost
-set(dlpack_header "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/3rdparty/dlpack/include/dlpack/dlpack.h")
-file(READ "${dlpack_header}" FILE_CONTENTS)
-if(NOT FILE_CONTENTS MATCHES ".*kDLMACA.*")
-  string(REPLACE "} DLDeviceType;" "  /*! \\brief MetaX MACA GPU */\n  kDLMACA = 22,\n  /*! \\brief Pinned MetaX MACA host memory */\n  kDLMACAHost = 23,\n} DLDeviceType;" NEW_CONTENTS "${FILE_CONTENTS}")
-  file(WRITE "${dlpack_header}" "${NEW_CONTENTS}")
-endif()
-set(ffi_core_pyi "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/core.pyi")
-file(READ "${ffi_core_pyi}" FILE_CONTENTS)
-if(NOT FILE_CONTENTS MATCHES ".*kDLMACA.*")
-  string(REPLACE "kDLTrn = 17" "kDLTrn = 17\n    kDLMACA = 22\n    kDLMACAHost = 23" NEW_CONTENTS "${FILE_CONTENTS}")
-  file(WRITE "${ffi_core_pyi}" "${NEW_CONTENTS}")
-endif()
-set(ffi_container_tensor "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/include/tvm/ffi/container/tensor.h")
-file(READ "${ffi_container_tensor}" FILE_CONTENTS)
-if(NOT FILE_CONTENTS MATCHES ".*kDLMACA.*")
-  string(REPLACE "device.device_type == kDLROCMHost;" "device.device_type == kDLROCMHost ||\n         device.device_type == kDLMACA || device.device_type == kDLMACAHost;" NEW_CONTENTS "${FILE_CONTENTS}")
-  file(WRITE "${ffi_container_tensor}" "${NEW_CONTENTS}")
-endif()
-set(ffi_cython_base_pxi "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/cython/base.pxi")
-file(READ "${ffi_cython_base_pxi}" FILE_CONTENTS)
-if(NOT FILE_CONTENTS MATCHES ".*kDLMACA.*")
-  string(REPLACE "kDLTrn = 18" "kDLTrn = 18\n        kDLMACA = 22\n        kDLMACAHost = 23" NEW_CONTENTS "${FILE_CONTENTS}")
-  file(WRITE "${ffi_cython_base_pxi}" "${NEW_CONTENTS}")
-endif()
-set(ffi_cython_device_pxi "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/cython/device.pxi")
-file(READ "${ffi_cython_device_pxi}" FILE_CONTENTS)
-if(NOT FILE_CONTENTS MATCHES ".*kDLMACA.*")
-  string(REPLACE "kDLTrn = 18" "kDLTrn = 18\n    kDLMACA = 22\n    kDLMACAHost = 23" NEW_CONTENTS "${FILE_CONTENTS}")
-  string(REPLACE "DLDeviceType.kDLTrn: \"trn\"," "DLDeviceType.kDLTrn: \"trn\",\n      DLDeviceType.kDLMACA: \"maca\",\n      DLDeviceType.kDLMACAHost: \"maca_host\"," NEW_CONTENTS "${NEW_CONTENTS}")
-  string(REPLACE "\"trn\": DLDeviceType.kDLTrn," "\"trn\": DLDeviceType.kDLTrn,\n        \"maca\": DLDeviceType.kDLMACA," NEW_CONTENTS "${NEW_CONTENTS}")
-  file(WRITE "${ffi_cython_device_pxi}" "${NEW_CONTENTS}")
-endif()
+# MACA support depends on DLPack/tvm-ffi knowing the vendor device types.
+# Keep this as an explicit source-level patch instead of mutating submodules
+# during CMake configure.
+function(tvm_maca_check_required_symbol file_path required_symbol)
+  if(NOT EXISTS "${file_path}")
+    message(FATAL_ERROR
+      "MACA support requires ${file_path}, but the file was not found. "
+      "Initialize submodules before configuring TVM.")
+  endif()
+
+  file(READ "${file_path}" _tvm_maca_file_contents)
+  if(NOT _tvm_maca_file_contents MATCHES "${required_symbol}")
+    message(FATAL_ERROR
+      "MACA support requires ${file_path} to define ${required_symbol}. "
+      "Apply the MACA DLPack/tvm-ffi device-type patch before configuring TVM.")
+  endif()
+endfunction()
+
+function(tvm_maca_check_required_symbols)
+  set(_tvm_maca_dlpack_header
+      "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/3rdparty/dlpack/include/dlpack/dlpack.h")
+  set(_tvm_maca_ffi_core_pyi
+      "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/core.pyi")
+  set(_tvm_maca_ffi_container_tensor
+      "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/include/tvm/ffi/container/tensor.h")
+  set(_tvm_maca_ffi_cython_base_pxi
+      "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/cython/base.pxi")
+  set(_tvm_maca_ffi_cython_device_pxi
+      "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-ffi/python/tvm_ffi/cython/device.pxi")
+
+  tvm_maca_check_required_symbol("${_tvm_maca_dlpack_header}" "kDLMACA")
+  tvm_maca_check_required_symbol("${_tvm_maca_dlpack_header}" "kDLMACAHost")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_core_pyi}" "kDLMACA")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_core_pyi}" "kDLMACAHost")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_container_tensor}" "kDLMACA")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_container_tensor}" "kDLMACAHost")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_cython_base_pxi}" "kDLMACA")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_cython_base_pxi}" "kDLMACAHost")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_cython_device_pxi}" "kDLMACA")
+  tvm_maca_check_required_symbol("${_tvm_maca_ffi_cython_device_pxi}" "kDLMACAHost")
+endfunction()
 
 # MACA Module
 find_maca(${USE_MACA})
@@ -64,6 +72,7 @@ if(USE_MACA)
   if(NOT MACA_FOUND)
     message(FATAL_ERROR "Cannot find MACA, USE_MACA=" ${USE_MACA})
   endif()
+  tvm_maca_check_required_symbols()
   message(STATUS "Build with MACA support")
   tvm_file_glob(GLOB RUNTIME_MACA_SRCS src/backend/maca/runtime/*.cc)
   set(_maca_libs ${MACA_MACAMCC_LIBRARY})

python/tvm/contrib/mxcc.py

@@ -0,0 +1,25 @@
+# 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.
+"""Compatibility wrapper for the MACA compiler helper."""
+
+from tvm.support import mxcc as _mxcc
+from tvm.support.mxcc import *  # pylint: disable=wildcard-import,unused-wildcard-import
+
+
+def get_maca_arch(maca_path="/opt/maca"):
+    """Return the MACA architecture using the support module callback."""
+    return _mxcc.get_maca_arch(maca_path)

python/tvm/relax/backend/contrib/maca.py

@@ -0,0 +1,60 @@
+# 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.
+"""Relax entry points for MACA external library offload."""
+
+from typing import Optional, Sequence, Tuple
+
+from tvm.ir import IRModule
+
+_SUPPORTED_LIBRARIES = ("mcblas", "mcdnn", "mcfft")
+
+
+def supported_libraries() -> Tuple[str, ...]:
+    """Return MACA external libraries known by this frontend."""
+    return _SUPPORTED_LIBRARIES
+
+
+def _normalize_library(library: str) -> str:
+    library = library.lower()
+    if library not in _SUPPORTED_LIBRARIES:
+        supported = ", ".join(_SUPPORTED_LIBRARIES)
+        raise ValueError(f"Unknown MACA external library '{library}'. Supported libraries: {supported}")
+    return library
+
+
+def is_available(library: Optional[str] = None) -> bool:
+    """Return whether MACA Relax external library offload is available."""
+    if library is not None:
+        _normalize_library(library)
+    return False
+
+
+def partition_for_maca(mod: IRModule, libraries: Optional[Sequence[str]] = None) -> IRModule:
+    """Partition a Relax module for MACA external libraries.
+
+    The public entry point is present so callers can reliably detect the feature.
+    The actual mcBLAS/mcDNN/mcFFT codegen and runtime integration are not present
+    in this tree yet, so invoking partitioning fails explicitly instead of later
+    surfacing as a missing module, missing global function, or silent no-op.
+    """
+    if libraries is not None:
+        for library in libraries:
+            _normalize_library(library)
+    raise NotImplementedError(
+        "MACA Relax external library offload is not implemented. "
+        "Expected mcBLAS/mcDNN/mcFFT codegen and runtime integration are absent."
+    )

python/tvm/support/mxcc.py

@@ -28,6 +28,50 @@
 from tvm.support import utils
 
 
+def _normalize_maca_arch(arch):
+    """Normalize MACA architecture names to mxcc offload-arch form."""
+    if arch is None:
+        return None
+    arch = str(arch).strip().lower()
+    if not arch:
+        return None
+    if arch.startswith("xcore"):
+        return arch
+    if re.fullmatch(r"\d+[a-zA-Z]*", arch):
+        return "xcore" + arch
+    raise ValueError(f"Invalid MACA architecture: {arch}")
+
+
+def _maca_compute_version_from_arch(arch):
+    """Convert xcore architecture names to compute version strings."""
+    arch = _normalize_maca_arch(arch)
+    if arch is None:
+        return None
+    match = re.match(r"xcore(\d+)", arch)
+    if not match:
+        raise ValueError(f"Invalid MACA architecture: {arch}")
+    digits = match.group(1)
+    if len(digits) >= 4:
+        major = int(digits[:-2])
+        minor = int(digits[-2:])
+    elif len(digits) == 3:
+        major = int(digits[:1])
+        minor = int(digits[1:])
+    else:
+        major = int(digits)
+        minor = 0
+    return f"{major}.{minor}"
+
+
+def _target_maca_arch(target):
+    """Extract a normalized MACA architecture from a Target object."""
+    target = target or tvm.target.Target.current()
+    if target is None:
+        return None
+    mcpu = getattr(target, "mcpu", None)
+    return _normalize_maca_arch(mcpu)
+
+
 def compile_maca(code, target_format="mcbin", arch=None, options=None, path_target=None):  # pylint: disable=unused-argument
     """Compile maca code with MXCC from env.
 
@@ -83,6 +127,9 @@ def compile_maca(code, target_format="mcbin", arch=None, options=None, path_targ
         cmd.extend(["-emit-llvm", "-maca-device-only"])
     else:
         cmd.append("-fatbin")
+    arch = _normalize_maca_arch(arch)
+    if arch:
+        cmd.append(f"-offload-arch={arch}")
     cmd.append("-O3")
     if kernels_output_dir is not None:
         cmd += ["-lineinfo"]
@@ -249,9 +296,9 @@ def find_maca_path():
 
 
 @tvm_ffi.register_global_func
-def tvm_callback_maca_compile(code, target):  # pylint: disable=unused-argument
+def tvm_callback_maca_compile(code, target):
     """use mxcc to generate fatbin code for better optimization"""
-    dev_obj = compile_maca(code, target_format="mcbin")
+    dev_obj = compile_maca(code, target_format="mcbin", arch=_target_maca_arch(target))
     return dev_obj
 
 
@@ -276,12 +323,7 @@ def get_target_compute_version(target=None):
     # 2. Target.current()
     target = target or tvm.target.Target.current()
     if target and target.mcpu:
-        arch = target.mcpu[5:]
-        major = arch[:2]
-        minor = arch[2:]
-        if minor == "00":
-            minor = "0"
-        return major + "." + minor
+        return _maca_compute_version_from_arch(target.mcpu)
 
     # 3. GPU compute version
     if tvm.maca(0).exist:

src/s_tir/meta_schedule/space_generator/space_generator.cc

@@ -19,13 +19,59 @@
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/runtime/logging.h>
 
+#include <cctype>
+#include <stdexcept>
+#include <string>
+
 #include "../../../target/canonicalizer/llvm/arm_aprofile.h"
 #include "../utils.h"
 
 namespace tvm {
 namespace s_tir {
 namespace meta_schedule {
 
+namespace {
+
+bool MACATargetSupportsWMMA(const Target& target) {
+  ffi::Optional<ffi::String> opt_mcpu = target->GetAttr<ffi::String>("mcpu");
+  if (!opt_mcpu) {
+    return false;
+  }
+
+  std::string arch = opt_mcpu.value();
+  for (char& ch : arch) {
+    ch = static_cast<char>(std::tolower(static_cast<unsigned char>(ch)));
+  }
+  if (support::StartsWith(arch, "xcore")) {
+    arch = arch.substr(5);
+  }
+
+  size_t digit_count = 0;
+  while (digit_count < arch.size() && std::isdigit(static_cast<unsigned char>(arch[digit_count]))) {
+    ++digit_count;
+  }
+  if (digit_count == 0) {
+    LOG(WARNING) << "ValueError: Unable to parse MACA target mcpu: " << opt_mcpu.value();
+    return false;
+  }
+
+  std::string version = arch.substr(0, digit_count);
+  try {
+    int arch_version = std::stoi(version);
+    // Keep this as an explicit allowlist until MACA exposes a queryable WMMA capability.
+    return arch_version == 1000;
+  } catch (const std::invalid_argument& e) {
+    LOG(WARNING) << "ValueError: Unable to parse MACA target mcpu: " << opt_mcpu.value()
+                 << ". Details: " << e.what();
+  } catch (const std::out_of_range& e) {
+    LOG(WARNING) << "ValueError: MACA target mcpu is out of range: " << opt_mcpu.value()
+                 << ". Details: " << e.what();
+  }
+  return false;
+}
+
+}  // namespace
+
 ffi::String GetRuleKindFromTarget(const Target& target) {
   if (target->kind->name == "llvm") {
     static auto target_has_feature_fn_ptr =
@@ -78,7 +124,10 @@ ffi::String GetRuleKindFromTarget(const Target& target) {
     return "cuda";
   }
   if (target->kind->name == "maca") {
-    return "maca-wmma";
+    if (MACATargetSupportsWMMA(target)) {
+      return "maca-wmma";
+    }
+    return "maca";
   }
 
   if (IsGPUTarget(target->kind->name)) {
@@ -227,6 +276,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
       .def_method("s_tir.meta_schedule.SpaceGeneratorGenerateDesignSpace",
                   &SpaceGeneratorNode::GenerateDesignSpace)
       .def("s_tir.meta_schedule.SpaceGeneratorPySpaceGenerator", SpaceGenerator::PySpaceGenerator)
+      .def("s_tir.meta_schedule.GetRuleKindFromTarget", GetRuleKindFromTarget)
       .def_method("s_tir.meta_schedule.SpaceGeneratorClone", &SpaceGeneratorNode::Clone);
 }