[GPU] Apply warp reduce to intermediate warp reduce results. (#11082)
Instead of doing vector.reduction after obtaining warp reduction
results, we apply another warp reduce to obtain the final reduction
results.
This shouldn't give much perf boost on it's own, but will give
flexibility to enable warp reduction on fp16 and int8 down
the line.
diff --git a/compiler/src/iree/compiler/Codegen/Common/VectorReductionToGPU.cpp b/compiler/src/iree/compiler/Codegen/Common/VectorReductionToGPU.cpp
index c6ed1a0..fe2b7c6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/VectorReductionToGPU.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/VectorReductionToGPU.cpp
@@ -56,6 +56,45 @@
return laneVal;
}
+// List of identity elements by operation.
+// https://en.wikipedia.org/wiki/Identity_element
+static Attribute getCombiningKindIdentity(OpBuilder &builder,
+ vector::CombiningKind combiningKind,
+ Type type) {
+ switch (combiningKind) {
+ case vector::CombiningKind::ADD:
+ return builder.getZeroAttr(type);
+ case vector::CombiningKind::MUL: {
+ if (type.isIntOrIndex()) {
+ return builder.getIntegerAttr(type, 1);
+ }
+ return builder.getFloatAttr(type, 1);
+ }
+ case vector::CombiningKind::MINUI:
+ case vector::CombiningKind::MINSI:
+ return builder.getIntegerAttr(type, std::numeric_limits<int64_t>::max());
+ case vector::CombiningKind::MAXUI:
+ case vector::CombiningKind::MAXSI:
+ return builder.getIntegerAttr(type, std::numeric_limits<int64_t>::min());
+ case vector::CombiningKind::AND:
+ return builder.getIntegerAttr(type, 1);
+ case vector::CombiningKind::OR:
+ case vector::CombiningKind::XOR:
+ return builder.getZeroAttr(type);
+ case vector::CombiningKind::MINF: {
+ auto posInfApFloat = APFloat::getInf(
+ type.cast<FloatType>().getFloatSemantics(), /*Negative=*/false);
+ return builder.getFloatAttr(type, posInfApFloat);
+ }
+ case vector::CombiningKind::MAXF: {
+ auto negInfApFloat = APFloat::getInf(
+ type.cast<FloatType>().getFloatSemantics(), /*Negative=*/true);
+ return builder.getFloatAttr(type, negInfApFloat);
+ }
+ }
+ return Attribute();
+}
+
/// Emit reduction across a group for a given input.
static Value groupReduction(Location loc, OpBuilder &builder, Value input,
vector::CombiningKind kind, uint32_t size,
@@ -67,6 +106,9 @@
// if we have more than one warp, reduce across warps.
if (size > warpSize) {
uint32_t numWarp = size / warpSize;
+ assert(numWarp <= warpSize &&
+ "Only support 1 level, need to implement recursive/loop for this "
+ "case.");
MemRefType memrefType =
MemRefType::get(numWarp, input.getType(), {},
gpu::GPUDialect::getWorkgroupAddressSpace());
@@ -79,11 +121,21 @@
SmallVector<Value> indices = {warpId};
builder.create<memref::StoreOp>(loc, laneVal, alloc, indices);
builder.create<gpu::BarrierOp>(loc);
- VectorType vecType = VectorType::get(numWarp, input.getType());
- Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
- Value vec =
- builder.create<vector::TransferReadOp>(loc, vecType, alloc, zero);
- laneVal = builder.create<vector::ReductionOp>(loc, kind, vec);
+ // Further reduce the outputs from each warps with a single warp reduce.
+ Value laneId = builder.create<arith::RemUIOp>(loc, threadX, cstWarpSize);
+ Value loadVal = builder.create<memref::LoadOp>(loc, alloc, laneId);
+ Value cstNumWarp = builder.create<arith::ConstantIndexOp>(loc, numWarp);
+ // Pad with identity element if numel < warpSize for valid warp reduction.
+ Value useIdentityElement = builder.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sge, laneId, cstNumWarp);
+ Attribute identityAttr =
+ getCombiningKindIdentity(builder, kind, loadVal.getType());
+ assert(identityAttr && "Unknown identity value for the reduction");
+ // TODO: Avoid reduction across all lanes if numWarp <= warpSize/2.
+ Value identity = builder.create<arith::ConstantOp>(loc, identityAttr);
+ Value selectedInput = builder.create<arith::SelectOp>(
+ loc, useIdentityElement, identity, loadVal);
+ laneVal = warpReduction(loc, builder, selectedInput, kind, warpSize);
}
return laneVal;
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/warp_reduction.mlir b/compiler/src/iree/compiler/Codegen/Common/test/warp_reduction.mlir
index 5057569..35f9cba 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/warp_reduction.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/warp_reduction.mlir
@@ -125,7 +125,9 @@
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i32
// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : i32
// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : i32
+// CHECK-DAG: %[[C2I:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C32I:.*]] = arith.constant 32 : index
+// CHECK-DAG: %[[IDENTITY:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[TID:.*]] = gpu.thread_id x
// CHECK-DAG: %[[VCST:.*]] = arith.constant dense<0.000000e+00> : vector<1xf32>
// CHECK: %[[F:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[V0:.*]] = %[[VCST]]) -> (vector<1xf32>) {
@@ -147,10 +149,22 @@
// CHECK: %[[WID:.*]] = arith.divui %[[TID]], %[[C32I]] : index
// CHECK: memref.store %[[S9]], %[[A]][%[[WID]]] : memref<2xf32, 3>
// CHECK: gpu.barrier
-// CHECK: %[[VL:.*]] = vector.transfer_read %[[A]][%[[C0]]], %{{.*}} {in_bounds = [true]} : memref<2xf32, 3>, vector<2xf32>
-// CHECK: %[[S10:.*]] = vector.reduction <add>, %[[VL]] : vector<2xf32> into f32
-// CHECK: %[[S11:.*]] = arith.addf %[[S10]], %[[E]] : f32
-// CHECK: %[[B:.*]] = vector.broadcast %[[S11]] : f32 to vector<1xf32>
+// CHECK: %[[LANE_ID:.*]] = arith.remui %[[TID]], %[[C32I]] : index
+// CHECK: %[[LOAD_VAL:.*]] = memref.load %[[A]][%[[LANE_ID]]] : memref<2xf32, 3>
+// CHECK: %[[USE_IDENTITY:.*]] = arith.cmpi sge, %[[LANE_ID]], %[[C2I]] : index
+// CHECK: %[[LANE_VAL:.*]] = arith.select %[[USE_IDENTITY]], %[[IDENTITY]], %[[LOAD_VAL]] : f32
+// CHECK: %[[S10:.*]], %{{.*}} = gpu.shuffle xor %[[LANE_VAL]], %[[C1]], %[[C32]] : f32
+// CHECK: %[[S11:.*]] = arith.addf %[[LANE_VAL]], %[[S10]] : f32
+// CHECK: %[[S12:.*]], %{{.*}} = gpu.shuffle xor %[[S11]], %[[C2]], %[[C32]] : f32
+// CHECK: %[[S13:.*]] = arith.addf %[[S11]], %[[S12]] : f32
+// CHECK: %[[S14:.*]], %{{.*}} = gpu.shuffle xor %[[S13]], %[[C4]], %[[C32]] : f32
+// CHECK: %[[S15:.*]] = arith.addf %[[S13]], %[[S14]] : f32
+// CHECK: %[[S16:.*]], %{{.*}} = gpu.shuffle xor %[[S15]], %[[C8]], %[[C32]] : f32
+// CHECK: %[[S17:.*]] = arith.addf %[[S15]], %[[S16]] : f32
+// CHECK: %[[S18:.*]], %{{.*}} = gpu.shuffle xor %[[S17]], %[[C16]], %[[C32]] : f32
+// CHECK: %[[S19:.*]] = arith.addf %[[S17]], %[[S18]] : f32
+// CHECK: %[[S20:.*]] = arith.addf %[[S19]], %[[E]] : f32
+// CHECK: %[[B:.*]] = vector.broadcast %[[S20]] : f32 to vector<1xf32>
// CHECK: scf.yield %[[B]] : vector<1xf32>
// CHECK: }
// CHECK: %[[DIV:.*]] = arith.divf %[[F]], %{{.*}} : vector<1xf32>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
index d391415..62cdc33 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
@@ -867,8 +867,9 @@
// CHECK-COUNT-5: nvvm.shfl.sync bfly
// CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm.ptr<f32, 3>
// CHECK: nvvm.barrier0
-// CHECK: llvm.load {{.*}} : !llvm.ptr<vector<8xf32>, 3>
-// CHECK: "llvm.intr.vector.reduce.fadd"(%{{.*}}, %{{.*}}) {reassoc = false} : (f32, vector<8xf32>) -> f32
+// CHECK: llvm.load {{.*}} : !llvm.ptr<f32, 3>
+// CHECK: llvm.icmp "sge" {{.*}}, {{.*}} : i64
+// CHECK-COUNT-5: nvvm.shfl.sync bfly
// -----
@@ -929,8 +930,9 @@
// CHECK-COUNT-5: nvvm.shfl.sync bfly
// CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm.ptr<f32, 3>
// CHECK: nvvm.barrier0
-// CHECK: llvm.load {{.*}} : !llvm.ptr<vector<8xf32>, 3>
-// CHECK: "llvm.intr.vector.reduce.fadd"(%{{.*}}, %{{.*}}) {reassoc = false} : (f32, vector<8xf32>) -> f32
+// CHECK: llvm.load {{.*}} : !llvm.ptr<f32, 3>
+// CHECK: llvm.icmp "sge" {{.*}}, {{.*}} : i64
+// CHECK-COUNT-5: nvvm.shfl.sync bfly
// CHECK: llvm.fdiv %{{.*}}, %{{.*}} : vector<4xf32>
// CHECK: llvm.store %{{.*}}, %{{.*}} {alignment = 4 : i64} : !llvm.ptr<vector<4xf32>>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline.mlir
index 9555eeb..368d490 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline.mlir
@@ -50,6 +50,9 @@
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : i32
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : i32
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i32
+// CHECK-DAG: %[[C16I:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32I:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[IDENTITY:.+]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[CF:.+]] = arith.constant 1.000000e+00 : f32
// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<4xf32>
// CHECK-DAG: %[[TID:.+]] = gpu.thread_id x
@@ -73,13 +76,25 @@
// CHECK: %[[WID:.+]] = arith.divui %{{.*}}, %{{.*}} : index
// CHECK: memref.store %[[R6]], %[[ALLOC]][%[[WID]]] : memref<16xf32, 3>
// CHECK: gpu.barrier
-// CHECK: %[[B:.+]] = vector.transfer_read %[[ALLOC]][%[[C0]]], %{{.*}} {in_bounds = [true]} : memref<16xf32, 3>, vector<16xf32>
-// CHECK: %[[R7:.+]] = vector.reduction <add>, %[[B]] : vector<16xf32> into f32
-// CHECK: %[[R8:.+]] = arith.addf %[[R7]], %[[CF]] : f32
-// CHECK: %[[R9:.+]] = vector.broadcast %[[R8]] : f32 to vector<1xf32>
+// CHECK: %[[LANE_ID:.+]] = arith.remui %[[TID]], %[[C32I]] : index
+// CHECK: %[[LOAD_VAL:.+]] = memref.load %[[ALLOC]][%[[LANE_ID]]] : memref<16xf32, 3>
+// CHECK: %[[USE_IDENTITY:.+]] = arith.cmpi sge, %[[LANE_ID]], %[[C16I]] : index
+// CHECK: %[[LANE_VAL:.+]] = arith.select %[[USE_IDENTITY]], %[[IDENTITY]], %[[LOAD_VAL]] : f32
+// CHECK: %[[S5:.+]], %{{.*}} = gpu.shuffle xor %[[LANE_VAL]], %[[C1]], %[[C32]] : f32
+// CHECK: %[[R7:.+]] = arith.addf %[[LANE_VAL]], %[[S5]] : f32
+// CHECK: %[[S6:.+]], %{{.*}} = gpu.shuffle xor %[[R7]], %[[C2]], %[[C32]] : f32
+// CHECK: %[[R8:.+]] = arith.addf %[[R7]], %[[S6]] : f32
+// CHECK: %[[S7:.+]], %{{.*}} = gpu.shuffle xor %[[R8]], %[[C4]], %[[C32]] : f32
+// CHECK: %[[R9:.+]] = arith.addf %[[R8]], %[[S7]] : f32
+// CHECK: %[[S8:.+]], %{{.*}} = gpu.shuffle xor %[[R9]], %[[C8]], %[[C32]] : f32
+// CHECK: %[[R10:.+]] = arith.addf %[[R9]], %[[S8]] : f32
+// CHECK: %[[S9:.+]], %{{.*}} = gpu.shuffle xor %[[R10]], %[[C16]], %[[C32]] : f32
+// CHECK: %[[R11:.+]] = arith.addf %[[R10]], %[[S9]] : f32
+// CHECK: %[[R12:.+]] = arith.addf %[[R11]], %[[CF]] : f32
+// CHECK: %[[R13:.+]] = vector.broadcast %[[R12]] : f32 to vector<1xf32>
// CHECK: %[[TID0:.+]] = arith.cmpi eq, %[[TID]], %[[C0]] : index
// CHECK: scf.if %[[TID0]] {
-// CHECK: vector.transfer_write %[[R9]], %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<1xf32>, memref<512xf32>
+// CHECK: vector.transfer_write %[[R13]], %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<1xf32>, memref<512xf32>
// CHECK: }
// -----
@@ -152,11 +167,23 @@
// CHECK: arith.addf
// CHECK: memref.store {{.*}} : memref<16xf32, 3>
// CHECK: gpu.barrier
-// CHECK: vector.transfer_read
-// CHECK: vector.reduction
+// CHECK: arith.remui
+// CHECK: memref.load
+// CHECK: arith.cmpi
+// CHECK: arith.select
+// CHECK: gpu.shuffle xor
+// CHECK: arith.addf
+// CHECK: gpu.shuffle xor
+// CHECK: arith.addf
+// CHECK: gpu.shuffle xor
+// CHECK: arith.addf
+// CHECK: gpu.shuffle xor
+// CHECK: arith.addf
+// CHECK: gpu.shuffle xor
+// CHECK: arith.addf
// CHECK: arith.addf
// CHECK: vector.broadcast %{{.*}} : f32 to vector<4xf32>
-// CHECK: %15 = arith.divf {{.*}} : vector<4xf32>
+// CHECK: %22 = arith.divf {{.*}} : vector<4xf32>
// CHECK: scf.for
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, memref<512x10240xf32>
// CHECK: }
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
index f7096b0..5214c2b 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -678,6 +678,16 @@
{64, uint64_t(groupSize)}, {64, uint64_t(maxWorkgroupSize)})
.getZExtValue();
}
+ // Current warp reduction pattern is a two step butterfly warp reduce.
+ // First, do warp reductions along multiple subgroups.
+ // Second, reduce results from multiple subgroups using single warp reduce.
+ // The final warp reduce requires numSubgroupUsed > subgroupSize to work.
+ // TODO(raikonenfnu): Add flexible num of warp reduce to handle more configs.
+ // TT::CPU and TT::ARM_Valhall is not going through warp reduce.
+ const int64_t numSubgroupsUsed = groupSize / subgroupSize;
+ if (numSubgroupsUsed > subgroupSize) {
+ return failure();
+ }
std::array<int64_t, 3> workgroupSize = {groupSize, 1, 1};
// Tile all the parallel dimension to 1.
SmallVector<unsigned> partitionedLoops =
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_reduction_subgroup.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_reduction_subgroup.mlir
index 4de4227..84b8d7b 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_reduction_subgroup.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/pipeline_reduction_subgroup.mlir
@@ -76,9 +76,18 @@
// CHECK: spirv.ControlBarrier <Workgroup>, <Workgroup>, <AcquireRelease|WorkgroupMemory>
-// CHECK-COUNT-8: spirv.Load "Workgroup" %{{.+}} : f32
-// CHECK-COUNT-7: spirv.FAdd %{{.+}}, %{{.+}} : f32
-// CHECK: spirv.FAdd %{{.+}}, %[[F0]] : f32
+// CHECK: %[[LOAD_VAL:.+]] = spirv.Load "Workgroup" {{.+}} : f32
+// CHECK: %[[USE_IDENTITY:.+]] = spirv.SGreaterThanEqual {{.+}}, %[[C8]] : i32
+// CHECK: %[[LANE_VAL:.+]] = spirv.Select %[[USE_IDENTITY]], %[[F0]], %[[LOAD_VAL]] : i1, f32
+// CHECK: %[[S4:.+]] = spirv.GroupNonUniformShuffleXor <Subgroup> %[[LANE_VAL]], %[[C1]] : f32, i32
+// CHECK: %[[ADD7:.+]] = spirv.FAdd %[[LANE_VAL]], %[[S4]] : f32
+// CHECK: %[[S5:.+]] = spirv.GroupNonUniformShuffleXor <Subgroup> %[[ADD7]], %[[C2]] : f32, i32
+// CHECK: %[[ADD8:.+]] = spirv.FAdd %[[ADD7]], %[[S5]] : f32
+// CHECK: %[[S6:.+]] = spirv.GroupNonUniformShuffleXor <Subgroup> %[[ADD8]], %[[C4]] : f32, i32
+// CHECK: %[[ADD9:.+]] = spirv.FAdd %[[ADD8]], %[[S6]] : f32
+// CHECK: %[[S7:.+]] = spirv.GroupNonUniformShuffleXor <Subgroup> %[[ADD9]], %[[C8]] : f32, i32
+// CHECK: %[[ADD10:.+]] = spirv.FAdd %[[ADD9]], %[[S7]] : f32
+// CHECK: spirv.FAdd %[[ADD10]], %[[F0]] : f32
// CHECK: %[[EQ:.+]] = spirv.IEqual %{{.+}}, %[[C0]] : i32
// CHECK: spirv.mlir.selection {
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/reduction_lowering.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/reduction_lowering.mlir
index 8eb1090..ce2a94f 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/reduction_lowering.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/reduction_lowering.mlir
@@ -46,6 +46,9 @@
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : i32
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : i32
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i32
+// CHECK-DAG: %[[C4I:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C32I:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[IDENTITY:.+]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[CF:.+]] = arith.constant 1.000000e+00 : f32
// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<4xf32>
// CHECK-DAG: %[[TID:.+]] = gpu.thread_id x
@@ -69,11 +72,23 @@
// CHECK: %[[WID:.+]] = arith.divui %{{.*}}, %{{.*}} : index
// CHECK: memref.store %[[R6]], %[[ALLOC]][%[[WID]]] : memref<4xf32, 3>
// CHECK: gpu.barrier
-// CHECK: %[[B:.+]] = vector.transfer_read %[[ALLOC]][%[[C0]]], %{{.*}} {in_bounds = [true]} : memref<4xf32, 3>, vector<4xf32>
-// CHECK: %[[R7:.+]] = vector.reduction <add>, %[[B]] : vector<4xf32> into f32
-// CHECK: %[[R8:.+]] = arith.addf %[[R7]], %[[CF]] : f32
-// CHECK: %[[R9:.+]] = vector.splat %[[R8]] : vector<1xf32>
+// CHECK: %[[LANE_ID:.+]] = arith.remui %[[TID]], %[[C32I]] : index
+// CHECK: %[[LOAD_VAL:.+]] = memref.load %[[ALLOC]][%[[LANE_ID]]] : memref<4xf32, 3>
+// CHECK: %[[USE_IDENTITY:.+]] = arith.cmpi sge, %[[LANE_ID]], %[[C4I]] : index
+// CHECK: %[[LANE_VAL:.+]] = arith.select %[[USE_IDENTITY]], %[[IDENTITY]], %[[LOAD_VAL]] : f32
+// CHECK: %[[S5:.+]], %{{.*}} = gpu.shuffle xor %[[LANE_VAL]], %[[C1]], %[[C32]] : f32
+// CHECK: %[[R7:.+]] = arith.addf %[[LANE_VAL]], %[[S5]] : f32
+// CHECK: %[[S6:.+]], %{{.*}} = gpu.shuffle xor %[[R7]], %[[C2]], %[[C32]] : f32
+// CHECK: %[[R8:.+]] = arith.addf %[[R7]], %[[S6]] : f32
+// CHECK: %[[S7:.+]], %{{.*}} = gpu.shuffle xor %[[R8]], %[[C4]], %[[C32]] : f32
+// CHECK: %[[R9:.+]] = arith.addf %[[R8]], %[[S7]] : f32
+// CHECK: %[[S8:.+]], %{{.*}} = gpu.shuffle xor %[[R9]], %[[C8]], %[[C32]] : f32
+// CHECK: %[[R10:.+]] = arith.addf %[[R9]], %[[S8]] : f32
+// CHECK: %[[S9:.+]], %{{.*}} = gpu.shuffle xor %[[R10]], %[[C16]], %[[C32]] : f32
+// CHECK: %[[R11:.+]] = arith.addf %[[R10]], %[[S9]] : f32
+// CHECK: %[[R12:.+]] = arith.addf %[[R11]], %[[CF]] : f32
+// CHECK: %[[R13:.+]] = vector.splat %[[R12]] : vector<1xf32>
// CHECK: %[[TID0:.+]] = arith.cmpi eq, %[[TID]], %[[C0]] : index
// CHECK: scf.if %[[TID0]] {
-// CHECK: vector.transfer_write %[[R9]], %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<1xf32>, memref<512xf32>
+// CHECK: vector.transfer_write %[[R13]], %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<1xf32>, memref<512xf32>
// CHECK: }
diff --git a/tests/e2e/regression/large_reduction.mlir b/tests/e2e/regression/large_reduction.mlir
index fde85c3..15e6fe4 100644
--- a/tests/e2e/regression/large_reduction.mlir
+++ b/tests/e2e/regression/large_reduction.mlir
@@ -46,6 +46,6 @@
%2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32
} -> tensor<2xf32>
- check.expect_eq_const(%result, dense<40.96> : tensor<2xf32>) : tensor<2xf32>
+ check.expect_almost_eq_const(%result, dense<40.96> : tensor<2xf32>) : tensor<2xf32>
return
}
