Feat/dev datatype by bitzyz · Pull Request #22 · InfiniTensor/InfiniOps

bitzyz · 2026-03-16T01:51:09Z

✅Moved device-specific cast.h/common.h files from common/ to respective device folders (cuda/, cpu/, cambricon/)
✅Introduced ValueTag and TypeTag templates for compile-time type-safe runtime dispatch
✅Enhanced DataType struct with device-aware type mappings
✅Updated all kernel files to use the new dispatch mechanism

bitzyz · 2026-03-16T02:33:23Z

bitzyz · 2026-03-16T02:36:52Z

bitzyz · 2026-03-16T02:45:24Z

寒武纪平台编译正常

bitzyz · 2026-03-16T03:15:17Z

沐曦平台：

…spective device folders.

… templates in dispatch framework.

bitzyz · 2026-03-16T06:14:10Z

天数平台：

voltjia · 2026-03-18T08:21:08Z

src/cambricon/gemm/cnblas.h


 #include "base/gemm.h"
-#include "cambricon/common.h"
+#include "../common.h"


voltjia · 2026-03-18T08:21:22Z

src/cpu/add/add.h


 #include <utility>

+#include "../cast.h"


voltjia · 2026-03-18T08:22:16Z

src/cpu/causal_softmax/causal_softmax.h


 #include <cmath>

+#include "../cast.h"


voltjia · 2026-03-18T08:22:28Z

src/cpu/gemm/gemm.h


 #include <utility>

+#include "../cast.h"


voltjia · 2026-03-18T08:22:33Z

src/cpu/rms_norm/rms_norm.h


 #include <cmath>

+#include "../cast.h"


voltjia · 2026-03-18T08:56:26Z

src/dispatcher.h

                         : false));

    if (!handled) {
-      // TODO(lzm): change to logging.


这些 TODO 咋给去掉了。

voltjia · 2026-03-18T08:57:41Z

src/data_type.h

 };

-template <DataType dtype>
+// Forward declaration.


这几处新加的注释可以去掉。如果决定不去掉需要用 Markdown。

voltjia · 2026-03-18T08:59:20Z

src/data_type.h

+using TypeMapTypeDevice = typename TypeMap<dtype, device>::type;

-template <typename T>
+template <typename T, Device::Type D>


按照 style guide，这个地方的 D 应该改成 d，因为它应该遵循变量命名，而不是类型命名。下同。

voltjia · 2026-03-18T09:01:19Z

tests/test_rms_norm.py


 def _torch_rms_norm(input, weight, *, eps=1e-6, out=None):
    return torch.nn.functional.rms_norm(input, input.shape[-1:], weight=weight, eps=eps)
+


这个地方是不是多加了空格，有过 ruff format && ruff check 嘛？

voltjia · 2026-03-18T09:02:27Z

README.md

-| `-DWITH_NVIDIA=[ON\|OFF]`              | Compile the NVIDIA implementation  | n
-| `-DWITH_METAX=[ON\|OFF]`               | Compile the MetaX implementation   | n
-| `-DGENERATE_PYTHON_BINDINGS=[ON\|OFF]` | Generate Python bindings           | n
+| Option                                 | Functionality                       | Default


最右列要不再右移一格，这样离最近的多一个空格。

Ziminli · 2026-03-18T09:43:03Z

src/cuda/add/kernel.h

  void operator()(const Tensor input, const Tensor other,
                  Tensor out) const override {
-    DispatchFunc<AllTypes>(
+    DispatchFunc<Backend::device_value, AllTypes>(


啊，这样能编过的么？所以这里的意图是准备只在这个平台分发这些类型吗？按理来说之前其他平台倒是都不需要这样写，毕竟因为只会编译自己所在的平台目录下的代码嘛。目前如果想做多类型混合分发的话，应该是类似这样：

DispatchFunc<FloatTypes, List<Device::Type::kCpu, Device::Type::kNvidia>>( {static_cast<int64_t>(DataType::kFloat32), static_cast<int64_t>(Device::Type::kNvidia)}, 0, [](auto list_tag) { constexpr DataType DT = static_cast<DataType>(ListGet<0>(list_tag)); constexpr Device::Type Dev = static_cast<Device::Type>(ListGet<1>(list_tag)); using T = TypeMapType<DT>; // Remaining logic ... }, "MixedDispatch", List<>{});

之后我会加一个方便一点的高层封装，就可以这样：

DispatchFunc<FloatTypes, List<Device::Type::kCpu, Device::Type::kNvidia>>( {static_cast<int64_t>(DataType::kFloat32), static_cast<int64_t>(Device::Type::kNvidia)}, [](auto list_tag) { constexpr DataType DT = static_cast<DataType>(ListGet<0>(list_tag)); constexpr Device::Type Dev = static_cast<Device::Type>(ListGet<1>(list_tag)); using T = TypeMapType<DT>; // Remaining logic ... }, "MixedDispatch");

本 PR 其他同类地方同理。

Ziminli · 2026-03-18T10:02:56Z

src/dispatcher.h

原来如此，那这个文件可以还原，原本的 dispatcher 就是支持任意类型和多类型混合分发的。

Ziminli · 2026-03-18T10:03:28Z

src/data_type.h

+template <>
+struct TypeMap<DataType::kFloat16, Device::Type::kNvidia> {
+  using type = half;
+};
+template <>
+struct DataTypeMap<half, Device::Type::kNvidia> {
+  static constexpr DataType value = DataType::kFloat16;
+};
+template <>
+struct TypeMap<DataType::kBFloat16, Device::Type::kNvidia> {
+  using type = __nv_bfloat16;
+};
+template <>
+struct DataTypeMap<__nv_bfloat16, Device::Type::kNvidia> {
+  static constexpr DataType value = DataType::kBFloat16;
+};


这些应该也是为了配合 dispatcher 的改动？应该可以还原了。

bitzyz self-assigned this Mar 16, 2026

bitzyz force-pushed the feat/dev-datatype branch 2 times, most recently from d479f8c to 7cba646 Compare March 16, 2026 02:28

bitzyz requested a review from voltjia March 16, 2026 02:45

bitzyz force-pushed the feat/dev-datatype branch from 7cba646 to db13a72 Compare March 16, 2026 03:15

bitzyz added 2 commits March 16, 2026 03:23

feat: Move device-specific files from the common folder into their re…

74fda5b

…spective device folders.

feat: Add Device tag in DataType struct && using ValueTag and TypeTag…

1c72cfe

… templates in dispatch framework.

bitzyz force-pushed the feat/dev-datatype branch from db13a72 to 1c72cfe Compare March 16, 2026 03:26

voltjia requested changes Mar 18, 2026

View reviewed changes

Ziminli requested changes Mar 18, 2026

View reviewed changes

Conversation

bitzyz commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bitzyz commented Mar 16, 2026

Uh oh!

bitzyz commented Mar 16, 2026

Uh oh!

bitzyz commented Mar 16, 2026

Uh oh!

bitzyz commented Mar 16, 2026

Uh oh!

bitzyz commented Mar 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ziminli Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bitzyz commented Mar 16, 2026 •

edited

Loading

Ziminli Mar 18, 2026 •

edited

Loading