Skip to content

Commit 626f533

Browse files
sahas3sahas3
authored
[TorchToLinalg] slice: use dimSize for INT64_MAX end sentinel (#4623)
PyTorch encodes "slice to the end" as end = INT64_MAX. The previous code ran that value through castIntToIndex and a chain of clamp selects, which materializes a large index constant that downstream targets with 32-bit index cannot represent. When the end operand is the constant INT64_MAX, use dimSize directly and skip the clamp chain; the existing logic is unchanged for all other end values.
1 parent bddd9e6 commit 626f533

2 files changed

Lines changed: 65 additions & 11 deletions

File tree

‎lib/Conversion/TorchToLinalg/DataMovement.cpp‎

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
2727
#include "llvm/ADT/APInt.h"
2828

29+
#include <limits>
2930
#include <numeric>
3031

3132
using namespace mlir;
@@ -82,17 +83,26 @@ LogicalResult prepareArgumentsForSlicingOp(OpTy op, OpAdaptor adaptor,
8283
if (isa<Torch::NoneType>(torchTypeEnd.getType())) {
8384
end = dimSize;
8485
} else {
85-
end = castIntToIndex(rewriter, loc, end);
86-
Value endcmp = arith::CmpIOp::create(rewriter, loc,
87-
arith::CmpIPredicate::slt, end, zero);
88-
Value endadd = arith::AddIOp::create(rewriter, loc, end, dimSize);
89-
end = arith::SelectOp::create(rewriter, loc, endcmp, endadd, end);
90-
endcmp = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
91-
end, zero);
92-
end = arith::SelectOp::create(rewriter, loc, endcmp, negone, end);
93-
endcmp = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sgt,
94-
end, dimSize);
95-
end = arith::SelectOp::create(rewriter, loc, endcmp, dimSize, end);
86+
// If end is INT64_MAX (PyTorch sentinel for "slice to the end"), use
87+
// dimSize directly to avoid materializing a large index constant that
88+
// downstream targets with 32-bit index cannot represent.
89+
int64_t endConst;
90+
if (matchPattern(torchTypeEnd, m_TorchConstantInt(&endConst)) &&
91+
endConst == std::numeric_limits<int64_t>::max()) {
92+
end = dimSize;
93+
} else {
94+
end = castIntToIndex(rewriter, loc, end);
95+
Value endcmp = arith::CmpIOp::create(
96+
rewriter, loc, arith::CmpIPredicate::slt, end, zero);
97+
Value endadd = arith::AddIOp::create(rewriter, loc, end, dimSize);
98+
end = arith::SelectOp::create(rewriter, loc, endcmp, endadd, end);
99+
endcmp = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
100+
end, zero);
101+
end = arith::SelectOp::create(rewriter, loc, endcmp, negone, end);
102+
endcmp = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sgt,
103+
end, dimSize);
104+
end = arith::SelectOp::create(rewriter, loc, endcmp, dimSize, end);
105+
}
96106
}
97107

98108
// Slice logic: resultSize = floordiv(end - start + step - 1, step)

‎test/Conversion/TorchToLinalg/datamovement.mlir‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,3 +133,47 @@ func.func @aten.permute$identity_permutation(%arg0: !torch.vtensor<[64,32,16,8,4
133133
%1 = torch.aten.permute %arg0, %0 : !torch.vtensor<[64,32,16,8,4],f32>, !torch.list<int> -> !torch.vtensor<[64,32,16,8,4],f32>
134134
return %1 : !torch.vtensor<[64,32,16,8,4],f32>
135135
}
136+
137+
// -----
138+
139+
// CHECK-LABEL: func.func @torch.aten.slice$end_int64_max_dynamic(
140+
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor<[4,?],f32>
141+
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
142+
// CHECK-DAG: %[[T:.*]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[4,?],f32> -> tensor<4x?xf32>
143+
// CHECK: %[[DIM:.*]] = tensor.dim %[[T]], %[[C1]] : tensor<4x?xf32>
144+
// COM: No INT64_MAX sentinel is materialized: the end resolves to %[[DIM]], so
145+
// COM: the slice extent is computed as (dim - start) directly off the queried
146+
// COM: dim, with no min(INT64_MAX, dim) clamp in between.
147+
// CHECK-NOT: 9223372036854775807
148+
// CHECK: arith.subi %[[DIM]], %{{.*}} : index
149+
// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[T]][0, %{{.*}}] [4, %{{.*}}] [1, 1] : tensor<4x?xf32> to tensor<4x?xf32>
150+
// CHECK: torch_c.from_builtin_tensor %[[SLICE]]
151+
func.func @torch.aten.slice$end_int64_max_dynamic(%arg0: !torch.vtensor<[4,?],f32>) -> !torch.vtensor<[4,?],f32> {
152+
%int1 = torch.constant.int 1
153+
%int2 = torch.constant.int 2
154+
%intmax = torch.constant.int 9223372036854775807
155+
%0 = torch.aten.slice.Tensor %arg0, %int1, %int2, %intmax, %int1 : !torch.vtensor<[4,?],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[4,?],f32>
156+
return %0 : !torch.vtensor<[4,?],f32>
157+
}
158+
159+
// -----
160+
161+
// Same INT64_MAX end sentinel, but with step=2: the end still resolves to the
162+
// dim (no sentinel constant), and the strided size is (dim - start) floor-div
163+
// step.
164+
// CHECK-LABEL: func.func @torch.aten.slice$end_int64_max_dynamic_step2(
165+
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor<[4,?],f32>
166+
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
167+
// CHECK-DAG: %[[T:.*]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[4,?],f32> -> tensor<4x?xf32>
168+
// CHECK: %[[DIM:.*]] = tensor.dim %[[T]], %[[C1]] : tensor<4x?xf32>
169+
// CHECK-NOT: 9223372036854775807
170+
// CHECK: %[[LEN:.*]] = arith.subi %[[DIM]], %{{.*}} : index
171+
// CHECK: arith.floordivsi %{{.*}}, %{{.*}} : index
172+
// CHECK: tensor.extract_slice %[[T]][0, %{{.*}}] [4, %{{.*}}] [1, 2] : tensor<4x?xf32> to tensor<4x?xf32>
173+
func.func @torch.aten.slice$end_int64_max_dynamic_step2(%arg0: !torch.vtensor<[4,?],f32>) -> !torch.vtensor<[4,?],f32> {
174+
%int1 = torch.constant.int 1
175+
%int2 = torch.constant.int 2
176+
%intmax = torch.constant.int 9223372036854775807
177+
%0 = torch.aten.slice.Tensor %arg0, %int1, %int2, %intmax, %int2 : !torch.vtensor<[4,?],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[4,?],f32>
178+
return %0 : !torch.vtensor<[4,?],f32>
179+
}

0 commit comments

Comments
 (0)