Simplify paddingMode lowering

Evaluate paddingMode at compile time
llvm · Nov 17, 2024 · 3395df8 · 3395df8
1 parent 3ae0446
commit 3395df8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -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>(

diff --git a/lib/Conversion/TorchToLinalg/Uncategorized.cpp b/lib/Conversion/TorchToLinalg/Uncategorized.cpp
@@ -2574,22 +2574,23 @@ class ConvertAtenGridSamplerOp : public OpConversionPattern<AtenGridSamplerOp> {
       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 @@ class ConvertAtenGridSamplerOp : public OpConversionPattern<AtenGridSamplerOp> {
           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>(