diff --git a/mlir-tensorrt/compiler/lib/Conversion/StablehloToTensorRT/StablehloToTensorRT.cpp b/mlir-tensorrt/compiler/lib/Conversion/StablehloToTensorRT/StablehloToTensorRT.cpp index 1476e0100..63e9ee322 100644 --- a/mlir-tensorrt/compiler/lib/Conversion/StablehloToTensorRT/StablehloToTensorRT.cpp +++ b/mlir-tensorrt/compiler/lib/Conversion/StablehloToTensorRT/StablehloToTensorRT.cpp @@ -2248,7 +2248,7 @@ struct PadConverter : public ConvertHloOpToTensorRTPattern { rewriter.replaceOpWithNewOp( op, /*input=*/adaptor.getOperand(), - /*offset=*/*edgePaddingLow, + /*offset=*/sliceOffset, /*size=*/size, /*stride=*/stride, /*slice_mode=*/tensorrt::SliceMode::kFILL, diff --git a/mlir-tensorrt/test/Conversion/StablehloToTensorRT/stablehlo-to-tensorrt.mlir b/mlir-tensorrt/test/Conversion/StablehloToTensorRT/stablehlo-to-tensorrt.mlir index 35d262b79..6dd22dbaf 100644 --- a/mlir-tensorrt/test/Conversion/StablehloToTensorRT/stablehlo-to-tensorrt.mlir +++ b/mlir-tensorrt/test/Conversion/StablehloToTensorRT/stablehlo-to-tensorrt.mlir @@ -1124,6 +1124,23 @@ func.func @hlo_pad_static(%arg0: tensor<10x48x48x32xf32>) -> tensor<10x48x48x48x // ----- +func.func @hlo_pad_static_low_high(%arg0: tensor<10x48x48x32xf32>) -> tensor<10x48x48x64xf32> { + %0 = "stablehlo.constant"() {value = dense<0.0> : tensor} : () -> tensor + %1 = "stablehlo.pad"(%arg0, %0) { + edge_padding_high = array, + edge_padding_low = array, + interior_padding = array + } : (tensor<10x48x48x32xf32>, tensor) -> tensor<10x48x48x64xf32> + func.return %1 : tensor<10x48x48x64xf32> +} + +// CHECK-LABEL: @hlo_pad_static +// CHECK-SAME: (%[[arg0:.+]]: tensor<10x48x48x32xf32> +// CHECK: %[[fill:.+]] = tensorrt.constant dense<0.0{{.*}}> : tensor +// CHECK: tensorrt.slice %[[arg0]][0, 0, 0, -16][10, 48, 48, 64][1, 1, 1, 1] fill(%[[fill]] : tensor) {mode = #tensorrt.slice_mode} : tensor<10x48x48x32xf32> to tensor<10x48x48x64xf32> + +// ----- + func.func @hlo_pad_dynamic_non_sliced_dim(%arg0: tensor) -> tensor { %0 = "stablehlo.constant"() {value = dense<0.0> : tensor} : () -> tensor %1 = "stablehlo.pad"(%arg0, %0) {