Fixups

Author: MacDue

mlir/include/mlir/Dialect/ArmSVE/IR/ArmSVE.td

@@ -45,7 +45,7 @@ def SVEPredicate : ScalableVectorOfRankAndLengthAndType<
 // Generalizations of SVBool and SVEPredicate to ranks >= 1.
 // These are masks with a single trailing scalable dimension.
 def SVBoolMask : TrailingScalableVectorOfSizeAndType<[16], [I1]>;
-def SVEMask : TrailingScalableVectorOfSizeAndType<[16, 8, 4, 2, 1], [I1]>;
+def SVEPredicateMask : TrailingScalableVectorOfSizeAndType<[16, 8, 4, 2, 1], [I1]>;
 
 //===----------------------------------------------------------------------===//
 // ArmSVE op definitions
@@ -243,14 +243,13 @@ def UmmlaOp : ArmSVE_Op<"ummla",
     "$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($dst)";
 }
 
-
-class SvboolTypeContraint<string lhsArg, string rhsArg> : TypesMatchWith<
+class SvboolTypeConstraint<string lhsArg, string rhsArg> : TypesMatchWith<
       "expected corresponding svbool type widened to [16]xi1",
       lhsArg, rhsArg,
       "VectorType(VectorType::Builder(::llvm::cast<VectorType>($_self)).setDim(::llvm::cast<VectorType>($_self).getRank() - 1, 16))">;
 
 def ConvertFromSvboolOp : ArmSVE_Op<"convert_from_svbool",
-                            [Pure, SvboolTypeContraint<"result", "source">]>
+                            [Pure, SvboolTypeConstraint<"result", "source">]>
 {
   let summary = "Convert a svbool type to a SVE predicate type";
   let description = [{
@@ -260,45 +259,61 @@ def ConvertFromSvboolOp : ArmSVE_Op<"convert_from_svbool",
 
     Example 1: Convert a 1-D svbool mask to a SVE predicate.
     ```mlir
-    %svbool = vector.load %memref[%c0] : memref<?xi1>, vector<[16]xi1>
-    %mask = arm_sve.convert_from_svbool %svbool : vector<[4]xi1>
+    %source = vector.load %memref[%c0] : memref<?xi1>, vector<[16]xi1>
+    %result = arm_sve.convert_from_svbool %source : vector<[4]xi1>
     ```
 
     Example 2: Convert a 2-D svbool mask to a mask of SVE predicates.
     ```mlir
-    %svbool = vector.load %memref[%c0, %c0] : memref<2x?xi1>, vector<2x[16]xi1>
-    %mask = arm_sve.convert_from_svbool %svbool : vector<2x[8]xi1>
+    %source = vector.load %memref[%c0, %c0] : memref<2x?xi1>, vector<2x[16]xi1>
+    %result = arm_sve.convert_from_svbool %source : vector<2x[8]xi1>
     ```
+
+    ---
+
+    A `svbool` is the smallest SVE predicate type that has a in-memory
+    representation (and maps to a full predicate register). In MLIR `svbool` is
+    represented as `vector<[16]xi1>`. Smaller SVE predicate types
+    (`vector<[1|2|4|8]xi1>`) must be stored as `svbool` then converted back to
+    a predicate after loading.
   }];
   let arguments = (ins SVBoolMask:$source);
-  let results = (outs SVEMask:$result);
+  let results = (outs SVEPredicateMask:$result);
   let assemblyFormat = "$source attr-dict `:` type($result)";
 }
 
 def ConvertToSvboolOp : ArmSVE_Op<"convert_to_svbool",
-                            [Pure, SvboolTypeContraint<"source", "result">]>
+                            [Pure, SvboolTypeConstraint<"source", "result">]>
 {
-  let summary = "Convert a predicate type to a svbool type";
+  let summary = "Convert a SVE predicate type to a svbool type";
   let description = [{
     Converts SVE predicate types (or vectors of predicate types, e.g.
     `vector<4x[4]xi1>`) to svbool types. Note: Only the trailing dimension can
     be scalable.
 
     Example 1: Convert a 1-D SVE predicate to a svbool mask.
     ```mlir
-    %mask = vector.create_mask %dim_size : vector<[4]xi1>
-    %svbool = arm_sve.convert_to_svbool %mask : vector<[4]xi1>
+    %source = vector.create_mask %dim_size : vector<[4]xi1>
+    %result = arm_sve.convert_to_svbool %source : vector<[4]xi1>
     // => Results in vector<[16]xi1>
     ```
 
     Example 2: Convert a 2-D mask of SVE predicates to a svbool mask.
     ```mlir
-    %mask = vector.create_mask %c2, %dim_size : vector<2x[2]xi1>
-    %svbool = arm_sve.convert_to_svbool %mask : vector<2x[2]xi1>
+    %source = vector.create_mask %c2, %dim_size : vector<2x[2]xi1>
+    %result = arm_sve.convert_to_svbool %source : vector<2x[2]xi1>
     // => Results in vector<2x[16]xi1>
     ```
+
+    ---
+
+    A `svbool` is the smallest SVE predicate type that has a in-memory
+    representation (and maps to a full predicate register). In MLIR `svbool` is
+    represented as `vector<[16]xi1>`. Smaller SVE predicate types
+    (`vector<[1|2|4|8]xi1>`) must be converted to a `svbool` before they can be
+    stored.
   }];
-  let arguments = (ins SVEMask:$source);
+  let arguments = (ins SVEPredicateMask:$source);
   let results = (outs SVBoolMask:$result);
   let assemblyFormat = "$source attr-dict `:` type($source)";
 }

mlir/include/mlir/IR/CommonTypeConstraints.td

@@ -38,10 +38,17 @@ def IsScalableVectorTypePred : CPred<[{::llvm::isa<::mlir::VectorType>($_self) &
                                    ::llvm::cast<VectorType>($_self).isScalable()}]>;
 
 // Whether a type is a scalable VectorType, with a single trailing scalable dimension.
-def IsTrailingScalableVectorTypePred : And<[CPred<"::llvm::isa<::mlir::VectorType>($_self)">,
-                            CPred<"::llvm::cast<::mlir::VectorType>($_self).getRank() > 0">,
-                            CPred<"::llvm::cast<::mlir::VectorType>($_self).getScalableDims().back()">,
-                            CPred<"!llvm::is_contained(::llvm::cast<::mlir::VectorType>($_self).getScalableDims().drop_back(), true)">]>;
+// Examples:
+// Valid:
+//   - vector<[4]xf32>, vector<2x3x[2]xi64>, vector<32x[8]xi32>
+// Invalid
+//   - vector<[4]x8xi32>, vector<[2]x[2]xf64>, vector<2x[8]x4xi32>
+def IsOnlyTrailingDimScalablePred : And<[
+  CPred<"::llvm::isa<::mlir::VectorType>($_self)">,
+  CPred<"::llvm::cast<::mlir::VectorType>($_self).getRank() > 0">,
+  CPred<"::llvm::cast<::mlir::VectorType>($_self).getScalableDims().back()">,
+  CPred<"!llvm::is_contained(::llvm::cast<::mlir::VectorType>($_self).getScalableDims().drop_back(), true)">
+]>;
 
 // Whether a type is a VectorType and all dimensions are scalable.
 def allDimsScalableVectorTypePred : And<[
@@ -410,8 +417,10 @@ class ScalableVectorOf<list<Type> allowedTypes> :
   ShapedContainerType<allowedTypes, IsScalableVectorTypePred,
           "scalable vector", "::mlir::VectorType">;
 
+// Any vector with a single trailing scalable dimension, with an element type in
+// the `allowedTypes` list.
 class TrailingScalableVectorOf<list<Type> allowedTypes> :
-  ShapedContainerType<allowedTypes, IsTrailingScalableVectorTypePred,
+  ShapedContainerType<allowedTypes, IsOnlyTrailingDimScalablePred,
           "trailing scalable vector", "::mlir::VectorType">;
 
 // Whether the number of elements of a vector is from the given

mlir/lib/Dialect/ArmSVE/Transforms/LegalizeForLLVMExport.cpp

@@ -70,6 +70,25 @@ using ScalableMaskedDivFOpLowering =
 
 namespace {
 
+/// Unrolls a conversion to/from equivalent vector types, to allow using a
+/// conversion intrinsic that only supports 1-D vector types.
+///
+/// Example:
+/// ```
+/// %result = arm_sve.convert_to_svbool %source : vector<2x[4]xi1>
+/// ```
+/// is rewritten into:
+/// ```
+/// %cst = arith.constant dense<false> : vector<2x[16]xi1>
+/// %1 = vector.extract %source[0] : vector<[4]xi1> from vector<2x[4]xi1>
+/// %2 = "arm_sve.intr.convert.to.svbool"(%1)
+///                : (vector<[4]xi1>) -> vector<[16]xi1>
+/// %3 = vector.insert %2, %cst [0] : vector<[16]xi1> into vector<2x[16]xi1>
+/// %4 = vector.extract %source[1] : vector<[4]xi1> from vector<2x[4]xi1>
+/// %5 = "arm_sve.intr.convert.to.svbool"(%4)
+///                : (vector<[4]xi1>) -> vector<[16]xi1>
+/// %result = vector.insert %5, %3 [1] : vector<[16]xi1> into vector<2x[16]xi1>
+/// ```
 template <typename Op, typename IntrOp>
 struct SvboolConversionOpLowering : public ConvertOpToLLVMPattern<Op> {
   using ConvertOpToLLVMPattern<Op>::ConvertOpToLLVMPattern;
@@ -86,9 +105,13 @@ struct SvboolConversionOpLowering : public ConvertOpToLLVMPattern<Op> {
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resultType, rewriter.getZeroAttr(resultType));
 
+    // We want to iterate over the input vector in steps of the trailing
+    // dimension. So this creates tile shape where all leading dimensions are 1,
+    // and the trailing dimension step is the size of the dimension.
     SmallVector<int64_t> tileShape(sourceType.getRank(), 1);
     tileShape.back() = sourceType.getShape().back();
 
+    // Iterate over all scalable mask/predicate slices of the source vector.
     for (SmallVector<int64_t> index :
          StaticTileOffsetRange(sourceType.getShape(), tileShape)) {
       auto extractOrInsertPosition = ArrayRef(index).drop_back();

mlir/test/Dialect/ArmSVE/legalize-for-llvm.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -convert-vector-to-llvm="enable-arm-sve" -convert-func-to-llvm -reconcile-unrealized-casts -canonicalize -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -convert-vector-to-llvm="enable-arm-sve" -convert-func-to-llvm -reconcile-unrealized-casts -split-input-file %s | FileCheck %s
 
 func.func @arm_sve_sdot(%a: vector<[16]xi8>,
                    %b: vector<[16]xi8>,