CCO SDMA micro optimizations by pemeliya · Pull Request #603 · ROCm/mori

pemeliya · 2026-08-25T16:29:03Z

ensured proper coalescing of ccoSdmaUnit stores
enforced global memory space for all internal loads / stores

before:

; /tf/mori/include/mori/cco/cco.hpp:1506
;     ccoSdmaQueueDeviceHandle handle = *shared;
	s_waitcnt vmcnt(0)                                         // 00000000BF84: BF8C0F70
	flat_load_dwordx4 v[4:7], v[16:17] offset:32               // 00000000BF88: DC5C0020 04000010
	flat_load_dwordx4 v[12:15], v[16:17]                       // 00000000BF90: DC5C0000 0C000010
	flat_load_dwordx4 v[8:11], v[16:17] offset:16              // 00000000BF98: DC5C0010 08000010
	flat_load_dwordx2 v[22:23], v[16:17] offset:48             // 00000000BFA0: DC540030 16000010
...........
; /tf/mori/include/mori/cco/cco.hpp:1379
;   *ccoSdmaUnitAt(handle, at) = img;
	v_mov_b32_e32 v22, v37                                     // 00000000C030: 7E2C0325
	flat_store_dwordx3 v[12:13], v[20:22]                      // 00000000C034: DC780000 0000140C
	flat_store_dwordx4 v[12:13], v[0:3] offset:12              // 00000000C03C: DC7C000C 0000000C
	flat_store_dword v[12:13], v37 offset:28                   // 00000000C044: DC70001C 0000250C

after:

;     ccoSdmaQueueDeviceHandle handle = *Iglobal(shared);
	s_waitcnt vmcnt(0)                                         // 00000000C064: BF8C0F70
	global_load_dwordx4 v[6:9], v[38:39], off offset:32        // 00000000C068: DC5C8020 067F0026
	global_load_dwordx2 v[0:1], v[38:39], off offset:48        // 00000000C070: DC548030 007F0026
	global_load_dwordx4 v[10:13], v[38:39], off offset:16      // 00000000C078: DC5C8010 0A7F0026
	global_load_dwordx4 v[14:17], v[38:39], off                // 00000000C080: DC5C8000 0E7F0026
.........       
; /tf/mori/include/mori/cco/cco.hpp:1402
; /tf/mori/include/mori/cco/cco.hpp:1411
;   dst->vec[0] = Uint4{p[0], p[1], p[2], p[3]};
	global_store_dwordx4 v[42:43], v[6:9], off                 // 00000000BFF8: DC7C8000 007F062A
; /tf/mori/include/mori/cco/cco.hpp:1412
;   dst->vec[1] = Uint4{p[4], p[5], p[6], 0};
	global_store_dwordx4 v[42:43], v[10:13], off offset:16     // 00000000C000: DC7C8010 007F0A2A

pemeliya · 2026-08-25T16:30:34Z

Benchmark kernel stats:
Before:

=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopWarp, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 47/27, scratch ld/st: 0/0
vgprs: 48 spills: 0
sgprs: 62 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0
=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopBlock, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 47/27, scratch ld/st: 0/0
vgprs: 42 spills: 0
sgprs: 67 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0
=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopThread, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 47/27, scratch ld/st: 0/0
vgprs: 46 spills: 0
sgprs: 54 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0

After:

=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopWarp, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 4/0, scratch ld/st: 0/0
vgprs: 52 spills: 0
sgprs: 62 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0
=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopBlock, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 4/0, scratch ld/st: 0/0
vgprs: 46 spills: 0
sgprs: 66 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0
=== void mori::cco::benchmark::sdma_put_bw<mori::cco::ccoCoopThread, 1u>(mori::cco::ccoWindowDevice*, mori::cco::ccoWindowDevice*, unsigned long, mori::cco::ccoDevComm, int, int, int) ===
flat ld/st: 4/0, scratch ld/st: 0/0
vgprs: 48 spills: 0
sgprs: 54 spills: 0, agprs: 0
private_seg: 0 / dyn_stack: 0

Copilot

Pull request overview

This PR applies low-level SDMA fast-path micro-optimizations in include/mori/cco/cco.hpp to improve instruction selection and memory coalescing for queue packet publishing on AMDGPU targets.

Changes:

Introduces an internal impl::global() helper to force global address space accesses for SDMA queue-related loads/stores and HIP atomics.
Reworks ccoSdmaUnit writes to use 16-byte vector stores (Uint4) for better coalescing when writing COPY/ATOMIC packets into the ring.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

jhchouuu

Tested on both arches — gfx950 (MI355X) and gfx1250 — with the baseline at the merge-base 624002c8 and identical build flags for both trees.

Correctness. The seven CCO SDMA tests (put, get, put_mt, block, edge, signal, warp_issue) in fork mode pass on both trees: 8/8 ranks on gfx950, 4/4 on gfx1250.

ISA. Both claims hold, and identically on the two arches (test_sdma_put):

	flat_load	flat_store	flat_atomic	global_load	global_store
base	100	84	12	4	0
PR	6	0	12	102	72

Packet stores go from three per unit to two. The baseline emits dwordx3@0 + dwordx4@offset:12 + dword@offset:28 — note offset 12 is not 16B-aligned — where this PR emits dwordx4@0 + dwordx4@offset:16. Same shape on gfx1250 with b96/b32/b128.

Perf. 8 B put, thread scope, quiet completion, 2 ranks, runs interleaved to cancel drift:

gfx950: median 6.076 → 5.970 us (-1.7%), and the PR is faster in 8 of 9 paired runs.
gfx1250: no measurable difference. Base-first ordering says "PR 0.5% slower", PR-first ordering says "PR 1.5% faster" — the sign follows run order, so an ordering artifact dominates. Pooling both blocks, so each binary ran first 9 times and second 9 times: 8.207 → 8.147 us, well inside the 7.98–8.60 spread.

Bandwidth is unchanged on both: 57.0 GB/s at 8 MB on gfx950, 980 GB/s at 8 MB on gfx1250.

The obvious gfx1250 worry — whether moving packet stores from flat to global could break "packet dwords land before the doorbell" — does not apply. ccoSdmaPublishStores() is a release fence at agent scope there, and s_waitcnt(0) on gfx9; neither relies on flat bumping an extra counter.

Comments inline, nothing blocking.

jhchouuu · 2026-09-01T09:06:40Z

+// Implicit cast to global memory space.
+template <typename T, typename T2 = typename std::remove_volatile<T>::type>
+__device__ __host__ inline static T2* global(T* ptr) {
+  return (T2*)(T2 CCO_GLOBAL_SPACE*)reinterpret_cast<uintptr_t>(ptr);


This reinterpret_cast<uintptr_t> looks like it could be simplified to a plain address-space cast. It cannot, and I would suggest a comment saying so — it is the most deletable-looking line in the PR, and deleting it silently undoes the entire change while every test still passes.

I tried it. Replacing the body with

return (T2*)(T2 CCO_GLOBAL_SPACE*)ptr;

and regenerating the device assembly for test_sdma_put (gfx950, same flags) puts everything back on flat:

as written direct cast

flat_load_dwordx2 / x4 6 / 0 68 / 32

flat_store_dwordx2 / x4 0 / 0 48 / 24

global_load + global_store 174 4

addrspacecast (addrspacecast p to as(1)) to as(0) is an inverse pair that gets folded away immediately, so the address space never survives to InferAddressSpaces. inttoptr cannot be folded, which is precisely why the round trip works. sgpr_count is 44 either way, so nothing is being paid for it.

yes it is tricky.. atomic_load/store builtin functions do not accept pointers with non-default memory space, that's why this workaround was necessary

jhchouuu · 2026-09-01T09:06:40Z

+#define CCO_GLOBAL_SPACE __attribute__((address_space(1)))
+// Implicit cast to global memory space.
+template <typename T, typename T2 = typename std::remove_volatile<T>::type>
+__device__ __host__ inline static T2* global(T* ptr) {


Three small things on the signature:

remove_volatile silently drops volatile. Nothing in ccoSdmaQueueDeviceHandle is volatile today, so this is inert — but rptr and doorbell are hardware-updated, and if anyone ever marks them volatile this helper will quietly discard it. A static_assert(!std::is_volatile<T>::value) would keep that honest.

__host__: address_space(1) means nothing on the host, and the attribute invites host callers into what is a device-only idiom.

T2 is a defaulted template parameter, so a caller can override it: impl::global<int, float>(p) compiles and type-puns (verified). Moving it to a using inside the body removes that.

Also static at namespace scope is redundant for a template.

jhchouuu · 2026-09-01T09:06:40Z

      do {
-        cachedHwReadIndex = __hip_atomic_load(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
+        cachedHwReadIndex =
+            __hip_atomic_load(impl::global(rptr), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);


The description says "enforced global memory space for all internal loads / stores", but three sites stay on flat:

L1252, the __hip_atomic_fetch_add(cachedWptr, ...) immediately above — one per put, and the hottest single memory op in this path. Both trees still show 12 flat_atomic_add_x2 in test_sdma_put, on gfx950 and gfx1250 alike.

L1620 / L1628 / L1636, **(deviceHandles + q) in ccoSdmaCommitThread/Warp/Block — these are the 6 flat_load_dwordx2 still present in the PR build.

L1827, the waitSignal spin load.

Leaving them may well be deliberate, but the reservation atomic in particular seems worth either converting or saying why not.

I see, then I guess I have missed these spots. I have just compiled few tests to check for flat load/stores, need to have a look at them once again

jhchouuu · 2026-09-01T09:06:40Z

                       : CCO_SDMA_UNIT * (n + sig);
 }

+using Uint4 = uint32_t __attribute__((ext_vector_type(4)));


Uint4 lands in the public mori::cco namespace of a public header, which is a fairly generic name to claim. include/mori/collective/allreduce/twoshot_sdma_kernel.hpp already defines its own Uint4 (function-local, so no clash today).

global() above is correctly tucked into mori::cco::impl — this could go the same way, or be named ccoUint4.

jhchouuu · 2026-09-01T09:06:40Z

-  *ccoSdmaUnitAt(handle, at) = img;
+  auto dst = impl::global(ccoSdmaUnitAt(queueBuf, at));
+  dst->vec[0] = Uint4{p[0], p[1], p[2], p[3]};
+  dst->vec[1] = Uint4{p[4], p[5], p[6], 0};


The trailing 0 used to carry its reason:

img.dw[7] = 0; // NOP (op=0), pads COPY out to the unit; engine skips it

It is an opcode the engine skips rather than padding, and that is not recoverable from the bare literal. Worth keeping the comment here — ccoSdmaWriteAtomic just below is a good contrast, since there the eighth dword is real packet content.

jhchouuu · 2026-09-01T09:06:40Z

 // Where a unit at monotonic byte offset `at` lands. Units are 32B-aligned in a
 // 32B-multiple ring, so this never needs to split a packet.
-inline __device__ ccoSdmaUnit* ccoSdmaUnitAt(ccoSdmaQueueDeviceHandle& handle, uint64_t at) {
+inline __device__ ccoSdmaUnit* ccoSdmaUnitAt(uint32_t* queueBuf, uint64_t at) {


Minor, but the alignment requirement changes character here. The old *ccoSdmaUnitAt(...) = img compiled to dwordx3@0 + dwordx4@offset:12 + dword@offset:28 — the compiler was not assuming 16 B alignment. The new Uint4 stores are dwordx4@0 and @offset:16, which do.

It holds today: every at is a multiple of CCO_SDMA_UNIT (32), and queueBuf comes from hsaKmtAllocMemory so it is page-aligned. But the failure mode moved from "slower" to "wrong", so it may be worth pinning down — a __builtin_assume(at % CCO_SDMA_UNIT == 0), or an assert under MORI_CCO_SDMA_DEBUG.

pemeliya · 2026-09-01T16:40:01Z

Thanks for taking a look! I am going to be away until 17.09, so I guess we can put this PR on hold.
Or, otherwise, feel free to have a look at them yourself if you think these make sense.

pemeliya added 2 commits August 24, 2026 09:31

flat -> global and vecs opts

ff1e276

simplifying

3b84fc1

fixing precommit

93e344e

pemeliya requested review from jhchouuu and a lite review from Copilot August 25, 2026 16:39

Copilot started reviewing on behalf of pemeliya August 25, 2026 16:40 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Comment thread include/mori/cco/cco.hpp

Undefine macro after usage

f77502a

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

jhchouuu self-assigned this Sep 1, 2026

jhchouuu reviewed Sep 1, 2026

View reviewed changes

Navigation Menu

Uh oh!

CCO SDMA micro optimizations - #603

CCO SDMA micro optimizations#603
pemeliya wants to merge 4 commits into
mainfrom
pemeliya/cco_improvements

pemeliya commented Aug 25, 2026 •

edited

Loading

Uh oh!

pemeliya commented Aug 25, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

jhchouuu left a comment

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

pemeliya Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

pemeliya Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

pemeliya commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	as written	direct cast
flat_load_dwordx2 / x4	6 / 0	68 / 32
flat_store_dwordx2 / x4	0 / 0	48 / 24
global_load + global_store	174	4

Uh oh!

Conversation

pemeliya commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pemeliya commented Aug 25, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pemeliya commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

pemeliya commented Aug 25, 2026 •

edited

Loading