fix format

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/adaptor/aten_adaptor.hpp

@@ -7,6 +7,7 @@
 
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
 #include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAGuard.h>
 #endif
 
 #ifdef ENABLE_NVIDIA_API

scripts/install.py

@@ -4,8 +4,8 @@
 import sys
 from set_env import (
     set_env,
-    ensure_metax_hpc_compiler_includes,
-    xmake_flags_need_metax_aten_torch_includes,
+    ensure_aten_torch_compiler_includes,
+    xmake_flags_need_aten_torch_compiler_includes,
 )
 
 PROJECT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -16,8 +16,8 @@ def run_cmd(cmd):
 
 
 def install(xmake_config_flags=""):
-    if xmake_flags_need_metax_aten_torch_includes(xmake_config_flags):
-        ensure_metax_hpc_compiler_includes()
+    if xmake_flags_need_aten_torch_compiler_includes(xmake_config_flags):
+        ensure_aten_torch_compiler_includes()
     run_cmd(f"xmake f {xmake_config_flags} -cv")
     run_cmd("xmake")
     run_cmd("xmake install")

scripts/set_env.py

@@ -2,25 +2,15 @@
 import platform
 
 
-def _maca_root_from_env():
-    return (
-        os.environ.get("MACA_PATH")
-        or os.environ.get("MACA_HOME")
-        or os.environ.get("MACA_ROOT")
-        or ""
-    ).strip()
-
-
-def metax_hpc_compiler_include_dirs():
-    """Directories needed so g++ finds cuda_runtime_api.h (cu-bridge) when compiling against PyTorch c10/cuda headers on MetaX/HPCC."""
-    maca = _maca_root_from_env()
-    if not maca:
-        return []
-    return [
-        os.path.join(maca, "tools", "cu-bridge", "include"),
-        os.path.join(maca, "include", "hcr"),
-        os.path.join(maca, "include"),
-    ]
+def _hpcc_toolkit_root() -> str:
+    """HPCC/MACA install root (cu-bridge, headers). Env vars first; else common container path."""
+    for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"):
+        v = os.environ.get(key, "").strip()
+        if v:
+            return v
+    if os.path.isdir("/opt/hpcc"):
+        return "/opt/hpcc"
+    return ""
 
 
 def _prepend_path_var(name, prefixes):
@@ -32,14 +22,16 @@ def _prepend_path_var(name, prefixes):
     os.environ[name] = f"{chunk}:{cur}" if cur else chunk
 
 
-def ensure_metax_hpc_compiler_includes():
-    """
-    Prepend HPCC/cu-bridge includes to CPATH, CPLUS_INCLUDE_PATH, and C_INCLUDE_PATH.
-    g++ uses CPLUS_INCLUDE_PATH for .cc files; C_INCLUDE_PATH alone is not enough.
-    """
-    dirs = metax_hpc_compiler_include_dirs()
-    if not dirs:
+def ensure_aten_torch_compiler_includes() -> None:
+    """If HPCC root is known, prepend cu-bridge + HPCC headers for g++ compiling ATen .cc (c10/cuda)."""
+    root = _hpcc_toolkit_root()
+    if not root:
         return
+    dirs = [
+        os.path.join(root, "tools", "cu-bridge", "include"),
+        os.path.join(root, "include", "hcr"),
+        os.path.join(root, "include"),
+    ]
     for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
         _prepend_path_var(var, dirs)
 
@@ -69,12 +61,20 @@ def _truthy_flag_value(v: str) -> bool:
     return v in ("y", "yes", "true", "1", "on")
 
 
-def xmake_flags_need_metax_aten_torch_includes(flags: str) -> bool:
-    """True when install.py-style args enable MetaX GPU and ATen (PyTorch) together."""
+# xmake.lua GPU / accelerator backends (any of these + aten may compile C++ against torch+cuda-style headers).
+_XMAKE_GPU_BACKEND_KEYS = frozenset(
+    {
+        "metax-gpu",
+    }
+)
+
+
+def xmake_flags_need_aten_torch_compiler_includes(flags: str) -> bool:
+    """True when ATen is enabled with any GPU/accelerator backend (install.py / xmake f ...)."""
     d = _parse_xmake_cli_flag_values(flags)
-    return _truthy_flag_value(d.get("metax-gpu", "n")) and _truthy_flag_value(
-        d.get("aten", "n")
-    )
+    if not _truthy_flag_value(d.get("aten", "n")):
+        return False
+    return any(_truthy_flag_value(d.get(k, "n")) for k in _XMAKE_GPU_BACKEND_KEYS)
 
 
 def set_env():

src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc

@@ -81,8 +81,8 @@ VarlenFlashPrepared prepare_varlen_flash_tensors(PlannedMeta *p) {
     t.max_seqlen_q = p->max_seqlen_q;
     t.max_seqlen_k = p->max_seqlen_k;
     t.alibi_slopes = p->alibi_slopes
-                         ? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes).contiguous())
-                         : std::nullopt;
+                       ? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes).contiguous())
+                       : std::nullopt;
     t.scale = p->scale;
     return t;
 }

src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.maca

@@ -1,8 +1,9 @@
 #include "../../../devices/metax/metax_common.h"
 #include "../../../devices/metax/metax_handle.h"
 #include "../../../devices/metax/metax_kernel_common.h"
+
 #include "binary_cross_entropy_with_logits_metax.h"
-#include <hc_runtime.h>
+
 #include <type_traits>
 
 namespace op::bce_with_logits::metax {

src/infiniop/ops/equal/metax/equal_metax.maca

@@ -13,27 +13,37 @@ infiniStatus_t Descriptor::create(
     Descriptor **desc_ptr,
     infiniopTensorDescriptor_t out_desc,
     std::vector<infiniopTensorDescriptor_t> input_desc_vec) {
+
     auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
+
     const auto &a_desc = input_desc_vec.at(0);
     auto compute_dtype = a_desc->dtype();
     auto out_dtype = out_desc->dtype();
+
     const auto &b_desc = input_desc_vec.at(1);
     const auto &c_shape = out_desc->shape();
     const auto &a_shape = a_desc->shape();
     const auto &b_shape = b_desc->shape();
+
     CHECK_DTYPE(compute_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16,
                 INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64);
+
     CHECK_DTYPE(out_dtype, INFINI_DTYPE_BOOL);
+
     CHECK_SAME_SHAPE(c_shape, a_shape, b_shape);
+
     CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, compute_dtype, out_desc, input_desc_vec)
+
     return INFINI_STATUS_SUCCESS;
 }
+
 infiniStatus_t Descriptor::calculate(
     void *workspace,
     size_t workspace_size,
     void *output,
     std::vector<const void *> inputs,
     void *stream) const {
+
     if (workspace_size < _workspace_size) {
         return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
     }