From bd3d85d6611106263bc61a48130163c8ad7b662c Mon Sep 17 00:00:00 2001 From: cliu1003 Date: Wed, 28 Oct 2026 17:44:24 +0800 Subject: [PATCH 1/2] fix(compiler): harden dynamic shape memory handling Materialize host-written scalar staging buffers before OptimizeMemRefs. Once OptimizeMemRefs coalesces such an allocation with a slot whose earlier device use may still be in flight, the host store can race it; moving hip-materialize-host-scalars first removes those allocations from the reuse candidate set. PromoteStridedHipOperands only creates alloc + memref.copy temporaries, so it can stay after OptimizeMemRefs, where it also covers the reinterpret_cast views that pass introduces. Only fold a same-type dynamic Reshape when the shape operand proves every runtime dimension is preserved. Type equality alone cannot tell [1,8,S,2S] from [1,8,2S,S]. When a dynamic Slice bound cannot be resolved with host shape arithmetic, read it back through the shared synchronized scalar path instead of falling back to the data dimension, which is only an upper bound on the extent. Co-Authored-By: Claude Opus 5.5 Made-with: Cursor Co-authored-by: Cursor --- .../OnnxToHip/ReshapeConversion.cpp | 74 ++++++++++++++++++- lib/Conversion/OnnxToHip/SliceConversion.cpp | 50 ++++++++----- lib/Dialect/Transforms/Pipelines.cpp | 45 ++++++----- .../Conversion/onnx-to-hip/test_reshape.mlir | 52 +++++++++++++ .../Conversion/onnx-to-hip/test_slice.mlir | 46 ++++++++++-- .../pipeline-static-shape-lockdown.mlir | 47 ++++++++---- test/lit/e2e/test_mlp_model.mlir | 2 + 7 files changed, 250 insertions(+), 66 deletions(-) diff --git a/lib/Conversion/OnnxToHip/ReshapeConversion.cpp b/lib/Conversion/OnnxToHip/ReshapeConversion.cpp index c4eb3dce5f..0dae8c0d1a 100644 --- a/lib/Conversion/OnnxToHip/ReshapeConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReshapeConversion.cpp @@ -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()) + value = cast.getIn(); + + auto dimOp = value.getDefiningOp(); + if (!dimOp || dimOp.getSource() != data) + return false; + + std::optional 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. 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("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()) + element = cast.getIn(); + + if (std::optional 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) //===----------------------------------------------------------------------===// @@ -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(); } diff --git a/lib/Conversion/OnnxToHip/SliceConversion.cpp b/lib/Conversion/OnnxToHip/SliceConversion.cpp index 2ad6343894..2fc38a6ac8 100644 --- a/lib/Conversion/OnnxToHip/SliceConversion.cpp +++ b/lib/Conversion/OnnxToHip/SliceConversion.cpp @@ -4,6 +4,7 @@ */ #include "OnnxToHipUtils.h" +#include "ReadbackScalar.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" @@ -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` @@ -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(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 @@ -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 dynSizes; for (int64_t i = 0; i < resultType.getRank(); ++i) { if (!resultType.isDynamicDim(i)) @@ -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, diff --git a/lib/Dialect/Transforms/Pipelines.cpp b/lib/Dialect/Transforms/Pipelines.cpp index afa7e09e3a..0ba4a0d812 100644 --- a/lib/Dialect/Transforms/Pipelines.cpp +++ b/lib/Dialect/Transforms/Pipelines.cpp @@ -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 @@ -295,33 +295,32 @@ void mlir::hip::buildOnnxToHipPipelineTail( // device memory. No-op on graphs that emit no linalg ops. pm.addNestedPass(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(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(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( 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(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); diff --git a/test/lit/Conversion/onnx-to-hip/test_reshape.mlir b/test/lit/Conversion/onnx-to-hip/test_reshape.mlir index 9e332a1ac5..dfc3a79de5 100644 --- a/test/lit/Conversion/onnx-to-hip/test_reshape.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_reshape.mlir @@ -64,6 +64,45 @@ module { return %result : tensor<1x128x4096xf16> } + // Equal dynamic ranked types do not imply an identity runtime reshape. + func.func @test_reshape_same_dynamic_type( + %data: tensor<1x8x?x?xf16>, %shape: tensor<4xi64>) + -> tensor<1x8x?x?xf16> { + %result = "onnx.Reshape"(%data, %shape) {allowzero = 0 : si64} + : (tensor<1x8x?x?xf16>, tensor<4xi64>) -> tensor<1x8x?x?xf16> + return %result : tensor<1x8x?x?xf16> + } + + // Shape(data) proves that every runtime output dimension equals the + // corresponding input dimension, so this dynamic Reshape is an identity. + func.func @test_reshape_dynamic_identity_from_shape( + %data: tensor<1x8x?x?xf16>) -> tensor<1x8x?x?xf16> { + %shape = "onnx.Shape"(%data) + : (tensor<1x8x?x?xf16>) -> tensor<4xi64> + %result = "onnx.Reshape"(%data, %shape) {allowzero = 0 : si64} + : (tensor<1x8x?x?xf16>, tensor<4xi64>) -> tensor<1x8x?x?xf16> + return %result : tensor<1x8x?x?xf16> + } + + // Equal dynamic types still require the Reshape when the shape operand + // exchanges two runtime dimensions. + func.func @test_reshape_dynamic_reordered_dims( + %data: tensor<1x8x?x?xf16>) -> tensor<1x8x?x?xf16> { + %c1 = arith.constant 1 : i64 + %c8 = arith.constant 8 : i64 + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %d2 = tensor.dim %data, %c2 : tensor<1x8x?x?xf16> + %d3 = tensor.dim %data, %c3 : tensor<1x8x?x?xf16> + %d2_i64 = arith.index_cast %d2 : index to i64 + %d3_i64 = arith.index_cast %d3 : index to i64 + %shape = tensor.from_elements %c1, %c8, %d3_i64, %d2_i64 + : tensor<4xi64> + %result = "onnx.Reshape"(%data, %shape) {allowzero = 0 : si64} + : (tensor<1x8x?x?xf16>, tensor<4xi64>) -> tensor<1x8x?x?xf16> + return %result : tensor<1x8x?x?xf16> + } + // --- Rank-0 scalar → rank-1 1-element (Qwen vision shape-arith chain) --- // onnx.Reshape on a tensor scalar to tensor<1xi64> can't be expressed // via tensor.expand_shape (source rank 0 can't supply a reassoc group). @@ -224,6 +263,19 @@ module { // CHECK-NOT: tensor.collapse_shape // CHECK: return +// CHECK-LABEL: func.func @test_reshape_same_dynamic_type +// CHECK: tensor.reshape +// CHECK-NOT: onnx.Reshape + +// CHECK-LABEL: func.func @test_reshape_dynamic_identity_from_shape +// CHECK-NOT: tensor.reshape +// CHECK-NOT: onnx.Reshape +// CHECK: return %arg1 + +// CHECK-LABEL: func.func @test_reshape_dynamic_reordered_dims +// CHECK: tensor.reshape +// CHECK-NOT: onnx.Reshape + // A runtime (non-constant) scalar must be read to the host with a synchronized // hip.readback_scalar, NOT a bare tensor.extract: the latter lowers to an // unsynchronized host load of a device buffer and reads stale memory on diff --git a/test/lit/Conversion/onnx-to-hip/test_slice.mlir b/test/lit/Conversion/onnx-to-hip/test_slice.mlir index bfe654c18f..1fd41ed4a9 100644 --- a/test/lit/Conversion/onnx-to-hip/test_slice.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_slice.mlir @@ -202,9 +202,8 @@ module { return %r : tensor } - // Test 9: an opaque `starts` (a graph input, not shape arithmetic) is not - // host-resolvable, so the extent falls back to the data dim upper bound - // rather than emitting an extent that cannot be justified. + // Test 9: opaque bounds are synchronized to the host so the allocation uses + // the actual Slice extent instead of the data-dim upper bound. func.func @test_slice_native_opaque_starts( %data: tensor, %starts: tensor<1xi64>, %ends: tensor<1xi64>) -> tensor { @@ -214,9 +213,9 @@ module { : (tensor, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor - // CHECK-NOT: arith.minsi - // CHECK: %[[DIM:.*]] = tensor.dim %arg1, %{{.*}} : tensor - // CHECK: tensor.empty(%[[DIM]]) : tensor + // CHECK-COUNT-2: hip.readback_scalar + // CHECK: arith.subi + // CHECK: tensor.empty(%{{.*}}) : tensor // CHECK: hip.slice return %r : tensor } @@ -448,4 +447,39 @@ module { // CHECK: hip.slice return %r : tensor } + + // Test 19: a dynamic Slice end may come from an arbitrary producer chain. + // Slice reads the final bound through the shared synchronized scalar path + // instead of interpreting every producer operation itself. + func.func @test_slice_native_reshape_wrapped_end( + %data: tensor, %mask: tensor<1x?xi64>) -> tensor { + // CHECK-LABEL: func.func @test_slice_native_reshape_wrapped_end + %starts = arith.constant dense<[0]> : tensor<1xi64> + %axes = arith.constant dense<[0]> : tensor<1xi64> + %shape = "onnx.Shape"(%mask) {start = 1 : si64, end = 2 : si64} + : (tensor<1x?xi64>) -> tensor<1xi64> + %squeeze_axes = arith.constant dense<[0]> : tensor<1xi64> + %length = "onnx.Squeeze"(%shape, %squeeze_axes) + : (tensor<1xi64>, tensor<1xi64>) -> tensor + %neg_one = arith.constant dense<-1> : tensor + %eight = arith.constant dense<8> : tensor + %one = arith.constant dense<1> : tensor + %minus_one = "onnx.Add"(%length, %neg_one) + : (tensor, tensor) -> tensor + %quotient = "onnx.Div"(%minus_one, %eight) + : (tensor, tensor) -> tensor + %end_scalar = "onnx.Add"(%one, %quotient) + : (tensor, tensor) -> tensor + %shape_spec = arith.constant dense<[1]> : tensor<1xi64> + %ends = "onnx.Reshape"(%end_scalar, %shape_spec) {allowzero = 0 : si64} + : (tensor, tensor<1xi64>) -> tensor<1xi64> + %r = "onnx.Slice"(%data, %starts, %ends, %axes) + : (tensor, tensor<1xi64>, tensor<1xi64>, + tensor<1xi64>) -> tensor + + // CHECK: hip.readback_scalar + // CHECK: tensor.empty(%{{.*}}) : tensor + // CHECK: hip.slice + return %r : tensor + } } diff --git a/test/lit/Pipeline/pipeline-static-shape-lockdown.mlir b/test/lit/Pipeline/pipeline-static-shape-lockdown.mlir index c2e20a738b..526e04ff4d 100644 --- a/test/lit/Pipeline/pipeline-static-shape-lockdown.mlir +++ b/test/lit/Pipeline/pipeline-static-shape-lockdown.mlir @@ -2,28 +2,22 @@ // Licensed under the MIT License. // //===----------------------------------------------------------------------===// -// Lockdown: a static-shape model with no host-fed scalar staging buffers -// must flow through the MaterializeHostScalars + PoolAllocs reorder -// (introduced for the dynseqlen regression) without picking up any host -// scratch. MaterializeHostScalars must reject every alloc here (no host -// I/O on the f32 path, no integer/index element type on the staging path), -// leaving PoolAllocs to handle the IR exactly as it did before the -// reorder. +// Lock down the pre-pooling pass order. Host-fed scalars must move to host +// scratch before OptimizeMemRefs considers ordinary allocations for reuse, +// while a static pure-GPU function must not pick up host scratch. // // If a future refactor relaxes MaterializeHostScalars's candidate filter -// (e.g. drops the `hasHostIO` requirement, accepts f32, or removes the -// arg-0 `!hip.context` check) this test fails immediately. +// (e.g. drops the `hasHostIO` requirement or removes the arg-0 +// `!hip.context` check) this test fails immediately. //===----------------------------------------------------------------------===// -// RUN: hip-mlir-opt --hip-materialize-host-scalars --hip-pool-allocs %s | FileCheck %s +// RUN: hip-mlir-opt --hip-materialize-host-scalars --hip-optimize-memrefs --hip-promote-strided-operands --hip-pool-allocs %s | FileCheck %s // ===== Static, pure-GPU: no host I/O at all. ===== // -// Two memref<8x8xf32> allocs flow through PoolAllocs as-is; no -// hip.get_host_scratch is emitted because: -// - element type f32 is not isIntOrIndex (filter rejects), -// - no memref.store / memref.load user (the `hasHostIO` flag stays -// false even if the type filter is relaxed in the future). +// Two memref<8x8xf32> allocs flow through PoolAllocs as-is. No +// hip.get_host_scratch is emitted because there is no memref.store / +// memref.load user, so the `hasHostIO` requirement rejects both allocations. // // CHECK-LABEL: func.func @static_pure_gpu_lockdown // CHECK-NOT: hip.get_host_scratch @@ -41,3 +35,26 @@ func.func @static_pure_gpu_lockdown( outs(%alloc1 : memref<8x8xf32>) return %alloc1 : memref<8x8xf32> } + +// The host-written scalar is replaced before OptimizeMemRefs runs, so it is a +// view of host scratch rather than a reusable memref.alloc slot previously +// occupied by an asynchronous HIP operation. +// +// CHECK-LABEL: func.func @host_scalar_excluded_from_reuse +// CHECK: %[[SCRATCH:.*]] = hip.get_host_scratch +// CHECK: %[[SHAPE:.*]] = memref.view %[[SCRATCH]] +// CHECK: memref.store %{{.*}}, %[[SHAPE]][] +// CHECK: hip.add{{.*}}ins(%[[SHAPE]] +func.func @host_scalar_excluded_from_reuse( + %ctx: !hip.context, %a: memref, %b: memref, + %host_value: i64) -> memref { + %gpu_tmp = memref.alloc() : memref + hip.add(%ctx) ins(%a, %b : memref, memref) + outs(%gpu_tmp : memref) + %host_scalar = memref.alloc() : memref + memref.store %host_value, %host_scalar[] : memref + %out = memref.alloc() : memref + hip.add(%ctx) ins(%host_scalar, %a : memref, memref) + outs(%out : memref) + return %out : memref +} diff --git a/test/lit/e2e/test_mlp_model.mlir b/test/lit/e2e/test_mlp_model.mlir index c43229bbee..e4506863e8 100644 --- a/test/lit/e2e/test_mlp_model.mlir +++ b/test/lit/e2e/test_mlp_model.mlir @@ -32,7 +32,9 @@ // DEFAULT-BUFFERIZE: lower-affine // DEFAULT-BUFFERIZE: cse // DEFAULT-BUFFERIZE: canonicalize +// DEFAULT-BUFFERIZE: hip-materialize-host-scalars // DEFAULT-BUFFERIZE: hip-optimize-memrefs +// DEFAULT-BUFFERIZE: hip-promote-strided-operands // DEFAULT-BUFFERIZE: hip-resolve-memref-dims // DEFAULT-BUFFERIZE: cse // DEFAULT-BUFFERIZE-NOT: hip-hoist-alloc-size-arith From bfa2b3a0285380a7392a561fe0dcd898b665d639 Mon Sep 17 00:00:00 2001 From: cliu1003 Date: Wed, 28 Oct 2026 17:01:46 +0800 Subject: [PATCH 2/2] fix(runtime): read host-mapped GQA seqlens_k without a stream sync hip-materialize-host-scalars now runs before OptimizeMemRefs, so a host-written seqlens_k can land in hipHostMalloc'd host scratch. GQA still read it with hipMemcpyAsync(DeviceToHost) + hipStreamSynchronize, draining all queued GPU work on the first GQA call of every forward pass. Route every seqlens_k host read in gqa.cpp through read_seqlens_k_to_host, which reads hipMemoryTypeHost pointers directly and keeps the synchronized copy for everything else. The fallback copy uses hipMemcpyDefault so host-accessible memory the pointer query does not classify is not rejected as an invalid DeviceToHost source. Co-Authored-By: Claude Opus 5.5 Made-with: Cursor Co-authored-by: Cursor --- lib/Runtime/real/gqa.cpp | 96 ++++++++++++++++++++-------------------- 1 file changed, 49 insertions(+), 47 deletions(-) diff --git a/lib/Runtime/real/gqa.cpp b/lib/Runtime/real/gqa.cpp index 2a21be000e..5fd7cd527a 100644 --- a/lib/Runtime/real/gqa.cpp +++ b/lib/Runtime/real/gqa.cpp @@ -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(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: @@ -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) { @@ -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 @@ -587,11 +611,8 @@ static int gqa_forward_fused( seqlens_k_val = seqlens_k_pre; } else if (B > 1) { std::vector 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) { @@ -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 @@ -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 @@ -1526,11 +1539,8 @@ static int gqa_forward_hipblaslt( seqlens_k_val = seqlens_k_pre; } else if (B > 1) { std::vector 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) { @@ -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; } }