Merge remote-tracking branch 'origin/repo-refactor' into r-refactor

Bob-Chen222 · Jun 5, 2024 · d785ece · d785ece
2 parents a2a3dd8 + 6fe5dba
commit d785ece
Show file tree

Hide file tree

Showing 27 changed files with 134 additions and 126 deletions.
diff --git a/lib/kernels/src/hip/loss_function_kernels.cpp b/lib/kernels/src/hip/loss_function_kernels.cpp
@@ -13,14 +13,12 @@
  * limitations under the License.
  */
 
-#include "flexflow/model.h"
-#include "utils/hip_helper.h"
+#include "kernels/loss_function_kernels.h"
+#include "device.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
 
-using namespace Legion;
-
 __global__ void
     sparse_categorical_crossentropy_loss_backward(float *logit_grad,
                                                   int const *label,
@@ -59,7 +57,8 @@ __global__ void identity_loss_backward(float *loss_grad,
   }
 }
 
-void Loss::sparse_categorical_crossentropy_loss_backward_kernel_wrapper(
+void sparse_categorical_crossentropy_loss_backward_kernel(
+    hipStream_t stream,
     float *logit_grad_ptr,
     float const *logit_ptr,
     int const *label_ptr,
@@ -69,12 +68,14 @@ void Loss::sparse_categorical_crossentropy_loss_backward_kernel_wrapper(
     int num_classes,
     int k,
     float scale_factor) {
-  hipStream_t stream;
+  // hipStream_t stream;
   checkCUDA(get_legion_stream(&stream));
   checkCUDA(hipMemcpy(logit_grad_ptr,
                       logit_ptr,
                       logit_volume * sizeof(float),
                       hipMemcpyDeviceToDevice));
+
+  // launch kernel in hip
   hipLaunchKernelGGL(sparse_categorical_crossentropy_loss_backward,
                      GET_BLOCKS(num_samples),
                      CUDA_NUM_THREADS,
@@ -94,17 +95,17 @@ void Loss::sparse_categorical_crossentropy_loss_backward_kernel_wrapper(
                      logit_grad_ptr,
                      logit_grad_volume,
                      0,
-                     scale_factor * k);
+                     scale_factor);
 }
 
-void Loss::categorical_crossentropy_loss_backward_kernel_wrapper(
-    float *logit_grad_ptr,
-    float const *logit_ptr,
-    float const *label_ptr,
-    size_t logit_volume,
-    size_t logit_grad_volume,
-    float scale_factor) {
-  hipStream_t stream;
+void categorical_crossentropy_loss_backward_kernel(hipStream_t stream,
+                                                   float *logit_grad_ptr,
+                                                   float const *logit_ptr,
+                                                   float const *label_ptr,
+                                                   size_t logit_volume,
+                                                   size_t logit_grad_volume,
+                                                   float scale_factor) {
+  // hipStream_t stream;
   checkCUDA(get_legion_stream(&stream));
   hipLaunchKernelGGL(categorical_crossentropy_loss_backward,
                      GET_BLOCKS(logit_volume),
@@ -115,6 +116,7 @@ void Loss::categorical_crossentropy_loss_backward_kernel_wrapper(
                      logit_ptr,
                      label_ptr,
                      logit_volume);
+
   // Scale logit gradients by loss->scale_factor
   hipLaunchKernelGGL(scale_kernel,
                      GET_BLOCKS(logit_grad_volume),
@@ -127,14 +129,14 @@ void Loss::categorical_crossentropy_loss_backward_kernel_wrapper(
                      scale_factor);
 }
 
-void Loss::mean_squared_error_avg_loss_backward_kernel_wrapper(
-    float *logit_grad_ptr,
-    float const *logit_ptr,
-    float const *label_ptr,
-    size_t logit_volume,
-    size_t logit_grad_volume,
-    float scale_factor) {
-  hipStream_t stream;
+void mean_squared_error_avg_loss_backward_kernel(hipStream_t stream,
+                                                 float *logit_grad_ptr,
+                                                 float const *logit_ptr,
+                                                 float const *label_ptr,
+                                                 size_t logit_volume,
+                                                 size_t logit_grad_volume,
+                                                 float scale_factor) {
+  // hipStream_t stream;
   checkCUDA(get_legion_stream(&stream));
   hipLaunchKernelGGL(mean_squared_error_avg_loss_backward,
                      GET_BLOCKS(logit_volume),
@@ -157,12 +159,13 @@ void Loss::mean_squared_error_avg_loss_backward_kernel_wrapper(
                      scale_factor);
 }
 
-void Loss::identity_loss_backward_kernel_wrapper(float *loss_grad_ptr,
-                                                 float const *loss_ptr,
-                                                 size_t loss_volume,
-                                                 size_t loss_grad_volume,
-                                                 float scale_factor) {
-  hipStream_t stream;
+void identity_loss_backward_kernel(hipStream_t stream,
+                                   float *loss_grad_ptr,
+                                   float const *loss_ptr,
+                                   size_t loss_volume,
+                                   size_t loss_grad_volume,
+                                   float scale_factor) {
+  // hipStream_t stream;
   checkCUDA(get_legion_stream(&stream));
   hipLaunchKernelGGL(identity_loss_backward,
                      GET_BLOCKS(loss_volume),

diff --git a/lib/kernels/src/hip/attention_kernels.cpp → ...kernels/src/hip/ops/attention_kernels.cpp b/lib/kernels/src/hip/attention_kernels.cpp → ...kernels/src/hip/ops/attention_kernels.cpp
diff --git a/lib/kernels/src/hip/batch_matmul_kernels.cpp → ...nels/src/hip/ops/batch_matmul_kernels.cpp b/lib/kernels/src/hip/batch_matmul_kernels.cpp → ...nels/src/hip/ops/batch_matmul_kernels.cpp
diff --git a/lib/kernels/src/hip/batch_norm_kernels.cpp → ...ernels/src/hip/ops/batch_norm_kernels.cpp b/lib/kernels/src/hip/batch_norm_kernels.cpp → ...ernels/src/hip/ops/batch_norm_kernels.cpp
diff --git a/lib/kernels/src/hip/cast_kernels.cpp → lib/kernels/src/hip/ops/cast_kernels.cpp b/lib/kernels/src/hip/cast_kernels.cpp → lib/kernels/src/hip/ops/cast_kernels.cpp
diff --git a/lib/kernels/src/hip/combine_kernels.cpp → lib/kernels/src/hip/ops/combine_kernels.cpp b/lib/kernels/src/hip/combine_kernels.cpp → lib/kernels/src/hip/ops/combine_kernels.cpp
diff --git a/lib/kernels/src/hip/concat_kernels.cpp → lib/kernels/src/hip/ops/concat_kernels.cpp b/lib/kernels/src/hip/concat_kernels.cpp → lib/kernels/src/hip/ops/concat_kernels.cpp
diff --git a/lib/kernels/src/hip/conv_2d_kernels.cpp → lib/kernels/src/hip/ops/conv_2d_kernels.cpp b/lib/kernels/src/hip/conv_2d_kernels.cpp → lib/kernels/src/hip/ops/conv_2d_kernels.cpp
diff --git a/lib/kernels/src/hip/conv_2d_kernels.h → lib/kernels/src/hip/ops/conv_2d_kernels.h b/lib/kernels/src/hip/conv_2d_kernels.h → lib/kernels/src/hip/ops/conv_2d_kernels.h
diff --git a/lib/kernels/src/hip/dropout_kernels.cpp → lib/kernels/src/hip/ops/dropout_kernels.cpp b/lib/kernels/src/hip/dropout_kernels.cpp → lib/kernels/src/hip/ops/dropout_kernels.cpp
@@ -14,14 +14,15 @@
  */
 
 #include "kernels/dropout_kernels.h"
-#include "kernels/hip_helper.h"
+#include "device.h"
+#include "kernels/ff_handle.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
 namespace Kernels {
 namespace Dropout {
 
-DropoutPerDeviceState init_kernel(PerDeviceFFHandle handler,
+DropoutPerDeviceState init_kernel(PerDeviceFFHandle handle,
                                   float rate,
                                   unsigned long long seed,
                                   ArrayShape output_shape,
@@ -47,8 +48,6 @@ DropoutPerDeviceState init_kernel(PerDeviceFFHandle handler,
     dropoutStates = allocator.allocate(totalSize);
     reserveSpace = ((char *)dropoutStates) + dropoutStateSize;
   }
-  // checkCUDA(hipMalloc(&dropoutStates, dropoutStateSize));
-  // checkCUDA(hipMalloc(&reserveSpace, reserveSpaceSize));
   checkCUDNN(miopenSetDropoutDescriptor(dropoutDesc,
                                         handle.dnn,
                                         rate,
@@ -74,7 +73,7 @@ void forward_kernel(hipStream_t stream,
                     DropoutPerDeviceState &m,
                     float const *input_ptr,
                     float *output_ptr) {
-  checkCUDNN(miopenSetStream(m->handle.dnn, stream));
+  checkCUDNN(miopenSetStream(m.handle.dnn, stream));
 
   checkCUDNN(miopenDropoutForward(m.handle.dnn,
                                   m.dropoutDesc,
@@ -91,7 +90,7 @@ void backward_kernel(hipStream_t stream,
                      DropoutPerDeviceState &m,
                      float const *output_grad_ptr,
                      float *input_grad_ptr) {
-  checkCUDNN(miopenSetStream(m->handle.dnn, stream));
+  checkCUDNN(miopenSetStream(m.handle.dnn, stream));
 
   checkCUDNN(miopenDropoutBackward(m.handle.dnn,
                                    m.dropoutDesc,

diff --git a/...ernels/src/hip/element_binary_kernels.cpp → ...ls/src/hip/ops/element_binary_kernels.cpp b/...ernels/src/hip/element_binary_kernels.cpp → ...ls/src/hip/ops/element_binary_kernels.cpp
diff --git a/...kernels/src/hip/element_unary_kernels.cpp → ...els/src/hip/ops/element_unary_kernels.cpp b/...kernels/src/hip/element_unary_kernels.cpp → ...els/src/hip/ops/element_unary_kernels.cpp
diff --git a/lib/kernels/src/hip/flat_kernels.cpp → lib/kernels/src/hip/ops/flat_kernels.cpp b/lib/kernels/src/hip/flat_kernels.cpp → lib/kernels/src/hip/ops/flat_kernels.cpp
@@ -14,11 +14,11 @@
  */
 
 #include "kernels/flat_kernels.h"
-#include "kernels/hip_helper.h"
+#include "device.h"
+#include "kernels/accessor.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
-
 namespace Kernels {
 namespace Flat {
 
@@ -31,7 +31,6 @@ void forward_kernel(hipStream_t stream,
                            (input.shape.num_elements()) * sizeof(float),
                            hipMemcpyDeviceToDevice,
                            stream));
-  // checkCUDA(hipDeviceSynchronize());
 }
 
 void backward_kernel(hipStream_t stream,
@@ -49,10 +48,6 @@ void backward_kernel(hipStream_t stream,
                      output_grad_ptr,
                      input.shape.num_elements(),
                      alpha);
-  // checkCUDA(hipMemcpyAsync(acc_input_grad.ptr, acc_output_grad.ptr,
-  //                           acc_input_grad.rect.volume() * sizeof(float),
-  //                           hipMemcpyDeviceToDevice));
-  // checkCUDA(hipDeviceSynchronize());
 }
 
 } // namespace Flat

diff --git a/lib/kernels/src/hip/gather_kernels.cpp → lib/kernels/src/hip/ops/gather_kernels.cpp b/lib/kernels/src/hip/gather_kernels.cpp → lib/kernels/src/hip/ops/gather_kernels.cpp
@@ -14,15 +14,11 @@
  */
 
 #include "kernels/gather_kernels.h"
+#include "device.h"
 #include "kernels/datatype_dispatch.h"
-#include "kernels/hip_helper.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
-
-GatherPerDeviceState::GatherPerDeviceState(FFHandler handler)
-    : PerDeviceOpState(handler){};
-
 namespace Kernels {
 namespace Gather {
 

diff --git a/lib/kernels/src/hip/layer_norm_kernels.cpp → ...ernels/src/hip/ops/layer_norm_kernels.cpp b/lib/kernels/src/hip/layer_norm_kernels.cpp → ...ernels/src/hip/ops/layer_norm_kernels.cpp
diff --git a/lib/kernels/src/hip/linear_kernels.cpp → lib/kernels/src/hip/ops/linear_kernels.cpp b/lib/kernels/src/hip/linear_kernels.cpp → lib/kernels/src/hip/ops/linear_kernels.cpp
diff --git a/lib/kernels/src/hip/partition_kernels.cpp → ...kernels/src/hip/ops/partition_kernels.cpp b/lib/kernels/src/hip/partition_kernels.cpp → ...kernels/src/hip/ops/partition_kernels.cpp
diff --git a/lib/kernels/src/hip/pool_2d_kernels.cpp → lib/kernels/src/hip/ops/pool_2d_kernels.cpp b/lib/kernels/src/hip/pool_2d_kernels.cpp → lib/kernels/src/hip/ops/pool_2d_kernels.cpp
diff --git a/lib/kernels/src/hip/reduce_kernels.cpp → lib/kernels/src/hip/ops/reduce_kernels.cpp b/lib/kernels/src/hip/reduce_kernels.cpp → lib/kernels/src/hip/ops/reduce_kernels.cpp
diff --git a/lib/kernels/src/hip/reduction_kernels.cpp → ...kernels/src/hip/ops/reduction_kernels.cpp b/lib/kernels/src/hip/reduction_kernels.cpp → ...kernels/src/hip/ops/reduction_kernels.cpp
diff --git a/lib/kernels/src/hip/replicate_kernels.cpp → ...kernels/src/hip/ops/replicate_kernels.cpp b/lib/kernels/src/hip/replicate_kernels.cpp → ...kernels/src/hip/ops/replicate_kernels.cpp
diff --git a/lib/kernels/src/hip/reshape_kernels.cpp → lib/kernels/src/hip/ops/reshape_kernels.cpp b/lib/kernels/src/hip/reshape_kernels.cpp → lib/kernels/src/hip/ops/reshape_kernels.cpp
diff --git a/lib/kernels/src/hip/reverse_kernels.cpp → lib/kernels/src/hip/ops/reverse_kernels.cpp b/lib/kernels/src/hip/reverse_kernels.cpp → lib/kernels/src/hip/ops/reverse_kernels.cpp
diff --git a/lib/kernels/src/hip/softmax_kernels.cpp → lib/kernels/src/hip/ops/softmax_kernels.cpp b/lib/kernels/src/hip/softmax_kernels.cpp → lib/kernels/src/hip/ops/softmax_kernels.cpp
@@ -14,40 +14,36 @@
  */
 
 #include "kernels/softmax_kernels.h"
-#include "kernels/hip_helper.h"
+#include "device.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
-// declare Legion names
-using Legion::Domain;
-
-SoftmaxPerDeviceState::SoftmaxPerDeviceState(FFHandler handler,
-                                             Softmax const *softmax,
-                                             Domain const &input_domain)
-    : PerDeviceOpState(handler) {
-  checkCUDNN(miopenCreateTensorDescriptor(&inputTensor));
-  checkCUDNN(cudnnSetTensorDescriptorFromDomain(inputTensor, input_domain));
-  dim = softmax->dim;
-  profiling = softmax->profiling;
-  std::strcpy(op_name, softmax->name);
-}
 
 namespace Kernels {
 namespace Softmax {
 
+SoftmaxPerDeviceState init_kernel(PerDeviceFFHandle const &handle, int dim) {
+  ffTensorDescriptor_t inputTensor;
+
+  checkCUDNN(miopenCreateTensorDescriptor(&inputTensor));
+
+  SoftmaxPerDeviceState per_device_state = {handle, inputTensor, dim};
+  return per_device_state;
+}
+
 void forward_kernel(hipStream_t stream,
-                    SoftmaxPerDeviceState const *m,
+                    SoftmaxPerDeviceState const &m,
                     float const *input_ptr,
                     float *output_ptr) {
-  checkCUDNN(miopenSetStream(m->handle.dnn, stream));
+  checkCUDNN(miopenSetStream(m.handle.dnn, stream));
 
   float alpha = 1.0f, beta = 0.0f;
-  checkCUDNN(miopenSoftmaxForward_V2(m->handle.dnn,
+  checkCUDNN(miopenSoftmaxForward_V2(m.handle.dnn,
                                      &alpha,
-                                     m->inputTensor,
+                                     m.inputTensor,
                                      input_ptr,
                                      &beta,
-                                     m->inputTensor,
+                                     m.inputTensor,
                                      output_ptr,
                                      MIOPEN_SOFTMAX_ACCURATE,
                                      MIOPEN_SOFTMAX_MODE_CHANNEL));

diff --git a/lib/kernels/src/hip/split_kernels.cpp → lib/kernels/src/hip/ops/split_kernels.cpp b/lib/kernels/src/hip/split_kernels.cpp → lib/kernels/src/hip/ops/split_kernels.cpp
@@ -14,12 +14,10 @@
  */
 
 #include "kernels/split_kernels.h"
-#include "kernels/hip_helper.h"
+#include "device.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
-// declare Legion names
-using Legion::coord_t;
 
 namespace Kernels {
 namespace Split {

diff --git a/lib/kernels/src/hip/topk_kernels.cpp → lib/kernels/src/hip/ops/topk_kernels.cpp b/lib/kernels/src/hip/topk_kernels.cpp → lib/kernels/src/hip/ops/topk_kernels.cpp
@@ -14,15 +14,10 @@
  */
 
 #include "kernels/topk_kernels.h"
-#include "kernels/hip_helper.h"
+#include "device.h"
 #include <hip/hip_runtime.h>
 
 namespace FlexFlow {
-// declare Legion names
-using Legion::coord_t;
-
-TopKPerDeviceState::TopKPerDeviceState(FFHandler handler)
-    : PerDeviceOpState(handler) {}
 
 namespace Kernels {
 namespace TopK {
@@ -36,6 +31,11 @@ struct Entry {
   T value;
 };
 
+TopKPerDeviceState init_kernel(bool sorted) {
+  TopKPerDeviceState per_device_state = {sorted};
+  return per_device_state;
+}
+
 template <typename T>
 struct LinearData {
   typedef Entry<T> Entry;
@@ -371,7 +371,7 @@ __global__ void topk_forward_kernel(T const *__restrict__ input,
 }
 
 void forward_kernel(hipStream_t stream,
-                    TopKPerDeviceState const *m,
+                    TopKPerDeviceState const &m,
                     float const *input_ptr,
                     float *output_ptr,
                     int *indices_ptr,
@@ -428,7 +428,7 @@ __global__ void topk_backward_kernel(T const *__restrict__ value_grad_ptr,
 }
 
 void backward_kernel(hipStream_t stream,
-                     TopKPerDeviceState const *m,
+                     TopKPerDeviceState const &m,
                      float const *value_grad_ptr,
                      int const *indices_ptr,
                      float *in_grad_ptr,