[mlir] Add pattern to handle trivial shape_cast in SPIR-V

[pzread]

Summary

Handle the trivial case of size-1 vector.shape_cast. The shape_cast is created when we drop the unit dims of 1 element vector.

[pzread]

Add more tests

[pzread]

Refactor

[kuhar]

LGTM but I'm not very familiar with how the unrealized conversion casts are handled in the pipeline, so you may want to wait for a second review before submitting.

[dcaballe]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

63–64

Shouldn't this case have been handled by a canonicalization/folding patterns before the lowering to SPIRV? e.g., by vector::populateVectorShapeCastLoweringPatterns?

[pzread]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

63–64

Currently IREE SPIRV pipeline doesn't have vector::populateVectorShapeCastLoweringPatterns to handle shape_cast. We can add that and it will also solve the problem.

I don't know which one is preferable. I'll let @antiagainst decide.

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

63–64

populateVectorShapeCastLoweringPatterns is more involved and can potentially generate quite a lot of data movement vector ops. It need more patterns to clean up everything for SPIR-V. It's hard to properly sequence it with all other vector patterns. So I'm fine to handle this special case directly when lowering to SPIR-V--it's pretty straightforward here.

(Some day we might need to take another look of how all those vector patterns are used in spirv https://github.com/openxla/iree/issues/11110; but that's a large effort to take on..)

66

We don't need this check anymore with dstType == srcType?

66

Actually ignore this. Misread it..

[pzread]

Oh, I didn't notice the windows build failed... But it seems to be due to an unrelated reason. I'll monitor the build status and revert if required.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

static int getNumBits(Type type) {
   // though in practice it will likely be stored as in a 4xi64 vector register.
   if (auto vectorType = dyn_cast<VectorType>(type))
     return vectorType.getNumElements() * vectorType.getElementTypeBitWidth();
   return type.getIntOrFloatBitWidth();
 }
 
 namespace {
 
+struct VectorShapeCast final : public OpConversionPattern<vector::ShapeCastOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::ShapeCastOp shapeCastOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type dstType = getTypeConverter()->convertType(shapeCastOp.getType());
+    if (!dstType)
+      return failure();
+
+    // If dstType is same as the source type or the vector size is 1, it can be
+    // directly replaced by the source.
+    if (dstType == adaptor.getSource().getType() ||
+        shapeCastOp.getResultVectorType().getNumElements() == 1) {
+      rewriter.replaceOp(shapeCastOp, adaptor.getSource());
+      return success();
+    }
+
+    // Lowering for size-n vectors when n > 1 hasn't been implemented.
+    return failure();
+  }
+};
+
 struct VectorBitcastConvert final
     : public OpConversionPattern<vector::BitCastOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::BitCastOp bitcastOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstType = getTypeConverter()->convertType(bitcastOp.getType());
spirv::CLFMaxOp, spirv::CLFMinOp, spirv::CLUMaxOp, spirv::CLUMinOp,          \
       spirv::CLSMaxOp, spirv::CLSMinOp
 
 #define GL_MAX_MIN_OPS                                                         \
   spirv::GLFMaxOp, spirv::GLFMinOp, spirv::GLUMaxOp, spirv::GLUMinOp,          \
       spirv::GLSMaxOp, spirv::GLSMinOp
 
 void mlir::populateVectorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                          RewritePatternSet &patterns) {
-  patterns.add<
-      VectorBitcastConvert, VectorBroadcastConvert,
+  patterns.add<VectorBitcastConvert, VectorBroadcastConvert,
                VectorExtractElementOpConvert, VectorExtractOpConvert,
-      VectorExtractStridedSliceOpConvert, VectorFmaOpConvert<spirv::GLFmaOp>,
+               VectorExtractStridedSliceOpConvert,
+               VectorFmaOpConvert<spirv::GLFmaOp>,
                VectorFmaOpConvert<spirv::CLFmaOp>, VectorInsertElementOpConvert,
                VectorInsertOpConvert, VectorReductionPattern<GL_MAX_MIN_OPS>,
-      VectorReductionPattern<CL_MAX_MIN_OPS>, VectorInsertStridedSliceOpConvert,
-      VectorShuffleOpConvert, VectorSplatPattern>(typeConverter,
-                                                  patterns.getContext());
+               VectorReductionPattern<CL_MAX_MIN_OPS>, VectorShapeCast,
+               VectorInsertStridedSliceOpConvert, VectorShuffleOpConvert,
+               VectorSplatPattern>(typeConverter, patterns.getContext());
 }
 
 void mlir::populateVectorReductionToSPIRVDotProductPatterns(
     RewritePatternSet &patterns) {
   patterns.add<VectorReductionToDotProd>(patterns.getContext());
 }

mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir

 //       CHECK:   %[[MIN0:.+]] = spirv.GL.UMin %[[S0]], %[[S1]]
 //       CHECK:   %[[MIN1:.+]] = spirv.GL.UMin %[[MIN0]], %[[S2]]
 //       CHECK:   %[[MIN2:.+]] = spirv.GL.UMin %[[MIN1]], %[[S]]
 //       CHECK:   return %[[MIN2]]
 func.func @reduction_minui(%v : vector<3xi32>, %s: i32) -> i32 {
   %reduce = vector.reduction <minui>, %v, %s : vector<3xi32> into i32
   return %reduce : i32
 }
+
+// -----
+
+// CHECK-LABEL: @shape_cast_same_type
+//  CHECK-SAME: (%[[ARG0:.*]]: vector<2xf32>)
+//       CHECK:   return %[[ARG0]]
+func.func @shape_cast_same_type(%arg0 : vector<2xf32>) -> vector<2xf32> {
+  %1 = vector.shape_cast %arg0 : vector<2xf32> to vector<2xf32>
+  return %arg0 : vector<2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @shape_cast_size1_vector
+//  CHECK-SAME: (%[[ARG0:.*]]: vector<f32>)
+//       CHECK:   %[[R0:.+]] = builtin.unrealized_conversion_cast %[[ARG0]] : vector<f32> to f32
+//       CHECK:   %[[R1:.+]] = builtin.unrealized_conversion_cast %[[R0]] : f32 to vector<1xf32>
+//       CHECK:   return %[[R1]]
+func.func @shape_cast_size1_vector(%arg0 : vector<f32>) -> vector<1xf32> {
+  %1 = vector.shape_cast %arg0 : vector<f32> to vector<1xf32>
+  return %1 : vector<1xf32>
+}