[Mlir-commits] [mlir] [mlir][ArmSME] More precisely model dataflow in ArmSME to SCF lowerings (PR #73922)

[Cullen Rhodes]

================
@@ -242,27 +247,25 @@ struct ConstantOpToArmSMELowering : public OpRewritePattern<arith::ConstantOp> {
 
     // Lower non-zero constants to a loop of 'arm_sme.move_vector_to_tile_slice'
     // ops that broadcast the constant to each tile slice.
-    OpBuilder::InsertionGuard g(rewriter);
     auto loc = constantOp.getLoc();
 
     // Unpack 1-d vector type from 2-d vector type.
-    auto tileSliceType =
-        VectorType::get(tileType.getShape().drop_front(), tileElementType,
-                        /*scalableDims=*/{true});
+    VectorType tileSliceType = VectorType::Builder(tileType).dropDim(0);
     auto denseAttr1D = DenseElementsAttr::get(
         tileSliceType, denseAttr.getSplatValue<Attribute>());
     auto constantOp1D = rewriter.create<arith::ConstantOp>(loc, denseAttr1D);
 
-    auto tile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
-
-    auto forOp = getLoopOverTileSlices(rewriter, loc, tileElementType);
-    auto tileSliceIndex = forOp.getInductionVar();
-
-    // Create 'arm_sme.move_vector_to_tile_slice' to write vector to tile slice.
-    rewriter.create<arm_sme::MoveVectorToTileSliceOp>(
-        loc, tileType, constantOp1D, tile, tileSliceIndex);
-
-    rewriter.replaceOp(constantOp, tile);
+    auto initTile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
+    auto forOp =
+        createLoopOverTileSlices(rewriter, loc, initTile, [&](auto forOp) {
+          auto tileSliceIndex = forOp.getInductionVar();
+          auto currentTile = forOp.getRegionIterArg(0);
+          // Create 'arm_sme.move_vector_to_tile_slice' to write vector to tile
+          // slice.
+          return rewriter.create<arm_sme::MoveVectorToTileSliceOp>(
+              loc, tileType, constantOp1D, currentTile, tileSliceIndex);
+        });
+    rewriter.replaceOp(constantOp, forOp.getResult(0));
----------------
c-rhodes wrote:

this is nice! Did clang-format format this? Indentation looks more than I'd expect

https://github.com/llvm/llvm-project/pull/73922