issue/111: Merge branch 'main' of github.com:PanZezhong1725/InfiniCore into...

issue/111: Merge branch 'main' of github.com:PanZezhong1725/InfiniCore into issue/111-rmsnorm-kunlun

src/utils/rearrange.cc

@@ -13,7 +13,7 @@ namespace utils {
 RearrangeMeta::RearrangeMeta(std::vector<ptrdiff_t> meta)
     : _meta(std::move(meta)) {}
 
-std::optional<RearrangeMeta> RearrangeMeta::create(
+Result<RearrangeMeta> RearrangeMeta::create(
     const size_t *shape,
     const ptrdiff_t *dst_strides_,
     const ptrdiff_t *src_strides_,
@@ -32,7 +32,9 @@ std::optional<RearrangeMeta> RearrangeMeta::create(
         // 剔除初始的 1 长维度
         if (shape[i] != 1) {
             auto sd = dst_strides_[i] * unit, ss = src_strides_[i] * unit;
-            // assert (sd != 0)
+            if (sd == 0) {
+                return INFINI_STATUS_BAD_TENSOR_STRIDES;
+            }
             dims.push_back(Dim{shape[i], sd, ss});
         }
     }
@@ -81,7 +83,7 @@ std::optional<RearrangeMeta> RearrangeMeta::create(
     for (ptrdiff_t i = ndim; i > 0; --i) {
         meta[1 + i - 1] *= meta[1 + i];
     }
-    return {RearrangeMeta(std::move(meta))};
+    return Result<RearrangeMeta>(meta);
 }
 
 size_t RearrangeMeta::ndim() const { return (_meta.size() - 2) / 3; }

src/utils/rearrange.h

 #ifndef __INFINIUTILS_REARRANGE_H__
 #define __INFINIUTILS_REARRANGE_H__
 
-#include <optional>
-#include <stddef.h>
+#include "result.hpp"
+#include <cstddef>
 #include <vector>
 
 namespace utils {
@@ -12,7 +12,7 @@ class RearrangeMeta {
     RearrangeMeta(std::vector<ptrdiff_t>);
 
 public:
-    static std::optional<RearrangeMeta> create(
+    static Result<RearrangeMeta> create(
         const size_t *shape,
         const ptrdiff_t *dst_strides,
         const ptrdiff_t *src_strides,

src/utils/result.hpp

+#ifndef __INFINIUTILS_RESULT_H__
+#define __INFINIUTILS_RESULT_H__
+
+#include "check.h"
+#include <infinicore.h>
+#include <variant>
+
+#define CHECK_RESULT(RESULT)    \
+    if (!RESULT) {              \
+        return RESULT.status(); \
+    }
+
+namespace utils {
+
+template <typename T, typename = std::enable_if_t<!std::is_same_v<T, infiniStatus_t>>>
+class Result {
+    std::variant<infiniStatus_t, T> _result;
+
+public:
+    explicit Result(T value) : _result(std::move(value)) {}
+    Result(infiniStatus_t status) : _result(status) {
+        if (status == INFINI_STATUS_SUCCESS) {
+            std::cerr << "Warning: Result created with success status but value is not set." << std::endl;
+            std::abort();
+        }
+    }
+
+    infiniStatus_t status() const {
+        return _result.index() == 0 ? std::get<0>(_result) : INFINI_STATUS_SUCCESS;
+    }
+
+    T take() {
+        return std::move(std::get<1>(_result));
+    }
+
+    operator bool() const {
+        return status() == INFINI_STATUS_SUCCESS;
+    }
+
+    T *operator->() {
+        return &std::get<1>(_result);
+    }
+
+    const T *operator->() const {
+        return &std::get<1>(_result);
+    }
+
+    T &operator*() {
+        return std::get<1>(_result);
+    }
+
+    const T &operator*() const {
+        return std::get<1>(_result);
+    }
+};
+
+} // namespace utils
+
+#endif // __INFINIUTILS_RESULT_H__

test/infiniop/libinfiniop/utils.py

@@ -171,6 +171,11 @@ def get_args():
         action="store_true",
         help="Run METAX GPU test",
     )
+    parser.add_argument(
+        "--moore",
+        action="store_true",
+        help="Run MTHREADS GPU test",
+    )
     parser.add_argument(
         "--kunlun",
         action="store_true",
@@ -443,6 +448,11 @@ def get_test_devices(args):
         import torch
 
         devices_to_test.append(InfiniDeviceEnum.METAX)
+    if args.moore:
+        import torch
+        import torch_musa
+
+        devices_to_test.append(InfiniDeviceEnum.MOORE)
     if args.kunlun:
         import torch_xmlir
 

xmake.lua

@@ -89,7 +89,8 @@ option("moore-gpu")
 option_end()
 
 if has_config("moore-gpu") then
-    add_defines("ENABLE_MUSA_API")
+    add_defines("ENABLE_MOORE_API")
+    includes("xmake/musa.lua")
 end
 
 -- 海光
@@ -154,6 +155,9 @@ target("infinirt")
     if has_config("metax-gpu") then
         add_deps("infinirt-metax")
     end
+    if has_config("moore-gpu") then
+        add_deps("infinirt-moore")
+    end
     if has_config("kunlun-xpu") then
         add_deps("infinirt-kunlun")
     end
@@ -197,6 +201,9 @@ target("infiniop")
     if has_config("metax-gpu") then
         add_deps("infiniop-metax")
     end
+    if has_config("moore-gpu") then
+        add_deps("infiniop-moore")
+    end
     if has_config("kunlun-xpu") then
         add_deps("infiniop-kunlun")
     end

xmake/ascend.lua

 add_defines("ENABLE_ASCEND_API")
-local ASCEND_HOME = os.getenv("ASCEND_HOME")
+local ASCEND_HOME = os.getenv("ASCEND_HOME") or os.getenv("ASCEND_TOOLKIT_HOME")
 local SOC_VERSION = os.getenv("SOC_VERSION")
 
 -- Add include dirs

xmake/musa.lua

+local MUSA_HOME = os.getenv("MUSA_INSTALL_PATH")
+add_includedirs(MUSA_HOME .. "/include")
+add_linkdirs(MUSA_HOME .. "/lib")
+add_links("musa", "musart", "mudnn", "mublas")
+
+rule("mu")
+    set_extensions(".mu")
+    on_load(function (target)
+        target:add("includedirs", "include")
+    end)
+
+    on_build_file(function (target, sourcefile)
+        local objectfile = target:objectfile(sourcefile)
+        os.mkdir(path.directory(objectfile))
+
+        local mcc = MUSA_HOME .. "/bin/mcc"
+        local includedirs = table.concat(target:get("includedirs"), " ")
+        local args = {"-c", sourcefile, "-o", objectfile, "-I" .. MUSA_HOME .. "/include", "-O3", "-fPIC", "-Wall", "-std=c++17", "-pthread"}
+        for _, includedir in ipairs(target:get("includedirs")) do
+            table.insert(args, "-I" .. includedir)
+        end
+
+        os.execv(mcc, args)
+        table.insert(target:objectfiles(), objectfile)
+    end)
+rule_end()
+
+target("infiniop-moore")
+    set_kind("static")
+    on_install(function (target) end)
+    set_languages("cxx17")
+    set_warnings("all", "error")
+    add_cxflags("-lstdc++", "-fPIC", "-Wno-comment")
+    add_files("../src/infiniop/devices/musa/*.cc", "../src/infiniop/ops/*/musa/*.cc")
+    add_files("../src/infiniop/ops/*/musa/*.mu", {rule = "mu"})
+target_end()
+
+target("infinirt-moore")
+    set_kind("static")
+    set_languages("cxx17")
+    on_install(function (target) end)
+    add_deps("infini-utils")
+    set_warnings("all", "error")
+    add_cxflags("-lstdc++", "-fPIC")
+    add_files("../src/infinirt/musa/*.cc")
+target_end()

xmake/test.lua

 target("infiniutils-test")
     set_kind("binary")
     add_deps("infini-utils")
-    on_install(function (target) end)
 
     set_warnings("all", "error")
     set_languages("cxx17")
     
     add_files(os.projectdir().."/src/utils-test/*.cc")
+    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+target_end()
 
 target("infiniop-test")
     set_kind("binary")
     add_deps("infini-utils")
-    on_install(function (target) end)
     set_default(false)
 
     local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")
@@ -31,4 +31,6 @@ target("infiniop-test")
     add_includedirs(os.projectdir().."/src/infiniop-test/include")
     add_files(os.projectdir().."/src/infiniop-test/src/*.cpp")
     add_files(os.projectdir().."/src/infiniop-test/src/ops/*.cpp")
+
+    set_installdir(INFINI_ROOT)
 target_end()