Merge pull request #627 from gongchensu/feat/add-dtype-support

Issue/626 - Add I32 and I64 dtype support to add operation (CPU) and …

src/infiniop/ops/add/cpu/add_cpu.cc

@@ -19,7 +19,7 @@ infiniStatus_t Descriptor::create(
     const auto &a_shape = a_desc->shape();
     const auto &b_shape = b_desc->shape();
 
-    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16);
+    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64);
 
     CHECK_SAME_SHAPE(c_shape, a_shape, b_shape);
 
@@ -45,6 +45,10 @@ infiniStatus_t Descriptor::calculate(
         return _device_info->calculate<AddOp, double>(_info, output, inputs, stream);
     case INFINI_DTYPE_BF16:
         return _device_info->calculate<AddOp, bf16_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_I32:
+        return _device_info->calculate<AddOp, int32_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_I64:
+        return _device_info->calculate<AddOp, int64_t>(_info, output, inputs, stream);
     default:
         return INFINI_STATUS_BAD_TENSOR_DTYPE;
     }

src/infiniop/ops/add/nvidia/add_nvidia.cu

@@ -22,7 +22,7 @@ infiniStatus_t Descriptor::create(
     const auto &a_shape = a_desc->shape();
     const auto &b_shape = b_desc->shape();
 
-    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16);
+    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64);
 
     CHECK_SAME_SHAPE(c_shape, a_shape, b_shape);
 
@@ -50,6 +50,10 @@ infiniStatus_t Descriptor::calculate(
         return _device_info->calculate<256, cuda::AddOp, cuda_bfloat16>(_info, workspace, output, inputs, stream);
     case INFINI_DTYPE_F32:
         return _device_info->calculate<256, cuda::AddOp, float>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_I32:
+        return _device_info->calculate<256, cuda::AddOp, int32_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_I64:
+        return _device_info->calculate<256, cuda::AddOp, int64_t>(_info, workspace, output, inputs, stream);
     case INFINI_DTYPE_F64:
         return _device_info->calculate<256, cuda::AddOp, double>(_info, workspace, output, inputs, stream);
     default:

test/infiniop/add.py

@@ -61,13 +61,15 @@ _TEST_CASES = [
 ]
 
 # Data types used for testing
-_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.F32, InfiniDtype.BF16]
+_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.F32, InfiniDtype.BF16, InfiniDtype.I32, InfiniDtype.I64]
 
 # Tolerance map for different data types
 _TOLERANCE_MAP = {
     InfiniDtype.F16: {"atol": 1e-3, "rtol": 1e-3},
     InfiniDtype.F32: {"atol": 1e-7, "rtol": 1e-7},
     InfiniDtype.BF16: {"atol": 1e-3, "rtol": 1e-3},
+    InfiniDtype.I32: {"atol": 0, "rtol": 0},
+    InfiniDtype.I64: {"atol": 0, "rtol": 0},
 }
 
 DEBUG = False

test/infiniop/libinfiniop/utils.py

@@ -70,9 +70,19 @@ class TestTensor(CTensor):
             else:
                 torch_shape.append(shape[i])
         if mode == "random":
-            self._torch_tensor = torch.rand(
-                torch_shape, dtype=to_torch_dtype(dt), device=torch_device_map[device]
-            )
+            # For integer types, use randint instead of rand
+            if dt in [InfiniDtype.I8, InfiniDtype.I16, InfiniDtype.I32, InfiniDtype.I64,
+                      InfiniDtype.U8, InfiniDtype.U16, InfiniDtype.U32, InfiniDtype.U64,
+                      InfiniDtype.BYTE, InfiniDtype.BOOL]:
+                randint_low = -2000000000 if randint_low is None else randint_low
+                randint_high = 2000000000 if randint_high is None else randint_high
+                self._torch_tensor = torch.randint(
+                    randint_low, randint_high, torch_shape, dtype=to_torch_dtype(dt), device=torch_device_map[device]
+                )
+            else:
+                self._torch_tensor = torch.rand(
+                    torch_shape, dtype=to_torch_dtype(dt), device=torch_device_map[device]
+                )
         elif mode == "zeros":
             self._torch_tensor = torch.zeros(
                 torch_shape, dtype=to_torch_dtype(dt), device=torch_device_map[device]