Add onnx GridSample support for border padding mode #3819
Conversation
| Value zeroInt = | ||
| b.create<arith::ConstantOp>(loc, b.getIntegerAttr(int64type, 0)); |
There was a problem hiding this comment.
It might be better to keep this constant outside the lambda function, and just get it from the capture (otherwise you will re-write the constant multiple times in the IR).
There was a problem hiding this comment.
Since paddingMode is checked at compile time, I have removed zeroInt completely.
| Value isZero = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, | ||
| paddingMode, zeroInt); |
There was a problem hiding this comment.
The padding mode is always a constant (coming from ONNX, at least). We should be able to get the paddingModeInt by matchPattern(paddingMode, m_TorchConstantInt(&paddingModeInt)). This way, we can avoid writing IR to check conditions that can be determined statically at compile time. Since the select op and cmp op will likely get folded anyway, this is not hugely problematic, but on principal I'd prefer to generate the simplest possible IR from the get-go.
There was a problem hiding this comment.
Yes, I'm evaluating it at compile time now.
| auto lambdaBorder = [&](OpBuilder &b, Location loc, Value x, | ||
| Value SizeSubOne) -> Value { | ||
| Value xMaxZero = b.create<arith::MaximumFOp>(loc, x, zeroFloat); | ||
| return b.create<arith::MinimumFOp>(loc, xMaxZero, SizeSubOne); | ||
| }; | ||
|
|
||
| auto lambdaPadding = [&](OpBuilder &b, Location loc, Value paddingMode, | ||
| Value x, Value SizeSubOne) -> Value { |
There was a problem hiding this comment.
x and sizeSubOne should be the only inputs for both of these lambda functions. The others can be gotten through the default capture. Just be sure to define these functions after you introduce the other arguments.
There was a problem hiding this comment.
If we want to capture the other arguments, the lambdas will have to be defined inside the GenericOp lambda later in the code, but I would want to avoid that as it would clutter the implementation there. There are a few other helper lambdas such as lambdaInterpolate and lambdaExtract which take arguments like this, so I decided to make it similar to those.
|
Thanks for the feedback! I'll work on the requested changes. |
Evaluate paddingMode at compile time
Adds support for padding_mode="border"
Clamps the grid coordinates between 0 and size - 1 like this when using this padding mode: