refactor: remove NativeScalar instead of using BackendTypeMap directly

Author: kilinchange

infini_train/include/core/backend_type_map.h

@@ -8,11 +8,6 @@ namespace infini_train::core {
 /**
  * Backend type mapping: DataType -> backend-native dispatch type
  *
- * NativeScalar   — maps framework low-precision scalar types (FP16/BF16) to
- *                  backend-native scalar types (__half / __nv_bfloat16).
- *                  Primary template intentionally undefined.
- *                  Each backend specializes only the types it supports.
- *
  * BackendTypeMap — maps DataType to the C++ type used by kernels/dispatch.
  *                  Primary template intentionally undefined — there is NO
  *                  default fallback to the framework TypeMap<DType>.
@@ -22,53 +17,20 @@ namespace infini_train::core {
  *                        call INFINI_REGISTER_STANDARD_BACKEND_TYPES(Dev)
  *                        once at file scope in the backend's dispatch header.
  *                    - Low-precision types (FP16, BF16):
- *                        specialize NativeScalar<Dev, infini_train::FP16/BF16>.
- *                        The generic partial specializations below then resolve
- *                        automatically via SFINAE-safe helper.
+ *                        directly specialize BackendTypeMap<Dev, kFLOAT16/kBFLOAT16>
+ *                        in the backend's dispatch header (the native scalar type
+ *                        differs per backend, e.g. __half on CUDA).
  *
  * If a backend does not register a dtype, HasMappedType_v returns false and
  * DispatchByTypeMap fires a clear static_assert at compile time.
  */
 
-// -----------------------------------------------------------------------------
-// NativeScalar: framework scalar -> backend native scalar
-// Primary template intentionally undefined.
-// -----------------------------------------------------------------------------
-template <Device::DeviceType Dev, typename Scalar> struct NativeScalar;
-
-template <Device::DeviceType Dev, typename Scalar> using NativeScalar_t = typename NativeScalar<Dev, Scalar>::type;
-
 // -----------------------------------------------------------------------------
 // BackendTypeMap: DataType -> backend dispatch type
 // Primary template intentionally undefined — no TypeMap<DType> fallback.
 // -----------------------------------------------------------------------------
 template <Device::DeviceType Dev, DataType DType> struct BackendTypeMap;
 
-// -----------------------------------------------------------------------------
-// SFINAE-safe helper for low-precision type routing.
-// When NativeScalar<Dev, Scalar> is undefined, this struct has no `type`
-// member, making HasMappedType_v<..., kFLOAT16/kBFLOAT16> return false and
-// triggering the static_assert in dispatch rather than an opaque hard error.
-// -----------------------------------------------------------------------------
-namespace detail {
-
-template <Device::DeviceType Dev, typename Scalar, typename = void>
-struct BackendLowPrecisionTypeHelper {}; // no `type` member when NativeScalar absent
-
-template <Device::DeviceType Dev, typename Scalar>
-struct BackendLowPrecisionTypeHelper<Dev, Scalar, std::void_t<typename NativeScalar<Dev, Scalar>::type>> {
-    using type = typename NativeScalar<Dev, Scalar>::type;
-};
-
-} // namespace detail
-
-// Low-precision partial specializations: generic over Dev, resolved via NativeScalar.
-template <Device::DeviceType Dev>
-struct BackendTypeMap<Dev, DataType::kFLOAT16> : detail::BackendLowPrecisionTypeHelper<Dev, infini_train::FP16> {};
-
-template <Device::DeviceType Dev>
-struct BackendTypeMap<Dev, DataType::kBFLOAT16> : detail::BackendLowPrecisionTypeHelper<Dev, infini_train::BF16> {};
-
 } // namespace infini_train::core
 
 // -----------------------------------------------------------------------------
@@ -80,7 +42,9 @@ struct BackendTypeMap<Dev, DataType::kBFLOAT16> : detail::BackendLowPrecisionTyp
 //
 //   INFINI_REGISTER_STANDARD_BACKEND_TYPES(Device::DeviceType::kCUDA)
 //
-// FP16 and BF16 are NOT registered here — they are handled via NativeScalar.
+// FP16 and BF16 are NOT registered here — backends must specialize
+// BackendTypeMap<DEV, kFLOAT16/kBFLOAT16> directly with their native scalar
+// type (e.g. __half / __nv_bfloat16 on CUDA).
 // -----------------------------------------------------------------------------
 #define INFINI_REGISTER_STANDARD_BACKEND_TYPES(DEV)                                                                    \
     namespace infini_train::core {                                                                                     \

infini_train/src/core/runtime/cpu/cpu_dispatch.h

@@ -7,22 +7,22 @@
 #include "infini_train/include/dtype_dispatch.h"
 
 // -----------------------------------------------------------------------------
-// CPU NativeScalar specializations: FP16 -> FP16, BF16 -> BF16
+// CPU low-precision BackendTypeMap specializations:
+//   FP16 -> infini_train::FP16, BF16 -> infini_train::BF16
 // CPU uses the framework wrapper types directly (host-side conversion).
 // -----------------------------------------------------------------------------
 namespace infini_train::core {
-template <> struct NativeScalar<Device::DeviceType::kCPU, infini_train::FP16> {
+template <> struct BackendTypeMap<Device::DeviceType::kCPU, DataType::kFLOAT16> {
     using type = infini_train::FP16;
 };
 
-template <> struct NativeScalar<Device::DeviceType::kCPU, infini_train::BF16> {
+template <> struct BackendTypeMap<Device::DeviceType::kCPU, DataType::kBFLOAT16> {
     using type = infini_train::BF16;
 };
 } // namespace infini_train::core
 
 // Register all standard (non-low-precision) dtypes for the CPU backend.
-// FP16/BF16 are handled above via NativeScalar specializations +
-// BackendTypeMap partial specializations in backend_type_map.h.
+// FP16/BF16 are registered explicitly above.
 INFINI_REGISTER_STANDARD_BACKEND_TYPES(infini_train::Device::DeviceType::kCPU)
 
 namespace infini_train::core::cpu {

infini_train/src/core/runtime/cuda/cuda_dispatch.h

@@ -10,21 +10,21 @@
 #include "infini_train/include/dtype_dispatch.h"
 
 // -----------------------------------------------------------------------------
-// CUDA NativeScalar specializations: FP16 -> __half, BF16 -> __nv_bfloat16
+// CUDA low-precision BackendTypeMap specializations:
+//   FP16 -> __half, BF16 -> __nv_bfloat16
 // -----------------------------------------------------------------------------
 namespace infini_train::core {
-template <> struct NativeScalar<Device::DeviceType::kCUDA, infini_train::FP16> {
+template <> struct BackendTypeMap<Device::DeviceType::kCUDA, DataType::kFLOAT16> {
     using type = __half;
 };
 
-template <> struct NativeScalar<Device::DeviceType::kCUDA, infini_train::BF16> {
+template <> struct BackendTypeMap<Device::DeviceType::kCUDA, DataType::kBFLOAT16> {
     using type = __nv_bfloat16;
 };
 } // namespace infini_train::core
 
 // Register all standard (non-low-precision) dtypes for the CUDA backend.
-// FP16/BF16 are handled above via NativeScalar specializations +
-// BackendTypeMap partial specializations in backend_type_map.h.
+// FP16/BF16 are registered explicitly above with their CUDA-native scalar types.
 INFINI_REGISTER_STANDARD_BACKEND_TYPES(infini_train::Device::DeviceType::kCUDA)
 
 namespace infini_train::core::cuda {
@@ -38,7 +38,7 @@ template <DataType DType> struct CudaTypeMap;
 
 // Register all supported dtypes by delegating to BackendTypeMap<kCUDA, DType>.
 // Standard types come from INFINI_REGISTER_STANDARD_BACKEND_TYPES above;
-// FP16/BF16 come from BackendTypeMap partial specializations + NativeScalar.
+// FP16/BF16 come from the explicit BackendTypeMap specializations above.
 #define INFINI_REGISTER_CUDA_TYPEMAP(DTYPE)                                                                            \
     template <>                                                                                                        \
     struct CudaTypeMap<DataType::DTYPE>                                                                                \