[pschlan-amd]

fix(optimize-memrefs): never share a slot with a host-accessed buffer

Problem

hip-optimize-memrefs reuses a dead buffer for a later allocation, and decides "dead" from op order in the block. Op order is stream order only for device ops. A memref.store or memref.load runs on the host as soon as the host reaches it, while device ops issued earlier may still be queued.

The standard onnxruntime-genai GQA preamble triggers this:

%sum = memref.alloc() : memref<1x1xi64>
hip.reduce_sum ins(%mask, ...) outs(%sum)        // device
hip.sub ins(%sum, %one) outs(%sub)               // device, still queued
hip.cast ins(%sub) outs(%seqlens_k)              // device
%n = arith.index_cast %dim : index to i64        // Shape(mask)[1]
%total = memref.alloc() : memref<i64>            // reused %sum
memref.store %n, %total[]                        // host
hip.cast ins(%total) outs(%total_sequence_length)

%sum is dead after hip.sub in op order, so %total got %sum's slot through a memref.reinterpret_cast. hip-materialize-host-scalars then moved that shared buffer into host-mapped scratch. Whichever side wrote last won:

• Host store lands between reduce_sum and sub. hip.sub reads Shape(mask)[1], so seqlens_k becomes max_length - 1. Attention then runs over the whole static KV buffer, including the unwritten rows. The output is silently wrong, with no error.
• Host store lands before reduce_sum. total_sequence_length takes the mask sum instead of the buffer length. Most fused GQA paths don't read it, so this usually stays hidden.

The failure depends on timing. In a one-layer GQA replica (one prefill step, then decode steps, against an fp32 reference) it hit the first decode step in roughly 1 of 5 runs. The wrong output exactly matched attention over all 4096 buffer rows, and the device seqlens_k read back as 4095 instead of 1536. Enabling HIPDNN_EP_DEBUG=1, or serialising with AMD_SERIALIZE_KERNEL=3 AMD_SERIALIZE_COPY=3, hid it.

Fix

An allocation with a memref.load or memref.store on any alias (as resolved by BufferViewFlowAnalysis) is left out of slot assignment. It neither reuses a slot nor provides one, and keeps its own allocation. Such buffers are the small shape and scalar staging buffers that hip-materialize-host-scalars later places in host scratch, where each already gets its own offset. The memory cost is a few bytes per buffer.

Device-only buffers are reused exactly as before.

Tests

• New lit case @no_reuse_host_store in test/lit/Dialect/hip-optimize-memrefs.mlir mirrors the preamble above. It checks that %total keeps its own memref.alloc and that no reinterpret_cast onto %sum appears.
• On this branch: full lit suite (check-hip-mlir-lit), 663 passed and 9 unsupported.
• On the GPU (gfx1151), measured on the branch the fix came from: the one-layer replica above failed 3–4 of 16 runs before and 0 of 16 after.

Impact

Numerics are unchanged whenever the race doesn't fire. Generated code changes only in that host-accessed allocations no longer alias device buffers. I verified the exact pattern only in the replica's IR. Because it comes from the standard genai preamble (seqlens_k = ReduceSum(mask) - 1, total_sequence_length = Shape(mask)[1]), other GQA models are likely affected too, depending on how the allocations line up after CSE.

Analyzed and fixed using AI assistance (Cursor / Claude Opus 5.5).

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[pschlan-amd]

/run-test test_mode=performance models=all

[amd-yanjunz]

ROCm GenAI CI finished with result: UNSTABLE.

• PR: fix(optimize-memrefs): never share a slot with a host-accessed buffer #1389
• Source comment: test_mode=performance models=all
• TEST_MODE: performance
• MODELZOO: LLM_ROCm_dynamic
• TEST_MODELS: (all)
• FORBID_MODELS: (none)
• NODE: (GenAI job default)
• CI_OPTIONS: (none)
• Jenkins build: http://xsjvitisaijenkins:8080/job/daily_hip_ep_genai/2964/
• Main baseline commit: 5e4585c
• Main baseline build: http://xsjvitisaijenkins:8080/job/daily_hip_ep_genai/2952/
• Main baseline: not usable, the baseline build has not finished yet
• Package build run: https://github.com/ROCm/hip-ep/actions/runs/41042256413
• NTS report: http://xsjnts.xilinx.com/ROCm/hipdnn-detail?database=development&collection=daily_hip_ep_genai&id=6ae49c912ca0a072895aa929
• Compare with latest daily baseline (fallback): http://xsjnts.xilinx.com/ROCm/compare?base_url=http%3A%2F%2Fxsjnts.xilinx.com%2FROCm%2Fhipdnn-detail%3Fids%3D6ae42f8865373de3d75aa894%26collection%3DHIP-EP-Dynamic-StrixHalo&base_job=&base=&curr_url=&curr_job=daily_hip_ep_genai&curr=2964

[fhanuman]

This is a generic pass, so we should not put any host- or device-specific logic in it. I made a similar comment on another PR (PR #1305), and we seem to be replicating the same solution in this PR. We need a backend-specific (CPU, GPU, or NPU) pass that runs late in the compiler flow. Please see the follow-up comment in the following PR.

#1126

[pschlan-amd]

Didn't realize there's already work in other PRs on this issue, hence closing here.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	1.12256E-07	64	0
GroupQueryAttention_seq256	59.2561	2621440	0
MatMulNBits_o_seq128	259.905	368640	0
QMoE_seq128	24.2583	368640	0

Threshold: 0.01 | Run: 8160 - Commit: d2e3d06

[github-actions]

Binary sizes

Baseline: run 40994390384 - 5e4585c (dll-sizes)

File	This (MB)	Main (MB)	Delta (MB)
hipgpu.dll	59.38	59.38	0.00
custom_kernels_gfx1170.dll	5.62	5.62	0.00
custom_kernels_gfx11-generic.dll	5.87	5.87	0.00

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.93	7.01	243	3	1159
GroupQueryAttention_seq128	11080.90	2.58172	0	5	167
matmul_down_seq128	740.48	3.25	118	3	267

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.97	47.12	250	3	15505

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.97	10.06	253	3	15544

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	89.1	82.6	1.25	12.34
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	117.5	49.6	1.15	5.40

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	116	49.8

VLM Benchmark Results

Model	TTFT (ms)	TPS	Output
gemma-4-E4B-it-fp16-int4-kquant-prune-logits	3905.3	7.3	This is a high-quality, outdoor photograph featuring a beautiful, fluffy white dog, likely a West Highland White Terrier or a si
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	4861.6	14.6	The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

TTFT / TPS vs main

Baseline: run 40994390384 - 5e4585c | Gate: TTFT +20% / TPS -20%

Suite	Model	TTFT (ms)	Main	Delta	TPS	Main	Delta
oga	gpt-oss-20b-webgpu-int4-rtn-block-32	89.1	91.0	-2.1%	82.6	82.7	-0.1%
oga	Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	117.5	116.4	+1.0%	49.6	49.6	+0.1%
vlm	gemma-4-E4B-it-fp16-int4-kquant-prune-logits	3905.3	3775.7	+3.4%	7.3	8.0	-8.8%
vlm	Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	4861.6	4880.8	-0.4%	14.6	14.6	+0.1%

Run: 8160 - Commit: d2e3d06