[mlir][ArmSME] Split lowering of arith.constant from vector.transfer_write

[c-rhodes]

Summary

An 'arith.constant dense<0>' is currently lowered to 'arm_sme.zero' as
part of the 'vector.transfer_write' lowering during '-vector-to-arm-sme'
conversion. This patch makes this lowering independent of the
'vector.transfer_write'. This can then be extended for further tile
types and non-zero constants.

[c-rhodes]

Update comment.

[awarzynski]

LG, ta!

mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp

if (llvm::isa<FloatType>(elemType))
     return val && val.isSplat() && val.getSplatValue<APFloat>().isZero();
   if (llvm::isa<IntegerType>(elemType))
     return val && val.isSplat() && val.getSplatValue<APInt>().isZero();
   return false;
 }
 
 namespace {
 
-/// Look at `vector.transfer_write` operations and convert suitable candidates
-/// to ArmSME operations, e.g.:
+/// Conversion pattern for vector.transfer_write. Currently only supports:
 ///
 ///    %cst = arith.constant dense<0> : vector<[16]x[16]xi8>
 ///    vector.transfer_write %cst, %arg0 : vector<[16]x[16]xi8>, memref<?x?xi8>
 ///
 /// is converted to:
 ///
 ///    %0 = arm_sme.zero : vector<[16]x[16]xi8>
 ///    arm_sme.tile_store %arg0[%c0, %c0], %0 : memref<?x?xi8>,
 ///    vector<[16]x[16]xi8>
 ///
+/// The conversion from arith.constant dense<0> to arm_sme.zero is done in
+/// ConstantOpToArmSMELowering.
 struct TransferWriteToArmSMELowering
     : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
                                 PatternRewriter &rewriter) const final {
     auto vType = writeOp.getVectorType();
     if (vType.getRank() != 2)
       return failure();
     if (vType.getShape() != ArrayRef<int64_t>({kMinNumElts, kMinNumElts}))
       return failure();
     if (vType.getElementType() != rewriter.getI8Type())
       return failure();
     if (vType.getScalableDims().size() != 2)
       return failure();
 
-    auto loc = writeOp.getLoc();
-
     if (!llvm::isa<MemRefType>(writeOp.getSource().getType()))
       return failure();
 
     auto constant = writeOp.getVector().getDefiningOp<arith::ConstantOp>();
     if (!constant)
       return failure();
 
     auto denseAttr = dyn_cast<DenseElementsAttr>(constant.getValueAttr());
     if (!denseAttr || !isSplatZero(vType.getElementType(), denseAttr))
       return failure();
 
-    auto zero = rewriter.create<arm_sme::ZeroOp>(loc, vType);
-
     rewriter.replaceOpWithNewOp<arm_sme::TileStoreOp>(
-        writeOp, zero, writeOp.getSource(), writeOp.getIndices());
+        writeOp, writeOp.getVector(), writeOp.getSource(),
+        writeOp.getIndices());
     return success();
   }
 };
 
 /// Conversion pattern for vector.load.
 struct VectorLoadToArmSMELowering : public OpRewritePattern<vector::LoadOp> {
   using OpRewritePattern<vector::LoadOp>::OpRewritePattern;
 
LogicalResult matchAndRewrite(vector::StoreOp store,
 
     rewriter.replaceOpWithNewOp<arm_sme::TileStoreOp>(
         store, store.getValueToStore(), store.getBase(), store.getIndices());
 
     return success();
   }
 };
 
+/// Conversion pattern for dense arith.constant.
+struct ConstantOpToArmSMELowering : public OpRewritePattern<arith::ConstantOp> {
+  using OpRewritePattern<arith::ConstantOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::ConstantOp constantOp,
+                                PatternRewriter &rewriter) const final {
+    auto vType = dyn_cast<VectorType>(constantOp.getType());
+    if (!vType)
+      return failure();
+    if (vType.getRank() != 2)
+      return failure();
+    if (vType.getShape() != ArrayRef<int64_t>({kMinNumElts, kMinNumElts}))
+      return failure();
+    if (vType.getElementType() != rewriter.getI8Type())
+      return failure();
+    if (vType.getScalableDims().size() != 2)
+      return failure();
+
+    auto denseAttr = dyn_cast<DenseElementsAttr>(constantOp.getValueAttr());
+    if (!denseAttr || !isSplatZero(vType.getElementType(), denseAttr))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<arm_sme::ZeroOp>(constantOp, vType);
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::populateVectorToArmSMEPatterns(RewritePatternSet &patterns,
                                           MLIRContext &ctx) {
   patterns.add<TransferWriteToArmSMELowering, VectorLoadToArmSMELowering,
-               VectorStoreToArmSMELowering>(&ctx);
+               VectorStoreToArmSMELowering, ConstantOpToArmSMELowering>(&ctx);
 }

mlir/test/Dialect/ArmSME/vector-ops-to-sme.mlir

-// RUN: mlir-opt %s -convert-vector-to-arm-sme -split-input-file | mlir-opt | FileCheck %s
-
+// RUN: mlir-opt %s -convert-vector-to-arm-sme -split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL:   func.func @transfer_write_2d_zero(
 // CHECK-SAME:      %[[ARG_0:.*]]: memref<?x?xi8>) {
 func.func @transfer_write_2d_zero(%arg0 : memref<?x?xi8>) {
 // CHECK:           %[[C_0:.*]] = arith.constant 0 : index
 // CHECK:           %[[ZERO:.*]] = arm_sme.zero : vector<[16]x[16]xi8>
 // CHECK:           arm_sme.tile_store %[[ZERO]], %[[ARG_0]][%[[C_0]], %[[C_0]]] : memref<?x?xi8>, vector<[16]x[16]xi8>
 // CHECK:           return
   %c0 = arith.constant 0 : index
   %cst = arith.constant dense<0> : vector<[16]x[16]xi8>
   vector.transfer_write %cst, %arg0[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[16]xi8>, memref<?x?xi8>
   return
 }
 
 // -----
 
+// CHECK-LABEL: @arith_constant_dense_2d_zero_i8
+// CHECK: %[[ZERO:.*]] = arm_sme.zero : vector<[16]x[16]xi8>
+func.func @arith_constant_dense_2d_zero_i8() {
+  %zero = arith.constant dense<0> : vector<[16]x[16]xi8>
+  "prevent.dce"(%zero) : (vector<[16]x[16]xi8>) -> ()
+  return
+}
+
+// -----
+
 // The following tests check the 'vector.transfer_write' -> 'arm_sme.intr.zero'
 // lowering only occurs for vector types of correct rank, shape, element size
 // and number of scalable dims.
 
 // CHECK-LABEL: @transfer_write_2d_zero__bad_type
 // CHECK: vector.transfer_write
 // CHECK-NOT: arm_sme.intr.zero
 func.func @transfer_write_2d_zero__bad_type(%arg0 : memref<?x?xi4>) {
func.func @transfer_write_2d_zero__non_memref_type(%arg0 : tensor<?x?xi8>) -> tensor<?x?xi8> {
   %0 = vector.transfer_write %cst, %arg0[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[16]xi8>, tensor<?x?xi8>
   return %0 : tensor<?x?xi8>
 }
 
 // -----
 
 // CHECK-LABEL: @transfer_write_2d_zero__non_zero_value
 // CHECK: vector.transfer_write
+// CHECK-NOT: arm_sme.zero
 // CHECK-NOT: arm_sme.tile_store
 func.func @transfer_write_2d_zero__non_zero_value(%arg0 : memref<?x?xi8>) {
   %c0 = arith.constant 0 : index
   %cst = arith.constant dense<1> : vector<[16]x[16]xi8>
   vector.transfer_write %cst, %arg0[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[16]xi8>, memref<?x?xi8>
   return
 }