[mlir][vector] Add scalable vectors support to OuterProductOp

[WanderAway]

Summary

This will probably be the first in a series of patches that tries to enable code generation for ARM SME (extension of SVE).

Since SME's core operation is the outer product instruction, I figured that it would probably be a good idea to enable the outer product operation to properly accept and generate scalable vectors.

[kaitingwang]

This patch catches the corner case where RHS is non-scalable vector and LHS is scalable vector, throw an error. When RHS is scalable vector, but LHS is non-scalable vector, this case is supported and testcase shows it. This patch looks good to me.

[WanderAway]

@jsetoain @nicolasvasilache @aartbik @dcaballe
Some additional feedback would be appreciated, thanks.

[dcaballe]

Adding support for SME very well welcome! Thanks!
Just a couple of comments for now.

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

2661

Vector types in MLIR may have multiple dimensions so there could be multiple scalable dims (see getNumScalableDims). It would be great if you could add a test with multiple scalar dimensions and make sure that that works.

2671

I would expect this to be implemented (if it's not already?) with a trait similar to SameTypeOperands so that we can use it for all the vector operations that do not expect the same operand types but they do expect the same kind of vector. Perhaps something like SameVectorKindOperand?

[WanderAway]

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

2671

Technically it is not necessary for both operands for the outer product op to be scalable or fixed:

mlir
<[2]xf64> x <2xf64> -> <[2]x2xf64>
<2xf64> x <[2]xf64> -> <2x[2]xf64>

However (if I understand correctly), the current implementation of scalable vectors only support the latter, with the scalable dimension on the inner dimension of the vector.

This can get kind of ugly with data layout issues etc. so I'm just emitting an error here in case the first case gets supported later on...

[WanderAway]

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

2661

As specified on line 2358 in OuterProductOp::verify(), the input operands for a outer product op are expected to be 1d vectors (otherwise this would turn into sort of a gemm or gemv operation).

[WanderAway]

Updated the test case to check for invalid multi-dimensional input vectors.

[jsetoain]

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

2671

Indeed, only innermost dimensions are allowed to be scalable at the moment. It should not be too traumatic to extend it to allow per-dimension scalability but, also, I'm not sure that's going to be needed. Re.: type traits, we don't have any to indicate both have to be either scalable or fixed length, but requesting matching types should enforce it. vector<4xi32> and vector<[4]xi32> are two different types.

[WanderAway]

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

2671

The problem is that currently the OuterProductOp can take vectors of varying sizes as input
e.g. OuterProduct(<2xf32>, <3xf32>) => <2x3xf32>

Plus the operand can be a scalar as well
e.g. OuterProduct(f32, <3xf32>) => <3xf32>

It turns out that a constraint for these (in addition to the "scalability" check) is surprisingly hard to write in tablegen/ODS, and I think I will update the verify() code in C++ instead.

[WanderAway]

Added verification of scalable operands.

[kaitingwang]

This comment was removed by kaitingwang.

[WanderAway]

This comment was removed by WanderAway.

[WanderAway]

@dcaballe do you mind taking another look? Thanks

[dcaballe]

Sorry for the delay! I'm not sure why your ping didn't land on my inbox.
Ok, I understand the complexity. Let's go with this as is for now. Thanks!

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

ParseResult OuterProductOp::parse(OpAsmParser &parser, OperationState &result) {
   if (operandsInfo.size() < 2)
     return parser.emitError(parser.getNameLoc(),
                             "expected at least 2 operands");
   VectorType vLHS = tLHS.dyn_cast<VectorType>();
   VectorType vRHS = tRHS.dyn_cast<VectorType>();
   if (!vLHS)
     return parser.emitError(parser.getNameLoc(),
                             "expected vector type for operand #1");
-  VectorType resType =
-      vRHS ? VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
-                             vLHS.getElementType())
-           : VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType());
+
+  unsigned numScalableDims = vLHS.getNumScalableDims();
+  VectorType resType;
+  if (vRHS) {
+    numScalableDims += vRHS.getNumScalableDims();
+    resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
+                              vLHS.getElementType(), numScalableDims);
+  } else {
+    // Scalar RHS operand
+    resType = VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType(),
+                              numScalableDims);
+  }
 
   if (!result.attributes.get(OuterProductOp::getKindAttrStrName())) {
     result.attributes.append(
         OuterProductOp::getKindAttrStrName(),
         CombiningKindAttr::get(result.getContext(),
                                OuterProductOp::getDefaultKind()));
   }
 
if (vRHS) {
     if (vRHS.getRank() != 1)
       return emitOpError("expected 1-d vector for operand #2");
     if (vRES.getRank() != 2)
       return emitOpError("expected 2-d vector result");
     if (vLHS.getDimSize(0) != vRES.getDimSize(0))
       return emitOpError("expected #1 operand dim to match result dim #1");
     if (vRHS.getDimSize(0) != vRES.getDimSize(1))
       return emitOpError("expected #2 operand dim to match result dim #2");
+    if (vRHS.isScalable() != vLHS.isScalable())
+      return emitOpError("expected either all or none of vector operands #1 "
+                         "and #2 to be scalable");
   } else {
     // An AXPY operation.
     if (vRES.getRank() != 1)
       return emitOpError("expected 1-d vector result");
     if (vLHS.getDimSize(0) != vRES.getDimSize(0))
       return emitOpError("expected #1 operand dim to match result dim #1");
   }