Address comments

Author: c-rhodes

mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp

@@ -415,7 +415,7 @@ struct FoldExtractFromVectorOfSMELikeCreateMasks
   }
 };
 
-// Shuffles arith extend ops after vector.extract op.
+// Rewrites: vector.extract(arith.extend) -> arith.extend(vector.extract).
 //
 // This transforms IR like:
 //   %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
@@ -460,13 +460,13 @@ struct SwapVectorExtractOfArithExtend
     Operation *newExtend =
         rewriter.create(loc, extendOpName, Value(newExtract), resultType);
 
-    rewriter.replaceOp(extractOp, newExtend->getResult(0));
+    rewriter.replaceOp(extractOp, newExtend);
 
     return success();
   }
 };
 
-// Shuffles arith extend ops after vector.scalable.extract op.
+// Same as above, but for vector.scalable.extract.
 //
 // This transforms IR like:
 //   %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
@@ -499,16 +499,15 @@ struct SwapVectorScalableExtractOfArithExtend
 
     // Create new extract from source of extend.
     VectorType extractResultVectorType =
-        VectorType::Builder(resultType)
-            .setElementType(extendSourceVectorType.getElementType());
+        resultType.clone(extendSourceVectorType.getElementType());
     Value newExtract = rewriter.create<vector::ScalableExtractOp>(
         loc, extractResultVectorType, extendSource, extractOp.getPos());
 
     // Extend new extract to original result type.
     Operation *newExtend =
         rewriter.create(loc, extendOpName, Value(newExtract), resultType);
 
-    rewriter.replaceOp(extractOp, newExtend->getResult(0));
+    rewriter.replaceOp(extractOp, newExtend);
 
     return success();
   }

mlir/test/Dialect/ArmSME/vector-legalization.mlir

@@ -322,23 +322,23 @@ func.func @extract_from_arith_ext(%src: vector<4x[8]xi8>) -> vector<[8]xi32> {
 // CHECK-SAME:                                       %[[SRC:[a-z0-9]+]]: vector<4x[8]xi8>,
 // CHECK-SAME:                                       %[[DIM:[a-z0-9]+]]: index
 // CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][%[[DIM]]] : vector<[8]xi8> from vector<4x[8]xi8>
-// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: %[[EXTEND:.*]] = arith.extui %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
 // CHECK: return %[[EXTEND]]
 func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: index) -> vector<[8]xi32> {
-  %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+  %0 = arith.extui %src : vector<4x[8]xi8> to vector<4x[8]xi32>
   %1 = vector.extract %0[%dim] : vector<[8]xi32> from vector<4x[8]xi32>
   return %1 : vector<[8]xi32>
 }
 
 // -----
 
 // CHECK-LABEL: @scalable_extract_from_arith_ext(
-// CHECK-SAME:                                   %[[SRC:.*]]: vector<[8]xi8>
-// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xi8> from vector<[8]xi8>
-// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[4]xi8> to vector<[4]xi32>
+// CHECK-SAME:                                   %[[SRC:.*]]: vector<[8]xf16>
+// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xf16> from vector<[8]xf16>
+// CHECK: %[[EXTEND:.*]] = arith.extf %[[EXTRACT]] : vector<[4]xf16> to vector<[4]xf32>
 // CHECK: return %[[EXTEND]]
-func.func @scalable_extract_from_arith_ext(%src: vector<[8]xi8>) -> vector<[4]xi32> {
-  %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
-  %1 = vector.scalable.extract %0[0] : vector<[4]xi32> from vector<[8]xi32>
-  return %1 : vector<[4]xi32>
+func.func @scalable_extract_from_arith_ext(%src: vector<[8]xf16>) -> vector<[4]xf32> {
+  %0 = arith.extf %src : vector<[8]xf16> to vector<[8]xf32>
+  %1 = vector.scalable.extract %0[0] : vector<[4]xf32> from vector<[8]xf32>
+  return %1 : vector<[4]xf32>
 }