feat: add folder for onnx.Slice.

Author: ttjost

src/Dialect/ONNX/ONNXOps.td.inc

@@ -8946,6 +8946,7 @@ def ONNXSliceOp:ONNX_Op<"Slice",
       return sh;
     }
   }];
+  let hasFolder = 1;
 }
 
 def ONNXSoftmaxOp:ONNX_Op<"Softmax",

src/Dialect/ONNX/ONNXOps/Tensor/Slice.cpp

@@ -201,3 +201,16 @@ LogicalResult ONNXSliceOp::inferShapes(
   ONNXSliceOpShapeHelper shapeHelper(getOperation(), {});
   return shapeHelper.computeShapeAndUpdateType(elementType);
 }
+
+//===----------------------------------------------------------------------===//
+// Folder
+//===----------------------------------------------------------------------===//
+OpFoldResult ONNXSliceOp::fold(FoldAdaptor adaptor) {
+
+  auto inputTy = llvm::dyn_cast<RankedTensorType>(getData().getType());
+  auto outputTy = llvm::dyn_cast<RankedTensorType>(getOutput().getType());
+  if (inputTy && inputTy == outputTy && inputTy.hasStaticShape()) {
+    return getData();
+  }
+  return nullptr;
+}

test/mlir/onnx/onnx_fold.mlir

@@ -43,4 +43,20 @@ func.func @test_reduceMeanIsNoopWithEmptyAxes(%arg0: tensor<4x512x256x8xf32>) ->
 // CHECK-LABEL: @test_reduceMeanIsNoopWithEmptyAxes
 // CHECK-SAME: (%[[VAL_0:.*]]: tensor<4x512x256x8xf32>) -> tensor<4x512x256x8xf32> {
 // CHECK:   return %[[VAL_0]] : tensor<4x512x256x8xf32>
+// CHECK: }
+
+// -----
+
+func.func @test_slice(%arg0: tensor<16x1x2500x4xf32>) -> tensor<16x1x2500x4xf32> {
+  %0 = onnx.Constant dense<0> : tensor<1xi64>
+  %1 = onnx.Constant dense<4> : tensor<1xi64>
+  %2 = onnx.Constant dense<3> : tensor<1xi64>
+  %3 = onnx.Constant dense<1> : tensor<1xi64>
+  %4 = "onnx.Slice"(%arg0, %0, %1, %2, %3) : (tensor<16x1x2500x4xf32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<16x1x2500x4xf32>
+  return %4 : tensor<16x1x2500x4xf32>
+}
+
+// CHECK-LABEL: @test_slice
+// CHECK-SAME: (%[[VAL_0:.*]]: tensor<16x1x2500x4xf32>) -> tensor<16x1x2500x4xf32> {
+// CHECK:   return %[[VAL_0]] : tensor<16x1x2500x4xf32>
 // CHECK: }

utils/gen_onnx_mlir.py

@@ -447,7 +447,7 @@
 ]
 
 # Op with fold function
-OpsWithFolder = ["Constant", "Squeeze", "SqueezeV11", "ReduceMean"]
+OpsWithFolder = ["Constant", "Squeeze", "SqueezeV11", "ReduceMean", "Slice"]
 
 # Op with ConstantLike trait
 OpsWithConstantLike = ["Constant"]