[onnx] Gridsampler addition of nearest mode (llvm#3320)

Added nearest neighbor selection for onnx.Gridsampler
kumardeepakamd · May 10, 2024 · adafd51 · adafd51
1 parent 4b24909
commit adafd51
Show file tree

Hide file tree

Showing 5 changed files with 169 additions and 64 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -123,12 +123,20 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
         if (gridShape[3] != 2)
           return rewriter.notifyMatchFailure(binder.op,
                                              "gridShape[3] expected to be 2");
-        std::string mode;
-        if (binder.customOpNameStringAttr(mode, "mode", "linear"))
+        std::string iModeString;
+        int64_t iModeInt;
+        if (binder.customOpNameStringAttr(iModeString, "mode", "linear"))
           return rewriter.notifyMatchFailure(binder.op, "mode bind failure");
-        if (mode != "linear" && mode != "bilinear")
+
+        if (iModeString == "linear" || iModeString == "bilinear") {
+          iModeInt = 0;
+        } else if (iModeString == "nearest") {
+          iModeInt = 1;
+        } else {
           return rewriter.notifyMatchFailure(
-              binder.op, "currently only mode : linear supported");
+              binder.op, "currently only mode : linear and nearest supported");
+        }
+
         std::string padding;
         if (binder.customOpNameStringAttr(padding, "padding_mode", "zeros"))
           return rewriter.notifyMatchFailure(binder.op,
@@ -143,7 +151,8 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
 
         Value interpolationMode = rewriter.create<Torch::ConstantIntOp>(
             binder.getLoc(), rewriter.getType<Torch::IntType>(),
-            rewriter.getIntegerAttr(rewriter.getIntegerType(64), 0));
+            rewriter.getIntegerAttr(rewriter.getIntegerType(64), iModeInt));
+
         Value paddingMode = rewriter.create<Torch::ConstantIntOp>(
             binder.getLoc(), rewriter.getType<Torch::IntType>(),
             rewriter.getIntegerAttr(rewriter.getIntegerType(64), 0));

diff --git a/lib/Conversion/TorchToLinalg/Uncategorized.cpp b/lib/Conversion/TorchToLinalg/Uncategorized.cpp
@@ -2524,14 +2524,38 @@ class ConvertAtenGridSamplerOp : public OpConversionPattern<AtenGridSamplerOp> {
       Value result = b.create<tensor::ExtractOp>(loc, input, index);
       return result;
     };
-    auto lambdaInter = [&](OpBuilder &b, Location loc, Value x, Value y,
-                           Value d) -> Value {
+
+    auto lambdaLinear = [&](OpBuilder &b, Location loc, Value x, Value y,
+                            Value d) -> Value {
       Value dm = b.create<arith::SubFOp>(loc, oneFloat, d);
       Value ra = b.create<arith::MulFOp>(loc, x, dm);
       Value rb = b.create<arith::MulFOp>(loc, y, d);
       Value res = b.create<arith::AddFOp>(loc, ra, rb);
       return res;
     };
+
+    auto lambdaNearest = [&](OpBuilder &b, Location loc, Value x, Value y,
+                             Value d) -> Value {
+      Value halfConst = rewriter.create<arith::ConstantOp>(
+          loc, rewriter.getFloatAttr(floatType, 0.5));
+      Value checkClosest =
+          b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, d, halfConst);
+      Value res = b.create<arith::SelectOp>(loc, checkClosest, x, y);
+      return res;
+    };
+
+    auto lambdaInterpolate = [&](OpBuilder &b, Location loc, Value iMode,
+                                 Value x, Value y, Value d) -> Value {
+      Value linear = lambdaLinear(b, loc, x, y, d);
+      Value nearest = lambdaNearest(b, loc, x, y, d);
+      Value zeroInt =
+          b.create<arith::ConstantOp>(loc, b.getIntegerAttr(int64type, 0));
+      Value checkMode = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+                                                iMode, zeroInt);
+      Value res = b.create<arith::SelectOp>(loc, checkMode, linear, nearest);
+      return res;
+    };
+
     auto resultType = getTypeConverter()
                           ->convertType(op.getResult().getType())
                           .cast<RankedTensorType>();
@@ -2545,6 +2569,7 @@ class ConvertAtenGridSamplerOp : public OpConversionPattern<AtenGridSamplerOp> {
     if (resultType.isDynamicDim(3))
       resultSize.push_back(rewriter.create<tensor::DimOp>(loc, grid, 2));
     Value alignCorners = adaptor.getAlignCorners();
+    Value interMode = adaptor.getInterpolationMode();
     Value resultFinal =
         rewriter.create<tensor::EmptyOp>(loc, resultType, resultSize);
     auto sGrid = rewriter.create<linalg::GenericOp>(
@@ -2633,10 +2658,12 @@ class ConvertAtenGridSamplerOp : public OpConversionPattern<AtenGridSamplerOp> {
           Value lw1a = b.create<arith::SIToFPOp>(loc, floatType, lower1);
           Value d1 = b.create<arith::SubFOp>(loc, result0, lw0a);
           Value d0 = b.create<arith::SubFOp>(loc, result1, lw1a);
-          Value resultScaled0 = lambdaInter(b, loc, result00b, result01b, d0);
-          Value resultScaled1 = lambdaInter(b, loc, result10b, result11b, d0);
-          Value resultScaled =
-              lambdaInter(b, loc, resultScaled0, resultScaled1, d1);
+          Value resultScaled0 =
+              lambdaInterpolate(b, loc, interMode, result00b, result01b, d0);
+          Value resultScaled1 =
+              lambdaInterpolate(b, loc, interMode, result10b, result11b, d0);
+          Value resultScaled = lambdaInterpolate(
+              b, loc, interMode, resultScaled0, resultScaled1, d1);
           b.create<linalg::YieldOp>(loc, resultScaled);
         });
     rewriter.replaceOp(op, sGrid.getResults());

diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/gridsampler.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/gridsampler.py
@@ -115,3 +115,41 @@ def GridSamplerBasic3_basic(module, tu: TestUtils):
         [[[[-0.3498, -0.8196], [-0.2127, 0.2138], [-0.6515, -0.0513]]]]
     ).type(torch.FloatTensor)
     module.forward(inp, grd)
+
+
+class GridSamplerBasic4(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [None, ([1, 1, 4, 4], torch.float32, True), ([1, 1, 3, 2], torch.float32, True)]
+    )
+    def forward(self, x, g):
+        interpolation_mode = (1,)
+        padding_mode = (0,)
+        align_corners = (False,)
+        tRes = torch.ops.aten.grid_sampler(
+            x, g, interpolation_mode[0], padding_mode[0], align_corners[0]
+        )
+        return tRes
+
+
+@register_test_case(module_factory=lambda: GridSamplerBasic4())
+def GridSamplerBasic4_basic(module, tu: TestUtils):
+    inp = torch.tensor(
+        [
+            [
+                [
+                    [0.4963, 0.7682, 0.0885, 0.1320],
+                    [0.3074, 0.6341, 0.4901, 0.8964],
+                    [0.4556, 0.6323, 0.3489, 0.4017],
+                    [0.0223, 0.1689, 0.2939, 0.5185],
+                ]
+            ]
+        ]
+    ).type(torch.FloatTensor)
+    grd = torch.tensor(
+        [[[[-0.3498, -0.8196], [-0.2127, 0.2138], [-0.6515, -0.0513]]]]
+    ).type(torch.FloatTensor)
+    module.forward(inp, grd)
diff --git a/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir b/test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
@@ -545,7 +545,7 @@ func.func @test_gelu_tanh_2(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtenso
 // CHECK: %[[B0:.*]] = torch.constant.bool false
 // CHECK: %[[A0:.*]] = torch.aten.grid_sampler %arg0, %arg1, %[[INT0]], %[[INT0_0]], %[[B0]] : !torch.vtensor<[5,10,10,4],f32>
 func.func @test_grid_sampler01(%arg0: !torch.vtensor<[5,10,10,4],f32>, %arg1: !torch.vtensor<[5,7,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
-  %0 = torch.operator "onnx.GridSample" (%arg0, %arg1) {torch.onnx.align_corners = 0 : si64, torch.onnx.padding_mode = "zeros"} : (!torch.vtensor<[5,10,10,4],f32>, !torch.vtensor<[5,7,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32>
+  %0 = torch.operator "onnx.GridSample" (%arg0, %arg1) {torch.onnx.align_corners = 0 : si64, torch.onnx.mode = "linear", torch.onnx.padding_mode = "zeros"} : (!torch.vtensor<[5,10,10,4],f32>, !torch.vtensor<[5,7,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32>
   return %0 : !torch.vtensor<[?,?,?,?],f32>
 }
 
@@ -563,6 +563,18 @@ func.func @test_grid_sampler02(%arg0: !torch.vtensor<[5,10,10,4],f32>, %arg1: !t
 
 // -----
 
+// CHECK-LABEL: @test_grid_sampler03
+// CHECK: %[[INT1:.*]] = torch.constant.int 1
+// CHECK: %[[INT0:.*]] = torch.constant.int 0
+// CHECK: %[[B0:.*]] = torch.constant.bool true
+// CHECK: %[[A0:.*]] = torch.aten.grid_sampler %arg0, %arg1, %[[INT1]], %[[INT0]], %[[B0]] : !torch.vtensor<[5,10,10,4],f32>
+func.func @test_grid_sampler03(%arg0: !torch.vtensor<[5,10,10,4],f32>, %arg1: !torch.vtensor<[5,7,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
+  %0 = torch.operator "onnx.GridSample" (%arg0, %arg1) {torch.onnx.align_corners = 1 : si64, torch.onnx.mode = "nearest", torch.onnx.padding_mode = "zeros"} : (!torch.vtensor<[5,10,10,4],f32>, !torch.vtensor<[5,7,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32>
+  return %0 : !torch.vtensor<[?,?,?,?],f32>
+}
+
+// -----
+
 // CHECK-LABEL: func.func @test_less_or_equal
 func.func @test_less_or_equal(%arg0: !torch.vtensor<[3,4,5],f32>, %arg1: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],i1> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 16 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
   // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[3,4,5],f32>

diff --git a/test/Conversion/TorchToLinalg/gridsampler.mlir b/test/Conversion/TorchToLinalg/gridsampler.mlir
@@ -2,27 +2,32 @@
 
 // CHECK: #map
 // CHECK-LABEL: func @grid_sampler
-// CHECK-DAG:  %[[TC0:.*]] = torch_c.to_builtin_tensor %[[ARG0:.*]] : !torch.vtensor<[4,10,10,4],f32> -> tensor<4x10x10x4xf32>
-// CHECK-DAG:  %[[TC1:.*]] = torch_c.to_builtin_tensor %[[ARG1:.*]] : !torch.vtensor<[4,6,8,2],f32> -> tensor<4x6x8x2xf32>
-// CHECK-DAG:  %[[FALSE:.*]] = torch.constant.bool false
-// CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
-// CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
-// CHECK-DAG:  %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG:  %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK-DAG:  %[[CST1:.*]] = arith.constant 1.000000e+00 : f32
-// CHECK-DAG:  %[[CST2:.*]] = arith.constant 2.000000e+00 : f32
-// CHECK-DAG:  %[[C2_3:.*]] = arith.constant 2 : index
-// CHECK-DAG:  %[[DIM:.*]] = tensor.dim %[[TC0]], %[[C2_3]] : tensor<4x10x10x4xf32>
-// CHECK-DAG:  %[[C3:.*]] = arith.constant 3 : index
-// CHECK-DAG:  %[[DIM_4:.*]] = tensor.dim %[[TC0]], %[[C3]] : tensor<4x10x10x4xf32>
-// CHECK-DAG:  %[[X2:.*]] = arith.subi %[[DIM:.*]], %[[C1]] : index
-// CHECK-DAG:  %[[X3:.*]] = arith.subi %[[DIM_4]], %[[C1:.*]] : index
-// CHECK-DAG:  %[[X4:.*]] = arith.index_cast %[[X2]] : index to i64
-// CHECK-DAG:  %[[X5:.*]] = arith.index_cast %[[X3]] : index to i64
-// CHECK-DAG:  %[[X6:.*]] = arith.sitofp %[[X4]] : i64 to f32
-// CHECK-DAG:  %[[X7:.*]] = arith.sitofp %[[X5]] : i64 to f32
-// CHECK-DAG:  %[[X8:.*]] = arith.divf %[[X6]], %[[CST2]] : f32
-// CHECK-DAG:  %[[X9:.*]] = arith.divf %[[X7]], %[[CST2]] : f32
+// CHECK-DAG: %[[TC0:.*]] = torch_c.to_builtin_tensor %[[ARG0:.*]] : !torch.vtensor<[4,10,10,4],f32> -> tensor<4x10x10x4xf32>
+// CHECK-DAG: %[[TC1:.*]] = torch_c.to_builtin_tensor %[[ARG1:.*]] : !torch.vtensor<[4,6,8,2],f32> -> tensor<4x6x8x2xf32>
+// CHECK-DAG: %[[FALSE:.*]] = torch.constant.bool false
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[CST1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[CST2:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK-DAG: %[[C2_3:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[DIM:.*]] = tensor.dim %[[TC0]], %[[C2_3]] : tensor<4x10x10x4xf32>
+// CHECK-DAG: %[[X73:.*]] = arith.cmpi eq, %[[X3:.*]], %[[C27:.*]] : i64
+// CHECK-DAG: %[[X74:.*]] = arith.select %[[X73:.*]], %[[X70:.*]], %[[X72:.*]] : f32
+// CHECK-DAG: %[[X75:.*]] = arith.subf %[[Xcst_1:.*]], %[[X57:.*]] : f32
+// CHECK-DAG: %[[X76:.*]] = arith.mulf %[[X66:.*]], %[[X75:.*]] : f32
+// CHECK-DAG: %[[X77:.*]] = arith.mulf %[[X74:.*]], %[[X57:.*]] : f32
+// CHECK-DAG: %[[X78:.*]] = arith.addf %[[X76:.*]], %[[X77:.*]] : f32
+// CHECK-DAG: %[[C28:.*]] = arith.constant 5.000000e-01 : f32
+// CHECK-DAG: %[[X79:.*]] = arith.cmpf olt, %[[X57:.*]], %[[X28:.*]] : f32
+// CHECK-DAG: %[[X80:.*]] = arith.select %[[X79:.*]], %[[X66:.*]], %[[X74:.*]] : f32
+// CHECK-DAG: %[[C29:.*]] = arith.constant 0 : i64
+// CHECK-DAG: %[[X81:.*]] = arith.cmpi eq, %[[X3:.*]], %[[C29:.*]] : i64
+// CHECK-DAG: %[[X82:.*]] = arith.select %[[X81:.*]], %[[X78:.*]], %[[X80:.*]] : f32
+// CHECK-DAG: linalg.yield %[[X82:.*]] : f32
+// CHECK-DAG: %[[X14:.*]] = torch_c.from_builtin_tensor %[[X13:.*]] : tensor<?x?x?x?xf32> -> !torch.vtensor<[?,?,?,?],f32>
+
 func.func @grid_sampler(%arg0: !torch.vtensor<[4,10,10,4],f32>, %arg1: !torch.vtensor<[4,6,8,2],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
   %true = torch.constant.bool 0
   %int0 = torch.constant.int 0
@@ -35,22 +40,25 @@ func.func @grid_sampler(%arg0: !torch.vtensor<[4,10,10,4],f32>, %arg1: !torch.vt
 
 // CHECK-LABEL: func @grid_sampler2
 // CHECK: #map
-// CHECK-DAG:  %[[X15:.*]] = arith.mulf %[[X13:.*]], %[[X8:.*]] : f32
-// CHECK-DAG:  %[[X16:.*]] = arith.mulf %[[X14:.*]], %[[X9:.*]] : f32
-// CHECK-DAG:  %[[X40:.*]] = arith.mulf %[[EXTRACTED:.*]], %[[X39:.*]] : f32
-// CHECK-DAG:  %[[X41:.*]] = arith.mulf %[[X31:.*]], %[[X37:.*]] : f32
-// CHECK-DAG:  %[[X42:.*]] = arith.addf %[[X40:.*]], %[[X41]] : f32
-// CHECK-DAG:   %[[X43:.*]] = arith.subf %[[CST_1:.*]], %[[X37]] : f32
-// CHECK-DAG:   %[[X45:.*]] = arith.mulf %[[X34:.*]], %[[X37]] : f32
-// CHECK-DAG:   %[[X46:.*]] = arith.addf %[[X44:.*]], %[[X45]] : f32
-// CHECK-DAG:   %[[X47:.*]] = arith.subf %[[CST_1]], %[[X38:.*]] : f32
-// CHECK-DAG:   %[[X48:.*]] = arith.mulf %[[X42]], %[[XX47:.*]] : f32
-// CHECK-DAG:   %[[X49:.*]] = arith.mulf %[[X46]], %[[XX38:.*]] : f32
-// CHECK-DAG:   %[[X50:.*]] = arith.addf %[[X48]], %[[X49]] : f32
-// CHECK-DAG:   linalg.yield %[[X50]] : f32
-// CHECK: } -> tensor<?x?x?x?xf32>
-// CHECK: %[[X12:.*]] = torch_c.from_builtin_tensor %[[X11:.*]] : tensor<?x?x?x?xf32> -> !torch.vtensor<[?,?,?,?],f32>
-// CHECK: return %[[X12]] : !torch.vtensor<[?,?,?,?],f32>
+// CHECK-DAG: %[[X70:.*]] = arith.addf %[[X68:.*]], %[[X69:.*]] : f32
+// CHECK-DAG: %[[X29:.*]] = arith.constant 5.000000e-01 : f32
+// CHECK-DAG: %[[X71:.*]] = arith.cmpf olt, %[[X58:.*]], %[[X29:.*]] : f32
+// CHECK-DAG: %[[X72:.*]] = arith.select %[[X71:.*]], %[[X52:.*]], %[[X54:.*]] : f32
+// CHECK-DAG: %[[X30:.*]] = arith.constant 0 : i64
+// CHECK-DAG: %[[X73:.*]] = arith.cmpi eq, %[[X3:.*]], %[[X30:.*]] : i64
+// CHECK-DAG: %[[X74:.*]] = arith.select %[[X73:.*]], %[[X70:.*]], %[[X72:.*]] : f32
+// CHECK-DAG: %[[X75:.*]] = arith.subf %[[X1:.*]], %[[X57:.*]] : f32
+// CHECK-DAG: %[[X76:.*]] = arith.mulf %[[X66:.*]], %[[X75:.*]] : f32
+// CHECK-DAG: %[[X77:.*]] = arith.mulf %[[X74:.*]], %[[X57:.*]] : f32
+// CHECK-DAG: %[[X78:.*]] = arith.addf %[[X76:.*]], %[[X77:.*]] : f32
+// CHECK-DAG: %[[X31:.*]] = arith.constant 5.000000e-01 : f32
+// CHECK-DAG: %[[X79:.*]] = arith.cmpf olt, %[[X57:.*]], %[[X31:.*]] : f32
+// CHECK-DAG: %[[X80:.*]] = arith.select %[[X79:.*]], %[[X66:.*]], %[[X74:.*]] : f32
+// CHECK-DAG: %[[X32:.*]] = arith.constant 0 : i64
+// CHECK-DAG: %[[X81:.*]] = arith.cmpi eq, %[[X3:.*]], %[[X32:.*]] : i64
+// CHECK-DAG: %[[X82:.*]] = arith.select %[[X81:.*]], %[[X78:.*]], %[[X80:.*]] : f32
+// CHECK-DAG: linalg.yield %[[X50:.*]] : f32
+// CHECK: return %[[X12:.*]] : !torch.vtensor<[?,?,?,?],f32>
 func.func @grid_sampler2(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
   %true = torch.constant.bool 0
   %int0 = torch.constant.int 0
@@ -64,25 +72,36 @@ func.func @grid_sampler2(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg1: !torch.vte
 // CHECK-LABEL: func @grid_sampler3
 // CHECK: #map
 // CHECK-DAG:  %[[X15:.*]] = arith.mulf %[[X13:.*]], %[[X8:.*]] : f32
-// CHECK-DAG:  %[[X16:.*]] = arith.mulf %[[X14:.*]], %[[X9:.*]] : f32
-// CHECK-DAG:  %[[X40:.*]] = arith.mulf %[[EXTRACTED:.*]], %[[X39:.*]] : f32
-// CHECK-DAG:  %[[X41:.*]] = arith.mulf %[[X31:.*]], %[[X37:.*]] : f32
-// CHECK-DAG:  %[[X42:.*]] = arith.addf %[[X40:.*]], %[[X41]] : f32
-// CHECK-DAG:   %[[X43:.*]] = arith.subf %[[CST_1:.*]], %[[X37]] : f32
-// CHECK-DAG:   %[[X45:.*]] = arith.mulf %[[X34:.*]], %[[X37]] : f32
-// CHECK-DAG:   %[[X46:.*]] = arith.addf %[[X44:.*]], %[[X45]] : f32
-// CHECK-DAG:   %[[X47:.*]] = arith.subf %[[CST_1]], %[[X38:.*]] : f32
-// CHECK-DAG:   %[[X48:.*]] = arith.mulf %[[X42]], %[[XX47:.*]] : f32
-// CHECK-DAG:   %[[X49:.*]] = arith.mulf %[[X46]], %[[XX38:.*]] : f32
-// CHECK-DAG:   %[[X50:.*]] = arith.addf %[[X48]], %[[X49]] : f32
-// CHECK-DAG:   linalg.yield %[[X50]] : f32
-// CHECK: } -> tensor<?x?x?x?xf32>
-// CHECK: %[[X12:.*]] = torch_c.from_builtin_tensor %[[X11:.*]] : tensor<?x?x?x?xf32> -> !torch.vtensor<[?,?,?,?],f32>
-// CHECK: return %[[X12]] : !torch.vtensor<[?,?,?,?],f32>
+// CHECK-DAG:      %[[Y60:.*]] = arith.mulf %[[X48:.*]], %[[X59:.*]] : f32
+// CHECK-DAG:      %[[Y61:.*]] = arith.mulf %[[X50:.*]], %[[X58:.*]] : f32
+// CHECK-DAG:      %[[Y62:.*]] = arith.addf %[[X60:.*]], %[[X61:.*]] : f32
+// CHECK-DAG:      %[[Y28:.*]] = arith.constant 5.000000e-01 : f32
+// CHECK-DAG:      %[[Y64:.*]] = arith.select %[[X63:.*]], %[[X48:.*]], %[[X50:.*]] : f32
+// CHECK-DAG:      %[[Y29:.*]] = arith.constant 0 : i6
+// CHECK-DAG:      %[[Y65:.*]] = arith.cmpi eq, %[[X3:.*]], %[[X28:.*]] : i64
+// CHECK-DAG:      %[[Y66:.*]] = arith.select %[[X65:.*]], %[[X62:.*]], %[[X64:.*]] : f32
+// CHECK-DAG:      %[[Y67:.*]] = arith.subf %[[X1:.*]], %[[X58:.*]] : f32
+// CHECK-DAG:      %[[Y68:.*]] = arith.mulf %[[X52:.*]], %[[X67:.*]] : f32
+// CHECK-DAG:      %[[Y69:.*]] = arith.mulf %[[X54:.*]], %[[X58:.*]] : f32
+// CHECK-DAG:      %[[Y70:.*]] = arith.addf %[[X68:.*]], %[[X69:.*]] : f32
+// CHECK-DAG:      %[[Y30:.*]] = arith.constant 5.000000e-01 : f32
+// CHECK-DAG:      %[[Y31:.*]] = arith.constant 0 : i64
+// CHECK: return %[[X12:.*]] : !torch.vtensor<[?,?,?,?],f32>
 func.func @grid_sampler3(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
   %false = torch.constant.bool 1
   %int0 = torch.constant.int 0
   %int1 = torch.constant.int 0
   %4 = torch.aten.grid_sampler %arg0, %arg1, %int0, %int1, %false : !torch.vtensor<[?,?,?,?],f32>, !torch.vtensor<[?,?,?,?],f32>, !torch.int, !torch.int, !torch.bool -> !torch.vtensor<[?,?,?,?],f32>
   return %4 : !torch.vtensor<[?,?,?,?],f32>
 }
+
+// -----
+
+// CHECK-LABEL: func @grid_sampler4
+func.func @grid_sampler4(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
+  %false = torch.constant.bool 1
+  %int0 = torch.constant.int 0
+  %int1 = torch.constant.int 1
+  %4 = torch.aten.grid_sampler %arg0, %arg1, %int0, %int1, %false : !torch.vtensor<[?,?,?,?],f32>, !torch.vtensor<[?,?,?,?],f32>, !torch.int, !torch.int, !torch.bool -> !torch.vtensor<[?,?,?,?],f32>
+  return %4 : !torch.vtensor<[?,?,?,?],f32>
+}