From 3395df830c1b54db197a917bec2286d5fb0edda5 Mon Sep 17 00:00:00 2001
From: Atri Sarkar <atrisarkar91@gmail.com>
Date: Sun, 17 Nov 2024 21:53:29 +0530
Subject: [PATCH] Simplify paddingMode lowering

Evaluate paddingMode at compile time
---
 .../TorchOnnxToTorch/DefaultDomainGtoP.cpp    |  4 +---
 .../TorchToLinalg/Uncategorized.cpp           | 21 ++++++++++---------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
index b4324a7f7a39..55ddc65b3a03 100644
--- 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
index 5ecf58947de7..957937180236 100644
--- 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>(