Change datatype for linear kernels away from void * in .cc

lib/kernels/include/kernels/array_shape.h

@@ -42,7 +42,7 @@ struct ArrayShape {
 
   ArrayShape reversed_dim_order() const;
   ArrayShape sub_shape(std::optional<legion_dim_t> start,
-                       std::optional<legion_dim_t> end);
+                       std::optional<legion_dim_t> end) const;
 
 public:
   LegionTensorDims dims;

lib/kernels/include/kernels/element_unary_kernels.h

@@ -9,9 +9,6 @@
 
 namespace FlexFlow {
 
-using ElementUnaryUnifiedAttrs =
-    std::variant<ElementUnaryAttrs, ElementScalarUnaryAttrs>;
-
 struct ElementUnaryPerDeviceState {
   ffTensorDescriptor_t inputTensor, outputTensor;
   req<ffActivationDescriptor_t> actiDesc;

lib/kernels/include/kernels/gather_kernels.h

@@ -2,36 +2,34 @@
 #define _FLEXFLOW_OPS_KERNELS_GATHER_KERNELS_H
 
 #include "accessor.h"
-#include "device.h"
+#include "kernels/device.h"
 
 namespace FlexFlow {
 
 struct GatherPerDeviceState {
-  int legion_dim;
-  req<DataType> index_data_type;
+  PerDeviceFFHandle handle;
+  legion_dim_t legion_dim;
 };
+
 FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(GatherPerDeviceState,
-                                             legion_dim,
-                                             index_data_type);
+                                             handle,
+                                             legion_dim);
 
 namespace Kernels {
 namespace Gather {
+
 void forward_kernel(ffStream_t stream,
                     GatherPerDeviceState const &m,
                     GenericTensorAccessorR const &input,
                     GenericTensorAccessorR const &index,
-                    GenericTensorAccessorW const &output,
-                    size_t stride,
-                    size_t input_dim_size,
-                    size_t output_dim_size);
+                    GenericTensorAccessorW const &output);
+
 void backward_kernel(ffStream_t stream,
                      GatherPerDeviceState const &m,
                      GenericTensorAccessorR const &output_grad,
                      GenericTensorAccessorR const &index,
-                     GenericTensorAccessorW const &input_grad,
-                     size_t stride,
-                     size_t input_dim_size,
-                     size_t output_dim_size);
+                     GenericTensorAccessorW const &input_grad);
+
 } // namespace Gather
 } // namespace Kernels
 } // namespace FlexFlow

lib/kernels/include/kernels/linear_kernels.h

@@ -50,22 +50,22 @@ bool use_activation(Activation activation);
 
 void forward_kernel(ffStream_t stream,
                     LinearPerDeviceState const &m,
-                    void const *input_ptr,
-                    void *output_ptr,
-                    void const *filter_ptr,
-                    void const *bias_ptr,
+                    float const *input_ptr,
+                    float *output_ptr,
+                    float const *filter_ptr,
+                    float const *bias_ptr,
                     int in_dim,
                     int out_dim,
                     int batch_size);
 void backward_kernel(ffStream_t stream,
                      LinearPerDeviceState const &m,
-                     void const *input_ptr,
-                     void *input_grad_ptr,
-                     void const *output_ptr,
-                     void *output_grad_ptr,
-                     void const *kernel_ptr,
-                     void *kernel_grad_ptr,
-                     void *bias_ptr,
+                     float const *input_ptr,
+                     float *input_grad_ptr,
+                     float const *output_ptr,
+                     float *output_grad_ptr,
+                     float const *kernel_ptr,
+                     float *kernel_grad_ptr,
+                     float *bias_ptr,
                      int in_dim,
                      int out_dim,
                      int batch_size);

lib/kernels/src/cuda/ops/gather_kernels.cu

@@ -25,10 +25,10 @@ template <typename IndexType>
 __global__ void gather_forward(float const *input,
                                IndexType const *index,
                                float *output,
-                               size_t output_size,
-                               size_t stride,
-                               size_t input_dim_size,
-                               size_t output_dim_size) {
+                               coord_t output_size,
+                               coord_t stride,
+                               coord_t input_dim_size,
+                               coord_t output_dim_size) {
   CUDA_KERNEL_LOOP(o, output_size) {
     // output tensor shape: [*, output_dim_size, stride]
     // output tensor stride: [output_dim_size * stride, stride, 1]
@@ -39,10 +39,10 @@ __global__ void gather_forward(float const *input,
     // [outer_index, index[0], left_over]
     // Therefore, input_index = outer_index * (stride * input_dim_size)
     //                        + index[0] * stride + left_over;
-    size_t outer_index = o / (stride * output_dim_size);
+    coord_t outer_index = o / (stride * output_dim_size);
     // coord_t index_2 = (o / stride) % dim_size
-    size_t left_over = o % stride;
-    size_t input_idx =
+    coord_t left_over = o % stride;
+    coord_t input_idx =
         outer_index * (stride * input_dim_size) + index[o] * stride + left_over;
     output[o] = input[input_idx];
   }
@@ -52,10 +52,10 @@ template <typename IndexType>
 __global__ void gather_backward(float const *output_grad,
                                 IndexType const *index,
                                 float *input_grad,
-                                size_t output_size,
-                                size_t stride,
-                                size_t input_dim_size,
-                                size_t output_dim_size) {
+                                coord_t output_size,
+                                coord_t stride,
+                                coord_t input_dim_size,
+                                coord_t output_dim_size) {
   CUDA_KERNEL_LOOP(o, output_size) {
     // output tensor shape: [*, output_dim_size, stride]
     // output tensor stride: [output_dim_size * stride, stride, 1]
@@ -66,10 +66,10 @@ __global__ void gather_backward(float const *output_grad,
     // [outer_index, index[0], left_over]
     // Therefore, input_index = outer_index * (stride * input_dim_size)
     //                        + index[0] * stride + left_over;
-    size_t outer_index = o / (stride * output_dim_size);
+    coord_t outer_index = o / (stride * output_dim_size);
     // coord_t index_2 = (o / stride) % dim_size
-    size_t left_over = o % stride;
-    size_t input_idx =
+    coord_t left_over = o % stride;
+    coord_t input_idx =
         outer_index * (stride * input_dim_size) + index[o] * stride + left_over;
 
     atomicAdd(&input_grad[input_idx], output_grad[o]);
@@ -78,100 +78,97 @@ __global__ void gather_backward(float const *output_grad,
 
 template <DataType IndexType>
 struct ForwardKernel {
-  void operator()(cudaStream_t stream,
-                  GatherPerDeviceState const &m,
+  void operator()(ffStream_t stream,
                   GenericTensorAccessorR const &input,
                   GenericTensorAccessorR const &index,
                   GenericTensorAccessorW const &output,
-                  size_t stride,
-                  size_t input_dim_size,
-                  size_t output_dim_size) {
-    /*size_t stride = 1;
-    for (int i = 0; i < m->legion_dim; i++) {
-      stride *= (output.domain.hi()[i] - output.domain.lo()[i] + 1);
-    }
-    size_t dim_size =
-        output.domain.hi()[m->legion_dim] - output.domain.lo()[m->legion_dim] +
-    1;
-*/
-    gather_forward<real_type<IndexType>>
-        <<<GET_BLOCKS(output.shape.get_volume()),
-           CUDA_NUM_THREADS,
-           0,
-           stream>>>(input.get<DataType::FLOAT>(),
-                     index.get<IndexType>(),
-                     output.get<DataType::FLOAT>(),
-                     output.shape.get_volume(),
-                     stride,
-                     input_dim_size,
-                     output_dim_size);
+                  coord_t output_size,
+                  coord_t stride,
+                  coord_t input_dim_size,
+                  coord_t output_dim_size) {
+    gather_forward<<<GET_BLOCKS(output_size), CUDA_NUM_THREADS, 0, stream>>>(
+        input.get_float_ptr(),
+        index.get<IndexType>(),
+        output.get_float_ptr(),
+        output_size,
+        stride,
+        input_dim_size,
+        output_dim_size);
   }
 };
 
-void forward_kernel(cudaStream_t stream,
+template <DataType IndexType>
+struct BackwardKernel {
+  void operator()(ffStream_t stream,
+                  GenericTensorAccessorR const &output_grad,
+                  GenericTensorAccessorR const &index,
+                  GenericTensorAccessorW const &input_grad,
+                  coord_t output_size,
+                  coord_t stride,
+                  coord_t input_dim_size,
+                  coord_t output_dim_size) {
+    gather_backward<<<GET_BLOCKS(output_size), CUDA_NUM_THREADS, 0, stream>>>(
+        output_grad.get_float_ptr(),
+        index.get<IndexType>(),
+        input_grad.get_float_ptr(),
+        output_size,
+        stride,
+        input_dim_size,
+        output_dim_size);
+  }
+};
+
+void forward_kernel(ffStream_t stream,
                     GatherPerDeviceState const &m,
                     GenericTensorAccessorR const &input,
                     GenericTensorAccessorR const &index,
-                    GenericTensorAccessorW const &output,
-                    size_t stride,
-                    size_t input_dim_size,
-                    size_t output_dim_size) {
-  DataTypeDispatch1<ForwardKernel>{}(m.index_data_type,
+                    GenericTensorAccessorW const &output) {
+  checkCUDA(get_legion_stream(&stream));
+
+  coord_t stride =
+      output.shape
+          .sub_shape(std::nullopt, legion_dim_t{m.legion_dim.value() + 1})
+          .get_volume();
+  coord_t output_dim_size = output.shape[m.legion_dim];
+  coord_t input_dim_size = input.shape[m.legion_dim];
+
+  assert(index.data_type == DataType::INT32 ||
+         index.data_type == DataType::INT64);
+
+  DataTypeDispatch1<ForwardKernel>{}(index.data_type,
                                      stream,
-                                     m,
                                      input,
                                      index,
                                      output,
+                                     output.shape.get_volume(),
                                      stride,
                                      input_dim_size,
                                      output_dim_size);
 }
 
-template <DataType IndexType>
-struct BackwardKernel {
-  void operator()(cudaStream_t stream,
-                  GatherPerDeviceState const &m,
-                  GenericTensorAccessorR const &output_grad,
-                  GenericTensorAccessorR const &index,
-                  GenericTensorAccessorW const &input_grad,
-                  size_t stride,
-                  size_t input_dim_size,
-                  size_t output_dim_size) {
-    /*size_t stride = 1;
-    for (int i = 0; i < m->legion_dim; i++) {
-      stride *= (output_grad.domain.hi()[i] - output_grad.domain.lo()[i] + 1);
-    }
-    size_t dim_size = output_grad.domain.hi()[m->legion_dim] -
-                      output_grad.domain.lo()[m->legion_dim] + 1;
-    */
-    gather_backward<real_type<IndexType>>
-        <<<GET_BLOCKS(output_grad.shape.get_volume()),
-           CUDA_NUM_THREADS,
-           0,
-           stream>>>(output_grad.get<DataType::FLOAT>(),
-                     index.get<IndexType>(),
-                     input_grad.get<DataType::FLOAT>(),
-                     output_grad.shape.get_volume(),
-                     stride,
-                     input_dim_size,
-                     output_dim_size);
-  }
-};
-
-void backward_kernel(cudaStream_t stream,
+void backward_kernel(ffStream_t stream,
                      GatherPerDeviceState const &m,
                      GenericTensorAccessorR const &output_grad,
                      GenericTensorAccessorR const &index,
-                     GenericTensorAccessorW const &input_grad,
-                     size_t stride,
-                     size_t input_dim_size,
-                     size_t output_dim_size) {
-  DataTypeDispatch1<BackwardKernel>{}(m.index_data_type,
+                     GenericTensorAccessorW const &input_grad) {
+  checkCUDA(get_legion_stream(&stream));
+
+  coord_t stride =
+      output_grad.shape
+          .sub_shape(std::nullopt, legion_dim_t{m.legion_dim.value() + 1})
+          .get_volume();
+  coord_t output_dim_size = output_grad.shape[m.legion_dim];
+  coord_t input_dim_size = input_grad.shape[m.legion_dim];
+
+  assert(index.data_type == DataType::INT32 ||
+         index.data_type == DataType::INT64);
+
+  DataTypeDispatch1<BackwardKernel>{}(index.data_type,
                                       stream,
-                                      m,
                                       output_grad,
                                       index,
                                       input_grad,
+                                      output_grad.shape.get_volume(),
                                       stride,
                                       input_dim_size,
                                       output_dim_size);