[DispatchCreation] Hoist scalar tensor.extract and tensor.extract_slice (#24552)
This PR aims to :
- Enable hoisting tensor.extract ops that read from a scalar tensor
already outside the dispatch.
- Loosen the requirement to hoist tensor.extract_slice ops.
### Context
I encountered this error when working on a modified version of LFM2.5. I
attach a reproducer which captures the idea.
A causal mask path stayed fused inside a QK matmul dispatch, which
produced large vectors due to tile size propagation in an mmt4d ukernel.
```
mask_slice_qk_repro.mlir:45:18: error: One or more operations with large vector sizes (32768 bytes) were found:
%scores_3d = torch.aten.bmm %q, %k : !torch.vtensor<[16,?,64],f32>, !torch.vtensor<[16,64,256],f32> -> !torch.vtensor<[16,?,256],f32>
^
<unknown>:0: note: %cst = arith.constant dense<0xFF800000> : vector<16x16x16x16xf32>
<unknown>:0: note: %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16x16x16xf32>
<unknown>:0: note: %cst_1 = arith.constant dense<0> : vector<16x16x16x16xi8>
```
This originated due to the mask's slicing offset coming from a
tensor.extract and subsequently producing a scalar metadata chain that
remained inside the dispatch.
### Proposed fix
- `HoistUniformScalarComputePass` can accept more "candidate ops" than
just arith ops. I just included tensor.extract since its what I ran
into. `isUniformScalarForDispatch` is still in charge to verify that the
candidate op is hoistable, so i added the logic to check the
tensor.extract ops.
- `IREE::Flow::isOffsetSizeAndStrideMappableToFlow` got split into two:
`isOffsetSizeAndStrideStructurallyMappableToFlow` just checks if the
slice can be represented as one flat contiguous byte range, and
`isOffsetSizeAndStrideMappableToFlow` checks for that and the additional
tensor.extract provenance.
- `isHoistableOp` in HoistEncodingOps.cpp was rejecting extract slice
ops whose offset, size, and stride where produced by an extract op due
to calling `isOffsetSizeAndStrideMappableToFlow` on them. Now it calls
`isOffsetSizeAndStrideStructurallyMappableToFlow`.
### Additional Notes:
- I tried to not interfere with the codebase's original intentions.
- Since the extract and extract_slice make it out of the dispatch, the
large vectors never occur. I thought this was the right way to address
the root cause of the problem.
- Inspecting the mmt4d ukernel tile size propagation, it seems that the
problematic large vectors originated due to propagating a pack op tiling
config to the outer dims of an accumulator, which should not happen
afaiu. I could work on that separate issue if it is of interest.
<details>
<summary>mlir reproducer</summary>
```mlir
module @module {
func.func @forward(
%query: !torch.vtensor<[1,16,?,64],f32>,
%key: !torch.vtensor<[1,16,64,256],f32>,
%mask: !torch.vtensor<[256,256],ui8>,
%positions: !torch.vtensor<[1],si64>)
-> !torch.vtensor<[1,16,?,256],f32>
attributes {torch.assume_strict_symbolic_shapes} {
%s = torch.symbolic_int "s" {min_val = 1, max_val = 256} : !torch.int
torch.bind_symbolic_shape %query, [%s], affine_map<()[s0] -> (1, 16, s0, 64)> : !torch.vtensor<[1,16,?,64],f32>
%int0 = torch.constant.int 0
%int1 = torch.constant.int 1
%int2 = torch.constant.int 2
%int16 = torch.constant.int 16
%int64 = torch.constant.int 64
%int256 = torch.constant.int 256
%int-1 = torch.constant.int -1
%seq_len = torch.aten.size.int %query, %int2 : !torch.vtensor<[1,16,?,64],f32>, !torch.int -> !torch.int
%pos_tensor = torch.aten.select.int %positions, %int0, %int-1 : !torch.vtensor<[1],si64>, !torch.int, !torch.int -> !torch.vtensor<[],si64>
%pos = torch.aten.item %pos_tensor : !torch.vtensor<[],si64> -> !torch.int
%end = torch.aten.add.int %pos, %seq_len : !torch.int, !torch.int -> !torch.int
%bool_dtype = torch.constant.int 11
%mask_bool = torch.prims.convert_element_type %mask, %bool_dtype : !torch.vtensor<[256,256],ui8>, !torch.int -> !torch.vtensor<[256,256],i1>
%mask_slice = torch.aten.slice.Tensor %mask_bool, %int0, %pos, %end, %int1 : !torch.vtensor<[256,256],i1>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?,256],i1>
torch.bind_symbolic_shape %mask_slice, [%s], affine_map<()[s0] -> (s0, 256)> : !torch.vtensor<[?,256],i1>
%float-Inf = torch.constant.float 0xFFF0000000000000
%float0 = torch.constant.float 0.000000e+00
%f32_dtype = torch.constant.int 6
%none = torch.constant.none
%cpu = torch.constant.device "cpu"
%neg_inf = torch.aten.scalar_tensor %float-Inf, %f32_dtype, %none, %cpu, %none : !torch.float, !torch.int, !torch.none, !torch.Device, !torch.none -> !torch.vtensor<[],f32>
%mask_bias = torch.aten.where.ScalarSelf %mask_slice, %float0, %neg_inf : !torch.vtensor<[?,256],i1>, !torch.float, !torch.vtensor<[],f32> -> !torch.vtensor<[?,256],f32>
torch.bind_symbolic_shape %mask_bias, [%s], affine_map<()[s0] -> (s0, 256)> : !torch.vtensor<[?,256],f32>
%q_shape = torch.prim.ListConstruct %int16, %seq_len, %int64 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
%q = torch.aten.view %query, %q_shape : !torch.vtensor<[1,16,?,64],f32>, !torch.list<int> -> !torch.vtensor<[16,?,64],f32>
torch.bind_symbolic_shape %q, [%s], affine_map<()[s0] -> (16, s0, 64)> : !torch.vtensor<[16,?,64],f32>
%k_shape = torch.prim.ListConstruct %int16, %int64, %int256 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
%k = torch.aten.view %key, %k_shape : !torch.vtensor<[1,16,64,256],f32>, !torch.list<int> -> !torch.vtensor<[16,64,256],f32>
%scores_3d = torch.aten.bmm %q, %k : !torch.vtensor<[16,?,64],f32>, !torch.vtensor<[16,64,256],f32> -> !torch.vtensor<[16,?,256],f32>
torch.bind_symbolic_shape %scores_3d, [%s], affine_map<()[s0] -> (16, s0, 256)> : !torch.vtensor<[16,?,256],f32>
%scores_shape = torch.prim.ListConstruct %int1, %int16, %seq_len, %int256 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%scores = torch.aten.view %scores_3d, %scores_shape : !torch.vtensor<[16,?,256],f32>, !torch.list<int> -> !torch.vtensor<[1,16,?,256],f32>
torch.bind_symbolic_shape %scores, [%s], affine_map<()[s0] -> (1, 16, s0, 256)> : !torch.vtensor<[1,16,?,256],f32>
%result = torch.aten.add.Tensor %scores, %mask_bias, %int1 : !torch.vtensor<[1,16,?,256],f32>, !torch.vtensor<[?,256],f32>, !torch.int -> !torch.vtensor<[1,16,?,256],f32>
torch.bind_symbolic_shape %result, [%s], affine_map<()[s0] -> (1, 16, s0, 256)> : !torch.vtensor<[1,16,?,256],f32>
return %result : !torch.vtensor<[1,16,?,256],f32>
}
}
```
</details>
Compile command:
```
iree-compile \
mask_slice_qk_repro.mlir \
-o mask_slice_qk_repro.vmfb \
--iree-input-type=auto \
--iree-hal-target-device=local \
--iree-opt-data-tiling=true \
--iree-llvmcpu-enable-ukernels=all \
--iree-hal-local-target-device-backends=llvm-cpu \
--iree-hal-local-host-device-backends=llvm-cpu \
--iree-llvmcpu-target-cpu-features=host
```
Assisted by Codex 5.5
---------
Signed-off-by: Juan Ignacio Pisula <pisula@roofline.ai>diff --git a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
index aef20cd..5f8520b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp
@@ -8,6 +8,7 @@
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
+#include "llvm/ADT/STLExtras.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
@@ -67,12 +68,11 @@
return hasExtract;
}
-bool isOffsetSizeAndStrideMappableToFlow(ArrayRef<OpFoldResult> offsets,
- ArrayRef<OpFoldResult> sizes,
- ArrayRef<OpFoldResult> strides,
- ArrayRef<int64_t> baseShape) {
+bool isOffsetSizeAndStrideStructurallyMappableToFlow(
+ ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+ ArrayRef<OpFoldResult> strides, ArrayRef<int64_t> baseShape) {
if (offsets.size() != baseShape.size()) {
- // Unhanded rank-reducing case.
+ // Unhandled rank-reducing case.
return false;
}
auto getVal = [](OpFoldResult valueOrAttr, int64_t dynamicVal) -> int64_t {
@@ -90,16 +90,6 @@
OpFoldResult size = sizes[dim - 1];
OpFoldResult stride = strides[dim - 1];
- // The offsets and sizes dont have to be static for all dimensions. When
- // `fullSlices` is true, the offset and sizes can be dynamic. But in some
- // cases, the dynamic offset/size value is obtained by computing from
- // another tensor which lives on the device. To avoid host-round tripping
- // try to infer when a value is extracted from a tensor.
- if (producedByValueExtract(offset) || producedByValueExtract(stride) ||
- producedByValueExtract(size)) {
- return false;
- }
-
int64_t staticOffset = getVal(offset, ShapedType::kDynamic);
int64_t staticSize = getVal(size, ShapedType::kDynamic);
int64_t staticStride = getVal(stride, ShapedType::kDynamic);
@@ -125,6 +115,29 @@
return true;
}
+bool isOffsetSizeAndStrideMappableToFlow(ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes,
+ ArrayRef<OpFoldResult> strides,
+ ArrayRef<int64_t> baseShape) {
+ if (!isOffsetSizeAndStrideStructurallyMappableToFlow(offsets, sizes, strides,
+ baseShape)) {
+ return false;
+ }
+
+ // The offsets and sizes don't have to be static for all dimensions. When
+ // `fullSlices` is true, the offset and sizes can be dynamic. But in some
+ // cases, the dynamic offset/size value is obtained by computing from another
+ // tensor which lives on the device. To avoid host-round tripping try to infer
+ // when a value is extracted from a tensor.
+ for (auto [offset, size, stride] : llvm::zip_equal(offsets, sizes, strides)) {
+ if (producedByValueExtract(offset) || producedByValueExtract(size) ||
+ producedByValueExtract(stride)) {
+ return false;
+ }
+ }
+ return true;
+}
+
LogicalResult
convertInsertSliceOpToFlowUpdateOp(RewriterBase &rewriter,
tensor::InsertSliceOp insertOp) {
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h
index 0713698..592277e 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h
+++ b/compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h
@@ -14,7 +14,14 @@
namespace mlir::iree_compiler::IREE::Flow {
/// Indicates whether the given offsets/sizes/strides representing a slice from
-/// baseShape is a contiguous slice, and this is mappable to Flow ops.
+/// baseShape is a contiguous slice structurally representable by Flow ops.
+bool isOffsetSizeAndStrideStructurallyMappableToFlow(
+ llvm::ArrayRef<OpFoldResult> offsets, llvm::ArrayRef<OpFoldResult> sizes,
+ llvm::ArrayRef<OpFoldResult> strides, llvm::ArrayRef<int64_t> baseShape);
+
+/// Indicates whether the given offsets/sizes/strides representing a slice from
+/// baseShape is mappable to Flow ops without tensor-load-derived dynamic
+/// offsets/sizes/strides.
bool isOffsetSizeAndStrideMappableToFlow(llvm::ArrayRef<OpFoldResult> offsets,
llvm::ArrayRef<OpFoldResult> sizes,
llvm::ArrayRef<OpFoldResult> strides,
diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
index d225a32..4796602 100644
--- a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
@@ -159,15 +159,16 @@
}
/// Returns true if the op is hoistable outside dispatches, which indicates that
-/// the ops can be either mappable to Flow ops or get hoisted to globals.
+/// the op can be moved across the dispatch boundary by this pass or hoisted to
+/// globals.
static bool isHoistableOp(Operation *op) {
if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();
SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();
ArrayRef<int64_t> srcShape = sliceOp.getSourceType().getShape();
- return IREE::Flow::isOffsetSizeAndStrideMappableToFlow(offsets, sizes,
- strides, srcShape);
+ return IREE::Flow::isOffsetSizeAndStrideStructurallyMappableToFlow(
+ offsets, sizes, strides, srcShape);
}
// ConstExprHoistingPolicy has an assumption that any root op is not hoistable
// because they are already hoisted. This is not the case when the parent op
diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp
index e0c6096..382627e 100644
--- a/compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp
@@ -9,6 +9,7 @@
#include "iree/compiler/DispatchCreation/Passes.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
@@ -31,11 +32,36 @@
return v.getParentRegion()->getParentOp() != dispatch;
};
- // Check that all operands are defined outside the dispatch and all operand
- // and result types are scalars or vectors.
- return llvm::all_of(op->getOperands(), isOutsideDispatch) &&
- llvm::all_of(op->getOperandTypes(), isIntOrIndex) &&
- llvm::all_of(op->getResultTypes(), isIntOrIndex);
+ // Arith candidates are uniform when all operands are defined outside the
+ // dispatch and all operand/result types are int/index scalars.
+ if (isa<arith::ArithDialect>(op->getDialect())) {
+ return llvm::all_of(op->getOperands(), isOutsideDispatch) &&
+ llvm::all_of(op->getOperandTypes(), isIntOrIndex) &&
+ llvm::all_of(op->getResultTypes(), isIntOrIndex);
+ }
+
+ // A tensor.extract is only uniform scalar metadata when it extracts the sole
+ // int/index element from a tensor and all extraction dependencies are already
+ // outside this dispatch.
+ if (auto extractOp = dyn_cast<tensor::ExtractOp>(op)) {
+ auto tensorType =
+ dyn_cast<RankedTensorType>(extractOp.getTensor().getType());
+ return tensorType && tensorType.hasStaticShape() &&
+ tensorType.getNumElements() == 1 &&
+ extractOp.getResult().getType().isIntOrIndex() &&
+ isOutsideDispatch(extractOp.getTensor()) &&
+ llvm::all_of(extractOp.getIndices(), isOutsideDispatch);
+ }
+
+ return false;
+}
+
+// Restrict to arith ops and scalar extracts from one-element
+// tensors to avoid unexpected hoisting of flow/stream/hal.dispatch
+// workgroups count/id ops.
+static bool isHoistCandidate(Operation *op) {
+ return isa<arith::ArithDialect>(op->getDialect()) ||
+ isa<tensor::ExtractOp>(op);
}
namespace {
@@ -52,11 +78,9 @@
funcOp.walk([&](IREE::Flow::DispatchRegionOp dispatch) {
for (Block &body : dispatch.getBody()) {
SmallVector<Operation *> ops;
- // Restrict to arith ops to avoid unexpected hoisting of
- // flow/stream/hal.dispatch.workgroups.count/id ops.
// TODO: Add an op trait to tie count/id ops to region ops.
for (Operation &op : body.getOperations()) {
- if (isa<arith::ArithDialect>(op.getDialect())) {
+ if (isHoistCandidate(&op)) {
ops.push_back(&op);
}
}
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
index 2d075f5..95a1a7a 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
@@ -239,6 +239,88 @@
// -----
+// Do not hoist a `set_encoding` whose source is a structurally flow-mappable
+// slice if the dynamic offset is still produced by a tensor.extract inside the
+// dispatch. HoistUniformScalarCompute is responsible for moving such scalar
+// metadata out first.
+
+#encoding = #iree_encoding.testing<>
+util.func private @get_tensor(tensor<64x256xi1, #encoding>) -> tensor<64x256xi1>
+util.func public @no_hoist_slice_with_dispatch_local_extract_offset(%input: tensor<256x256xi1>, %offsets: tensor<1xi64>) -> tensor<64x256xi1> {
+ %0 = flow.dispatch.region -> (tensor<64x256xi1>) {
+ %c0 = arith.constant 0 : index
+ %offset_i64 = tensor.extract %offsets[%c0] : tensor<1xi64>
+ %offset = arith.index_cast %offset_i64 : i64 to index
+ %1 = tensor.extract_slice %input[%offset, 0] [64, 256] [1, 1] : tensor<256x256xi1> to tensor<64x256xi1>
+ %2 = iree_encoding.set_encoding %1 : tensor<64x256xi1> -> tensor<64x256xi1, #encoding>
+ %3 = util.call @get_tensor(%2) : (tensor<64x256xi1, #encoding>) -> tensor<64x256xi1>
+ flow.return %3 : tensor<64x256xi1>
+ }
+ util.return %0 : tensor<64x256xi1>
+}
+// CHECK-LABEL: util.func public @no_hoist_slice_with_dispatch_local_extract_offset(
+// CHECK: flow.dispatch.region
+// CHECK: %[[OFFSET_I64:.+]] = tensor.extract
+// CHECK: %[[OFFSET:.+]] = arith.index_cast %[[OFFSET_I64]]
+// CHECK: %[[SRC:.+]] = tensor.extract_slice %{{.*}}[%[[OFFSET]], 0] [64, 256] [1, 1]
+// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[SRC]]
+// CHECK: flow.return
+
+// -----
+
+// Hoist a `set_encoding` whose source is a structurally flow-mappable slice,
+// even if a dynamic offset is derived from a scalar tensor.extract that has
+// already been hoisted outside the dispatch.
+
+#encoding = #iree_encoding.testing<>
+util.func private @get_tensor(tensor<64x256xi1, #encoding>) -> tensor<64x256xi1>
+util.func public @hoist_slice_with_external_extract_offset(%input: tensor<256x256xi1>, %offsets: tensor<1xi64>) -> tensor<64x256xi1> {
+ %c0 = arith.constant 0 : index
+ %offset_i64 = tensor.extract %offsets[%c0] : tensor<1xi64>
+ %offset = arith.index_cast %offset_i64 : i64 to index
+ %0 = flow.dispatch.region -> (tensor<64x256xi1>) {
+ %1 = tensor.extract_slice %input[%offset, 0] [64, 256] [1, 1] : tensor<256x256xi1> to tensor<64x256xi1>
+ %2 = iree_encoding.set_encoding %1 : tensor<64x256xi1> -> tensor<64x256xi1, #encoding>
+ %3 = util.call @get_tensor(%2) : (tensor<64x256xi1, #encoding>) -> tensor<64x256xi1>
+ flow.return %3 : tensor<64x256xi1>
+ }
+ util.return %0 : tensor<64x256xi1>
+}
+// CHECK-LABEL: util.func public @hoist_slice_with_external_extract_offset(
+// CHECK-SAME: %[[INPUT:.+]]: tensor<256x256xi1>, %[[OFFSETS:.+]]: tensor<1xi64>
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[OFFSET_I64:.+]] = tensor.extract %[[OFFSETS]][%[[C0]]] : tensor<1xi64>
+// CHECK: %[[OFFSET:.+]] = arith.index_cast %[[OFFSET_I64]] : i64 to index
+// CHECK: %[[SRC:.+]] = tensor.extract_slice %[[INPUT]][%[[OFFSET]], 0] [64, 256] [1, 1]
+// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[SRC]]
+// CHECK: flow.dispatch.region
+
+// -----
+
+// Do not hoist a `set_encoding` whose source is an `extract_slice` that is NOT
+// structurally flow-mappable: dim 1 has size 12 (not 1) while dim 2 is
+// partially sliced, so the slice is not a single contiguous byte range.
+
+#encoding = #iree_encoding.testing<>
+util.func private @get_tensor(tensor<12x100x64xf16, #encoding>) -> tensor<12x100x64xf16>
+util.func public @no_hoist_non_flow_mappable_slice(%input: tensor<1x12x2080x64xf16>) -> tensor<12x100x64xf16> {
+ %0 = flow.dispatch.region -> (tensor<12x100x64xf16>) {
+ %1 = tensor.extract_slice %input[0, 0, 0, 0] [1, 12, 100, 64] [1, 1, 1, 1]
+ : tensor<1x12x2080x64xf16> to tensor<12x100x64xf16>
+ %2 = iree_encoding.set_encoding %1 : tensor<12x100x64xf16> -> tensor<12x100x64xf16, #encoding>
+ %3 = util.call @get_tensor(%2) : (tensor<12x100x64xf16, #encoding>) -> tensor<12x100x64xf16>
+ flow.return %3 : tensor<12x100x64xf16>
+ }
+ util.return %0 : tensor<12x100x64xf16>
+}
+// CHECK-LABEL: util.func public @no_hoist_non_flow_mappable_slice(
+// CHECK: flow.dispatch.region
+// CHECK: %[[SRC:.+]] = tensor.extract_slice
+// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[SRC]]
+// CHECK: flow.return
+
+// -----
+
// Avoid hoisting `set_encoding` operations that have pad encodings
#encoding = #iree_encoding.padding<[0, ?]>
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir b/compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir
index a29880f..365a1fd 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir
@@ -68,3 +68,48 @@
// CHECK: arith.muli
// CHECK: tensor.empty
// CHECK: flow.return
+
+// -----
+
+util.func public @hoist_one_element_tensor_extract(%arg0: tensor<1xi64>, %m: index) -> tensor<?xf32> {
+ %0 = flow.dispatch.region -> (tensor<?xf32>{%m}) {
+ %c0 = arith.constant 0 : index
+ %extracted = tensor.extract %arg0[%c0] : tensor<1xi64>
+ %idx = arith.index_cast %extracted : i64 to index
+ %sum = arith.addi %idx, %m : index
+ %1 = tensor.empty(%sum) : tensor<?xf32>
+ flow.return %1 : tensor<?xf32>
+ }
+ util.return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: @hoist_one_element_tensor_extract
+// CHECK-SAME: %[[ARG0:.+]]: tensor<1xi64>, %[[M:.+]]: index
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[C0]]] : tensor<1xi64>
+// CHECK: %[[IDX:.+]] = arith.index_cast %[[EXTRACT]] : i64 to index
+// CHECK: %[[SUM:.+]] = arith.addi %[[IDX]], %[[M]] : index
+// CHECK: flow.dispatch.region
+// CHECK: tensor.empty(%[[SUM]])
+// CHECK: flow.return
+
+// -----
+
+util.func public @dont_hoist_extract_from_multi_element_tensor(%arg0: tensor<4xi64>, %i: index, %m: index) -> tensor<?xf32> {
+ %0 = flow.dispatch.region -> (tensor<?xf32>{%m}) {
+ %extracted = tensor.extract %arg0[%i] : tensor<4xi64>
+ %idx = arith.index_cast %extracted : i64 to index
+ %sum = arith.addi %idx, %m : index
+ %1 = tensor.empty(%sum) : tensor<?xf32>
+ flow.return %1 : tensor<?xf32>
+ }
+ util.return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: @dont_hoist_extract_from_multi_element_tensor
+// CHECK: flow.dispatch.region
+// CHECK: %[[EXTRACT:.+]] = tensor.extract
+// CHECK: %[[IDX:.+]] = arith.index_cast %[[EXTRACT]]
+// CHECK: %[[SUM:.+]] = arith.addi %[[IDX]]
+// CHECK: tensor.empty(%[[SUM]])
+// CHECK: flow.return
diff --git a/tests/e2e/stablehlo_ops/CMakeLists.txt b/tests/e2e/stablehlo_ops/CMakeLists.txt
index 8c00921..9ff2976 100644
--- a/tests/e2e/stablehlo_ops/CMakeLists.txt
+++ b/tests/e2e/stablehlo_ops/CMakeLists.txt
@@ -693,7 +693,7 @@
"pad.mlir"
# "philox.mlir" # TODO(#12509): WebGPU SPIR-V broken
# "pow.mlir" # TODO(#11321): error: value nan cannot be represented as 'f32'
- "reduce.mlir"
+ # "reduce.mlir" # TODO(#24650): error: use of variable address space 'immediate'
"reduce_window.mlir"
"remainder.mlir"
"reshape.mlir"
