[mlir][vector] Add cast op between scalable and fixed-length vectors

[jsetoain]

Summary

This patch adds the vector.scalable_cast operation that allows
casting fixed-length vectors to scalable vectors and vice-versa.

This functionality is needed in order to interface VLS (vector length
specific) code with scalable functions and intrinsics.

This change has been discussed in:

https://discourse.llvm.org/t/rfc-interfacing-between-fixed-length-and-scalable-vectors-for-vls-vector-code-on-scalable-vector-architectures/63074

[zhanghb97]

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

2061

According to the discussion, from the syntax and semantic point of view, maybe we should allow a small vector to be inserted into a large one.
Take the invalid test case in this patch as an example:

%0 = vector.scalable_cast %vin : vector<8xf32> to vector<[16]xf32>

Although the size of the scalable vector depends on the target hardware, the vector<8xf32> should be smaller than vector<[16]xf32> in most cases, so this seems safe to use. Should we classify the cases like this as legal?

[jsetoain]

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

2061

Indeed. But you can already insert smaller vectors into larger ones for fixed-length vectors, so I can't see a real reason or advantage to having the ability to go directly from small fixed to large scalable. If you are casting between fixed-length and scalable length, you already know the size of the physical underlaying vector; you can use that and then cast. There's an example of that in my last comment to that thread. If I'm missing something, please let me know in the thread :-)

[nicolasvasilache]

It is unclear to me why we want different semantics than LLVM.

vector.scalable_insert/extract

that work only on 1-D vectors seem less surprising.

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

let extraClassDeclaration = [{
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
+def Vector_ScalableCastOp :
+  Vector_Op<"scalable_cast", [NoSideEffect,
+    AllElementTypesMatch<["source", "result"]>,
+    PredOpTrait<"is a cast between scalable and fixed-length vectors", CPred<
+      "getSourceVectorType().isScalable() != getResultVectorType().isScalable()"
+      >>,
+    PredOpTrait<"fixed number of elements is multiple of scalable number of elements",
+      Or<[
+        And<[CPred<"getSourceVectorType().isScalable()">,
+             Neg<CPred<"getResultVectorType().getNumElements() % getSourceVectorType().getNumElements()">>]>,
+        And<[CPred<"getResultVectorType().isScalable()">,
+             Neg<CPred<"getSourceVectorType().getNumElements() % getResultVectorType().getNumElements()">>]>
+         ]>>
+  ]>,
+    Arguments<(ins VectorOfRank<[1]>:$source)>,
+    Results<(outs VectorOfRank<[1]>:$result)> {
+  let summary = "scalable_cast casts between scalable and fixed-length vectors";
+  let description = [{
+    The scalable_cast operation performs the cast between a scalable vector and
+    a fixed-length vector by assuming a vector scale.
+
+    This operation allows to interface VLS (Vector-Length Specific) code with
+    scalable vector functions and operations, such as the ones found in the
+    ArmSVE dialect.
+
+    The scalable_cast operation doesn't allow a cast between element types. If
+    needed, an element-type cast should be performed before or after the cast
+    between scalable and fixed-length types.
+
+    For the cast to be valid, the number of dimensions of the fixed-length
+    vector must be a multiple of the scalable vector.
+
+    Example:
+
+    ```mlir
+    // Casting from a scalable vector to a fixed-length vector assuming a
+    // vector scale of 2
+    %1 = vector.scalable_cast %0 : vector<[32]xf32> to vector<64xf32>
+
+    // Casting from a fixed-length vector to a scalable vector assuming a
+    // vector scale of 16
+    %3 = vector.scalable_cast %2 : vector<128xf32> to vector<[8]xf32>
+
+    // Casting from a fixed-length vector to a scalable vector assuming a
+    // vector scale of 1
+    %5 = vector.scalable_cast %4 : vector<8xbf16> to vector<[8]xbf16>
+
+    ```
+  }];
+  let extraClassDeclaration = [{
+    VectorType getSourceVectorType() {
+      return getSource().getType().cast<VectorType>();
+    }
+    VectorType getResultVectorType() {
+      return getResult().getType().cast<VectorType>();
+    }
+  }];
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
+  let hasFolder = 1;
+}
+
 def Vector_BitCastOp :
   Vector_Op<"bitcast", [NoSideEffect, AllRanksMatch<["source", "result"]>]>,
     Arguments<(ins AnyVectorOfAnyRank:$source)>,
     Results<(outs AnyVectorOfAnyRank:$result)>{
   let summary = "bitcast casts between vectors";
   let description = [{
     The bitcast operation casts between vectors of the same rank, the minor 1-D
     vector size is casted to a vector with a different element type but same

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) {
       Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC);
       desc = rewriter.create<InsertOp>(loc, fma, desc, i);
     }
     rewriter.replaceOp(op, desc);
     return success();
   }
 };
 
+class VectorScalableCastOpRewritePattern
+    : public OpRewritePattern<ScalableCastOp> {
+public:
+  using OpRewritePattern<ScalableCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ScalableCastOp op,
+                                PatternRewriter &rewriter) const override {
+    auto resultType = op.getResultVectorType();
+    if (resultType.isScalable()) {
+      // Fixed-length to scalable cast
+      Value insVec = rewriter.create<arith::ConstantOp>(
+          op.getLoc(), resultType, rewriter.getZeroAttr(resultType));
+      rewriter.replaceOpWithNewOp<LLVM::vector_insert>(
+          op, op.getSource(), insVec, rewriter.getI64IntegerAttr(0));
+    } else {
+      // Scalable to fixed-length cast
+      rewriter.replaceOpWithNewOp<LLVM::vector_extract>(
+          op, resultType, op.getSource(), rewriter.getI64IntegerAttr(0));
+    }
+    return success();
+  }
+};
+
 /// Returns the strides if the memory underlying `memRefType` has a contiguous
 /// static layout.
 static llvm::Optional<SmallVector<int64_t, 4>>
 computeContiguousStrides(MemRefType memRefType) {
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   if (failed(getStridesAndOffset(memRefType, strides, offset)))
     return None;

 
 } // namespace
 
 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns,
     bool reassociateFPReductions, bool force32BitVectorIndices) {
   MLIRContext *ctx = converter.getDialect()->getContext();
-  patterns.add<VectorFMAOpNDRewritePattern>(ctx);
+  patterns.add<VectorFMAOpNDRewritePattern, VectorScalableCastOpRewritePattern>(
+      ctx);
   populateVectorInsertExtractStridedSliceTransforms(patterns);
   patterns.add<VectorReductionOpConversion>(converter, reassociateFPReductions);
   patterns.add<VectorCreateMaskOpRewritePattern>(ctx, force32BitVectorIndices);
   patterns
       .add<VectorBitCastOpConversion, VectorShuffleOpConversion,
            VectorExtractElementOpConversion, VectorExtractOpConversion,
            VectorFMAOp1DConversion, VectorInsertElementOpConversion,
            VectorInsertOpConversion, VectorPrintOpConversion,

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 
 void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   // Pattern to rewrite a ShapeCastOp(ConstantOp) -> ConstantOp.
   results.add<ShapeCastConstantFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// VectorScalableCastOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ScalableCastOp::fold(ArrayRef<Attribute> operands) {
+  if (auto otherOp = getSource().getDefiningOp<ScalableCastOp>())
+    if (getResult().getType() == otherOp.getSource().getType())
+      return otherOp.getSource();
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // VectorBitCastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult BitCastOp::verify() {
   auto sourceVectorType = getSourceVectorType();
   auto resultVectorType = getResultVectorType();
 
   for (int64_t i = 0, e = sourceVectorType.getRank() - 1; i < e; i++) {

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: @scalable_cast_fixed_to_scalable
+// CHECK-SAME: %[[A:arg[0-9]+]]: vector<8xi32>
+func.func @scalable_cast_fixed_to_scalable(%a: vector<8xi32>)
+        -> vector<[4]xi32> {
+  // CHECK-NEXT: %[[C0:.*]] = arith.constant dense<0> : vector<[4]xi32>
+  // CHECK-NEXT: %[[R:[0-9]+]] = llvm.intr.experimental.vector.insert
+  // CHECK-SAME: %[[A]], %[[C0]][0] : vector<8xi32> into vector<[4]xi32>
+  // CHECK-NEXT: return %[[R]] : vector<[4]xi32>
+  %0 = vector.scalable_cast %a : vector<8xi32> to vector<[4]xi32>
+  return %0 : vector<[4]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @scalable_cast_scalable_to_fixed
+// CHECK-SAME: %[[A:arg[0-9]+]]: vector<[4]xf32>
+func.func @scalable_cast_scalable_to_fixed(%a: vector<[4]xf32>)
+        -> vector<8xf32> {
+  // CHECK-NEXT: %[[R:[0-9]+]] = llvm.intr.experimental.vector.extract
+  // CHECK-SAME: %[[A]][0] : vector<8xf32> from vector<[4]xf32>
+  // CHECK-NEXT: return %[[R]] : vector<8xf32>
+  %0 = vector.scalable_cast %a : vector<[4]xf32> to vector<8xf32>
+  return %0 : vector<8xf32>
+}

mlir/test/Dialect/Vector/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: func @dont_reduce_one_element_vector
 //       CHECK: vector.reduction
 func.func @dont_reduce_one_element_vector(%a : vector<4xf32>) -> f32 {
   %s = vector.reduction <add>, %a : vector<4xf32> into f32
   return %s : f32
 }
+
+// -----
+
+// CHECK-LABEL: func @scalable_cast_fold_complementary
+//  CHECK-SAME: (%[[V:.+]]: vector<8xf32>)
+//       CHECK-NEXT: return %[[V]] : vector<8xf32>
+func.func @scalable_cast_fold_complementary(%a : vector<8xf32>) -> vector<8xf32> {
+  %0 = vector.scalable_cast %a : vector<8xf32> to vector<[4]xf32>
+  %1 = vector.scalable_cast %0 : vector<[4]xf32> to vector<8xf32>
+  return %1 : 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 @scalable_cast_element_type_mismatch(%vin : vector<8xf32>) {
+  // expected-error@+1 {{op failed to verify that all of {source, result} have same element type}}
+  %0 = vector.scalable_cast %vin : vector<8xf32> to vector<[4]xbf16>
+  return
+}
+
+// -----
+
+func.func @scalable_cast_non_scalable_cast(%vin : vector<8xf32>) {
+  // expected-error@+1 {{op failed to verify that is a cast between scalable and fixed-length vector}}
+  %0 = vector.scalable_cast %vin : vector<8xf32> to vector<8xf32>
+  return
+}
+
+// -----
+
+func.func @scalable_cast_invalid_cast(%vin : vector<8xf32>) {
+  // expected-error@+1 {{op failed to verify that fixed number of elements is multiple of scalable number of elements}}
+  %0 = vector.scalable_cast %vin : vector<8xf32> to vector<[16]xf32>
+  return
+}

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: @vector_scalable_casts
+func.func @vector_scalable_casts(%0: vector<8xf32>,
+                                 %1: vector<[4]xf32>) {
+  // CHECK-NEXT: %{{.*}} = vector.scalable_cast %{{.*}}: vector<8xf32> to vector<[8]xf32>
+  %2 = vector.scalable_cast %0 : vector<8xf32> to vector<[8]xf32>
+  // CHECK-NEXT: %{{.*}} = vector.scalable_cast %{{.*}}: vector<8xf32> to vector<[4]xf32>
+  %3 = vector.scalable_cast %0 : vector<8xf32> to vector<[4]xf32>
+  // CHECK-NEXT: %{{.*}} = vector.scalable_cast %{{.*}}: vector<[4]xf32> to vector<4xf32>
+  %4 = vector.scalable_cast %1 : vector<[4]xf32> to vector<4xf32>
+  // CHECK-NEXT: %{{.*}} = vector.scalable_cast %{{.*}}: vector<[4]xf32> to vector<8xf32>
+  %5 = vector.scalable_cast %1 : vector<[4]xf32> to vector<8xf32>
+  return
+}