[gfx906] Collection of fixes for MI50/MI60 (non-MFMA) GPUs

[dbsanfte]

Summary

This PR is an aggregation of fixes discovered while working with ComposableKernel on gfx906 (MI50/MI60) GPUs. These GPUs don't have MFMA instructions, so they rely on the DeviceGemmDl path which has some edge cases that aren't well-tested.

Note: This is a draft PR that will be updated as we discover more issues.

---

Fix 1: Buffer Load OOB Crash with Large K and Small M

Problem

DeviceGemmDl crashes on gfx906 when K >= 1472 with small M (e.g., M=1 decode case in LLM inference).

The crash occurs in gridwise_gemm_dl_v1r3.hpp during block_sync_lds() after an invalid buffer load.

Root Cause

CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK was disabled by default (set to 0).

Without the offset trick:

1. Buffer loads execute unconditionally
2. Bounds check happens after the load returns
3. If the load address is unmapped, the GPU crashes before the bounds check matters

With the offset trick enabled:

1. Invalid coordinates get 0x80000000 added to offset
2. This flags the load as OOB to hardware
3. Hardware safely returns zero instead of accessing unmapped memory

Solution

1. include/ck/ck.hpp: Enable CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK by default
2. include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v5r1.hpp: Use coordinate_has_valid_offset() instead of coordinate_has_valid_offset_assuming_visible_index_is_valid() for full bounds validation

Verification

INT8 GEMM tests pass for:

• M=1 (single-row decode)
• K up to 14336
• FFN projection dimensions for Qwen2.5 and Llama3

---

Fix 2: GridwiseGemmDlMultipleD Element Op Type Mismatch (FloatAcc != FloatC)

Problem

When FloatAcc differs from FloatC (e.g., INT8×INT8→INT32 accumulator with FP32 output scaling), the CDE element op is invoked with wrong storage types.

The element op contract is: (E& e, const C& c, const D& d...) where:

• E = FloatC (final output type, e.g., float)
• C = FloatAcc (accumulator type, e.g., int32_t)

Root Cause

Original code at lines 615-618 used generate_tie() returning the same c_thread_buf for both E& and C&:

auto dst_data_refs = generate_tie(
    [&](auto) -> auto& { return c_thread_buf(Number<c_offset>{}); },
    Number<2>{});

This causes:

1. Type mismatch when FloatAcc != FloatC (element op expects float& for e, gets int32_t&)
2. Compile errors with strictly-typed element ops
3. Undefined behavior during ThreadwiseTensorSliceTransfer which type-puns FloatAcc bits as FloatC

This bug has existed since the file was created in December 2022 (PR #517).

Solution

include/ck/tensor_operation/gpu/grid/gridwise_gemm_dl_multiple_d.hpp:

1. Introduce separate e_thread_buf<FloatC> for element op output
2. Pass (E& e) from e_thread_buf and (const C& c) from c_thread_buf using tie()
3. Transfer e_thread_buf (not c_thread_buf) to global memory

Minimal Repro

See original PR #3565 for compile-time repro that demonstrates the type mismatch.

---

Environment

• GPU: gfx906 (MI50)
• ROCm: 7.1.1
• Use case: INT8×INT8→INT32 GEMM with FP32 output for LLM inference

[Copilot]

Pull request overview

This PR addresses critical bugs in ComposableKernel's DeviceGemmDl implementation for gfx906 (MI50/MI60) GPUs that lack MFMA instructions. The fixes resolve out-of-bounds memory access crashes and type safety issues affecting INT8 GEMM operations with mixed accumulator and output types.

Changes:

• Enable hardware-assisted OOB protection by default for buffer loads on gfx906
• Fix bounds checking logic to validate all coordinates, preventing invalid memory accesses
• Correct type mismatches in element operations when accumulator type differs from output type

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
include/ck/ck.hpp	Enables offset trick for OOB buffer load protection by default
include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v5r1.hpp	Strengthens bounds validation to check all coordinates including visible indices
include/ck/tensor_operation/gpu/grid/gridwise_gemm_dl_multiple_d.hpp	Introduces separate buffer for element op output with correct type (FloatC) to fix type mismatch when FloatAcc differs from FloatC

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The inline comments identify 'e' as output and 'c' as input, but the element operation signature is '(E& e, const C& c, const D& d...)' where 'c' should be marked as const. The current comment doesn't indicate that 'c' is const, which may confuse readers about whether the element op modifies the accumulator buffer.

Suggested change

	// - c (input): comes from c_thread_buf (FloatAcc type)
	// Use tie() to create a tuple of references to different buffers
	auto dst_data_refs =
	tie(e_thread_buf(
	Number<c_offset>{}), // E& e (output to FloatC buffer)
	c_thread_buf(
	Number<c_offset>{}) // C& c (input from FloatAcc buffer)
	// - c (const input): comes from c_thread_buf (FloatAcc type)
	// Use tie() to create a tuple of references to different buffers
	auto dst_data_refs =
	tie(e_thread_buf(
	Number<c_offset>{}), // E& e (output to FloatC buffer)
	c_thread_buf(
	Number<c_offset>{}) // const C& c (input from FloatAcc buffer)

[Copilot]

With the introduction of e_thread_buf for element op output, consider whether e_thread_buf also needs initialization. If the element op always writes to all elements before the transfer, initialization may not be necessary, but this assumption should be documented or e_thread_buf should be cleared for consistency.

Suggested change

	// Initialize C
	c_thread_buf.Clear();
	// Initialize C and element-op output buffer
	c_thread_buf.Clear();
	e_thread_buf.Clear();