[mlir][vector] Add vector.scalable.insert/extract ops

[jsetoain]

Summary

These new operations match the semantics of
llvm.experimental.vector.insert and llvm.experimental.vector.extract.

vector.scalable.insert and vector.scalable.extract allow,
respectively, insert vectors into scalable vectors, and extract vectors
from scalable vectors.

The discussion about the inclusion of these operations is here:
https://discourse.llvm.org/t/rfc-interfacing-between-fixed-length-and-scalable-vectors-for-vls-vector-code-on-scalable-vector-architectures

[jsetoain]

Remove comments

[jsetoain]

Changing the name and adjusting the semantics of insert/extract ops

[dcaballe]

Thanks, Javier! LGTM. Some minor comments.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

782–783

Not sure I'm getting this properly, how could a vector be rank-1 and n-D at the same time?

796

I'm assuming this will be aligned with the new LLVM doc, whatever that ends up being :)

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

749

We should bail out if the vector has multiple dimensions, right?

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

1835

Don't need to capture TMP?

[jsetoain]

Answer comments and rebase on main

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

782–783

Throughout the documentation, we use "rank" and "dimensionality" interchangeably, but that is incorrect. This: vector<3xf32> is a 3-D vector; this: vector<2x3x4xf32> is a rank-3 vector. I have thought about this wording quite a lot while I was writing it, because the meaning is nor respected everywhere (comments and documentation). To prevent confusion, I'd be happy to try to avoid this language entirely, but it is an issue in the documentation.

796

I wrote the documentation here in far more detail than LLVM's LangRef, but I'm saying the same thing. If you think I should just respect whatever LLVM LangRef says and avoid clarifications, I'm happy to simplify this :-)

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

749

This should be prevented by the operand constraints in the operation definition:

Arguments<(ins VectorOfRank<[1]>:$source,

ScalableVectorOfRank<[1]>:$dest,
I64Attr:$pos)>,

Or is there a reason to verify it inside the lowering as well?

[dcaballe]

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

782–783

If I'm not mistaken, in MLIR we would say that vector<3xf32> is a 1-D vector with shape or vector size = {3}. vector<2x3x4xf32> would be a 3-D vector (or vector with rank = 3) with shape or vector size = {2x3x4}. I would stick to that terminology for "consistency". Where did you see this being inconsistent?

796

It looks good! It's fine as long as they end up being aligned. I just mentioned that since the doc in LLVM is being rewritten.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

749

Oh, that should be good enough for now. Thanks!

[awarzynski]

@jsetoain thanks for working on this! As hinted on IREE's Discord, I'm starting to look at scalable vectors myself and this would really help us :)

I've left a few minor comments, but otherwise LGTM. I'll let somebody more experienced to make this formal though :)

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

782–783

I don't have a specific example, but agree with @jsetoain here - the nomenclature is a bit confusing and I wouldn't mind it being clarified. I can post a question on Discourse and propose something unless you have some other suggestion? Javier's documentation makes sense to me as is, though to me vectors have no rank 😂 . That's a property of tensors or matrices. But "rank-1" is intuitive enough. Ultimately, we should keep this consistent with what's used elsewhere in MLIR. I'm just not sure what that would be. @dcaballe, wdyt?

Some refs found online:

• TensorFlow docs

Note: The rank of a tensor is not the same as the rank of a matrix. The rank of a tensor is the number of indices required to uniquely select each element of the tensor. Rank is also known as "order", "degree", or "ndims."

• Wolfram Mathworld:

The rank R of a tensor is independent of the number of dimensions N of the underlying space.

796

[nit] I would add a note that this is meant to be kept in sync/consistent with the relevant ops from LangRef. WDYT?

mlir/test/Dialect/Vector/invalid.mlir

1620

How about a similar test when the source vector is scalable?

mlir/test/Dialect/Vector/ops.mlir

792

Could you add a test for non-zero starting index?

[awarzynski]

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

782–783

I brought this up on Discourse. That helped me realize that we are confusing the type and value domains. And that @dcaballe was right and I was wrong :)

So, we need to make sure that we are always referring to the indexing space here, so

• vector<3xf32> is a rank-1 vector
• vector<2x3x4xf32> is a rank-3 vector

Now, the "vector" dialect introduces n-D vectors. Using their terminology we would have:

• vector<3xf32> is a 1-D vector
• vector<2x3x4xf32> is a 3-D vector

@jsetoain

This: vector<3xf32> is a 3-D vector

In this statement you are referring to the value domain :)

I'd be happy to try to avoid this language entirely

If we accept the TF definition of a tensor/vector rank, then this document is starting to make sense to me. WDYT?

[dcaballe]

I thought this was already in! What is missing here? Only some doc fixes? Do you need any help with that, @jsetoain?

[jsetoain]

Address comments

mlir/test/Dialect/Vector/invalid.mlir

1620

I can certainly add that, but we'd be testing the exact same code. I don't have a preference one way or another, but if you think it's worth the extra tests, happy to oblige :-)

[jsetoain]

In D127875#3864124, @dcaballe wrote:

I thought this was already in! What is missing here? Only some doc fixes? Do you need any help with that, @jsetoain?

Thank you for offering, Diego, it's greatly appreciated :-) Apologies for the delay, this should get things back in motion!

[dcaballe]

Thanks, @jsetoain! LGTM. Please, wait for @awarzynski's approval/reply back.

[awarzynski]

LGTM, thanks for the updates!

[jsetoain]

Rebase on main

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

let extraClassDeclaration = [{
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
+def Vector_ScalableInsertOp :
+  Vector_Op<"scalable.insert", [NoSideEffect,
+       AllElementTypesMatch<["source", "dest"]>,
+       AllTypesMatch<["dest", "res"]>,
+       PredOpTrait<"position is a multiple of the source length.",
+        CPred<
+          "(getPos() % getSourceVectorType().getNumElements()) == 0"
+        >>]>,
+     Arguments<(ins VectorOfRank<[1]>:$source,
+                    ScalableVectorOfRank<[1]>:$dest,
+                    I64Attr:$pos)>,
+     Results<(outs ScalableVectorOfRank<[1]>:$res)> {
+  let summary = "insert subvector into scalable vector operation";
+  let description = [{
+    Takes rank-1 fixed-length or scalable source `n-D` subvector, a rank-1
+    destination `m-D` vector, a position `pos` within the destination vector.
+    If `source` is fixed-length, it inserts the elements `(0, ..., n-1)` of the
+    `n-D` subvector into the scalable vector, starting from the specified
+    position `pos`, and up to `(pos + n - 1)`. If source vector is scalable,
+    the insertion position will be scaled by the runtime vector scale
+    `vscale`, and the source elements `(0, ..., vscale * n - 1)` will be
+    inserted in the destination from position `(vscale * pos)`, and up to
+    `(vscale * pos + vscale * n - 1)`.
+
+    The insertion position must be a multiple of the minimum size of the source
+    vector. For the operation to be well defined, the source vector must fit in
+    the destination vector from the specified position. Since the destination
+    vector is scalable and its runtime length is unknown, it can't be verified
+    and/or guaranteed at compile time.
+
+    Example:
+
+    ```mlir
+    %2 = vector.scalable.insert %0, %1[8] : vector<4xf32> into vector<[16]xf32>
+    %5 = vector.scalable.insert %3, %4[0] : vector<8xf32> into vector<[4]xf32>
+    %8 = vector.scalable.insert %6, %7[0] : vector<[4]xf32> into vector<[8]xf32>
+    ```
+
+    Invalid example:
+    ```mlir
+    %2 = vector.scalable.insert %0, %1[5] : vector<4xf32> into vector<[16]xf32>
+    ```
+  }];
+
+  let assemblyFormat = [{
+    $source `,` $dest `[` $pos `]` attr-dict `:` type($source) `into` type($dest)
+  }];
+
+  let extraClassDeclaration = [{
+    VectorType getSourceVectorType() {
+      return getSource().getType().cast<VectorType>();
+    }
+    VectorType getDestVectorType() {
+      return getDest().getType().cast<VectorType>();
+    }
+  }];
+}
+
+def Vector_ScalableExtractOp :
+  Vector_Op<"scalable.extract", [NoSideEffect,
+       AllElementTypesMatch<["source", "res"]>,
+       PredOpTrait<"position is a multiple of the result length.",
+        CPred<
+          "(getPos() % getResultVectorType().getNumElements()) == 0"
+        >>]>,
+     Arguments<(ins ScalableVectorOfRank<[1]>:$source,
+                    I64Attr:$pos)>,
+     Results<(outs VectorOfRank<[1]>:$res)> {
+  let summary = "extract subvector from scalable vector operation";
+  let description = [{
+    Takes rank-1 `n-D` source vector and a position `pos` within the source
+    vector, and extracts the elements `(pos, ..., pos + m - 1)` from the `n-D`
+    source vector into a new rank-1 `m-D` vector.
+
+    The extraction position must be a multiple of the minimum size of the result
+    vector. For the operation to be well defined, the destination vector must
+    fit in the source vector from the specified position. Since the source
+    vector is scalable and its runtime length is unknown, it can't be verified
+    and/or guaranteed at compile time.
+
+    Example:
+
+    ```mlir
+    %1 = vector.scalable.extract %0[8] : vector<4xf32> from vector<[8]xf32>
+    %3 = vector.scalable.extract %2[0] : vector<[4]xf32> from vector<[8]xf32>
+    ```
+
+    Invalid example:
+    ```mlir
+    %1 = vector.scalable.extract %0[5] : vector<4xf32> from vector<[16]xf32>
+    ```
+  }];
+
+  let assemblyFormat = [{
+    $source `[` $pos `]` attr-dict `:` type($res) `from` type($source)
+  }];
+
+  let extraClassDeclaration = [{
+    VectorType getSourceVectorType() {
+      return getSource().getType().cast<VectorType>();
+    }
+    VectorType getResultVectorType() {
+      return getRes().getType().cast<VectorType>();
+    }
+  }];
+}
+
 def Vector_InsertMapOp :
   Vector_Op<"insert_map", [NoSideEffect, AllTypesMatch<["dest", "result"]>]>,
     Arguments<(ins AnyVector:$vector, AnyVector:$dest, Variadic<Index>:$ids)>,
     Results<(outs AnyVector:$result)> {
   let summary = "vector insert map operation";
   let description = [{
     Inserts a N-D vector and within a larger vector starting at id. The new
     vector created will have the same size as the destination operand vector.

mlir/include/mlir/IR/OpBase.td

 // `allowedRanks` list
 class IsVectorOfRankPred<list<int> allowedRanks> :
   And<[IsVectorTypePred,
        Or<!foreach(allowedlength, allowedRanks,
                    CPred<[{$_self.cast<::mlir::VectorType>().getRank()
                            == }]
                          # allowedlength>)>]>;
 
+// Whether the number of elements of a fixed-length vector is from the given
+// `allowedRanks` list
+class IsFixedVectorOfRankPred<list<int> allowedRanks> :
+  And<[IsFixedVectorTypePred,
+       Or<!foreach(allowedlength, allowedRanks,
+                   CPred<[{$_self.cast<::mlir::VectorType>().getRank()
+                           == }]
+                         # allowedlength>)>]>;
+
+// Whether the number of elements of a scalable vector is from the given
+// `allowedRanks` list
+class IsScalableVectorOfRankPred<list<int> allowedRanks> :
+  And<[IsScalableVectorTypePred,
+       Or<!foreach(allowedlength, allowedRanks,
+                   CPred<[{$_self.cast<::mlir::VectorType>().getRank()
+                           == }]
+                         # allowedlength>)>]>;
+
 // Any vector where the rank is from the given `allowedRanks` list
 class VectorOfRank<list<int> allowedRanks> : Type<
   IsVectorOfRankPred<allowedRanks>,
   " of ranks " # !interleave(allowedRanks, "/"), "::mlir::VectorType">;
 
+// Any fixed-length vector where the rank is from the given `allowedRanks` list
+class FixedVectorOfRank<list<int> allowedRanks> : Type<
+  IsFixedVectorOfRankPred<allowedRanks>,
+  " of ranks " # !interleave(allowedRanks, "/"), "::mlir::VectorType">;
+
+// Any scalable vector where the rank is from the given `allowedRanks` list
+class ScalableVectorOfRank<list<int> allowedRanks> : Type<
+  IsScalableVectorOfRankPred<allowedRanks>,
+  " of ranks " # !interleave(allowedRanks, "/"), "::mlir::VectorType">;
+
 // Any vector where the rank is from the given `allowedRanks` list and the type
 // is from the given `allowedTypes` list
 class VectorOfRankAndType<list<int> allowedRanks,
                           list<Type> allowedTypes> : Type<
   And<[VectorOf<allowedTypes>.predicate,
        VectorOfRank<allowedRanks>.predicate]>,
   VectorOf<allowedTypes>.summary # VectorOfRank<allowedRanks>.summary,
   "::mlir::VectorType">;

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

if (positionAttrs.size() > 1) {
           nMinusOnePositionAttrs);
     }
 
     rewriter.replaceOp(insertOp, inserted);
     return success();
   }
 };
 
+/// Lower vector.scalable.insert ops to LLVM vector.insert
+struct VectorScalableInsertOpLowering
+    : public ConvertOpToLLVMPattern<vector::ScalableInsertOp> {
+  using ConvertOpToLLVMPattern<
+      vector::ScalableInsertOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::ScalableInsertOp insOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<LLVM::vector_insert>(
+        insOp, adaptor.getSource(), adaptor.getDest(), adaptor.getPos());
+    return success();
+  }
+};
+
+/// Lower vector.scalable.extract ops to LLVM vector.extract
+struct VectorScalableExtractOpLowering
+    : public ConvertOpToLLVMPattern<vector::ScalableExtractOp> {
+  using ConvertOpToLLVMPattern<
+      vector::ScalableExtractOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::ScalableExtractOp extOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<LLVM::vector_extract>(
+        extOp, typeConverter->convertType(extOp.getResultVectorType()),
+        adaptor.getSource(), adaptor.getPos());
+    return success();
+  }
+};
+
 /// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1.
 ///
 /// Example:
 /// ```
 ///   %d = vector.fma %a, %b, %c : vector<2x4xf32>
 /// ```
 /// is rewritten into:
 /// ```
patterns
            VectorLoadStoreConversion<vector::LoadOp, vector::LoadOpAdaptor>,
            VectorLoadStoreConversion<vector::MaskedLoadOp,
                                      vector::MaskedLoadOpAdaptor>,
            VectorLoadStoreConversion<vector::StoreOp, vector::StoreOpAdaptor>,
            VectorLoadStoreConversion<vector::MaskedStoreOp,
                                      vector::MaskedStoreOpAdaptor>,
            VectorGatherOpConversion, VectorScatterOpConversion,
            VectorExpandLoadOpConversion, VectorCompressStoreOpConversion,
-           VectorSplatOpLowering, VectorSplatNdOpLowering>(converter);
+           VectorSplatOpLowering, VectorSplatNdOpLowering,
+           VectorScalableInsertOpLowering, VectorScalableExtractOpLowering>(
+          converter);
   // Transfer ops with rank > 1 are handled by VectorToSCF.
   populateVectorTransferLoweringPatterns(patterns, /*maxTransferRank=*/1);
 }
 
 void mlir::populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns) {
   patterns.add<VectorMatmulOpConversion>(converter);
   patterns.add<VectorFlatTransposeOpConversion>(converter);
 }

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

func.func @splat(%a: vector<4xf32>, %b: f32) -> vector<4xf32> {
   return %r : vector<4xf32>
 }
 // CHECK-NEXT: %[[UNDEF:[0-9]+]] = llvm.mlir.undef : vector<4xf32>
 // CHECK-NEXT: %[[ZERO:[0-9]+]] = llvm.mlir.constant(0 : i32) : i32
 // CHECK-NEXT: %[[V:[0-9]+]] = llvm.insertelement %[[ELT]], %[[UNDEF]][%[[ZERO]] : i32] : vector<4xf32>
 // CHECK-NEXT: %[[SPLAT:[0-9]+]] = llvm.shufflevector %[[V]], %[[UNDEF]] [0 : i32, 0 : i32, 0 : i32, 0 : i32]
 // CHECK-NEXT: %[[SCALE:[0-9]+]] = arith.mulf %[[A]], %[[SPLAT]] : vector<4xf32>
 // CHECK-NEXT: return %[[SCALE]] : vector<4xf32>
+
+// -----
+
+// CHECK-LABEL: @vector_scalable_insert
+// CHECK-SAME: %[[SUB:.*]]: vector<4xf32>, %[[SV:.*]]: vector<[4]xf32>
+func.func @vector_scalable_insert(%sub: vector<4xf32>, %dsv: vector<[4]xf32>) -> vector<[4]xf32> {
+  // CHECK-NEXT: %[[TMP:.*]] = llvm.intr.experimental.vector.insert %[[SUB]], %[[SV]][0] : vector<4xf32> into vector<[4]xf32>
+  %0 = vector.scalable.insert %sub, %dsv[0] : vector<4xf32> into vector<[4]xf32>
+  // CHECK-NEXT: llvm.intr.experimental.vector.insert %[[SUB]], %[[TMP]][4] : vector<4xf32> into vector<[4]xf32>
+  %1 = vector.scalable.insert %sub, %0[4] : vector<4xf32> into vector<[4]xf32>
+  return %1 : vector<[4]xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @vector_scalable_extract
+// CHECK-SAME: %[[VEC:.*]]: vector<[4]xf32>
+func.func @vector_scalable_extract(%vec: vector<[4]xf32>) -> vector<8xf32> {
+  // CHECK-NEXT: %[[TMP:.*]] = llvm.intr.experimental.vector.extract %[[VEC]][0] : vector<8xf32> from vector<[4]xf32>
+  %0 = vector.scalable.extract %vec[0] : vector<8xf32> from vector<[4]xf32>
+  return %0 : vector<8xf32>
+}

mlir/test/Dialect/Vector/invalid.mlir

 func.func @warp_mismatch_rank(%laneid: index) {
   // expected-error@+1 {{'vector.warp_execute_on_lane_0' op expected vector type for distributed operands.}}
   %2 = vector.warp_execute_on_lane_0(%laneid)[32] -> (i32) {
     %0 = arith.constant dense<2>: vector<128xi32>
     vector.yield %0 : vector<128xi32>
   }
   return
 }
+
+// -----
+
+func.func @vector_scalable_insert_unaligned(%subv: vector<4xi32>, %vec: vector<[16]xi32>) {
+  // expected-error@+1 {{op failed to verify that position is a multiple of the source length.}}
+  %0 = vector.scalable.insert %subv, %vec[2] : vector<4xi32> into vector<[16]xi32>
+}
+
+// -----
+
+func.func @vector_scalable_extract_unaligned(%vec: vector<[16]xf32>) {
+  // expected-error@+1 {{op failed to verify that position is a multiple of the result length.}}
+  %0 = vector.scalable.extract %vec[5] : vector<4xf32> from vector<[16]xf32>
+}

mlir/test/Dialect/Vector/ops.mlir

^bb0(%arg0 : vector<128xi32>) :
     %1 = arith.addi %arg0, %0 : vector<128xi32>
 //       CHECK:       vector.yield %{{.*}} : vector<128xi32>
     vector.yield %1 : vector<128xi32>
 //  CHECK-NEXT:     }
   }
   return %2 : vector<4xi32>
 }
 
+// CHECK-LABEL: func @vector_scalable_insert(
+// CHECK-SAME: %[[SUB0:.*]]: vector<4xi32>, %[[SUB1:.*]]: vector<8xi32>,
+// CHECK-SAME: %[[SUB2:.*]]: vector<[4]xi32>, %[[SV:.*]]: vector<[8]xi32>
+func.func @vector_scalable_insert(%sub0: vector<4xi32>, %sub1: vector<8xi32>,
+                                 %sub2: vector<[4]xi32>, %sv: vector<[8]xi32>) {
+  // CHECK-NEXT: vector.scalable.insert %[[SUB0]], %[[SV]][12] : vector<4xi32> into vector<[8]xi32>
+  %0 = vector.scalable.insert %sub0, %sv[12] : vector<4xi32> into vector<[8]xi32>
+  // CHECK-NEXT: vector.scalable.insert %[[SUB1]], %[[SV]][0] : vector<8xi32> into vector<[8]xi32>
+  %1 = vector.scalable.insert %sub1, %sv[0] : vector<8xi32> into vector<[8]xi32>
+  // CHECK-NEXT: vector.scalable.insert %[[SUB2]], %[[SV]][0] : vector<[4]xi32> into vector<[8]xi32>
+  %2 = vector.scalable.insert %sub2, %sv[0] : vector<[4]xi32> into vector<[8]xi32>
+  return
+ }
 
+// CHECK-LABEL: func @vector_scalable_extract(
+// CHECK-SAME: %[[SV:.*]]: vector<[8]xi32>
+func.func @vector_scalable_extract(%sv: vector<[8]xi32>) {
+  // CHECK-NEXT: vector.scalable.extract %[[SV]][0] : vector<16xi32> from vector<[8]xi32>
+  %0 = vector.scalable.extract %sv[0] : vector<16xi32> from vector<[8]xi32>
+  // CHECK-NEXT: vector.scalable.extract %[[SV]][0] : vector<[4]xi32> from vector<[8]xi32>
+  %1 = vector.scalable.extract %sv[0] : vector<[4]xi32> from vector<[8]xi32>
+  return
+ }