add layout op runtime function

[jjsjann123]

#5118 PR3: enable codegen for layout op
#5115 PR2: add layout op runtime function <- this PR
#5114 PR1: add layout op

Runtime function signature.

  template <
      typename T,
      typename Index_T,
      int BLOCK_ROW_OUTER,
      int BLOCK_ROW_INNER,
      int BLOCK_COL,
      int UNROLL_FACTOR>
  __device__ void preprocessGroupedMatmulInputSf(
      T* output,
      const T* input,
      const nvfuser_index_t row_idx,
      const nvfuser_index_t col_idx,
      const Index_T* input_offsets,
      const Index_T* output_offsets,
      const nvfuser_index_t col_size,
      const nvfuser_index_t group_size)

where:
  BLOCK_ROW_OUTER, BLOCK_ROW_INNER, BLOCK_COL will be translated from
    BlockScalingFactorLayout e.g. Block128x4 is translated to 32, 4, 4.

This function will be used by codegen for `PreprocessGroupedMatmulInputSf`
  `output` is expected to be the beginning of output buffer,
     indexing will be done inside the function template with help of `row_idx`,
     `col_idx`, `expert_offsets`, `output_offsets` and `col_size`
  Meanwhil, indexing on `input` would have been resolved during device lowering.

Todo for future PRs:
Add vectorization support.

[github-actions]

Review updated until commit 967b7f5

Description

• Add runtime function for grouped block layout

• Include block layout resource in kernel code

• Update kernel summary with layout op support

• Implement swizzled memory layout transformation

---

Changes walkthrough 📝

	Relevant files
Enhancement	kernel.cpp Handle layout op in kernel IR scanner                                        csrc/kernel.cpp Added handler for PreprocessGroupedMatmulInputSf op Updates kernel IR scanner to track layout op usage +4/-0      compiled_kernel.cpp Integrate block layout resource in kernel codegen                csrc/runtime/compiled_kernel.cpp Added has_block_layout parameter to _getStructuredCode Include block_layout_cu resource if layout op is used Pass layout op flag from kernel summary +8/-2      block_layout.cu Implement grouped block layout preprocessing                          runtime/block_layout.cu Implemented preprocessGroupedMatmulInputSf device function Added swizzle padding logic for memory layout Supports grouped matmul input preprocessing Includes TODO for vectorization support +102/-0  kernel.h Add layout op flag in kernel summary                                          csrc/kernel.h Added has_preprocess_grouped_matmul_input_sf flag Tracks layout op usage in kernel summary +3/-0
Configuration changes	CMakeLists.txt Include block layout in build configuration                            CMakeLists.txt Added block_layout.cu to runtime files list Ensures layout code is compiled into runtime +1/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review Possible Issue The loop in preprocessGroupedMatmulInputSf uses input[i] for loading values, but it should index into the correct position within the current expert's input data, likely using input + input_offsets[expert_id] + c_row_idx * col_size + col_idx + i or similar, depending on layout. for (int i = 0; i < UNROLL_FACTOR && col_idx + i < col_size; ++i) { 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]; } Performance Issue The function uses a loop for unrolling with no vectorization, which may lead to suboptimal memory throughput. The comment indicates awareness, but without vectorized loads/stores, performance could be significantly impacted, especially for small UNROLL_FACTOR values. // 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 = outputOffsetAfterSwizzlePadding< BLOCK_ROW_OUTER, BLOCK_ROW_INNER, BLOCK_COL>(c_row_idx, col_idx + i, padded_col_size); out_group_offset[index] = input[i]; } } Inconsistent Flag Usage The has_block_layout parameter in _getStructuredCode is set using has_preprocess_grouped_matmul_input_sf, suggesting a mismatch between the operation name and the resource inclusion flag, which could lead to confusion or incorrect code generation in the future. kernel()->summary().has_preprocess_grouped_matmul_input_sf);

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

!test

[naoyam]

Is there any not-so-tedious way to verify the result?

[naoyam]

Stamping