method renaming

Author: jjsjann123

csrc/kernel.cpp

@@ -272,8 +272,8 @@ class KernelIrScanner : private IrVisitor {
     summary_.has_argsort = true;
   }
 
-  void handle(GroupedBlockScalingFactorLayoutOp* aop) final {
-    summary_.has_grouped_block_sf_layout = true;
+  void handle(PreprocessGroupedMatmulInputSf* aop) final {
+    summary_.has_preprocess_grouped_matmul_input_sf = true;
   }
 
   void handle(TopKOp* top) final {

csrc/kernel.h

@@ -143,7 +143,7 @@ struct KernelSummary {
   bool has_argsort = false;
 
   //! Do we have any grouped_block_sf_layout op?
-  bool has_grouped_block_sf_layout = false;
+  bool has_preprocess_grouped_matmul_input_sf = false;
 
   //! Do we have any topk op?
   bool has_topk = false;

csrc/runtime/compiled_kernel.cpp

@@ -1445,7 +1445,7 @@ std::string CompiledKernel::getStructuredCode() const {
       kernel()->summary().has_argsort,
       kernel()->summary().has_topk,
       kernel()->summary().has_scan,
-      kernel()->summary().has_grouped_block_sf_layout);
+      kernel()->summary().has_preprocess_grouped_matmul_input_sf);
 }
 
 std::string CompiledKernel::disassembledKernelSASS() const {

runtime/block_layout.cu

@@ -12,7 +12,7 @@ namespace {
 
 // TODO: support vectorized store
 template <int BLOCK_ROW_OUTER, int BLOCK_ROW_INNER, int BLOCK_COL>
-__device__ nvfuser_index_t offsetAfterSwizzlePadding(
+__device__ nvfuser_index_t outputOffsetAfterSwizzlePadding(
     const nvfuser_index_t row_idx,
     const nvfuser_index_t col_idx,
     const nvfuser_index_t padded_col_size) {
@@ -64,36 +64,37 @@ template <
     int BLOCK_ROW_INNER,
     int BLOCK_COL,
     int UNROLL_FACTOR>
-__device__ void groupedBlockLayout(
+__device__ void preprocessGroupedMatmulInputSf(
     T* output,
     const T* input,
     const nvfuser_index_t row_idx,
     const nvfuser_index_t col_idx,
-    const Index_T* expert_offsets,
+    const Index_T* input_offsets,
     const Index_T* output_offsets,
     const nvfuser_index_t col_size,
     const nvfuser_index_t group_size) {
   // find corresponding expert_id
   int expert_id = 0;
   for (int i = 0; i < group_size; ++i) {
-    if (row_idx < expert_offsets[i + 1]) {
+    if (row_idx < input_offsets[i + 1]) {
       expert_id = i;
       break;
     }
   }
 
   // row idx for current group
-  nvfuser_index_t c_row_idx = row_idx - expert_offsets[expert_id];
+  nvfuser_index_t c_row_idx = row_idx - input_offsets[expert_id];
   // compute output group offset for current group
   nvfuser_index_t padded_col_size =
       (col_size + BLOCK_COL - 1) / BLOCK_COL * BLOCK_COL;
   T* out_group_offset = output + output_offsets[expert_id] * padded_col_size;
 
   // TODO: vectorized load/store instead of for loop
   for (int i = 0; i < UNROLL_FACTOR && col_idx + i < col_size; ++i) {
-    nvfuser_index_t index =
-        offsetAfterSwizzlePadding<BLOCK_ROW_OUTER, BLOCK_ROW_INNER, BLOCK_COL>(
-            c_row_idx, col_idx + i, padded_col_size);
+    nvfuser_index_t index = outputOffsetAfterSwizzlePadding<
+        BLOCK_ROW_OUTER,
+        BLOCK_ROW_INNER,
+        BLOCK_COL>(c_row_idx, col_idx + i, padded_col_size);
     out_group_offset[index] = input[i];
   }
 }