Skip to content

Commit c852e84

Browse files
ziliangzlziliangzl
authored
[TorchToLinalg] Clamp nearest-interpolate indices before tensor.extract (#4548)
Nearest resize lowering maps each output position to a source index using coordinate transform plus rounding. For half_pixel, the fractional coordinate can be slightly negative; floor (and the default empty nearest mode) then yields -1. Those values reached tensor.extract without bounding to the input extent, so extraction could read out of bounds and produce incorrect results. Clamp the rounded float index to [0, length-1] per spatial dimension before fptosi. Refresh FileCheck lines in test/Conversion/TorchToLinalg/resize.mlir. --------- Signed-off-by: Ziliang Zhang <zzl.coding@gmail.com>
1 parent 60487c9 commit c852e84

2 files changed

Lines changed: 76 additions & 21 deletions

File tree

‎lib/Conversion/TorchToLinalg/Uncategorized.cpp‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2775,24 +2775,26 @@ static Value nearestInterpolate(OpBuilder &b, Location loc,
27752775
nearestFP = arith::SelectOp::create(b, loc, cmp, floor, ceil);
27762776
} else if (nearestMode == "round_prefer_ceil") {
27772777
Value cstHalf = arith::ConstantOp::create(b, loc, b.getF32FloatAttr(0.5));
2778-
Value cstOne = arith::ConstantOp::create(b, loc, b.getF32FloatAttr(1));
27792778
Value floor = math::FloorOp::create(b, loc, proj);
27802779
Value ceil = math::CeilOp::create(b, loc, proj);
27812780
Value decimal = arith::SubFOp::create(b, loc, proj, floor);
27822781
Value cmp = arith::CmpFOp::create(b, loc, arith::CmpFPredicate::UGE,
27832782
decimal, cstHalf);
27842783
nearestFP = arith::SelectOp::create(b, loc, cmp, ceil, floor);
2785-
Value inputSizeMOne = arith::SubFOp::create(b, loc, inputSizeFP, cstOne);
2786-
// don't extract out of bounds
2787-
nearestFP = arith::MinimumFOp::create(b, loc, nearestFP, inputSizeMOne);
27882784
} else if (nearestMode == "ceil") {
2789-
Value cstOne = arith::ConstantOp::create(b, loc, b.getF32FloatAttr(1));
2790-
Value inputSizeMOne = arith::SubFOp::create(b, loc, inputSizeFP, cstOne);
27912785
nearestFP = math::CeilOp::create(b, loc, proj);
2792-
nearestFP = arith::MinimumFOp::create(b, loc, nearestFP, inputSizeMOne);
27932786
} else {
27942787
llvm_unreachable("Unsupported nearest mode");
27952788
}
2789+
// Clamp to valid input indices. ONNX half_pixel (and asymmetric) coords can
2790+
// lie slightly outside [0, length-1] before rounding; without clamping,
2791+
// tensor.extract uses out-of-range indices (garbage on some backends).
2792+
Value cstOne = arith::ConstantOp::create(b, loc, b.getF32FloatAttr(1.0));
2793+
Value cstZero = arith::ConstantOp::create(b, loc, b.getF32FloatAttr(0.0));
2794+
Value inputSizeMOne = arith::SubFOp::create(b, loc, inputSizeFP, cstOne);
2795+
nearestFP = arith::MaximumFOp::create(b, loc, nearestFP, cstZero);
2796+
nearestFP = arith::MinimumFOp::create(b, loc, nearestFP, inputSizeMOne);
2797+
27962798
Value nearestInt =
27972799
arith::FPToSIOp::create(b, loc, b.getI64Type(), nearestFP);
27982800
Value nearest =

‎test/Conversion/TorchToLinalg/resize.mlir‎

Lines changed: 67 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -76,17 +76,27 @@ func.func @test_resize_sizes_nearest(%arg0: !torch.vtensor<[1,1,2,4],f32>, %arg1
7676
// CHECK: %[[x23:.*]] = arith.index_cast %[[x13]] : index to i64
7777
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
7878
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
79-
// CHECK: %[[x26:.*]] = math.floor %[[x25]] : f32
80-
// CHECK: %[[x31:.*]] = arith.fptosi %[[x26]] : f32 to i64
79+
// CHECK: %[[floorH:.*]] = math.floor %[[x25]] : f32
80+
// CHECK: %[[c1H:.*]] = arith.constant 1.000000e+00 : f32
81+
// CHECK: %[[c0H:.*]] = arith.constant 0.000000e+00 : f32
82+
// CHECK: %[[inLm1H:.*]] = arith.subf %[[x15]], %[[c1H]] : f32
83+
// CHECK: %[[clH1:.*]] = arith.maximumf %[[floorH]], %[[c0H]] : f32
84+
// CHECK: %[[clH2:.*]] = arith.minimumf %[[clH1]], %[[inLm1H]] : f32
85+
// CHECK: %[[x31:.*]] = arith.fptosi %[[clH2]] : f32 to i64
8186
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
8287
// CHECK: %[[x16:.*]] = arith.sitofp %[[c4_i64:.*]] : i64 to f32
8388
// CHECK: %[[x20:.*]] = arith.sitofp %[[x7:.*]] : i64 to f32
8489
// CHECK: %[[x22:.*]] = arith.divf %[[x20]], %[[x16]] : f32
8590
// CHECK: %[[x26:.*]] = arith.index_cast %[[x14]] : index to i64
8691
// CHECK: %[[x27:.*]] = arith.sitofp %[[x26]] : i64 to f32
8792
// CHECK: %[[x28:.*]] = arith.divf %[[x27]], %[[x22]] : f32
88-
// CHECK: %[[x29:.*]] = math.floor %[[x28]] : f32
89-
// CHECK: %[[x33:.*]] = arith.fptosi %[[x29]] : f32 to i64
93+
// CHECK: %[[floorW:.*]] = math.floor %[[x28]] : f32
94+
// CHECK: %[[c1W:.*]] = arith.constant 1.000000e+00 : f32
95+
// CHECK: %[[c0W:.*]] = arith.constant 0.000000e+00 : f32
96+
// CHECK: %[[inLm1W:.*]] = arith.subf %[[x16]], %[[c1W]] : f32
97+
// CHECK: %[[clW1:.*]] = arith.maximumf %[[floorW]], %[[c0W]] : f32
98+
// CHECK: %[[clW2:.*]] = arith.minimumf %[[clW1]], %[[inLm1W]] : f32
99+
// CHECK: %[[x33:.*]] = arith.fptosi %[[clW2]] : f32 to i64
90100
// CHECK: %[[x34:.*]] = arith.index_cast %[[x33]] : i64 to index
91101
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]], %[[x34]]] : tensor<1x1x2x4xf32>
92102
// CHECK: linalg.yield %[[extracted]] : f32
@@ -121,8 +131,13 @@ func.func @test_resize_nearest_1d(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !to
121131
// CHECK: %[[x23:.*]] = arith.index_cast %[[x13]] : index to i64
122132
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
123133
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
124-
// CHECK: %[[x29:.*]] = math.floor %[[x25]] : f32
125-
// CHECK: %[[x31:.*]] = arith.fptosi %[[x29]] : f32 to i64
134+
// CHECK: %[[flo1d:.*]] = math.floor %[[x25]] : f32
135+
// CHECK: %[[c1_1d:.*]] = arith.constant 1.000000e+00 : f32
136+
// CHECK: %[[c0_1d:.*]] = arith.constant 0.000000e+00 : f32
137+
// CHECK: %[[inLm1_1d:.*]] = arith.subf %[[x15]], %[[c1_1d]] : f32
138+
// CHECK: %[[cl1d1:.*]] = arith.maximumf %[[flo1d]], %[[c0_1d]] : f32
139+
// CHECK: %[[cl1d2:.*]] = arith.minimumf %[[cl1d1]], %[[inLm1_1d]] : f32
140+
// CHECK: %[[x31:.*]] = arith.fptosi %[[cl1d2]] : f32 to i64
126141
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
127142
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]]] : tensor<?x?x?xf32>
128143
// CHECK: linalg.yield %[[extracted]] : f32
@@ -157,10 +172,41 @@ func.func @test_resize_nearest_3d(%arg0: !torch.vtensor<[?,?,?,?,?],f32>, %arg1:
157172
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
158173
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
159174
// CHECK: %[[floor:.*]] = math.floor %[[x25]] : f32
160-
// CHECK: %[[x31:.*]] = arith.fptosi %[[floor]] : f32 to i64
175+
// CHECK: %[[c1_3da:.*]] = arith.constant 1.000000e+00 : f32
176+
// CHECK: %[[c0_3da:.*]] = arith.constant 0.000000e+00 : f32
177+
// CHECK: %[[inLm1_3da:.*]] = arith.subf %[[x15]], %[[c1_3da]] : f32
178+
// CHECK: %[[cl3da1:.*]] = arith.maximumf %[[floor]], %[[c0_3da]] : f32
179+
// CHECK: %[[cl3da2:.*]] = arith.minimumf %[[cl3da1]], %[[inLm1_3da]] : f32
180+
// CHECK: %[[x31:.*]] = arith.fptosi %[[cl3da2]] : f32 to i64
161181
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
162-
// CHECK: %[[x34:.*]] = arith.index_cast %[[Wfptosi:.*]] : i64 to index
163-
// CHECK: %[[x35:.*]] = arith.index_cast %[[Dfptosi:.*]] : i64 to index
182+
// CHECK: %[[x16w:.*]] = arith.sitofp %[[c3_i64:.*]] : i64 to f32
183+
// CHECK: %[[x20w:.*]] = arith.sitofp %[[x7:.*]] : i64 to f32
184+
// CHECK: %[[x22w:.*]] = arith.divf %[[x20w]], %[[x16w]] : f32
185+
// CHECK: %[[x26w:.*]] = arith.index_cast %[[x14]] : index to i64
186+
// CHECK: %[[x27w:.*]] = arith.sitofp %[[x26w]] : i64 to f32
187+
// CHECK: %[[x28w:.*]] = arith.divf %[[x27w]], %[[x22w]] : f32
188+
// CHECK: %[[floorW2:.*]] = math.floor %[[x28w]] : f32
189+
// CHECK: %[[c1W2:.*]] = arith.constant 1.000000e+00 : f32
190+
// CHECK: %[[c0W2:.*]] = arith.constant 0.000000e+00 : f32
191+
// CHECK: %[[inLm1W2:.*]] = arith.subf %[[x16w]], %[[c1W2]] : f32
192+
// CHECK: %[[clW2a:.*]] = arith.maximumf %[[floorW2]], %[[c0W2]] : f32
193+
// CHECK: %[[clW2b:.*]] = arith.minimumf %[[clW2a]], %[[inLm1W2]] : f32
194+
// CHECK: %[[Wfptosi:.*]] = arith.fptosi %[[clW2b]] : f32 to i64
195+
// CHECK: %[[x34:.*]] = arith.index_cast %[[Wfptosi]] : i64 to index
196+
// CHECK: %[[x16d:.*]] = arith.sitofp %[[c4_i64:.*]] : i64 to f32
197+
// CHECK: %[[x20d:.*]] = arith.sitofp %[[x8:.*]] : i64 to f32
198+
// CHECK: %[[x22d:.*]] = arith.divf %[[x20d]], %[[x16d]] : f32
199+
// CHECK: %[[x26d:.*]] = arith.index_cast %[[index4]] : index to i64
200+
// CHECK: %[[x27d:.*]] = arith.sitofp %[[x26d]] : i64 to f32
201+
// CHECK: %[[x28d:.*]] = arith.divf %[[x27d]], %[[x22d]] : f32
202+
// CHECK: %[[floorD2:.*]] = math.floor %[[x28d]] : f32
203+
// CHECK: %[[c1D2:.*]] = arith.constant 1.000000e+00 : f32
204+
// CHECK: %[[c0D2:.*]] = arith.constant 0.000000e+00 : f32
205+
// CHECK: %[[inLm1D2:.*]] = arith.subf %[[x16d]], %[[c1D2]] : f32
206+
// CHECK: %[[clD2a:.*]] = arith.maximumf %[[floorD2]], %[[c0D2]] : f32
207+
// CHECK: %[[clD2b:.*]] = arith.minimumf %[[clD2a]], %[[inLm1D2]] : f32
208+
// CHECK: %[[Dfptosi:.*]] = arith.fptosi %[[clD2b]] : f32 to i64
209+
// CHECK: %[[x35:.*]] = arith.index_cast %[[Dfptosi]] : i64 to index
164210
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]], %[[x34]], %[[x35]]] : tensor<?x?x?x?x?xf32>
165211
// CHECK: linalg.yield %[[extracted]] : f32
166212
%none = torch.constant.none
@@ -200,11 +246,13 @@ func.func @test_resize_nearest_ceil(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !
200246
// CHECK: %[[add:.*]] = arith.addf %[[x24]], %[[cst]] : f32
201247
// CHECK: %[[x25:.*]] = arith.divf %[[add]], %[[x21]] : f32
202248
// CHECK: %[[sub:.*]] = arith.subf %[[x25]], %[[cst]] : f32
203-
// CHECK: %[[cst3:.*]] = arith.constant 1.000000e+00 : f32
204-
// CHECK: %[[nM1:.*]] = arith.subf %[[inputsizefp:.*]], %[[cst3]]
205249
// CHECK: %[[ceil:.*]] = math.ceil %[[sub]] : f32
206-
// CHECK: %[[minindex:.*]] = arith.minimumf %[[ceil]], %[[nM1]]
207-
// CHECK: %[[x31:.*]] = arith.fptosi %[[minindex]] : f32 to i64
250+
// CHECK: %[[c1ceil:.*]] = arith.constant 1.000000e+00 : f32
251+
// CHECK: %[[c0ceil:.*]] = arith.constant 0.000000e+00 : f32
252+
// CHECK: %[[inLm1ceil:.*]] = arith.subf %[[x15]], %[[c1ceil]] : f32
253+
// CHECK: %[[minCl1:.*]] = arith.maximumf %[[ceil]], %[[c0ceil]] : f32
254+
// CHECK: %[[minCl2:.*]] = arith.minimumf %[[minCl1]], %[[inLm1ceil]] : f32
255+
// CHECK: %[[x31:.*]] = arith.fptosi %[[minCl2]] : f32 to i64
208256
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
209257
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]]] : tensor<?x?x?xf32>
210258
// CHECK: linalg.yield %[[extracted]] : f32
@@ -286,7 +334,12 @@ func.func @test_resize_nearest_half_pixel_round_prefer_floor(%arg0: !torch.vtens
286334
// CHECK: %[[sub2:.*]] = arith.subf %[[sub]], %[[floor]] : f32
287335
// CHECK: %[[cmpf:.*]] = arith.cmpf ule, %[[sub2]], %[[cst3]] : f32
288336
// CHECK: %[[select:.*]] = arith.select %[[cmpf]], %[[floor]], %[[ceil]] : f32
289-
// CHECK: %[[x31:.*]] = arith.fptosi %[[select]] : f32 to i64
337+
// CHECK: %[[c1_clamp:.*]] = arith.constant 1.000000e+00 : f32
338+
// CHECK: %[[c0_clamp:.*]] = arith.constant 0.000000e+00 : f32
339+
// CHECK: %[[inLm1:.*]] = arith.subf %[[x15]], %[[c1_clamp]] : f32
340+
// CHECK: %[[selMax0:.*]] = arith.maximumf %[[select]], %[[c0_clamp]] : f32
341+
// CHECK: %[[selClamped:.*]] = arith.minimumf %[[selMax0]], %[[inLm1]] : f32
342+
// CHECK: %[[x31:.*]] = arith.fptosi %[[selClamped]] : f32 to i64
290343
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
291344
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]]] : tensor<?x?x?xf32>
292345
// CHECK: linalg.yield %[[extracted]] : f32

0 commit comments

Comments
 (0)