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/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; } } 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