[Mlir-commits] [mlir] [mlir][vector] Clarify the semantics of gather/scatter indexing (nfc) (PR #181357)

[Krzysztof Drewniak]

================
@@ -2082,23 +2082,27 @@ def Vector_GatherOp :
     3-D and the result is 2-D:
 
     ```mlir
-    func.func @gather_3D_to_2D(
-        %base: memref<?x10x?xf32>, %ofs_0: index, %ofs_1: index, %ofs_2: index,
-        %indices: vector<2x3xi32>, %mask: vector<2x3xi1>,
-        %fall_thru: vector<2x3xf32>) -> vector<2x3xf32> {
-            %result = vector.gather %base[%ofs_0, %ofs_1, %ofs_2]
-                                   [%indices], %mask, %fall_thru : [...]
-            return %result : vector<2x3xf32>
+      %base: memref<?x10x?xf32>, %ofs_0: index, %ofs_1: index, %ofs_2: index,
+      %indices: vector<2x3xi32>, %mask: vector<2x3xi1>,
+      %fall_thru: vector<2x3xf32>) -> vector<2x3xf32> {
+          %result = vector.gather %base[%ofs_0, %ofs_1, %ofs_2]
+                                 [%indices], %mask, %fall_thru : [...]
     }
     ```
 
     The indexing semantics are then,
 
     ```
-    result[i,j] := if mask[i,j] then base[i0, i1, i2 + indices[i,j]]
+    result[i,j] := if mask[i,j] then base[%ofs_0, %ofs_1, %ofs_2 + indices[i,j]]
                    else pass_thru[i,j]
     ```
-    The index into `base` only varies in the innermost ((k-1)-th) dimension.
+    Note, `indices` are element offsets - they are expressed in units of
+    elements (not bytes). Each element in `indices` represents a displacement
+    in units of elements from the starting element, i.e. `%base[%ofs_0, %ofs_1,
+    %ofs_2]` for the example above. Importantly, for MemRefs, `indices` are
+    %interpreted assuming an identity (contiguous) MemRef
+    layout and do not account for non-identity strides.
----------------
krzysz00 wrote:

So I'd like to be clear that I don't like the "you have to unfold the index according to the shape of the memref" semantics that are being proposed here - they make the lowering of what's supposed to be a low-level operation extremely non-trivial, and make rewrites (such as folding in a subview) a rather painful pile of index computations.

This is why I'm pushing for "the index just gets added to the last dimension, and if that goes oob, so be it" and the handling of, say, gathering along multiple dimensions involving either k vectors of indices (one per dimension) or a vector of k indices per vector element.

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