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
 }