Skip to content

Users/mkulikow/ck/data prefetch in mxgemm pipeline - #9752

Merged
JeniferC99 merged 2 commits into
developfrom
users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline
Jul 27, 2026
Merged

Users/mkulikow/ck/data prefetch in mxgemm pipeline#9752
JeniferC99 merged 2 commits into
developfrom
users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline

Conversation

@AMD-mkulikow

Copy link
Copy Markdown
Contributor

JIRA ID : AICK-1670

Motivation

Added example for data cache prefetch in mx gemm pipeline while also fixing some bugs in data cache prefetch pipeline

Technical Details

Add a standalone flatmm example (mx_flatmm_data_cache_prefetch) that runs
MX GEMM through the compute TDM v1 pipeline (GemmPipelineAgBgCrCompTDMV1)
with hardware data cache prefetch on gfx1250. Prefetch destination is
selectable per operand (A/B) between L1, L2 or None via the
DataCachePrefetchKind trait, exposed through -prefetch_a_l1 / -prefetch_b_l1
CLI flags, with an optional -compare mode against a no-prefetch run.

Guarded behind gfx125 in the 18_flatmm CMakeLists.

Also fix data cache prefetch being silently disabled in the scaled
operator() of GemmPipelineAgBgCrCompTDMV1. The scaled path defaulted
data_cache_prefetch_a/b to false and only set them under
UseClusterLaunch, so with cluster launch off the runtime guards folded
away every prefetch and no global_prefetch_b8 was emitted despite an
L1/L2 policy. Default them to true (matching the non-scaled operator());
emission stays gated by the compile-time UseDataCachePrefetch policy
check, so None still emits nothing.

Test Plan

Checked on simulators:
test name: tile_example_mx_flatmm_mxgemm_data_cache_prefetch

Test Result

fp4 for 512x512x4096:
no prefetch/L1 prefetch: 71,334 / 54,889 ( 23.1% increase )

Submission Checklist

@therock-pr-bot

therock-pr-bot Bot commented Jul 23, 2026

Copy link
Copy Markdown

❌ PR Check — Action Required

Check Status Details
🌿 Branch Name ✅ Pass
📝 PR Title/Description ✅ Pass
Forbidden Files ✅ Pass
🧪 Unit Test ❌ Fail Error: Source/code files changed without an accompanying unit test.
Expected: add at least one test file named like test_<name>.py / test_<name>.cpp (or <name>_test.*).
Current: code file(s) changed: projects/composablekernel/example/ck_tile/18_flatmm/mxgemm/mx_flatmm_data_cache_prefetch.cpp, projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp, projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v1.hpp, projects/composablekernel/test/prefetch_op/prefetch_op_util.hpp; no test file found
🔎 pre-commit ✅ Pass
🚫 Draft PR 🔜 To Be Enabled
🚩 Feature Flag 🔜 To Be Enabled
📊 Code Coverage 🔜 To Be Enabled

⚠️ 1 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

  • ❌ Unit Test

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

@therock-pr-bot

therock-pr-bot Bot commented Jul 23, 2026

Copy link
Copy Markdown

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

  • ❌ Unit Test

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

@AMD-mkulikow
AMD-mkulikow force-pushed the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch from 6981d65 to ed4ce98 Compare July 23, 2026 12:06
@AMD-mkulikow
AMD-mkulikow force-pushed the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch from ed4ce98 to eb24e03 Compare July 23, 2026 14:52
@AMD-mkulikow
AMD-mkulikow force-pushed the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch from eb24e03 to 8fd7608 Compare July 24, 2026 08:15
Add a standalone flatmm example (mx_flatmm_data_cache_prefetch) that runs
MX GEMM through the compute TDM v1 pipeline (GemmPipelineAgBgCrCompTDMV1)
with hardware data cache prefetch on gfx1250. Prefetch destination is
selectable per operand (A/B) between L1, L2 or None via the
DataCachePrefetchKind trait, exposed through -prefetch_a_l1 / -prefetch_b_l1
CLI flags, with an optional -compare mode against a no-prefetch run.

Guarded behind gfx125 in the 18_flatmm CMakeLists.

Also fix data cache prefetch being silently disabled in the scaled
operator() of GemmPipelineAgBgCrCompTDMV1. The scaled path defaulted
data_cache_prefetch_a/b to false and only set them under
UseClusterLaunch, so with cluster launch off the runtime guards folded
away every prefetch and no global_prefetch_b8 was emitted despite an
L1/L2 policy. Default them to true (matching the non-scaled operator());
emission stays gated by the compile-time UseDataCachePrefetch policy
check, so None still emits nothing.
@AMD-mkulikow
AMD-mkulikow force-pushed the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch from 8fd7608 to 1f13323 Compare July 27, 2026 13:26
… handling for gfx1250

gfx1250 data-cache prefetch (prefetch_for_tdm / prefetch_for_flat) computed
byte offsets and coverage using the wrong cacheline size and double-counted the
packing factor for sub-byte packed types (e.g. pk_fp4), causing redundant/
mis-addressed prefetches.

Changes in tile_window.hpp:
- getCachelineSize(): L1 cacheline is 128B on gfx1250 (was 32B).
- prefetch_for_tdm_covers_more_calls / prefetch_for_flat_covers_more_calls:
  compute bytes_per_x_step, cacheline_part_covered and x_len_bytes in physical
  bytes by dividing by Traits::PackedSize instead of multiplying. This lets the
  IsOverprefetched coverage check actually fire for packed types, removing the
  redundant per-iteration A/B prefetches.
- prefetch_for_tdm / prefetch_for_flat issuers: x is already in packed-element
  units, so col_prefetch_stride converts cacheline bytes to packed elements with
  sizeof only (no extra /PackedSize), and the prefetch offset uses x directly
  (drop the x * PackedSize double-count). No-op for non-packed types
  (PackedSize == 1), consistent with the tdm_load_to_lds address convention.

prefetch_op_util.hpp: match the 128B L1 cacheline size in the prefetch op test.
@AMD-mkulikow
AMD-mkulikow force-pushed the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch from 1f13323 to 187b646 Compare July 27, 2026 13:33
@JeniferC99
JeniferC99 merged commit 5fa0bbe into develop Jul 27, 2026
36 of 40 checks passed
@JeniferC99
JeniferC99 deleted the users/mkulikow/ck/data_prefetch_in_mxgemm_pipeline branch July 27, 2026 18:51
assistant-librarian Bot pushed a commit to ROCm/composable_kernel that referenced this pull request Jul 27, 2026
Users/mkulikow/ck/data prefetch in mxgemm pipeline

JIRA ID : AICK-1670

## Motivation

Added example for data cache prefetch in mx gemm pipeline while also
fixing some bugs in data cache prefetch pipeline

## Technical Details

Add a standalone flatmm example (mx_flatmm_data_cache_prefetch) that
runs
MX GEMM through the compute TDM v1 pipeline
(GemmPipelineAgBgCrCompTDMV1)
with hardware data cache prefetch on gfx1250. Prefetch destination is
selectable per operand (A/B) between L1, L2 or None via the
DataCachePrefetchKind trait, exposed through -prefetch_a_l1 /
-prefetch_b_l1
CLI flags, with an optional -compare mode against a no-prefetch run.

Guarded behind gfx125 in the 18_flatmm CMakeLists.

Also fix data cache prefetch being silently disabled in the scaled
operator() of GemmPipelineAgBgCrCompTDMV1. The scaled path defaulted
data_cache_prefetch_a/b to false and only set them under
UseClusterLaunch, so with cluster launch off the runtime guards folded
away every prefetch and no global_prefetch_b8 was emitted despite an
L1/L2 policy. Default them to true (matching the non-scaled operator());
emission stays gated by the compile-time UseDataCachePrefetch policy
check, so None still emits nothing.

## Test Plan

Checked on simulators:
test name: tile_example_mx_flatmm_mxgemm_data_cache_prefetch

## Test Result

fp4 for 512x512x4096:
no prefetch/L1 prefetch: 71,334 / 54,889 ( 23.1% increase )

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
shumway pushed a commit to ROCm/composable_kernel that referenced this pull request Aug 18, 2026
Users/mkulikow/ck/data prefetch in mxgemm pipeline

JIRA ID : AICK-1670

## Motivation

Added example for data cache prefetch in mx gemm pipeline while also
fixing some bugs in data cache prefetch pipeline

## Technical Details

Add a standalone flatmm example (mx_flatmm_data_cache_prefetch) that
runs
MX GEMM through the compute TDM v1 pipeline
(GemmPipelineAgBgCrCompTDMV1)
with hardware data cache prefetch on gfx1250. Prefetch destination is
selectable per operand (A/B) between L1, L2 or None via the
DataCachePrefetchKind trait, exposed through -prefetch_a_l1 /
-prefetch_b_l1
CLI flags, with an optional -compare mode against a no-prefetch run.

Guarded behind gfx125 in the 18_flatmm CMakeLists.

Also fix data cache prefetch being silently disabled in the scaled
operator() of GemmPipelineAgBgCrCompTDMV1. The scaled path defaulted
data_cache_prefetch_a/b to false and only set them under
UseClusterLaunch, so with cluster launch off the runtime guards folded
away every prefetch and no global_prefetch_b8 was emitted despite an
L1/L2 policy. Default them to true (matching the non-scaled operator());
emission stays gated by the compile-time UseDataCachePrefetch policy
check, so None still emits nothing.

## Test Plan

Checked on simulators:
test name: tile_example_mx_flatmm_mxgemm_data_cache_prefetch

## Test Result

fp4 for 512x512x4096:
no prefetch/L1 prefetch: 71,334 / 54,889 ( 23.1% increase )

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants