Add i32-type

Author: junjihashimoto

gpu.h

@@ -190,7 +190,8 @@ struct TensorPool {
 
 enum NumType {
   kf16, // (experimental)
-  kf32
+  kf32,
+  ki32
 };
 
 /**
@@ -202,6 +203,8 @@ inline size_t sizeBytes(const NumType &type) {
     return sizeof(uint16_t);
   case kf32:
     return sizeof(float);
+  case ki32:
+    return sizeof(int32_t);
   default:
     LOG(kDefLog, kError, "Invalid NumType in size calculation.");
     return 0;
@@ -217,6 +220,8 @@ inline std::string toString(NumType type) {
     return "f16";
   case kf32:
     return "f32";
+  case ki32:
+    return "i32";
   default:
     LOG(kDefLog, kError, "Invalid NumType in string conversion.");
     return "unknown";
@@ -635,6 +640,18 @@ inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype,
   return tensor;
 }
 
+inline Tensor createTensor(Context &ctx, const Shape &shape, NumType dtype,
+                           const int32_t *data) {
+  assert(dtype == ki32);
+  Tensor tensor =
+      createTensor(ctx.pool, ctx.device, shape, dtype,
+                   WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |
+                       WGPUBufferUsage_CopySrc);
+  wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data,
+                       tensor.data.size);
+  return tensor;
+}
+
 /**
  * @brief Overload of the tensor factory function to instantiate a tensor on
  * the GPU with a given shape, data type. This overload also takes initial