Skip to content

Repository navigation

fix(compiler): correct memory handling for dynamic shapes #1272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cliu1003 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

fix(compiler): correct memory handling for dynamic shapes #1272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bd3d85d
fix(compiler): harden dynamic shape memory handling
cliu1003 Oct 28, 2026
bfa2b3a
fix(runtime): read host-mapped GQA seqlens_k without a stream sync
cliu1003 Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 72 additions & 2 deletions lib/Conversion/OnnxToHip/ReshapeConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -416,6 +416,71 @@ buildExpandShapeOutputShape(mlir::PatternRewriter &rewriter, mlir::Location loc,
return outputShape;
}

/// Return whether \p value is the runtime extent of dimension \p expectedDim
/// of \p data. ReshapeShapeFold represents dynamic shape entries as an
/// index_cast of tensor.dim, so ignore that element-type conversion.
static bool isInputDim(mlir::Value value, mlir::Value data,
int64_t expectedDim) {
if (auto cast = value.getDefiningOp<mlir::arith::IndexCastOp>())
value = cast.getIn();

auto dimOp = value.getDefiningOp<mlir::tensor::DimOp>();
if (!dimOp || dimOp.getSource() != data)
return false;

std::optional<int64_t> dimIndex = mlir::getConstantIntValue(dimOp.getIndex());
return dimIndex && *dimIndex == expectedDim;
}

/// Prove that a same-type dynamic Reshape preserves every runtime dimension.
///
/// Type equality alone is insufficient because two different runtime shapes
/// can both have type tensor<?x...?xT>. The shape operand must be a
/// host-visible tensor.from_elements whose entries either copy the
/// corresponding input dimension (ONNX zero with allowzero=0), equal its
/// static extent, or come directly from tensor.dim(data, i). Unknown entries,
/// including -1, conservatively keep the Reshape.
static bool isProvablyIdentityReshape(mlir::Operation *op, mlir::Value data,
mlir::RankedTensorType inputType,
mlir::RankedTensorType outputType) {
if (inputType != outputType)
return false;

if (op->getNumOperands() < 2)
return false;

mlir::Value shapeOperand = op->getOperand(1);
int64_t rank = inputType.getRank();

bool allowzero = false;
if (auto attr = op->getAttrOfType<mlir::IntegerAttr>("allowzero"))
allowzero = attr.getValue().getSExtValue() != 0;

for (int64_t i = 0; i < rank; ++i) {
mlir::Value element = getShapeOperandElement(shapeOperand, i, rank);
if (!element)
return false;

if (auto cast = element.getDefiningOp<mlir::arith::IndexCastOp>())
element = cast.getIn();

if (std::optional<int64_t> value = mlir::getConstantIntValue(element)) {
if (*value == 0 && !allowzero)
continue;

if (!inputType.isDynamicDim(i) && *value == inputType.getDimSize(i))
continue;

return false;
}

if (!isInputDim(element, data, i))
return false;
}

return true;
}

//===----------------------------------------------------------------------===//
// Reshape -> standard tensor ops (zero-cost metadata reinterpretation)
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -447,8 +512,13 @@ struct ReshapeToStdTensor : public mlir::RewritePattern {
int64_t inputRank = inputType.getRank();
int64_t outputRank = outputType.getRank();

// No-op: same type
if (inputType == outputType) {
// Equal static types are trivially an identity. For dynamic types, type
// equality alone is insufficient: [1,8,S,2S] and [1,8,2S,S] both have type
// tensor<1x8x?x?xT>, so require the shape operand to prove every runtime
// dimension is unchanged.
if (inputType == outputType &&
(inputType.hasStaticShape() ||
isProvablyIdentityReshape(op, data, inputType, outputType))) {
rewriter.replaceOp(op, data);
return mlir::success();
}
Expand Down
50 changes: 30 additions & 20 deletions lib/Conversion/OnnxToHip/SliceConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
*/

#include "OnnxToHipUtils.h"
#include "ReadbackScalar.h"

#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
Expand Down Expand Up @@ -37,9 +38,9 @@ namespace {
// long after that landed). Worth knowing when reading the extent logic
// below: that runtime D2Hs starts/ends/axes/steps and synchronizes the
// stream on every call, so this op already costs a host sync per
// execution. Dynamic output dims are computed from
// the slice bounds when those are host-resolvable, and otherwise fall back
// to `tensor.dim` on `data` (an upper bound — Slice cannot widen any axis).
// execution. Dynamic output dims are computed from the slice bounds using
// host-visible shape arithmetic where possible and synchronized scalar
// readback otherwise.
//
// That fallback is only an upper bound, and an upper bound is not a safe
// stand-in for the extent: the extent this pattern puts on the `tensor.empty`
Expand Down Expand Up @@ -276,6 +277,25 @@ static mlir::Value resolveHostIndex(mlir::OpBuilder &b, mlir::Location loc,
return mlir::arith::MulIOp::create(b, loc, l, r);
}

/// Read element \p idx of a Slice bound tensor to a synchronized host `index`.
/// Constants and host `tensor.from_elements` values fold without device
/// traffic; other producers use the shared scalar-readback path.
static mlir::Value readSliceBoundToIndex(mlir::OpBuilder &b, mlir::Location loc,
mlir::Value ctx, mlir::Value bound,
int64_t idx) {
if (!bound || idx < 0)
return {};
auto type = mlir::dyn_cast<mlir::RankedTensorType>(bound.getType());
if (!type || type.getRank() != 1 ||
!type.getElementType().isSignlessInteger())
return {};
if (type.hasStaticShape() && idx >= type.getNumElements())
return {};

mlir::Value scalar = readbackShapeEntryToHost(b, loc, ctx, bound, idx);
return mlir::arith::IndexCastOp::create(b, loc, b.getIndexType(), scalar);
}

/// Emit the number of elements ONNX Slice produces on one axis, given runtime
/// `start`/`end` bounds and a positive compile-time \p step. Mirrors the
/// compile-time arithmetic in SliceDecompose; `arith` ops are needed here only
Expand Down Expand Up @@ -582,23 +602,9 @@ struct SliceToHip : public mlir::RewritePattern {
if (stepsVec.size() != axesVec.size())
haveAxes = false;

// Compute each dynamic output dim from the slice bounds where possible;
// see the file header for why the data dim is not an acceptable substitute.
// Ops materialized by a resolution that then fails are pure and get DCE'd.
//
// Where the bounds do not resolve, the fallback keeps the upper bound and
// therefore keeps the over-sized shape. It could be exact instead:
// CompressConversion solves the same shrinking-extent problem by reading
// the true count back from the device with `hip.ReadbackDimOp`. That is not
// done here, but not because a readback is unaffordable on this op --
// `wrap_slice` already D2Hs its own bounds and syncs the stream every call,
// so the sync is paid regardless. It is that a readback would add a second
// sync point, in the shape computation ahead of the slice, to serve a case
// no model in scope reaches: on Gemma-4 every sliced axis resolves from its
// bounds, and host arithmetic is strictly better than a readback wherever
// it is available. If a model does land here, revisit it -- the cost is
// lower than it looks. Hence the debug line below: the fallback is a known
// performance cliff, and it should be findable without an RGP capture.
// Compute each dynamic output dim from the slice bounds. Preserve the
// existing host-arithmetic fast path, then use the shared synchronized
// scalar readback for producer chains that it cannot interpret.
llvm::SmallVector<mlir::Value> dynSizes;
for (int64_t i = 0; i < resultType.getRank(); ++i) {
if (!resultType.isDynamicDim(i))
Expand Down Expand Up @@ -640,6 +646,10 @@ struct SliceToHip : public mlir::RewritePattern {
}
mlir::Value s = resolveHostIndex(rewriter, loc, starts, k, 0);
mlir::Value e = resolveHostIndex(rewriter, loc, ends, k, 0);
if (!s)
s = readSliceBoundToIndex(rewriter, loc, context, starts, k);
if (!e)
e = readSliceBoundToIndex(rewriter, loc, context, ends, k);
if (s && e)
extent = buildAllocCapacity(
rewriter, loc, dim,
Expand Down
45 changes: 22 additions & 23 deletions lib/Dialect/Transforms/Pipelines.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ void mlir::hip::buildOnnxToHipPipelineTail(
// translation.
//
// Placement is load-bearing: it MUST run BEFORE MaterializeHostScalars
// (6b). The rank-0 fill bufferizes to `memref.alloc + linalg.fill`;
// (6a). The rank-0 fill bufferizes to `memref.alloc + linalg.fill`;
// only AFTER this pass does it become `memref.alloc + memref.store`,
// which is the shape MaterializeHostScalars recognises as a tiny
// host-fed scalar and redirects to host-mapped scratch. If this ran
Expand All @@ -295,33 +295,32 @@ void mlir::hip::buildOnnxToHipPipelineTail(
// device memory. No-op on graphs that emit no linalg ops.
pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());

// 6. HIP-specific buffer optimizations
// 6a. Correctness/ABI: redirect tiny host-fed memref.alloc ops (bufferized
// `tensor.from_elements` shape arithmetic) to runtime-owned host-mapped
// scratch before buffer reuse. This removes them from the memref.alloc
// set, so OptimizeMemRefs cannot coalesce a host-written scalar with a
// slot whose preceding device use may still be asynchronous.
//
// PromoteStridedHipOperands only introduces alloc + memref.copy
// temporaries, not host-stored scalar allocations, so it does not need
// to precede this pass. See MaterializeHostScalars.cpp;
// pipeline-static-shape-lockdown.mlir locks down this ordering.
pm.addNestedPass<func::FuncOp>(mlir::hip::createMaterializeHostScalarsPass());

// 6b. Reuse the remaining ordinary allocations. Keep this pass generic:
// MaterializeHostScalars has already removed the recognized host-written
// scalar staging allocations from its memref.alloc candidate set.
pm.addNestedPass<func::FuncOp>(mlir::hip::createOptimizeMemRefsPass());

// The remaining pre-pooling passes have two distinct roles:
// - Correctness/ABI passes that must introduce or redirect allocations
// before pooling (6a and 6b).
// - Shape preconditions used only to improve pool quality (grouped in 6c).

// 6a. Correctness/ABI: materialize identity-layout temporaries for
// non-identity-layout DPS-input memrefs. The interface-based selection
// covers plugin ops whose bare-pointer runtime ABI carries no descriptor
// offset or strides. Runs after linalg cleanup and OptimizeMemRefs, and
// before pooling so the temporaries become pool views rather than
// per-inference allocations.
// 6c. Correctness/ABI: materialize identity-layout temporaries for
// non-identity-layout DPS-input memrefs. This must run after
// OptimizeMemRefs, which may introduce memref.reinterpret_cast views
// whose offset or strides cannot be represented by the bare-pointer
// runtime ABI. The new temporaries are pooled below.
pm.addNestedPass<func::FuncOp>(
mlir::hip::createPromoteStridedHipOperandsPass());

// 6b. Correctness/ABI: redirect tiny host-fed memref.alloc ops (bufferized
// `tensor.from_elements` shape arithmetic) from the GPU pool to
// runtime-owned host-mapped scratch. Must run after 6a so its candidate
// analysis sees the final allocation set, and before pooling so
// host-written buffers never become views into device memory. See
// MaterializeHostScalars.cpp; pipeline-static-shape-lockdown.mlir locks
// down this ordering.
pm.addNestedPass<func::FuncOp>(mlir::hip::createMaterializeHostScalarsPass());

// 6c. Pool quality: normalize dynamic allocation sizes after all view and
// 6d. Pool quality: normalize dynamic allocation sizes after all view and
// allocation rewrites, immediately before hip-pool-allocs.
addPoolAllocsShapePreconditionPasses(pm);

Expand Down
96 changes: 49 additions & 47 deletions lib/Runtime/real/gqa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,35 @@ static bool gqa_cache_seqlens_enabled() {
// they already implement.
static constexpr int32_t kSeqlensKNotRead = -2;

// Copy `count` seqlens_k entries to `dst`. hip-materialize-host-scalars can
// place a host-written seqlens_k in hipHostMalloc'd host scratch; the host
// store has already retired, so it is read directly instead of with a copy
// plus a stream sync that would drain all queued GPU work. Anything else
// (including pointers the query does not recognise) takes the synchronized
// readback.
static hipError_t read_seqlens_k_to_host(hipStream_t stream,
const void *seqlens_k_ptr,
int32_t *dst, int64_t count) {
const size_t bytes = static_cast<size_t>(count) * sizeof(int32_t);
hipPointerAttribute_t attr{};
if (hipPointerGetAttributes(&attr, seqlens_k_ptr) != hipSuccess) {
// Clear the last error a failed query leaves for unrecognised pointers.
(void)hipGetLastError();
} else if (attr.type == hipMemoryTypeHost) {
std::memcpy(dst, seqlens_k_ptr, bytes);
return hipSuccess;
}

// hipMemcpyDefault, not DeviceToHost: if the query above could not classify
// host-accessible memory, DeviceToHost would reject it with
// `invalid argument`.
hipError_t err =
hipMemcpyAsync(dst, seqlens_k_ptr, bytes, hipMemcpyDefault, stream);
if (err != hipSuccess)
return err;
return hipStreamSynchronize(stream);
}

// Read seqlens_k_val from device (or the per-Compute() cache when
// HIPDNN_EP_GQA_CACHE_SEQLENS=1) once per call before the fused/decomposed
// dispatch decision. Two purposes:
Expand Down Expand Up @@ -154,10 +183,8 @@ static int32_t read_seqlens_k_for_dispatch(hipStream_t stream,
return state->seqlens_k_cached_val;

int32_t seqlens_k_val = 0;
if (hipMemcpyAsync(&seqlens_k_val, seqlens_k_ptr, sizeof(int32_t),
hipMemcpyDeviceToHost, stream) != hipSuccess)
return kSeqlensKNotRead;
if (hipStreamSynchronize(stream) != hipSuccess)
if (read_seqlens_k_to_host(stream, seqlens_k_ptr, &seqlens_k_val, 1) !=
hipSuccess)
return kSeqlensKNotRead;

if (gqa_cache_seqlens_enabled() && state) {
Expand Down Expand Up @@ -385,12 +412,9 @@ static int gqa_forward_fused(
int32_t seqlens_k_val = 0;
if (seqlens_k_pre != kSeqlensKNotRead) {
seqlens_k_val = seqlens_k_pre;
} else {
if (hipMemcpyAsync(&seqlens_k_val, seqlens_k_ptr, sizeof(int32_t),
hipMemcpyDeviceToHost, stream) != hipSuccess)
return -1;
if (hipStreamSynchronize(stream) != hipSuccess)
return -1;
} else if (read_seqlens_k_to_host(stream, seqlens_k_ptr, &seqlens_k_val,
1) != hipSuccess) {
return -1;
}
if (seqlens_k_val < 0) {
past_len = 0; // ORT prefill sentinel
Expand Down Expand Up @@ -587,11 +611,8 @@ static int gqa_forward_fused(
seqlens_k_val = seqlens_k_pre;
} else if (B > 1) {
std::vector<int32_t> seqlens_k_host(B);
if (hipMemcpyAsync(seqlens_k_host.data(), seqlens_k_ptr,
B * sizeof(int32_t), hipMemcpyDeviceToHost,
stream) != hipSuccess)
return -1;
if (hipStreamSynchronize(stream) != hipSuccess)
if (read_seqlens_k_to_host(stream, seqlens_k_ptr, seqlens_k_host.data(),
B) != hipSuccess)
return -1;
seqlens_k_val = seqlens_k_host[0];
for (int64_t b = 1; b < B; ++b) {
Expand All @@ -603,12 +624,9 @@ static int gqa_forward_fused(
return -1;
}
}
} else {
if (hipMemcpyAsync(&seqlens_k_val, seqlens_k_ptr, sizeof(int32_t),
hipMemcpyDeviceToHost, stream) != hipSuccess)
return -1;
if (hipStreamSynchronize(stream) != hipSuccess)
return -1;
} else if (read_seqlens_k_to_host(stream, seqlens_k_ptr, &seqlens_k_val,
1) != hipSuccess) {
return -1;
}
if (seqlens_k_val < 0) {
total_seq = sq; // ORT prefill sentinel
Expand Down Expand Up @@ -1350,14 +1368,9 @@ static int gqa_forward_hipblaslt(
int32_t seqlens_k_val = 0;
if (seqlens_k_pre != kSeqlensKNotRead) {
seqlens_k_val = seqlens_k_pre;
} else {
if (hipMemcpyAsync(&seqlens_k_val, seqlens_k_ptr, sizeof(int32_t),
hipMemcpyDeviceToHost, stream) != hipSuccess) {
return -1;
}
if (hipStreamSynchronize(stream) != hipSuccess) {
return -1;
}
} else if (read_seqlens_k_to_host(stream, seqlens_k_ptr, &seqlens_k_val,
1) != hipSuccess) {
return -1;
}

// ORT prefill sentinel: when there is no past KV yet, the producer
Expand Down Expand Up @@ -1526,11 +1539,8 @@ static int gqa_forward_hipblaslt(
seqlens_k_val = seqlens_k_pre;
} else if (B > 1) {
std::vector<int32_t> seqlens_k_host(B);
if (hipMemcpyAsync(seqlens_k_host.data(), seqlens_k_ptr,
B * sizeof(int32_t), hipMemcpyDeviceToHost,
stream) != hipSuccess)
return -1;
if (hipStreamSynchronize(stream) != hipSuccess)
if (read_seqlens_k_to_host(stream, seqlens_k_ptr, seqlens_k_host.data(),
B) != hipSuccess)
return -1;
seqlens_k_val = seqlens_k_host[0];
for (int64_t b = 1; b < B; ++b) {
Expand All @@ -1547,24 +1557,16 @@ static int gqa_forward_hipblaslt(
// (D2H or sync failure). Rare path; not cached because the same failure
// mode would have prevented the helper from caching too.
//
// Both failures are reported rather than returned bare: this is the one
// The failure is reported rather than returned bare: this is the one
// place the decomposed path can fail before it has done any work, and a
// silent -1 here surfaces only as a zero-filled attention output, which
// looks like a kernel bug rather than a seqlens_k that could not be read.
hipError_t cp =
hipMemcpyAsync(&seqlens_k_val, seqlens_k_ptr, sizeof(int32_t),
hipMemcpyDeviceToHost, stream);
if (cp != hipSuccess) {
hipError_t rd =
read_seqlens_k_to_host(stream, seqlens_k_ptr, &seqlens_k_val, 1);
if (rd != hipSuccess) {
fprintf(stderr,
"gqa_forward_hipblaslt: could not read seqlens_k from %p: %s\n",
seqlens_k_ptr, hipGetErrorString(cp));
return -1;
}
hipError_t sy = hipStreamSynchronize(stream);
if (sy != hipSuccess) {
fprintf(stderr,
"gqa_forward_hipblaslt: sync after seqlens_k read failed: %s\n",
hipGetErrorString(sy));
seqlens_k_ptr, hipGetErrorString(rd));
return -1;
}
}
Expand Down
Loading
Loading