Add vector.scan op

[harsh]

Summary

This patch adds the vector.scan op which computes the
scan for a given n-d vector. It requires specifying the operator,
the identity element and whether the scan is inclusive or
exclusive.

TEST: Added test in ops.mlir

[aartbik]

Some fly by comments. Can you say a bit more on the motivation for adding this op?
Will you also provide a lowering to LLVM IR?

mlir/lib/Dialect/Vector/VectorOps.cpp

4194

since you have a custom verifier, this revision also needs an invalid.mlir case that exercises all custom verification errors

4194

period at end

[harsh]

Hi Aart,

Thanks for taking a look at this patch. Let me provide some motivation for this op. I have been working closely with the IREE folks on try to implement a scan operation that can run on the
cpu and gpu. Since scan doesn't fit into the linalg structured op framework, we decided to prototype this in a linalg_ext dialect. This dialect currently resides in iree but the plan is
to upstream this to mlir-core. Here is the op definition https://github.com/google/iree/blob/main/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td#L250
though I plan to update that and this patch based on our recent findings. Having defined this op, we needed to find a way to lower this to the gpu dialect and thought that the jump
from the linalg_ext abstraction to the gpu abstraction would be too large. Thus, we decided on creating a vector.scan op which would provide a more progressive lowering. The plan is to
first add a lowering to the gpu dialect (where the scan would be lowered to gpu.shuffle ops) that is performant and then provide a cpu lowering. Would love to hear any thoughts
you have on this and any ideas on how we could do this better.

Thanks,
Harsh

[aartbik]

In D117171#3245189, @harsh wrote:

Having defined this op, we needed to find a way to lower this to the gpu dialect and thought that the jump
from the linalg_ext abstraction to the gpu abstraction would be too large. Thus, we decided on creating a vector.scan op which would provide a more progressive lowering. The plan is to
first add a lowering to the gpu dialect (where the scan would be lowered to gpu.shuffle ops) that is performant and then provide a cpu lowering. Would love to hear any thoughts
you have on this and any ideas on how we could do this better.

For most vector ops, we have an actual lowering to LLVM IR, just so we can run integration test on the new functionality. This also acts as extra "documentation" on what the op does.
This first lowering does not need to be super efficient, but if you see chance to add the lowering to this or a follow up CL, that would make the intended use of this op more clear to me.

[harsh]

In D117171#3252772, @aartbik wrote:

In D117171#3245189, @harsh wrote:

Having defined this op, we needed to find a way to lower this to the gpu dialect and thought that the jump
from the linalg_ext abstraction to the gpu abstraction would be too large. Thus, we decided on creating a vector.scan op which would provide a more progressive lowering. The plan is to
first add a lowering to the gpu dialect (where the scan would be lowered to gpu.shuffle ops) that is performant and then provide a cpu lowering. Would love to hear any thoughts
you have on this and any ideas on how we could do this better.

For most vector ops, we have an actual lowering to LLVM IR, just so we can run integration test on the new functionality. This also acts as extra "documentation" on what the op does.
This first lowering does not need to be super efficient, but if you see chance to add the lowering to this or a follow up CL, that would make the intended use of this op more clear to me.

Sounds good. I will add an initial lowering to LLVM IR / Standard Dialect that should showcase the use of this op in the subsequent patch. Thanks!

mlir/lib/Dialect/Vector/VectorOps.cpp

4194

thanks will add an invalid.mlir

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

4194

is this still TBD?

[harsh]

Added lowering to arith ops, tests and integration tests.

[harsh]

@aartbik - based on discussions with @ThomasRaoux , we felt it made more sense to do a vector to vector conversion instead of going directly to llvm. I have added the patch here along with some tests. Please take a look and let me know what you think. Thanks!

mlir/lib/Dialect/Vector/VectorOps.cpp

4194

I have just uploaded a new patch. Thanks!

[ThomasRaoux]

Thanks for adding an integration test! Few more comments.

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

2434

vector source and result need to have the exact same type?

2456

source and destination types should be the same? It would be good to have them only once and add it as restriction to make the op look simpler.

2457

The IR doesn't look correct, should it be %1#1?

mlir/lib/Dialect/Vector/VectorOps.cpp

4196

you should be able to specify that in tablegen? (related to the comments there)

4204

has to be greater or equal than?

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2382–2421 ↗

(On Diff #403523)

Move this function outside of the struct and add some documentation comment on top.

2423 ↗

(On Diff #403523)

Value can be nullptr so you don't need the Optional here.

Also as above move this function outside of the struct.

2453 ↗

(On Diff #403523)

nit, use: elType.isIntOrIndex()

2468 ↗

(On Diff #403523)

you can just initialize the offsets directly:
SmallVector<int64_t> offsets(dstRank, 0);
Same for sizes.

2505–2506 ↗

(On Diff #403523)

We should check precondition earlier not after generating a bunch of IR.

2525 ↗

(On Diff #403523)

remove

[harsh]

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

2434

Yes thanks! I changed it to use TCresIsSameAsOpBase

2456

I added the constraint but couldn't make the assemblyFormat of the op simpler. If I remove type($dest) I get

error: type of result #0, named 'dest', is not buildable and a buildable type cannot be inferred
$kind `,` $source `,` $initial_value attr-dict `:` `(` type($source) `,` type($initial_value) `)` `to` `(` `,` type($accumulated_value) `)`
^
/home/harsh/llvm-project/mlir/include/mlir/Dialect/Vector/VectorOps.td:2436:1: note: in custom assembly format for this operation
def Vector_ScanOp :
^
note: suggest adding a type constraint to the operation or adding a 'type($dest)' directive to the custom assembly format
$kind `,` $source `,` $initial_value attr-dict `:` `(` type($source) `,` type($initial_value) `)` `to` `(` `,` type($accumulated_value) `)`
^

which is surprising since I already added the type constraint. Any ideas on how to get around this?

2457

Yes, I removed the line.

mlir/lib/Dialect/Vector/VectorOps.cpp

4196

thanks yes I added it in the tablegen file.

4204

If the rank is 4, then doesn't that mean the reduction dimension can only be 0, 1, 2 or 3? So doesn't reduction_dim < rank ?

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2505–2506 ↗

(On Diff #403523)

So should I check if the kind matches early on and leave this check here? Or did you have something else in mind?

[harsh]

Updated based on Thomas' comments

[ThomasRaoux]

Looks good overall, I think you should be able to fix the syntax then it should be good to go.

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

2456

I think you want to use AllTypesMatch instead, here is an example:
https://github.com/llvm/llvm-project/blob/7d176844d05bf11b32c4379f1422379d3d979a1e/mlir/test/mlir-tblgen/op-result.td#L129

mlir/lib/Dialect/Vector/VectorOps.cpp

4204

I got mixed, my comment was about the "has to be" it should be "has to be less than"

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2357 ↗

(On Diff #403834)

This is a bit of a nit but could the two genOperator helper be merged into one and check the type based on Value x instead? This would be a bit less verbose.

2505–2506 ↗

(On Diff #403523)

I would have a precondition earlier and change this into an assert

[harsh]

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

2456

Thanks that what I was missing. Now the asm looks nice :)

mlir/lib/Dialect/Vector/VectorOps.cpp

4204

got it. I replaced "has to be <" with "has to be less than".

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2357 ↗

(On Diff #403834)

Yes, I have added this in the new patch.

2505–2506 ↗

(On Diff #403523)

sure sounds good

[harsh]

Updated based on second round of comments

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

let description = [{
     Right now we only support one contiguous set of scalable dimensions, all of
     them grouped and scaled with the value returned by 'vscale'.
   }];
   let results = (outs Index:$res);
   let assemblyFormat = "attr-dict";
   let verifier = ?;
 }
 
+//===----------------------------------------------------------------------===//
+// VectorScanOp
+//===----------------------------------------------------------------------===//
+
+def Vector_ScanOp :
+  Vector_Op<"scan", [NoSideEffect,
+     PredOpTrait<"source operand and result have same element type",
+                 TCresVTEtIsSameAsOpBase<0, 0>>]>,
+    Arguments<(ins Vector_CombiningKindAttr:$kind,
+                   AnyVector:$source,
+                   AnyType:$identity,
+                   I64ArrayAttr:$reduction_dims,
+                   BoolAttr:$inclusive)>,
+    Results<(outs AnyVector:$dest)> {
+  let summary = "Scan operation";
+  let description = [{
+    Performs an inclusive/exclusive scan on an n-D vector along a single
+    dimension converting it into an n-D vector using the given
+    operation (add/mul/min/max for int/fp and and/or/xor for
+    int only) and a specified value for the identity element. The identity
+    element is only used in the exclusive scan.
+
+    Example:
+
+    ```mlir
+    %1 = vector.scan "add", %0, %identity [1, 0] {inclusive = false} :
+      (vector<4x8x16x32xf32>, f32) to vector<4x8x16x32xf32>
+    %2 = vector.scan "add", %1, %identity [0] {inclusive = true} :
+      (vector<4x16xf32>, f32) to vector<4x16xf32>
+    ```
+  }];
+  let builders = [
+    OpBuilder<(ins "Value":$source, "Value":$identity,
+                   "CombiningKind":$kind,
+                   "ArrayAttr":$reduction_dims,
+                   CArg<"bool", "true">:$inclusive)>
+  ];
+  let extraClassDeclaration = [{
+    static StringRef getKindAttrName() { return "kind"; }
+    static StringRef getReductionDimAttrName() { return "reduction_dims"; }
+
+    VectorType getSourceType() {
+      return source().getType().cast<VectorType>();
+    }
+    VectorType getDestType() {
+      return dest().getType().cast<VectorType>();
+    }
+  }];
+  let assemblyFormat =
+    "$kind `,` $source `,` $identity $reduction_dims attr-dict `:` `(` type($source) `,` type($identity) `)` `to` type($dest)";
+}
+
 #endif // VECTOR_OPS

mlir/lib/Dialect/Vector/VectorOps.cpp

 
 } // namespace
 
 void CreateMaskOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<CreateMaskFolder>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// ScanOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(ScanOp op) {
+  VectorType srcType = op.getSourceType();
+  VectorType dstType = op.getDestType();
+  if (srcType != dstType)
+    return op.emitError("src and dst types must match");
+
+  // Check reduction dimensions < rank
+  int64_t rank = srcType.getRank();
+  for (auto i : op.reduction_dims().getAsRange<IntegerAttr>()) {
+    if (i.getInt() >= rank)
+      return op.emitOpError("reduction dimension ")
+             << i << " has to be < " << rank;
+  }
+
+  // Check identity element is integer or float
+  auto identityElementType = op.identity().getType();
+  if (!(identityElementType.isa<FloatType>() ||
+        identityElementType.isa<IntegerType>())) {
+    return op.emitOpError(
+        "expected identity element type to be float or integer");
+  }
+  return success();
+}
+
 void mlir::vector::populateVectorToVectorCanonicalizationPatterns(
     RewritePatternSet &patterns) {
   patterns
       .add<CreateMaskFolder, MaskedLoadFolder, MaskedStoreFolder, GatherFolder,
            ScatterFolder, ExpandLoadFolder, CompressStoreFolder,
            StridedSliceConstantMaskFolder, TransposeFolder>(
           patterns.getContext());
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Vector/VectorOps.cpp.inc"

mlir/test/Dialect/Vector/ops.mlir

 }
 
 // CHECK-LABEL: @get_vector_scale
 func @get_vector_scale() -> index {
   // CHECK: vector.vscale
   %0 = vector.vscale
   return %0 : index
 }
+
+// CHECK-LABEL: @vector_scan
+func @vector_scan(%0: vector<4x8x16x32xf32>) -> vector<4x8x16x32xf32> {
+  %1 = arith.constant 0.0 : f32
+  %2 = vector.scan <add>, %0, %1 [1, 2] {inclusive = true} :
+    (vector<4x8x16x32xf32>, f32) to vector<4x8x16x32xf32>
+  return %2 : vector<4x8x16x32xf32>
+}