[SPIRV] Add pattern for breaking down vector.bitcast (#14151)
The pattern added here is to handle bitcast ops that failed to unroll to
a supported vector size. For example, this can happen in cases like the
following after SPIRVVectorize
```mlir
%64 = arith.trunci %33 : vector<4xi32> to vector<4xi8>
%65 = vector.insert_strided_slice %64, %63 {offsets = [28], strides = [1]} : vector<4xi8> into vector<32xi8>
%66 = vector.extract_strided_slice %49 {offsets = [0], sizes = [16], strides = [1]} : vector<32xi8> to vector<16xi8>
%67 = vector.bitcast %66 : vector<16xi8> to vector<4xi32>
```
Where the intermediate `vector<32xi8>` can't be broken down without
first breaking down the bitcast.diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp
index fe5b880..f57d306 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVBreakDownLargeVector.cpp
@@ -27,6 +27,10 @@
patterns, [](vector::ExtractStridedSliceOp op) {
return op.getSourceVectorType().getNumElements() > 4;
});
+ vector::populateBreakDownVectorBitCastOpPatterns(
+ patterns, [](vector::BitCastOp op) {
+ return op.getSourceVectorType().getNumElements() > 4;
+ });
vector::InsertOp::getCanonicalizationPatterns(patterns, context);
vector::ExtractOp::getCanonicalizationPatterns(patterns, context);
if (failed(applyPatternsAndFoldGreedily(getOperation(),
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 abec761..4c7b43a 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
@@ -16,3 +16,18 @@
%0 = vector.extract_strided_slice %input {offsets = [1], sizes = [2], strides = [1]} : vector<4xf16> to vector<2xf16>
return %0: vector<2xf16>
}
+
+// -----
+
+// CHECK-LABEL: func @bitcast_16_elements
+func.func @bitcast_16_elements(%input: vector<16xi8>) -> vector<4xi32> {
+ // CHECK-DAG: %[[CST_I32:.*]] = arith.constant dense<0> : vector<4xi32>
+ // CHECK-DAG: arith.constant dense<0> : vector<4xi8>
+ // CHECK-COUNT-4: vector.extract
+ // CHECK-COUNT-4: vector.insert
+ // CHECK: vector.bitcast %{{.*}} : vector<4xi8> to vector<1xi32>
+ // CHECK: vector.insert_strided_slice {{.*}}, %[[CST_I32]]
+ // CHECK-COUNT-3: vector.bitcast
+ %0 = vector.bitcast %input : vector<16xi8> to vector<4xi32>
+ return %0 : vector<4xi32>
+}