Skip to content

Commit 2f10c15

Browse files
tpatabantpataban
committed
Lower tt.gather to linalg
Currently, `tt.gather` is not lowered to any middle layer dialects and the operator remains unhandled. This commit registers a lowering pass in TritonArithToLinalg that lowers `tt.gather` into a `linalg.generic` region that expresses the gather operation purely in terms of `arith::IndexCastOp`, `linalg::IndexOp`, and `tensor::ExtractOp`. Signed-off-by: Tharindu Patabandi <tpataban@qti.qualcomm.com>
1 parent 1477e98 commit 2f10c15

3 files changed

Lines changed: 111 additions & 0 deletions

File tree

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,6 +1026,71 @@ 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.
1033+
struct GatherConverter : public OpConversionPattern<triton::GatherOp> {
1034+
using OpConversionPattern<triton::GatherOp>::OpConversionPattern;
1035+
1036+
LogicalResult
1037+
matchAndRewrite(triton::GatherOp op, OpAdaptor adaptor,
1038+
ConversionPatternRewriter &rewriter) const override {
1039+
auto loc = op.getLoc();
1040+
Value src = adaptor.getSrc();
1041+
Value indices = adaptor.getIndices();
1042+
1043+
auto srcType = cast<RankedTensorType>(src.getType());
1044+
auto idxType = cast<RankedTensorType>(indices.getType());
1045+
auto resType = cast<RankedTensorType>(op.getResult().getType());
1046+
1047+
if (!srcType.hasStaticShape() || !idxType.hasStaticShape())
1048+
return rewriter.notifyMatchFailure(
1049+
op, "gather: only static src/indices shapes are supported");
1050+
1051+
int64_t axis = op.getAxis();
1052+
int64_t rank = idxType.getRank();
1053+
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));
1059+
1060+
Value init = tensor::EmptyOp::create(rewriter, loc, resType.getShape(),
1061+
resType.getElementType());
1062+
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+
}
1083+
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());
1090+
return success();
1091+
}
1092+
};
1093+
10291094
struct MulHiUIOpConverter : public OpConversionPattern<triton::MulhiUIOp> {
10301095
using OpConversionPattern<triton::MulhiUIOp>::OpConversionPattern;
10311096

‎lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns(
7979
patterns.add<DenseConstantConverter>(patterns.getContext());
8080
patterns.add<CumSumConverter>(patterns.getContext());
8181
patterns.add<ReshapeConverter>(patterns.getContext());
82+
patterns.add<GatherConverter>(patterns.getContext());
8283

8384
populateExternElementwiseOpToMLIROps(patterns);
8485

‎test/Conversion/TritonArithToLinalg/gather.mlir‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// RUN: triton-shared-opt --triton-arith-to-linalg %s | FileCheck %s
2+
3+
module {
4+
tt.func public @gather_axis1(%src: tensor<8x16xf32>, %indices: tensor<8x16xi32>) -> tensor<8x16xf32> {
5+
%0 = tt.gather %src[%indices] {axis = 1 : i32} : (tensor<8x16xf32>, tensor<8x16xi32>) -> tensor<8x16xf32>
6+
tt.return %0 : tensor<8x16xf32>
7+
}
8+
}
9+
10+
// CHECK-LABEL: func.func @gather_axis1(
11+
// CHECK-SAME: %[[SRC:.*]]: tensor<8x16xf32>,
12+
// 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
17+
// 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>
22+
// CHECK: }
23+
24+
// -----
25+
26+
module {
27+
tt.func public @gather_axis0(%src: tensor<8x16xf32>, %indices: tensor<8x16xi32>) -> tensor<8x16xf32> {
28+
%0 = tt.gather %src[%indices] {axis = 0 : i32} : (tensor<8x16xf32>, tensor<8x16xi32>) -> tensor<8x16xf32>
29+
tt.return %0 : tensor<8x16xf32>
30+
}
31+
}
32+
33+
// CHECK-LABEL: func.func @gather_axis0(
34+
// CHECK-SAME: %[[SRC:.*]]: tensor<8x16xf32>,
35+
// 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
40+
// 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>
45+
// CHECK: }

0 commit comments

Comments
 (0)