[Mlir-commits] [mlir] [mlir][vector] Improve vector.gather description (PR #153278)

[James Newling]

================
@@ -2058,39 +2058,52 @@ def Vector_GatherOp :
     Results<(outs AnyVectorOfNonZeroRank:$result)> {
 
   let summary = [{
-    gathers elements from memory or ranked tensor into a vector as defined by an
-    index vector and a mask vector
+    Gathers elements from memory or ranked tensor into a vector as defined by an
+    index vector and a mask vector.
   }];
 
   let description = [{
     The gather operation returns an n-D vector whose elements are either loaded
-    from memory or ranked tensor, or taken from a pass-through vector, depending
+    from a k-D memref or tensor, or taken from an n-D pass-through vector, depending
     on the values of an n-D mask vector.
-    If a mask bit is set, the corresponding result element is defined by the base
-    with indices and the n-D index vector (each index is a 1-D offset on the base).
-    Otherwise, the corresponding element is taken from the n-D pass-through vector.
-    Informally the semantics are:
+
+    If a mask bit is set, the corresponding result element is taken from `base`
+    at an index defined by k indices and n-D `index_vec`. Otherwise, the element
+    is taken from the pass-through vector. As an example, suppose that `base` is
+    3-D and the result is 2-D:
+
+    ```mlir
+    func.func @gather_3D_to_2D(
+        %base: memref<?x10x?xf32>, %i0: index, %i1: index, %i2: index,
+        %index_vec: vector<2x3xi32>, %mask: vector<2x3xi1>,
+        %fall_thru: vector<2x3xf32>) -> vector<2x3xf32> {
+            %result = vector.gather %base[%i0, %i1, %i2]
+                                   [%index_vec], %mask, %fall_thru : [...]
+            return %result : vector<2x3xf32>
+    }
     ```
-    result[0] := if mask[0] then base[index[0]] else pass_thru[0]
-    result[1] := if mask[1] then base[index[1]] else pass_thru[1]
-    etc.
+
+    The indexing semantics are then,
+
+    ```
+    result[i,j] := if mask[i,j] then base[i0, i1, i2 + index_vec[i,j]]
+                   else pass_thru[i,j]
----------------
newling wrote:

> > I'll add that memrefs can be strided (see this [test](https://github.com/llvm/llvm-project/blob/16ad20291dad174e441076c4c10d899b333fc0e7/mlir/test/Dialect/Vector/vector-gather-lowering.mlir#L151) so should strides be included?
> 
> Hm, not at the `vector.gather` nor `Vector` level, no. Are they?
> 

I meant, in that test (copied below)

```mlir
%0 = vector.gather %base[%c0][%v], %mask, %pass_thru : memref<4xf32, strided<[2]>>, vector<1xindex>, vector<1xi1>, vector<1xf32> into vector<1xf32>
```

if the explanation of vector.gather was pointer based:

```bash
    result[i,j] := if mask[i,j] then base[i0, i1, i2] + index_vec[i,j]
                   else pass_thru[i,j]
```                   

then we should probably include a stride in the above. i.e. it should `+ index_vec[i,j]` should be `+ stride * index_vec[i,j]`

This is one reason why I think the pointer-based definition of gather a bit less clear. With the tensor-based definition (just add index_vec[i,j] to another index, not a pointer) this question doesn't come up. 

https://github.com/llvm/llvm-project/pull/153278