llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Simplify paddingMode lowering 

Evaluate paddingMode at compile time
pull/3819/head
Atri Sarkar 2024-11-17 21:53:29 +05:30
parent 3ae0446aa2
commit 3395df830c
2 changed files with 12 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
View File

@ -163,9 +163,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
rewriter.getIntegerAttr(rewriter.getIntegerType(64), iModeInt));
Value paddingMode = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getType<Torch::IntType>(),
rewriter.getIntegerAttr(rewriter.getIntegerType(64),
paddingModeInt));
binder.getLoc(), paddingModeInt);
bool alignMode = align;
Value alignCorners = rewriter.create<Torch::ConstantBoolOp>(

21
lib/Conversion/TorchToLinalg/Uncategorized.cpp
View File

@ -2574,22 +2574,23 @@ public:
return b.create<arith::MinimumFOp>(loc, xMaxZero, SizeSubOne);
};
auto lambdaPadding = [&](OpBuilder &b, Location loc, Value paddingMode,
auto lambdaPadding = [&](OpBuilder &b, Location loc, int64_t paddingMode,
Value x, Value SizeSubOne) -> Value {
Value border = lambdaBorder(b, loc, x, SizeSubOne);
Value zeroInt =
b.create<arith::ConstantOp>(loc, b.getIntegerAttr(int64type, 0));
Value isZero = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
paddingMode, zeroInt);
// Border
if (paddingMode == 1) {
return lambdaBorder(b, loc, x, SizeSubOne);
}
return b.create<arith::SelectOp>(loc, isZero, x, border);
return x;
};
auto resultType = cast<RankedTensorType>(
getTypeConverter()->convertType(op.getResult().getType()));
Value alignCorners = adaptor.getAlignCorners();
Value interMode = adaptor.getInterpolationMode();
Value paddingMode = adaptor.getPaddingMode();
int64_t paddingModeInt;
matchPattern(op.getPaddingMode(), m_TorchConstantInt(&paddingModeInt));
SmallVector<Value> dynamicSizes{};
if (resultType.isDynamicDim(0))
@ -2623,9 +2624,9 @@ public:
Value unnorm1 =
b.create<arith::AddFOp>(loc, gPlusMul1, gr1HalfSelect);
Value result0 =
lambdaPadding(b, loc, paddingMode, unnorm0, innerDim0d);
lambdaPadding(b, loc, paddingModeInt, unnorm0, innerDim0d);
Value result1 =
lambdaPadding(b, loc, paddingMode, unnorm1, innerDim1d);
lambdaPadding(b, loc, paddingModeInt, unnorm1, innerDim1d);
Value checkLowerBound0 = b.create<arith::CmpFOp>(
loc, arith::CmpFPredicate::OLT, result0, zeroFloat);
Value checkLowerBound1 = b.create<arith::CmpFOp>(