Skip to content

Commit 3a1f1cb

Browse files
tpatabantpataban
committed
Lower tt.gather into tensor.gather
Signed-off-by: Tharindu Patabandi <tpataban@qti.qualcomm.com>
1 parent 2f10c15 commit 3a1f1cb

3 files changed

Lines changed: 132 additions & 53 deletions

File tree

‎include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp‎

Lines changed: 62 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,10 +1026,16 @@ struct JoinConverter : public OpConversionPattern<triton::JoinOp> {
10261026
}
10271027
};
10281028

1029-
// Lowers `tt.gather` to a `linalg.generic` that reads `indices` through the
1030-
// identity affine map and, for each output position, extracts the gathered
1031-
// element out of `src` via `tensor.extract` at a dynamically-computed
1032-
// coordinate.
1029+
// Lowers `tt.gather` to `tensor.gather`. `tt.gather`'s `indices` has the same
1030+
// shape as the output, with only the `axis` coordinate varying per element,
1031+
// whereas `tensor.gather` expects a coordinate tuple per output element
1032+
// covering every gathered dimension. We bridge the two by gathering along
1033+
// *all* source dimensions: a coordinate tensor of shape
1034+
// `indicesShape ++ [rank]` is built where the `axis` slot holds `indices`
1035+
// (cast to index) and every other slot `d` holds that element's own index
1036+
// along dimension `d` (an identity/iota broadcast). With every dimension
1037+
// listed in `gather_dims`, the rank-reduced result shape collapses to exactly
1038+
// `indicesShape`, matching `tt.gather`'s output shape.
10331039
struct GatherConverter : public OpConversionPattern<triton::GatherOp> {
10341040
using OpConversionPattern<triton::GatherOp>::OpConversionPattern;
10351041

@@ -1050,43 +1056,62 @@ struct GatherConverter : public OpConversionPattern<triton::GatherOp> {
10501056

10511057
int64_t axis = op.getAxis();
10521058
int64_t rank = idxType.getRank();
1059+
ArrayRef<int64_t> idxShape = idxType.getShape();
1060+
Type indexElemType = rewriter.getIndexType();
1061+
1062+
SmallVector<int64_t> coordsShape(idxShape.begin(), idxShape.end());
1063+
coordsShape.push_back(rank);
1064+
Value coords =
1065+
tensor::EmptyOp::create(rewriter, loc, coordsShape, indexElemType);
1066+
1067+
auto idxIndexType = RankedTensorType::get(idxShape, indexElemType);
1068+
Value axisComponent =
1069+
arith::IndexCastOp::create(rewriter, loc, idxIndexType, indices);
1070+
1071+
SmallVector<OpFoldResult> offsets(rank + 1, rewriter.getIndexAttr(0));
1072+
SmallVector<OpFoldResult> strides(rank + 1, rewriter.getIndexAttr(1));
1073+
SmallVector<OpFoldResult> sizes;
1074+
sizes.reserve(rank + 1);
1075+
for (int64_t dim : idxShape)
1076+
sizes.push_back(rewriter.getIndexAttr(dim));
1077+
sizes.push_back(rewriter.getIndexAttr(1));
1078+
1079+
SmallVector<AffineMap> iotaIndexingMaps(
1080+
1, rewriter.getMultiDimIdentityMap(rank));
1081+
1082+
for (int64_t dim = 0; dim < rank; ++dim) {
1083+
Value component;
1084+
if (dim == axis) {
1085+
component = axisComponent;
1086+
} else {
1087+
Value iotaInit =
1088+
tensor::EmptyOp::create(rewriter, loc, idxShape, indexElemType);
1089+
auto iotaOp = linalg::GenericOp::create(
1090+
rewriter, loc, TypeRange{iotaInit.getType()}, ValueRange{},
1091+
ValueRange{iotaInit}, iotaIndexingMaps,
1092+
getNParallelLoopsAttrs(rank),
1093+
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange) {
1094+
Value idx =
1095+
linalg::IndexOp::create(nestedBuilder, nestedLoc, dim);
1096+
linalg::YieldOp::create(nestedBuilder, nestedLoc, idx);
1097+
});
1098+
component = iotaOp.getResult(0);
1099+
}
10531100

1054-
// indices and the output share the same shape, so both use the identity
1055-
// map; src is captured below rather than driven through an ins() operand.
1056-
SmallVector<AffineMap> indexingMaps(
1057-
/*indices*/ 1 + /*output*/ 1,
1058-
rewriter.getMultiDimIdentityMap(rank));
1101+
offsets.back() = rewriter.getIndexAttr(dim);
1102+
coords = tensor::InsertSliceOp::create(rewriter, loc, component, coords,
1103+
offsets, sizes, strides);
1104+
}
10591105

1060-
Value init = tensor::EmptyOp::create(rewriter, loc, resType.getShape(),
1061-
resType.getElementType());
1106+
SmallVector<int64_t> gatherDims;
1107+
gatherDims.reserve(rank);
1108+
for (int64_t dim = 0; dim < rank; ++dim)
1109+
gatherDims.push_back(dim);
10621110

1063-
auto linalgOp = linalg::GenericOp::create(
1064-
rewriter, loc, op->getResultTypes(), ValueRange{indices},
1065-
ValueRange{init}, indexingMaps, getNParallelLoopsAttrs(rank),
1066-
[&](OpBuilder &nestedBuilder, Location nestedLoc,
1067-
ValueRange blockArgs) {
1068-
Value idxScalar = blockArgs[0];
1069-
Value idxAsIndex = arith::IndexCastOp::create(
1070-
nestedBuilder, nestedLoc, nestedBuilder.getIndexType(),
1071-
idxScalar);
1072-
1073-
SmallVector<Value> coords;
1074-
coords.reserve(rank);
1075-
for (int64_t dim = 0; dim < rank; ++dim) {
1076-
if (dim == axis) {
1077-
coords.push_back(idxAsIndex);
1078-
} else {
1079-
coords.push_back(
1080-
linalg::IndexOp::create(nestedBuilder, nestedLoc, dim));
1081-
}
1082-
}
1111+
auto gatherOp = tensor::GatherOp::create(
1112+
rewriter, loc, resType, src, coords, ArrayRef<int64_t>(gatherDims));
10831113

1084-
Value gathered =
1085-
tensor::ExtractOp::create(nestedBuilder, nestedLoc, src, coords);
1086-
linalg::YieldOp::create(nestedBuilder, nestedLoc, gathered);
1087-
});
1088-
1089-
rewriter.replaceOp(op, linalgOp->getResults());
1114+
rewriter.replaceOp(op, gatherOp);
10901115
return success();
10911116
}
10921117
};

‎python/examples/test_tensor_gather.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates, Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
import pytest
5+
import torch
6+
7+
import triton
8+
import triton.language as tl
9+
10+
from triton.backends.triton_shared.driver import CPUDriver
11+
12+
13+
@triton.jit
14+
def gather2d_kernel(
15+
src_ptr, idx_ptr, out_ptr, AXIS: tl.constexpr, M: tl.constexpr, N: tl.constexpr
16+
):
17+
rows = tl.arange(0, M)[:, None]
18+
cols = tl.arange(0, N)[None, :]
19+
offs = rows * N + cols
20+
21+
src = tl.load(src_ptr + offs)
22+
idx = tl.load(idx_ptr + offs)
23+
out = tl.gather(src, idx, AXIS)
24+
tl.store(out_ptr + offs, out)
25+
26+
27+
def gather2d(src, idx, axis):
28+
M, N = src.shape
29+
out = torch.empty_like(src)
30+
gather2d_kernel[1,](src, idx, out, AXIS=axis, M=M, N=N)
31+
return out
32+
33+
34+
@pytest.mark.parametrize("axis", [0, 1])
35+
def test_gather_2d(axis, device):
36+
if device == "cpu":
37+
triton.runtime.driver.set_active(CPUDriver())
38+
39+
torch.manual_seed(0)
40+
M, N = 8, 16
41+
src = torch.randn(M, N, device=device)
42+
dim_size = M if axis == 0 else N
43+
idx = torch.randint(0, dim_size, (M, N), dtype=torch.int32, device=device)
44+
45+
out = gather2d(src, idx, axis)
46+
ref = torch.gather(src, axis, idx.to(torch.int64))
47+
48+
torch.testing.assert_close(out, ref)

‎test/Conversion/TritonArithToLinalg/gather.mlir‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,18 @@ module {
1010
// CHECK-LABEL: func.func @gather_axis1(
1111
// CHECK-SAME: %[[SRC:.*]]: tensor<8x16xf32>,
1212
// CHECK-SAME: %[[INDICES:.*]]: tensor<8x16xi32>) -> tensor<8x16xf32> {
13-
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<8x16xf32>
14-
// CHECK: %[[GENERIC:.*]] = linalg.generic {indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel"]} ins(%[[INDICES]] : tensor<8x16xi32>) outs(%[[EMPTY]] : tensor<8x16xf32>) {
15-
// CHECK: ^bb0(%[[IDX:.*]]: i32, %{{.*}}: f32):
16-
// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[IDX]] : i32 to index
13+
// CHECK: %[[COORDS_INIT:.*]] = tensor.empty() : tensor<8x16x2xindex>
14+
// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[INDICES]] : tensor<8x16xi32> to tensor<8x16xindex>
15+
// CHECK: %[[IOTA_INIT0:.*]] = tensor.empty() : tensor<8x16xindex>
16+
// CHECK: %[[IOTA0:.*]] = linalg.generic {indexing_maps = [#{{.*}}], iterator_types = ["parallel", "parallel"]} outs(%[[IOTA_INIT0]] : tensor<8x16xindex>) {
17+
// CHECK: ^bb0(%{{.*}}: index):
1718
// CHECK: %[[ROW:.*]] = linalg.index 0 : index
18-
// CHECK: %[[EXTRACTED:.*]] = tensor.extract %[[SRC]]{{\[}}%[[ROW]], %[[IDX_CAST]]] : tensor<8x16xf32>
19-
// CHECK: linalg.yield %[[EXTRACTED]] : f32
20-
// CHECK: } -> tensor<8x16xf32>
21-
// CHECK: return %[[GENERIC]] : tensor<8x16xf32>
19+
// CHECK: linalg.yield %[[ROW]] : index
20+
// CHECK: } -> tensor<8x16xindex>
21+
// CHECK: %[[COORDS0:.*]] = tensor.insert_slice %[[IOTA0]] into %[[COORDS_INIT]][0, 0, 0] [8, 16, 1] [1, 1, 1] : tensor<8x16xindex> into tensor<8x16x2xindex>
22+
// CHECK: %[[COORDS1:.*]] = tensor.insert_slice %[[IDX_CAST]] into %[[COORDS0]][0, 0, 1] [8, 16, 1] [1, 1, 1] : tensor<8x16xindex> into tensor<8x16x2xindex>
23+
// CHECK: %[[GATHER:.*]] = tensor.gather %[[SRC]]{{\[}}%[[COORDS1]]] gather_dims([0, 1]) : (tensor<8x16xf32>, tensor<8x16x2xindex>) -> tensor<8x16xf32>
24+
// CHECK: return %[[GATHER]] : tensor<8x16xf32>
2225
// CHECK: }
2326

2427
// -----
@@ -33,13 +36,16 @@ module {
3336
// CHECK-LABEL: func.func @gather_axis0(
3437
// CHECK-SAME: %[[SRC:.*]]: tensor<8x16xf32>,
3538
// CHECK-SAME: %[[INDICES:.*]]: tensor<8x16xi32>) -> tensor<8x16xf32> {
36-
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<8x16xf32>
37-
// CHECK: %[[GENERIC:.*]] = linalg.generic {indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel"]} ins(%[[INDICES]] : tensor<8x16xi32>) outs(%[[EMPTY]] : tensor<8x16xf32>) {
38-
// CHECK: ^bb0(%[[IDX:.*]]: i32, %{{.*}}: f32):
39-
// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[IDX]] : i32 to index
39+
// CHECK: %[[COORDS_INIT:.*]] = tensor.empty() : tensor<8x16x2xindex>
40+
// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[INDICES]] : tensor<8x16xi32> to tensor<8x16xindex>
41+
// CHECK: %[[COORDS0:.*]] = tensor.insert_slice %[[IDX_CAST]] into %[[COORDS_INIT]][0, 0, 0] [8, 16, 1] [1, 1, 1] : tensor<8x16xindex> into tensor<8x16x2xindex>
42+
// CHECK: %[[IOTA_INIT1:.*]] = tensor.empty() : tensor<8x16xindex>
43+
// CHECK: %[[IOTA1:.*]] = linalg.generic {indexing_maps = [#{{.*}}], iterator_types = ["parallel", "parallel"]} outs(%[[IOTA_INIT1]] : tensor<8x16xindex>) {
44+
// CHECK: ^bb0(%{{.*}}: index):
4045
// CHECK: %[[COL:.*]] = linalg.index 1 : index
41-
// CHECK: %[[EXTRACTED:.*]] = tensor.extract %[[SRC]]{{\[}}%[[IDX_CAST]], %[[COL]]] : tensor<8x16xf32>
42-
// CHECK: linalg.yield %[[EXTRACTED]] : f32
43-
// CHECK: } -> tensor<8x16xf32>
44-
// CHECK: return %[[GENERIC]] : tensor<8x16xf32>
46+
// CHECK: linalg.yield %[[COL]] : index
47+
// CHECK: } -> tensor<8x16xindex>
48+
// CHECK: %[[COORDS1:.*]] = tensor.insert_slice %[[IOTA1]] into %[[COORDS0]][0, 0, 1] [8, 16, 1] [1, 1, 1] : tensor<8x16xindex> into tensor<8x16x2xindex>
49+
// CHECK: %[[GATHER:.*]] = tensor.gather %[[SRC]]{{\[}}%[[COORDS1]]] gather_dims([0, 1]) : (tensor<8x16xf32>, tensor<8x16x2xindex>) -> tensor<8x16xf32>
50+
// CHECK: return %[[GATHER]] : tensor<8x16xf32>
4551
// CHECK: }

0 commit comments

Comments
 (0)