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

[Andrzej Warzyński]

================
@@ -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]
----------------
banach-space wrote:

Sorry for not replying earlier, I was OOO.

> 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]

I view this differently. To me, `base[i0, i1, i2] + index_vec[i,j]` is the `vector` abstraction and that's all we care about here.  Later, these `vector` level indices are interpreted at either the `memref` or `tensor` abstraction levels - that's when "stride" would matter. Put differently, I agree that the actual meaning of this will depend on what `base` is:
```mlir
base[i0, i1, i2] + index_vec[i,j]
```

However, to me that should not be a concern at the `vector` level.

Anyway, this is a side point - it's obviously totally fine to see things differently. Your change is a much appreciated improvement, lets leave it as is.

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