[mlir][ArmSME] Rename slice move operations to insert/extract_tile_slice (#106755)

This renames:

- `arm_sme.move_tile_slice_to_vector` to `arm_sme.extract_tile_slice`
- `arm_sme.move_vector_to_tile_slice` to `arm_sme.insert_tile_slice`

The new names are more consistent with the rest of MLIR and should be
easier to understand. The current names (to me personally) are hard to
parse and easy to mix up when skimming through code.

Additionally, the syntax for `insert_tile_slice` has changed from:

```mlir
%4 = arm_sme.insert_tile_slice %0, %1, %2
  : vector<[16]xi8> into vector<[16]x[16]xi8>
```

To:

```mlir
%4 = arm_sme.insert_tile_slice %0, %1[%2]
  : vector<[16]xi8> into vector<[16]x[16]xi8>
```

This is for consistency with `extract_tile_slice`, but also helps with
readability as it makes it clear which operand is the index.

Author: MacDue

mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td

@@ -592,7 +592,7 @@ def StoreTileSliceOp : ArmSME_Op<"store_tile_slice", [
   }];
 }
 
-def MoveVectorToTileSliceOp : ArmSME_Op<"move_vector_to_tile_slice", [
+def InsertTileSliceOp : ArmSME_Op<"insert_tile_slice", [
     ArmSMETileOpInterface, Pure,
     AllTypesMatch<["tile", "result"]>,
     TypesMatchWith<
@@ -603,25 +603,25 @@ def MoveVectorToTileSliceOp : ArmSME_Op<"move_vector_to_tile_slice", [
         "::llvm::cast<mlir::VectorType>($_self).getElementType(),"
         "/*scalableDims=*/{true})">,
 ]> {
-  let summary = "Move 1-D scalable vector to slice of 2-D tile";
+  let summary = "Insert 1-D scalable vector into slice of 2-D tile";
   let description = [{
-    The vector to tile slice operation moves a 1-D scalable vector to a slice
-    of a 2-D scalable vector tile at the given index. The type of the 1-D
-    scalable vector to be moved must match the type of the tile slice. A tile
-    slice is a 1-D vector of horizontally or vertically contiguous elements
-    within a ZA tile. The updated tile is returned as the result.
+    Inserts a 1-D scalable vector to a slice of a 2-D scalable vector tile at
+    the given index. The type of the 1-D scalable vector to be inserted must
+    match the type of the tile slice. A tile slice is a 1-D vector of
+    horizontally or vertically contiguous elements within a ZA tile. The updated
+    tile is returned as the result.
 
     An optional tile slice layout attribute specifies whether the tile slice is
     horizontal (default) or vertical.
 
-    Example 1: Move a vector<[16]xi8> into tile horizontally (default) at given index.
+    Example 1: Insert `vector<[16]xi8>` into tile horizontally at the given index.
     ```mlir
-    %tile_update = arm_sme.move_vector_to_tile_slice %vector, %tile, %tile_slice_index : vector<[16]xi8> into vector<[16]x[16]xi8>
+    %tile_update = arm_sme.insert_tile_slice %vector, %tile[%tile_slice_index] : vector<[16]xi8> into vector<[16]x[16]xi8>
     ```
 
-    Example 2: Move a vector<[2]xf64> into tile vertically at given index.
+    Example 2: Insert `vector<[2]xf64>` into tile vertically at the given index.
     ```mlir
-    %tile_update = arm_sme.move_vector_to_tile_slice %vector, %tile, %tile_slice_index layout<vertical> : vector<[2]xf64> into vector<[2]x[2]xf64>
+    %tile_update = arm_sme.insert_tile_slice %vector, %tile[%tile_slice_index] layout<vertical> : vector<[2]xf64> into vector<[2]x[2]xf64>
     ```
   }];
   let arguments = (ins
@@ -636,35 +636,35 @@ def MoveVectorToTileSliceOp : ArmSME_Op<"move_vector_to_tile_slice", [
   }];
 
   let assemblyFormat = [{
-    $vector `,` $tile `,` $tile_slice_index (`layout` `` $layout^)?
+    $vector `,` $tile `[` $tile_slice_index `]` (`layout` `` $layout^)?
       attr-dict `:` type($vector) `into` type($result)
   }];
 }
 
-def MoveTileSliceToVectorOp : ArmSME_Op<"move_tile_slice_to_vector", [
+def ExtractTileSliceOp : ArmSME_Op<"extract_tile_slice", [
     ArmSMETileOpInterface, Pure,
     TypesMatchWith<
       "type of 'result' matches type of 'tile' slice",
       "tile", "result",
       "VectorType(VectorType::Builder(::llvm::cast<mlir::VectorType>($_self)).dropDim(0))">,
 ]> {
-  let summary = "Move slice of a 2-D tile to a 1-D scalable vector";
+  let summary = "Extract 1-D scalable vector from slice of 2-D tile";
   let description = [{
-    The tile slice to vector operation extracts a 1-D scalable slice from a 2-D
-    scalable tile at the given index. A tile slice is a 1-D vector of
-    horizontally or vertically contiguous elements within a ZA tile.
+    Extracts a 1-D scalable slice from a 2-D scalable tile at the given index.
+    A tile slice is a 1-D vector of horizontally or vertically contiguous
+    elements within a ZA tile.
 
     An optional tile slice layout attribute specifies whether the tile slice is
     horizontal (default) or vertical.
 
     Example 1: Extract `vector<[16]xi8>` from tile horizontally at the given index.
     ```mlir
-    %slice = arm_sme.move_tile_slice_to_vector %tile[%tile_slice_index] : vector<[16]xi8> from vector<[16]x[16]xi8>
+    %slice = arm_sme.extract_tile_slice %tile[%tile_slice_index] : vector<[16]xi8> from vector<[16]x[16]xi8>
     ```
 
     Example 2: Extract `vector<[2]xf64>` from tile vertically at the given index.
     ```mlir
-    %slice = arm_sme.move_tile_slice_to_vector %tile[%tile_slice_index] layout<vertical> : vector<[2]xf64> from vector<[2]x[2]xf64>
+    %slice = arm_sme.extract_tile_slice %tile[%tile_slice_index] layout<vertical> : vector<[2]xf64> from vector<[2]x[2]xf64>
     ```
   }];
 

mlir/lib/Conversion/ArithToArmSME/ArithToArmSME.cpp

@@ -64,7 +64,7 @@ struct ConstantOpToArmSMELowering : public OpRewritePattern<arith::ConstantOp> {
       return success();
     }
 
-    // Lower non-zero constants to a loop of 'arm_sme.move_vector_to_tile_slice'
+    // Lower non-zero constants to a loop of 'arm_sme.insert_tile_slice'
     // ops that broadcast the constant to each tile slice.
     auto loc = constantOp.getLoc();
 
@@ -79,9 +79,9 @@ struct ConstantOpToArmSMELowering : public OpRewritePattern<arith::ConstantOp> {
     auto initTile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
     auto makeLoopBody = [&](OpBuilder &b, Location loc, Value tileSliceIndex,
                             Value currentTile) {
-      // Create 'arm_sme.move_vector_to_tile_slice' to write vector to tile
+      // Create 'arm_sme.insert_tile_slice' to write vector to tile
       // slice.
-      auto nextTile = b.create<arm_sme::MoveVectorToTileSliceOp>(
+      auto nextTile = b.create<arm_sme::InsertTileSliceOp>(
           loc, tileType, constantOp1D, currentTile, tileSliceIndex);
       return nextTile.getResult();
     };

mlir/lib/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.cpp

@@ -575,23 +575,23 @@ struct StoreTileSliceConversion
   }
 };
 
-/// Lower `arm_sme.move_vector_to_tile_slice` to SME intrinsics.
-struct MoveVectorToTileSliceConversion
-    : public ConvertArmSMEOpToLLVMPattern<arm_sme::MoveVectorToTileSliceOp> {
+/// Lower `arm_sme.insert_tile_slice` to SME intrinsics.
+struct InsertTileSliceConversion
+    : public ConvertArmSMEOpToLLVMPattern<arm_sme::InsertTileSliceOp> {
   using ConvertArmSMEOpToLLVMPattern::ConvertArmSMEOpToLLVMPattern;
 
   LogicalResult
-  matchAndRewrite(arm_sme::MoveVectorToTileSliceOp moveVectorToTileSliceOp,
-                  arm_sme::MoveVectorToTileSliceOp::Adaptor adaptor,
+  matchAndRewrite(arm_sme::InsertTileSliceOp insertTileSliceOp,
+                  arm_sme::InsertTileSliceOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto loc = moveVectorToTileSliceOp.getLoc();
-    auto tileType = moveVectorToTileSliceOp.getTileType();
+    auto loc = insertTileSliceOp.getLoc();
+    auto tileType = insertTileSliceOp.getTileType();
 
-    auto tileId = getTileIdOrError(moveVectorToTileSliceOp);
+    auto tileId = getTileIdOrError(insertTileSliceOp);
     if (!tileId)
       return failure();
 
-    auto tileSlice = moveVectorToTileSliceOp.getTileSliceIndex();
+    auto tileSlice = insertTileSliceOp.getTileSliceIndex();
 
     // Cast tile slice from index to i32 for intrinsic.
     auto tileSliceI32 = rewriter.create<arith::IndexCastUIOp>(
@@ -606,42 +606,40 @@ struct MoveVectorToTileSliceConversion
     auto allActiveMask = rewriter.create<vector::SplatOp>(loc, predTy, one);
 
     // Create 'arm_sme.intr.write.(horiz|vert)' to write vector to tile slice.
-    switch (moveVectorToTileSliceOp.getLayout()) {
+    switch (insertTileSliceOp.getLayout()) {
     case arm_sme::TileSliceLayout::Horizontal:
       rewriter.create<arm_sme::aarch64_sme_write_horiz>(
           loc, tileId, tileSliceI32, allActiveMask,
-          moveVectorToTileSliceOp.getVector());
+          insertTileSliceOp.getVector());
       break;
     case arm_sme::TileSliceLayout::Vertical:
       rewriter.create<arm_sme::aarch64_sme_write_vert>(
           loc, tileId, tileSliceI32, allActiveMask,
-          moveVectorToTileSliceOp.getVector());
+          insertTileSliceOp.getVector());
       break;
     }
 
-    // Intrinsic has no result, replace 'arm_sme.move_vector_to_tile_slice' with
+    // Intrinsic has no result, replace 'arm_sme.insert_tile_slice' with
     // the input tile to preserve dataflow.
-    rewriter.replaceOp(moveVectorToTileSliceOp,
-                       moveVectorToTileSliceOp.getTile());
+    rewriter.replaceOp(insertTileSliceOp, insertTileSliceOp.getTile());
 
     return success();
   }
 };
 
-/// Lower `arm_sme.move_tile_slice_to_vector` to SME intrinsics.
-struct MoveTileSliceToVectorConversion
-    : public ConvertArmSMEOpToLLVMPattern<arm_sme::MoveTileSliceToVectorOp> {
+/// Lower `arm_sme.extract_tile_slice` to SME intrinsics.
+struct ExtractTileSliceConversion
+    : public ConvertArmSMEOpToLLVMPattern<arm_sme::ExtractTileSliceOp> {
   using ConvertArmSMEOpToLLVMPattern::ConvertArmSMEOpToLLVMPattern;
 
   LogicalResult
-  matchAndRewrite(arm_sme::MoveTileSliceToVectorOp moveTileSliceToVector,
-                  OpAdaptor,
+  matchAndRewrite(arm_sme::ExtractTileSliceOp extractTileSlice, OpAdaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto loc = moveTileSliceToVector.getLoc();
-    auto sliceType = moveTileSliceToVector.getSliceType();
-    auto sliceIndex = moveTileSliceToVector.getTileSliceIndex();
+    auto loc = extractTileSlice.getLoc();
+    auto sliceType = extractTileSlice.getSliceType();
+    auto sliceIndex = extractTileSlice.getTileSliceIndex();
 
-    auto tileId = getTileIdOrError(moveTileSliceToVector);
+    auto tileId = getTileIdOrError(extractTileSlice);
     if (!tileId)
       return failure();
 
@@ -659,16 +657,16 @@ struct MoveTileSliceToVectorConversion
         loc, rewriter.getI32Type(), sliceIndex);
 
     // Create 'arm_sme.intr.read.(horiz|vert)' to extract the tile slice.
-    switch (moveTileSliceToVector.getLayout()) {
+    switch (extractTileSlice.getLayout()) {
     case arm_sme::TileSliceLayout::Horizontal:
       rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_read_horiz>(
-          moveTileSliceToVector, sliceType, zeroVector, allTruePredicate,
-          tileId, sliceIndexI32);
+          extractTileSlice, sliceType, zeroVector, allTruePredicate, tileId,
+          sliceIndexI32);
       break;
     case arm_sme::TileSliceLayout::Vertical:
       rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_read_vert>(
-          moveTileSliceToVector, sliceType, zeroVector, allTruePredicate,
-          tileId, sliceIndexI32);
+          extractTileSlice, sliceType, zeroVector, allTruePredicate, tileId,
+          sliceIndexI32);
       break;
     }
 
@@ -985,8 +983,8 @@ void mlir::populateArmSMEToLLVMConversionPatterns(LLVMTypeConverter &converter,
   });
 
   addArmSMEConversionPatterns<
-      LoadTileSliceConversion, MoveTileSliceToVectorConversion,
-      MoveVectorToTileSliceConversion, StoreTileSliceConversion,
+      LoadTileSliceConversion, ExtractTileSliceConversion,
+      InsertTileSliceConversion, StoreTileSliceConversion,
       StreamingVLOpConversion, OuterProductOpConversion,
       OuterProductWideningOpConversion<arm_sme::FMopa2WayOp,
                                        arm_sme::aarch64_sme_mopa_wide>,

mlir/lib/Conversion/ArmSMEToSCF/ArmSMEToSCF.cpp

@@ -245,8 +245,8 @@ struct TileLoadOpConversion : public OpRewritePattern<arm_sme::TileLoadOp> {
 ///      : memref<?x?xi32>, vector<[4]xi1>,
 ///        vector<[4]xi32> into vector<[4]xi32>
 ///    // Insert slice into tile
-///    %tile_update = arm_sme.move_vector_to_tile_slice
-///      %slice, %iter_tile, %tile_slice_idx :
+///    %tile_update = arm_sme.insert_tile_slice
+///      %slice, %iter_tile[%tile_slice_idx] :
 ///      vector<[4]xi32> into vector<[4]x[4]xi32>
 ///    scf.yield %tile_update : vector<[4]x[4]xi32>
 ///  }
@@ -332,11 +332,11 @@ struct TileLoadOpWithMaskAndPadNonZeroConversion
         loc, tileSliceType, tileLoadOp.getBase(), memrefIndices, maskOp1D,
         /*passthru=*/pad1DOp);
 
-    // Create 'arm_sme.move_vector_to_tile_slice' to move slice into tile.
-    auto moveSlice = rewriter.create<arm_sme::MoveVectorToTileSliceOp>(
+    // Create 'arm_sme.insert_tile_slice' to insert slice into tile.
+    auto insertSlice = rewriter.create<arm_sme::InsertTileSliceOp>(
         loc, tileType, loadSlice->getResult(0), currentTile, tileSliceIndex,
         tileLoadOp.getLayout());
-    rewriter.create<scf::YieldOp>(loc, moveSlice.getResult());
+    rewriter.create<scf::YieldOp>(loc, insertSlice.getResult());
 
     rewriter.setInsertionPointAfter(forOp);