From 7b9cff1c84f77b178a4f202038598c7bd3e25785 Mon Sep 17 00:00:00 2001 From: James Newling Date: Tue, 12 Aug 2025 13:21:46 -0700 Subject: [PATCH 1/4] tighten the example --- .../mlir/Dialect/Vector/IR/VectorOps.td | 28 +++++++++++-------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td index 30c1d97ba58f1..db5de0c70d0d0 100644 --- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td @@ -2058,23 +2058,27 @@ 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. The indexing semantics are then, + ``` - 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. + result[i,j] := if mask[i,j] then + base[indices[0], indices[1], indices[2] + index_vec[i,j]] + else + pass_thru[i,j] ``` + The index into `base` only varies in the dimension k-1. 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 @@ -2082,8 +2086,8 @@ def Vector_GatherOp : 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. + during progressive lowering to bring other memory operations closer to hardware + ISA support for a gather. Examples: From 6b431e30736df7cdea1b2c8e52f7df18c82704a5 Mon Sep 17 00:00:00 2001 From: James Newling Date: Tue, 12 Aug 2025 14:46:17 -0700 Subject: [PATCH 2/4] add IR snippet for 3-D to 2-D, remove comment about gradual lowering --- .../mlir/Dialect/Vector/IR/VectorOps.td | 33 ++++++++++++------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td index db5de0c70d0d0..5a5a3200f7652 100644 --- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td @@ -2068,33 +2068,42 @@ def Vector_GatherOp : on the values of an n-D mask vector. 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. The indexing semantics are then, + 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> + } + ``` + + The indexing semantics are then, ``` - result[i,j] := if mask[i,j] then - base[indices[0], indices[1], indices[2] + index_vec[i,j]] - else - pass_thru[i,j] + 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 dimension k-1. + The index into `base` only varies in the innermost dimension, k-1. 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 progressive 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> ``` From 52133ed2a384f3426e5dc24c034dc49d15ea12f4 Mon Sep 17 00:00:00 2001 From: James Newling Date: Tue, 12 Aug 2025 14:48:34 -0700 Subject: [PATCH 3/4] k-1 clarify --- mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td index 5a5a3200f7652..156b5d6080bd6 100644 --- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td @@ -2089,7 +2089,7 @@ def Vector_GatherOp : 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 dimension, k-1. + 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 From 8e9936b251943c6aa9db310f78ca1dd5d5a20003 Mon Sep 17 00:00:00 2001 From: James Newling Date: Tue, 12 Aug 2025 17:57:14 -0700 Subject: [PATCH 4/4] indentation fix --- mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td index 156b5d6080bd6..c7b83674fb009 100644 --- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td @@ -2077,9 +2077,9 @@ def Vector_GatherOp : %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] + %result = vector.gather %base[%i0, %i1, %i2] [%index_vec], %mask, %fall_thru : [...] - return %result : vector<2x3xf32> + return %result : vector<2x3xf32> } ```