[mlir][tensor] Add gather/scatter op definitions to the tensor dialect.

[nicolasvasilache]

Summary

Gather/Scatter are examined from first principles in light of our recent progress on tensor-based codegen
and in-place bufferization.

In the future, lowering of these abstractions to operate inplace on buffers
will likely require a more powerful buffer representation than strided memref.

General context: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707
Relevant TL;DR parts of the proposal:

• gather: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-gatherop-and-friends-10
• need for more expressive types: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-bufferization-copy-view-and-the-need-for-more-expressive-types-12
• jagged buffer discussion: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-first-class-jagged-buffer-13

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Update.

[jpienaar]

Note this is failing with

tools/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.cpp.inc.d
error: expected literal, variable, directive, or optional group

$input `[` $indices `]` () $unique
                         ^

/var/lib/buildkite-agent/builds/llvm-project/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td:413:1: note: in custom assembly format for this operation
def Tensor_GatherOp : Tensor_Op<"gather", [
^

[jpienaar]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

413

Summary should be single line with no trailing punctuation (https://mlir.llvm.org/docs/OpDefinitions/#operation-documentation)

425

   ```mlir
``` ?

448

Coordinates being the input side?

465

New line before, tag with mlir again to ensure markdown generation does whats expected

1302

This sentence feels out of place.

1322

Not sure if I follow, wouldn't the result be consumed by the parent parallel op?

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1322

The parent ParallelCombiningOpInterface takes ops on tensors that do not return and combines them into the result, so there is no result to consume for this specific op
(more context in: https://discourse.llvm.org/t/rfc-parallel-abstraction-for-tensors-and-buffers/62607).

I can rephrase to make it more explicit, thanks!

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1322

Actually I can also take this out of scope for now: getting the op right in the sequential case is the first step.

[nicolasvasilache]

Address comments.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

425

thanks! also updated a few other places I had missed in the past.

448

rephrased

[bixia]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

434

to allow?

441

So 1 here is the size of the slice in dimension 1? What kind of syntax is this? The other two dimensions seem to be using numpy slice syntax. Would it be better to replace this with c:c+1 where c is the coordinate value in %input, and adjust the text description?

457–458

I think this is AoS not SoA, an array of coordinate tuples(structure).

The description doesn't mention the other dimensions beside the most minor dimension becomes part of the dimensions in the resulting tensor. It would be helpful to add this information.

473

"allow gather to view into the input data" doesn't make sense to me. Do you actually want to say:

Nested buffer support would allow us to show that the gather output is an array of 4-element slices?

1301

insert?

[nicolasvasilache]

Address comments and update op semantics according to: https://discourse.llvm.org/t/rfc-adding-gather-scatter-ops/64757

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

441

This was confusing indeed, I reworked this part, please lmk what you think.

Please also look at https://discourse.llvm.org/t/rfc-adding-gather-scatter-ops/64757.

457–458

You're right, thanks, updated.

I spelled out the types without dropping the 1s for now and tried to explain better, please let me know if this requires further clarifications.

473

I rephrased that part, please let me know if this still does not make sense.

[nicolasvasilache]

s/extract/insert

[nicolasvasilache]

Update scatter and add tests.

[nicolasvasilache]

More fixes + tests.

[nicolasvasilache]

Use AnySignlessIntegerOrIndex

[nicolasvasilache]

Drop spurious auto-introduced includes.

[ftynse]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

413

let summary, let description, etc. here and below

424

Please also specify that the trailing dimension of the index tensor contains the coordinates and it is expected to have its size equal to the number of dimensions being gathered.

425

Nit: otherwise _ is interpreted by markdown

511

Would it simplify deeper implementation to make this attribute mandatory and just have syntactic sugar that sets it to gather_dims(0, 1,2 ...) when it is missing from the textual form?

524

Is this expected to connect to InferTypeOpInterface?

1360–1362

If I am not misreading it, the RFC discussion concluded that uniqueness is always a requirement for scatter and that it won't support repeated indices at this point.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

566 ↗

(On Diff #456900)

(Small)DenseSet should be enough here, nothing relies on the order of insertion. More nitpicky, we the expected size of the array, it is likely more wasteful to create a parallel data structure than to just do binary search through the vector (which we enforced to be sorted in the verifier) several times.

578 ↗

(On Diff #456900)

Nit: can't we just append to resultShape after reserving resultShape() + sourceType.getRank() space in it?

588 ↗

(On Diff #456900)

I would plainly disallow empty gather_dims. It is confusing that gather_dims() is equivalent to gather_dims(0, 1, 2, ...) and may be interpreted as a (misnamed) list of dimensions not to gather.

2398–2409 ↗

(On Diff #456900)

This is common with gather, could we factor out the impl?

[nicolasvasilache]

Address comments.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

413

hmmm I manage to copy-paste the 1 bad op and reuse as a template .. fixed the other places, thanks for catching!

524

Ideally yes, but I am reluctant to embark on this path until some prototyping is available for memref<memref>. Can we do this in a followup?
Added a TODO in the meantime.

1360–1362

How would you phrase this differently ?
The presence of the unique attribute is enforced but we can't statically enforce true uniqueness.
UB seems the proper way to express this?

[ftynse]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1360–1362

UB is the proper way. I was thinking that we can just drop the attribute. Or is there some anticipation that it will become optional?

[nicolasvasilache]

Change gather/scatter to non-optional.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

What's the best way to achieve this these days, still use a custom parser/printer or can we just use a simpler embedded directive for this in the asm format?

1360–1362

ah yes ok, dropped the attribute, kept the syntax as being unambiguously clear about the semantics is important here.

[ftynse]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

For the syntax, I'd go with an embedded custom directive for this attribute specifically.

In any case, I was really asking a question as you may have a deeper prototype somewhere and it's hard for me to say whether this will help the implementation or just add more complexity in the syntax.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

Dropping optional is a good suggestion and will indeed simplify things.

Gave a shot to the custom thing but unless we change the syntax to something like:

input[%indices: indices_type] gather_dims(...)

I don't see a good way of writing it atm without implementing the whole parser (type must be resolved before I am able to know to what dim to automatically expand the no gather_dims).

I'll drop the implicit syntax and require explicit everything for now to remove the Optional, the syntax bikeshedding is separable.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

I also tried something like `DefaultValuedAttr<DenseI64ArrayAttr, "llvm::to_vector(llvm::seq<int64_t>(0, getIndices().getType().cast<RankedRensorType>().getIndicesType().getRank() - 1))">:$gather_dims```.
Unfortunately, the generated GatherOp::populateDefaultAttrs does not seem to receive enough information to be able to write such a default.

Maybe @Mogball will have more insight here.

[nicolasvasilache]

Drop implicit form of gather_dims and scatter_dims: all must be explicit atm.

[Mogball]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

You can't make the "default" value depend on another op attribute or property.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

Thanks for confirming @Mogball , would the feature make sense ?
I.e. non-constexpr default values are generally useful but maybe this is not useful enough to justify adding the support.

[Mogball]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

511

This can be achieved with a derived attribute that uses the value of an optional attribute or some derived value

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let builders = [
     OpBuilder<(ins "ValueRange":$elements)>
   ];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// GatherOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_GatherOp : Tensor_Op<"gather", [
+    NoSideEffect
+  ]> {
+  string summary = "gather a subset of a tensor at specified indices.";
+  string description = [{
+    The `gather` operation extracts a subset of the elements from an 
+    input tensor at the given indices.
+
+    In its most general form, the tensor of indices specifies all the coordinates
+    of every element to extract (i.e. COO format). The indices are expected to be
+    confined to coordinate values that fit the range of the input tensor, 
+    otherwise the behavior is undefined.
+
+    Example:
+
+    ```
+        // For each 1x2 triple of coordinates in %indices, extract the 
+        // element (i.e. 0-D subset) at the coordinates triple in %input.
+        // This corresponds to implicit COO coordinates = [0, 1, 2]
+        //
+        %out = gather %input[%indices] : 
+          tensor<4x4x4xf32>[tensor<1x2x3xindex>] -> tensor<1x2xf32>
+    ```
+
+    A slice variant is provided that allows specifying whole slices of the input
+    tensor.
+
+    Example:
+
+    ```
+        // For each 5x6 singleton of coordinates in %indices, extract the 2-D
+        // slice [:, 1, :] at the coordinates singleton in %input.
+        //
+        %out = gather %input[%indices] coordinates = [1] : 
+          tensor<4x4x4xf32>[tensor<5x6x1xindex>] -> tensor<5x6x4x4xf32>
+    ```
+
+    An optional `unique` unit attribute may be specified to indicate that the
+    coordinates are statically guaranteed to be unique at runtime. Incorrectly
+    setting the `unique` attribute when the coordinates are not truly unique is
+    undefined behavior.
+
+    Only full slices are meant to be supported by this op, if one desires 
+    partial slices one should compose this op with other tensor ops.
+    This is to avoid a slippery slope of complexity that would make the op 
+    unusable in practice.
+
+    At the tensor-level, the index tensor is specified in an SoA form (i.e. 
+    coordinate tuple is the most minor). It is the responsibility of further 
+    lowerings and bufferiation to implement various concrete layouts.
+
+    Note: As currently specified, the operation must lower to an abstraction that
+    performs copies to the output tensor. This is because the buffer type system
+    is currently not rich enough to allow multiple non-contiguous views in the 
+    same type. This is visible more clearly in a notional buffer version of the op:
+    ```
+        // memref<?x4xf32> is a contiguous buffer of ?x4 elements, gather from
+        // random input slices must copy to the contiguous output.
+        %out = gather %input[%indices] coordinates = [1] : 
+          memref<4x4xf32>[memref<?x1xindex>] -> memref<?x4xf32>
+
+        // Nested buffer support would allow gather to view into the input data.
+        %out = gather %input[%indices] coordinates = [1] : 
+          memref<4x4xf32>[memref<?x1xindex>] -> memref<? x memref<4xf32>>
+    ```
+  }];
+
+  let arguments = (ins AnyRankedTensor:$input, 
+                       AnyRankedTensor:$indices,
+                       OptionalAttr<IndexAttr>:$coordinates,
+                       UnitAttr:$unique);
+  let results = (outs AnyRankedTensor:$result);
+
+  let assemblyFormat = [{
+    $input `[` $indices `]` () $unique
+    `:` type($input) `[` type($indices) `]` `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // GenerateOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_GenerateOp : Tensor_Op<"generate",
     [RecursiveSideEffects,
      DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      SingleBlockImplicitTerminator<"mlir::tensor::YieldOp">]> {
   string summary = "Creates a dynamically sized tensor from elements";
   string description = [{
     This operation creates a dynamically sized tensor with elements of any type.
     It expects one index operand per dynamic extent of the result tensor.
 
     The body region defines the tensor's elements. It takes index operands as
     its region arguments that span the index space. The element at the given
     position is yielded with the `yield` operation (see `YieldOp`). There is
     no defined ordering to the invocations of the body. It is conceptually
     a "parallel map" operation.
 
     Example:
 
     ```mlir
       %tnsr = tensor.generate %m, %n {
       ^bb0(%i : index, %j : index, %k : index):
         ...
         yield %elem : f32
       } : tensor<?x3x?f32>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamicExtents);
   let results = (outs AnyRankedTensor:$result);
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat = "$dynamicExtents $body attr-dict `:` type($result)";
 
   let builders = [
     // Build op and populate its body per callback function.
     OpBuilder<(ins "Type":$resultTy, "ValueRange":$dynamicExtents,
       "function_ref<void(OpBuilder &, Location, ValueRange)>")>,
   ];
 
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
def Tensor_ParallelInsertSliceOp : Tensor_Op<"parallel_insert_slice", [
   ];
   
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// ParallelScatterOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_ParallelScatterOp : Tensor_Op<"parallel_scatter", [
+    NoSideEffect
+  ]> {
+  string summary = [{
+    Specify the tensor scatter update of a single thread of a parent 
+    ParallelCombiningOpInterface op, at specified indices.
+  }];
+  string description = [{
+    The `parallel_scatter` operation yields a subset tensor value to its parent 
+    ParallelCombiningOpInterface. These subset tensor values are aggregated to
+    in some unspecified order into a full tensor value returned by the parent 
+    parallel iterating op.
+    The `parallel_scatter` is one such op allowed in the 
+    ParallelCombiningOpInterface op.
+
+    Conflicting scatter indices result in undefined semantics, in that the indices
+    written to by multiple parallel updates might contain data from any of the 
+    updates, or even a malformed bit pattern.
+
+    In the future, a combinator region or symbol may be added to allow atomic RMW
+    semantics.
+
+    If an index is updated exactly once, the value contained at that index
+    in the resulting tensor will be equal to the value at a corresponding index
+    of a slice that was used for the updated. If an index is not updated at all,
+    its value will be equal to the one in the original tensor.
+
+    This op does not create a new value, which allows maintaining a clean
+    separation between the subset and full tensor.
+
+    Note that we cannot mark this operation as pure (NoSideEffects), even
+    though it has no side effects, because it will get DCEd during
+    canonicalization.
+
+    In its most general form, the tensor of indices specifies all the coordinates
+    of every element to extract (i.e. COO format). The indices are expected to be
+    confined to coordinate values that fit the range of the dest tensor, 
+    otherwise the behavior is undefined.
+
+    Example:
+
+    ```
+        // For each 1x2 triple of coordinates in %indices, insert the 
+        // element (i.e. 0-D subset) at the coordinates triple in %dest.
+        // This corresponds to implicit COO coordinates = [0, 1, 2]
+        //
+        parallel_scatter %input into %dest[%indices] : 
+          tensor<1x2xf32> into tensor<4x4x4xf32>[tensor<1x2x3xindex>] 
+            -> tensor<4x4x4xf32>
+    ```
+
+    A slice variant is provided that allows specifying whole tensor slices.
+
+    Example:
+
+    ```
+        // For each 5x6 singleton of coordinates in %indices, insert the 2-D
+        // slice [:, 1, :] at the coordinates singleton into %dest.
+        //
+        parallel_scatter %input into %dest[%indices] coordinates = [1] : 
+          tensor<5x6x4x4xf32> into tensor<4x4x4xf32>[tensor<5x6x1xindex>] 
+            -> tensor<4x4x4xf32>
+    ```
+
+    Only full slices are meant to be supported by this op, if one desires 
+    partial slices one should compose this op with other tensor ops.
+    This is to avoid a slippery slope of complexity that would make the op 
+    unusable in practice.
+
+    At the tensor-level, the index tensor is specified in an SoA form (i.e. 
+    coordinate tuple is the most minor). It is the responsibility of further 
+    lowerings and bufferiation to implement various concrete layouts.
+
+    After buffer allocation, the "parallel_scatter" op is expected to lower
+    into a yet-undefined buffer view op that supports nested buffer types (see
+    the discussion about copy vs view in the documentation of GatherOp and 
+    ScatterOp: the copy behavior is deemed prohibitive for the `parallel_scatter`
+    op).
+  }];
+
+  let arguments = (ins AnyRankedTensor:$input, 
+                       AnyRankedTensor:$indices,
+                       OptionalAttr<IndexAttr>:$coordinates,
+                       UnitAttr:$unique);
+  let results = (outs AnyRankedTensor:$result);
+
+  let assemblyFormat = [{
+    $input `into` $dest `[` $indices `]` () $unique
+    `:` type($input) `into` type($dest) `[` type($indices) `]` `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ScatterOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_ScatterOp : Tensor_Op<"scatter", [
+    NoSideEffect
+  ]> {
+  string summary = "scatter a subset of a tensor at specified indices.";
+  string description = [{
+    The `scatter` operation inserts a subset of the elements from an 
+    input tensor into a destination tensor at the given indices.
+
+    In its most general form, the tensor of indices specifies all the coordinates
+    of every element to extract (i.e. COO format). The indices are expected to be
+    confined to coordinate values that fit the range of the dest tensor, 
+    otherwise the behavior is undefined.
+
+    Example:
+
+    ```
+        // For each 1x2 triple of coordinates in %indices, insert the 
+        // element (i.e. 0-D subset) at the coordinates triple in %dest.
+        // This corresponds to implicit COO coordinates = [0, 1, 2]
+        //
+        %out = scatter %input into %dest[%indices] : 
+          tensor<1x2xf32> into tensor<4x4x4xf32>[tensor<1x2x3xindex>] 
+            -> tensor<4x4x4xf32>
+    ```
+
+    A slice variant is provided that allows specifying whole tensor slices.
+
+    Example:
+
+    ```
+        // For each 5x6 singleton of coordinates in %indices, insert the 2-D
+        // slice [:, 1, :] at the coordinates singleton into %dest.
+        //
+        %out = scatter %input into %dest[%indices] coordinates = [1] : 
+          tensor<5x6x4x4xf32> into tensor<4x4x4xf32>[tensor<5x6x1xindex>] 
+            -> tensor<4x4x4xf32>
+    ```
+
+    An optional `unique` unit attribute may be specified to indicate that the
+    coordinates are statically guaranteed to be unique at runtime. Incorrectly
+    setting the `unique` attribute when the coordinates are not truly unique is
+    undefined behavior.
+
+    Only full slices are meant to be supported by this op, if one desires 
+    partial slices one should compose this op with other tensor ops.
+    This is to avoid a slippery slope of complexity that would make the op 
+    unusable in practice.
+
+    At the tensor-level, the index tensor is specified in an SoA form (i.e. 
+    coordinate tuple is the most minor). It is the responsibility of further 
+    lowerings and bufferiation to implement various concrete layouts.
+
+    Note: As currently specified, the operation must lower to an abstraction that
+    performs copies to the output tensor. This is because the buffer type system
+    is currently not rich enough to allow multiple non-contiguous views in the 
+    same type. This is visible more clearly in a notional buffer version of the op:
+    ```
+        // memref<?x4xf32> is a contiguous buffer of ?x4 elements, scatter into
+        // random input slices must copy to the contiguous dest.
+        some_side_effecting_op_writing_into %input, ...: memref<4x4xf32>
+        scatter %input into %dest[%indices] coordinates = [1] : 
+          memref<4x4xf32> into memref<?x4xf32>[memref<?x1xindex>]
+
+        // Nested buffer support in the producing op would allow writing directly
+        // into the dest buffer.
+        %v = some_nested_buffer_view_op %dest[%indices] coordinates = [1] : 
+          memref<? x memref<4xf32>>
+        some_side_effecting_op_writing_into %v, ...: memref<? x memref<4xf32>>
+    ```
+  }];
+
+  let arguments = (ins AnyRankedTensor:$input, 
+                       AnyRankedTensor:$indices,
+                       OptionalAttr<IndexAttr>:$coordinates,
+                       UnitAttr:$unique);
+  let results = (outs AnyRankedTensor:$result);
+
+  let assemblyFormat = [{
+    $input `into` $dest `[` $indices `]` () $unique
+    `:` type($input) `into` type($dest) `[` type($indices) `]` `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // SplatOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_SplatOp : Tensor_Op<"splat", [
     NoSideEffect,
     TypesMatchWith<"operand type matches element type of result",
                    "aggregate", "input",
                    "$_self.cast<TensorType>().getElementType()">