Batch Matmul Op

lib/kernels/include/kernels/batch_matmul_kernels.h

@@ -6,17 +6,26 @@
 
 namespace FlexFlow {
 
-class BatchMatmulPerDeviceState : public PerDeviceOpState {
-public:
-  BatchMatmulPerDeviceState(FFHandler handler);
-  int a_seq_length_dim, b_seq_length_dim;
+struct BMMPerDeviceState {
+  PerDeviceFFHandle handle;
+  Allocator allocator;
+  int a_seq_length_dim;
+  int b_seq_length_dim;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(BMMPerDeviceState,
+                          handle,);
+
 namespace Kernels {
 namespace BatchMatmul {
 
+BMMPerDeviceState init_kernel(PerDeviceFFHandle handle,
+                  Allocator allocator,
+                  int a_seq_length_dim,
+                  int b_seq_length_dim);
+
 void forward_kernel(ffStream_t stream,
-                    BatchMatmulPerDeviceState const *,
+                    BMMPerDeviceState const *meta,
                     float *o_ptr,
                     float const *a_ptr,
                     float const *b_ptr,
@@ -25,12 +34,10 @@ void forward_kernel(ffStream_t stream,
                     int n,
                     int k,
                     int batch,
-                    int a_seq_length_dim = -1,
-                    int b_seq_length_dim = -1,
                     int seq_length = -1);
 
 void backward_kernel(ffStream_t stream,
-                     BatchMatmulPerDeviceState const *,
+                    BMMPerDeviceState const *meta,
                      float const *o_ptr,
                      float const *o_grad_ptr,
                      float const *a_ptr,

lib/kernels/src/hip/batch_matmul_kernels.cpp

@@ -19,9 +19,6 @@
 
 namespace FlexFlow {
 
-BatchMatmulPerDeviceState::BatchMatmulPerDeviceState(FFHandler handler)
-    : PerDeviceOpState(handler) {}
-
 namespace Kernels {
 namespace BatchMatmul {
 
@@ -42,9 +39,9 @@ void forward_kernel(hipStream_t stream,
                     int k,
                     int batch,
                     hipStream_t stream,
-                    int a_seq_length_dim,
-                    int b_seq_length_dim,
                     int seq_length) {
+  int a_seq_length_dim = meta->a_seq_length_dim;
+  int b_seq_length_dim = meta->b_seq_length_dim;
   checkCUDA(hipblasSetStream(meta->handle.blas, stream));
   checkCUDNN(miopenSetStream(meta->handle.dnn, stream));
 

lib/op-attrs/include/op-attrs/ops/batch_matmul.h

@@ -12,8 +12,10 @@ struct BatchMatmulAttrs {
 };
 FF_VISITABLE_STRUCT(BatchMatmulAttrs, a_seq_length_dim, b_seq_length_dim);
 
-CHECK_VALID_OP_ATTR(BatchMatmulAttrs);
+int get_aSeqLengthDim(BatchMatmulAttrs const &attrs);
+int get_bSeqLengthDim(BatchMatmulAttrs const &attrs);
 
+CHECK_VALID_OP_ATTR(BatchMatmulAttrs);
 } // namespace FlexFlow
 
 #endif

lib/op-attrs/src/batch_matmul.cc

@@ -2,6 +2,14 @@
 
 namespace FlexFlow {
 
+int get_aSeqLengthDim(BatchMatmulAttrs const &attrs) {
+  return attrs.a_seq_length_dim;
+}
+
+int get_bSeqLengthDim(BatchMatmulAttrs const &attrs) {
+  return attrs.b_seq_length_dim;
+}
+
 /* bool BatchMatmulAttrs::is_valid( */
 /*     ParallelTensorShape const &lhs, ParallelTensorShape const &rhs) const {
  */