feat: introduce Scalar abstraction to support multiple scalar types

Author: kilinchange

infini_train/include/dtype_dispatch.h

@@ -17,7 +17,7 @@
 
 #define LOG_UNSUPPORTED_DTYPE(DTYPE, CONTEXT_IDENTIFIER)                                                               \
     LOG_LOC(FATAL, std::string(CONTEXT_IDENTIFIER)                                                                     \
-                       + ": Unsupported data type: " + kDataTypeToDesc.at(static_cast<infini_train::DataType>(dtype)))
+                       + ": Unsupported data type: " + kDataTypeToDesc.at(static_cast<infini_train::DataType>(DTYPE)))
 
 // Helper macros to count the number of arguments
 #define PP_NARG(...) PP_NARG_(__VA_ARGS__, PP_RSEQ_N())

infini_train/include/scalar.h

@@ -0,0 +1,59 @@
+#pragma once
+
+#include <cstdint>
+#include <type_traits>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/common/cpu/common_cpu.h"
+
+namespace infini_train {
+
+struct Scalar {
+    enum class Kind : uint8_t { kBool, kDouble, kInt64, kUInt64 };
+
+    Scalar() : kind(Kind::kInt64), i(0) {}
+    Scalar(bool v) : kind(Kind::kBool), u(v ? 1 : 0) {}
+
+    template <typename T, typename std::enable_if_t<std::is_floating_point_v<T>, int> = 0>
+    Scalar(T v) : kind(Kind::kDouble), d(static_cast<double>(v)) {}
+
+    template <typename T,
+              typename std::enable_if_t<std::is_integral_v<T> && std::is_signed_v<T> && !std::is_same_v<T, bool>, int>
+              = 0>
+    Scalar(T v) : kind(Kind::kInt64), i(static_cast<int64_t>(v)) {}
+
+    template <typename T,
+              typename std::enable_if_t<std::is_integral_v<T> && std::is_unsigned_v<T> && !std::is_same_v<T, bool>, int>
+              = 0>
+    Scalar(T v) : kind(Kind::kUInt64), u(static_cast<uint64_t>(v)) {}
+
+    Scalar(FP16 v) : kind(Kind::kDouble), d(static_cast<float>(v)) {}
+    Scalar(BF16 v) : kind(Kind::kDouble), d(static_cast<float>(v)) {}
+
+    template <typename T> T to() const {
+        switch (kind) {
+        case Kind::kBool:
+            return common::cpu::Cast<T>(u != 0);
+        case Kind::kDouble:
+            return common::cpu::Cast<T>(d);
+        case Kind::kInt64:
+            return common::cpu::Cast<T>(i);
+        case Kind::kUInt64:
+            return common::cpu::Cast<T>(u);
+        default:
+            LOG(FATAL) << "Unknown scalar kind";
+        }
+
+        std::abort();
+    }
+
+    Kind kind;
+    union {
+        double d;
+        int64_t i;
+        uint64_t u;
+    };
+};
+
+} // namespace infini_train

infini_train/include/tensor.h

@@ -12,6 +12,7 @@
 
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/scalar.h"
 
 namespace infini_train {
 namespace autograd {
@@ -78,8 +79,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     size_t NumElements() const;
     DataType Dtype() const;
 
-    // Fill tensor with a scalar value (accepts double, automatically converts to tensor's dtype)
-    void Fill(double value);
+    void Fill(Scalar value);
 
     Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> EigenMatrix();
     Eigen::Map<Eigen::Matrix<float, 1, Eigen::Dynamic, Eigen::RowMajor>> EigenVector();

infini_train/src/kernels/cpu/fill.cc

@@ -1,19 +1,18 @@
 #include "glog/logging.h"
 
-#include "infini_train/include/common/cpu/common_cpu.h"
 #include "infini_train/include/dispatcher.h"
 #include "infini_train/include/dtype_dispatch.h"
 #include "infini_train/include/tensor.h"
 
 #include "infini_train/src/core/runtime/cpu/cpu_dispatch.h"
 
 namespace infini_train::kernels::cpu {
-void Fill(std::shared_ptr<Tensor> tensor, double value) {
+void Fill(std::shared_ptr<Tensor> tensor, Scalar scalar) {
     core::cpu::DispatchCpuFunc<INFINI_ALL_TYPES>(
         tensor->Dtype(),
         [=]<typename T>() {
             auto data = reinterpret_cast<T *>(tensor->DataPtr());
-            T casted_value = common::cpu::Cast<T>(value);
+            const T casted_value = scalar.to<T>();
             std::fill(data, data + tensor->NumElements(), casted_value);
         },
         "CPU Fill");

infini_train/src/kernels/cuda/fill.cu

@@ -1,7 +1,6 @@
 #include <cstddef>
 #include <memory>
 
-#include "infini_train/include/common/cpu/common_cpu.h"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/dispatcher.h"
@@ -20,7 +19,7 @@ template <typename T> __global__ void FillKernel(T *data, T value, size_t size)
 }
 
 // TODO(dcj): refactor Fill kernel with elementwise template
-void Fill(std::shared_ptr<Tensor> tensor, double value) {
+void Fill(std::shared_ptr<Tensor> tensor, Scalar scalar) {
     const int num_tokens = tensor->NumElements();
     const int threads_per_block = 256;
     const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block;
@@ -32,7 +31,7 @@ void Fill(std::shared_ptr<Tensor> tensor, double value) {
     core::cuda::DispatchCudaFunc<INFINI_ALL_TYPES>(
         tensor->Dtype(),
         [=]<typename T>() {
-            T casted_value = common::cpu::Cast<T>(value);
+            const T casted_value = scalar.to<T>();
             FillKernel<T><<<num_blocks, threads_per_block, 0, cuda_stream>>>(static_cast<T *>(tensor->DataPtr()),
                                                                              casted_value, tensor->NumElements());
         },

infini_train/src/tensor.cc

@@ -18,12 +18,10 @@
 #include "infini_train/include/autograd/outer.h"
 #include "infini_train/include/autograd/reduction.h"
 #include "infini_train/include/autograd/transform.h"
-#include "infini_train/include/common/cpu/common_cpu.h"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/dispatcher.h"
-#include "infini_train/include/dtype_dispatch.h"
 #include "infini_train/include/nn/init.h"
 
 namespace infini_train {
@@ -104,7 +102,7 @@ size_t Tensor::NumElements() const { return num_elements_; }
 
 DataType Tensor::Dtype() const { return dtype_; }
 
-void Tensor::Fill(double value) {
+void Tensor::Fill(Scalar value) {
     auto device = GetDevice();
     core::DeviceGuard guard(device);
     auto kernel = Dispatcher::Instance().GetKernel({device.type(), "Fill"});