[SPIRV] Unroll wide elementwise vector operations (#25053) I hit this while compiling a JAX training step to WebGPU through IREE in the browser. A `vector<32xf32>` addition survived until SPIR-V conversion and caused compilation to fail. This uses MLIR's existing unrolling patterns in the vector-breakdown pass to split wide elementwise operations into four-element operations. Signed-off-by: anutosh491 <andersonbhat491@gmail.com>

diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp
index c64f7c3..137b735 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp

@@ -144,10 +144,29 @@
 
 struct SPIRVBreakDownLargeVectorPass final
     : impl::SPIRVBreakDownLargeVectorPassBase<SPIRVBreakDownLargeVectorPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithDialect, vector::VectorDialect>();
+  }
+
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
     patterns.add<BreakDownCastExtractExtend>(context, /*benefits=*/10);
+    vector::populateVectorUnrollPatterns(
+        patterns, vector::UnrollVectorOptions().setNativeShapeFn(
+                      [](Operation *op) -> std::optional<SmallVector<int64_t>> {
+                        if (!OpTrait::hasElementwiseMappableTraits(op) ||
+                            op->getNumResults() != 1) {
+                          return std::nullopt;
+                        }
+                        auto type =
+                            dyn_cast<VectorType>(op->getResult(0).getType());
+                        if (!type || type.getRank() != 1 ||
+                            type.getDimSize(0) <= 4 || type.getDimSize(0) % 4) {
+                          return std::nullopt;
+                        }
+                        return SmallVector<int64_t>{4};
+                      }));
     // Convert vector.extract_strided_slice into a chain of vector.extract and
     // then a chain of vector.insert ops. This helps to cancel with previous
     // vector.insert/extract ops, especially for fP16 cases where we have

diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/break_down_large_vector.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/break_down_large_vector.mlir
index 958a2d0..67c5264 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/test/break_down_large_vector.mlir
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/test/break_down_large_vector.mlir

@@ -84,3 +84,62 @@
 //       CHECK:   %[[AND3:.+]] = arith.andi %[[SHR3]], %[[MASK]] : i32
 //       CHECK:   %[[RES:.+]] = vector.from_elements %[[AND0]], %[[AND1]], %[[AND2]], %[[AND3]] : vector<4xi32>
 //       CHECK:   return %[[RES]] : vector<4xi32>
+
+// -----
+
+func.func @bitcast_extract_extend_8_elements(%input: vector<4xi32>) -> vector<8xi32> {
+  %bitcast = vector.bitcast %input : vector<4xi32> to vector<32xi4>
+  %extract = vector.extract_strided_slice %bitcast offsets = [8], sizes = [8], strides = [1] : vector<32xi4> to vector<8xi4>
+  %extend = arith.extui %extract : vector<8xi4> to vector<8xi32>
+  return %extend : vector<8xi32>
+}
+
+// CHECK-LABEL: func @bitcast_extract_extend_8_elements
+//  CHECK-SAME: (%[[INPUT:.+]]: vector<4xi32>)
+//       CHECK: %[[BASE:.+]] = vector.extract %[[INPUT]][1] : i32 from vector<4xi32>
+// CHECK-COUNT-8: arith.andi {{.*}} : i32
+//       CHECK: %[[RES:.+]] = vector.from_elements {{.*}} : vector<8xi32>
+//       CHECK: return %[[RES]] : vector<8xi32>
+
+// -----
+
+// CHECK-LABEL: func @large_elementwise_add
+// CHECK-NOT: arith.addf {{.*}}vector<32xf32>
+// CHECK: arith.addf {{.*}}vector<4xf32>
+// CHECK-NOT: arith.addf {{.*}}vector<32xf32>
+func.func @large_elementwise_add(%lhs: vector<32xf32>, %rhs: vector<32xf32>) -> vector<4xf32> {
+  %sum = arith.addf %lhs, %rhs : vector<32xf32>
+  %part = vector.extract_strided_slice %sum offsets = [8], sizes = [4], strides = [1] : vector<32xf32> to vector<4xf32>
+  return %part : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @native_elementwise_add
+// CHECK: arith.addf {{.*}}vector<4xf32>
+func.func @native_elementwise_add(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> vector<4xf32> {
+  %sum = arith.addf %lhs, %rhs : vector<4xf32>
+  return %sum : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @elementwise_add_16xi8
+// CHECK-COUNT-4: arith.addi {{.*}} : vector<4xi8>
+// CHECK-NOT: arith.addi
+// CHECK: return {{.*}} : vector<16xi8>
+func.func @elementwise_add_16xi8(%lhs: vector<16xi8>, %rhs: vector<16xi8>) -> vector<16xi8> {
+  %sum = arith.addi %lhs, %rhs : vector<16xi8>
+  return %sum : vector<16xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func @elementwise_add_8xf16
+// CHECK-COUNT-2: arith.addf {{.*}} : vector<4xf16>
+// CHECK-NOT: arith.addf
+// CHECK: return {{.*}} : vector<8xf16>
+func.func @elementwise_add_8xf16(%lhs: vector<8xf16>, %rhs: vector<8xf16>) -> vector<8xf16> {
+  %sum = arith.addf %lhs, %rhs : vector<8xf16>
+  return %sum : vector<8xf16>
+}