diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td index 30c1d97ba58f1..c7b83674fb009 100644 --- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td @@ -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, %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] ``` + The index into `base` only varies in the innermost ((k-1)-th) dimension. If a mask bit is set and the corresponding index is out-of-bounds for the given base, the behavior is undefined. If a mask bit is not set, the value comes from the pass-through vector regardless of the index, and the index is allowed to be out-of-bounds. - The gather operation can be used directly where applicable, or can be used - during progressively lowering to bring other memory operations closer to - hardware ISA support for a gather. - Examples: ```mlir + // 1-D memref gathered to 2-D vector. %0 = vector.gather %base[%c0][%v], %mask, %pass_thru : memref, vector<2x16xi32>, vector<2x16xi1>, vector<2x16xf32> into vector<2x16xf32> + // 2-D memref gathered to 1-D vector. %1 = vector.gather %base[%i, %j][%v], %mask, %pass_thru : memref<16x16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32> ```