[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:

@banach-space Thanks for the feedback, and apologies for landing this faster than necessary. Let me know if you think this can be improved further and I'll definitely make a follow-up PR. 

With respect to 

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

I find interpreting base as a pointer less clear. 

> This is a bit tricky/nuanced though, as Tensors have no notion of "base address" 😅

Exactly!

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? 

Another subtle difference is what 'out of bounds' means. Current lowering ends up as vector.loads of single elements

```
[...]
%foo = vector.load %base[%i, %j] : memref<100x100xf32>, vector<1xf32>
[...]
```

There is nothing in the vector.load definition about out-of-bounds, but I assume the natural definition there would be that if %j excedes 99 above, it's out of bounds a UB. Which I think is more inline with the current definition of adding index_vec[i,j] to i2. 


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