Add support for split (#2201)

* Add support for split

Signed-off-by: philass <[email protected]>

* Add lit tests

Signed-off-by: philass <[email protected]>

---------

Signed-off-by: philass <[email protected]>

Author: philass

src/Builder/OpBuildTable.inc

@@ -188,7 +188,7 @@ op_dialect_version_map_["SoftmaxCrossEntropyLoss"] = {13};
 op_dialect_version_map_["Softplus"] = {1};
 op_dialect_version_map_["Softsign"] = {1};
 op_dialect_version_map_["SpaceToDepth"] = {13};
-op_dialect_version_map_["Split"] = {13, 11};
+op_dialect_version_map_["Split"] = {18, 13, 11};
 op_dialect_version_map_["SplitToSequence"] = {11};
 op_dialect_version_map_["Sqrt"] = {13};
 op_dialect_version_map_["Squeeze"] = {13, 11};
@@ -586,6 +586,8 @@ import_handler_map_["SpaceToDepth"] =
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXSpaceToDepthOp>;
 import_handler_map_["Split"] = 
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXSplitOp>;
+import_handler_map_["SplitV13"] = 
+   &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXSplitV13Op>;
 import_handler_map_["SplitV11"] = 
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXSplitV11Op>;
 import_handler_map_["SplitToSequence"] = 

src/Dialect/ONNX/DialectBuilder.cpp

@@ -255,8 +255,8 @@ ValueRange OnnxBuilder::split(
   IntegerAttr axisAttr =
       IntegerAttr::get(b().getIntegerType(64, /*isSigned=*/true),
           APInt(64, axis, /*isSigned=*/true));
-  return createOpAndInferShapes<ONNXSplitOp>(
-      toTensors(outputTypes), toTensor(input), toTensor(split), axisAttr)
+  return createOpAndInferShapes<ONNXSplitOp>(toTensors(outputTypes),
+      toTensor(input), toTensor(split), axisAttr, IntegerAttr())
       .getResults();
 }
 

src/Dialect/ONNX/ONNXOps.td.inc

@@ -8556,6 +8556,53 @@ def ONNXSpaceToDepthOp:ONNX_Op<"SpaceToDepth",
 }
 
 def ONNXSplitOp:ONNX_Op<"Split",
+  [Pure, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>, DeclareOpInterfaceMethods<ShapeHelperOpInterface>]> {
+  let summary = "ONNX Split operation";
+  let description = [{
+  Split a tensor into a list of tensors, along the specified 'axis'.
+  Either input 'split' or the attribute 'num_outputs' should be specified, but not both.
+  If the attribute 'num_outputs' is specified, then the tensor is split into equal sized parts.
+  If the tensor is not evenly splittable into `num_outputs`, the last chunk will be smaller.
+  If the input 'split' is specified, it indicates the sizes of each output in the split.
+  }];
+  let arguments = (ins AnyTypeOf<[TensorOf<[UI8]>, TensorOf<[UI16]>, TensorOf<[UI32]>, TensorOf<[UI64]>, TensorOf<[I8]>, TensorOf<[I16]>, TensorOf<[I32]>, TensorOf<[I64]>, TensorOf<[BF16]>, TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>, TensorOf<[StringType]>, TensorOf<[I1]>, TensorOf<[Complex<F32>]>, TensorOf<[Complex<F64>]>]>:$input,
+    AnyTypeOf<[TensorOf<[I64]>, NoneType]>:$split,
+    DefaultValuedAttr<SI64Attr, "0">:$axis,
+    OptionalAttr<SI64Attr>:$num_outputs);
+  let results = (outs Variadic<AnyTypeOf<[TensorOf<[UI8]>, TensorOf<[UI16]>, TensorOf<[UI32]>, TensorOf<[UI64]>, TensorOf<[I8]>, TensorOf<[I16]>, TensorOf<[I32]>, TensorOf<[I64]>, TensorOf<[BF16]>, TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>, TensorOf<[StringType]>, TensorOf<[I1]>, TensorOf<[Complex<F32>]>, TensorOf<[Complex<F64>]>]>>:$outputs);
+  let builders = [
+    OpBuilder<(ins "Value":$input, "Value":$split, "IntegerAttr":$axis, "IntegerAttr":$num_outputs), [{
+      auto resultType = UnrankedTensorType::get(input.getType().cast<ShapedType>().getElementType());
+      build($_builder, $_state, resultType, input, split, axis, num_outputs);
+    }]>,
+    OpBuilder<(ins "ValueRange":$operands, "ArrayRef<NamedAttribute>":$attributes), [{
+      auto resultType = UnrankedTensorType::get(operands[0].getType().cast<ShapedType>().getElementType());
+      build($_builder, $_state, {resultType}, operands, attributes);
+    }]>
+    ];
+    let extraClassDeclaration = [{
+      static int getNumberOfOperands() {
+        return 2;
+      }
+      static int getNumberOfResults() {
+        return -1;
+      }
+      static std::vector<int> getTypeMap() {
+        return {20};
+      }
+    }];
+    let extraClassDefinition = [{
+      onnx_mlir::ONNXOpShapeHelper * $cppClass::getShapeHelper(mlir::Operation *op, llvm::ArrayRef<mlir::Value> oper, 
+          onnx_mlir::IndexExprBuilder *ieb, onnx_mlir::IndexExprScope *scope) {
+        onnx_mlir::ONNXOpShapeHelper *sh = new onnx_mlir::ONNXSplitOpShapeHelper(op, oper, ieb, scope);
+        assert(sh && "failed to allocate shape helper");
+        return sh;
+      }
+    }];
+    let hasVerifier = 1;
+}
+
+def ONNXSplitV13Op:ONNX_Op<"SplitV13",
   [Pure, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>, DeclareOpInterfaceMethods<ShapeHelperOpInterface>]> {
   let summary = "ONNX Split operation";
   let description = [{
@@ -8591,12 +8638,11 @@ def ONNXSplitOp:ONNX_Op<"Split",
     let extraClassDefinition = [{
       onnx_mlir::ONNXOpShapeHelper * $cppClass::getShapeHelper(mlir::Operation *op, llvm::ArrayRef<mlir::Value> oper, 
           onnx_mlir::IndexExprBuilder *ieb, onnx_mlir::IndexExprScope *scope) {
-        onnx_mlir::ONNXOpShapeHelper *sh = new onnx_mlir::ONNXSplitOpShapeHelper(op, oper, ieb, scope);
+        onnx_mlir::ONNXOpShapeHelper *sh = new onnx_mlir::ONNXSplitV13OpShapeHelper(op, oper, ieb, scope);
         assert(sh && "failed to allocate shape helper");
         return sh;
       }
     }];
-    let hasVerifier = 1;
 }
 
 def ONNXSplitV11Op:ONNX_Op<"SplitV11",

src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp

@@ -598,6 +598,7 @@ struct ONNXCommonSplitOpShapeHelper : public ONNXOpShapeHelper {
 // clang-format off
 using ONNXSplitOpShapeHelper = ONNXCommonSplitOpShapeHelper<mlir::ONNXSplitOp>;
 using ONNXSplitV11OpShapeHelper = ONNXCommonSplitOpShapeHelper<mlir::ONNXSplitV11Op>;
+using ONNXSplitV13OpShapeHelper = ONNXCommonSplitOpShapeHelper<mlir::ONNXSplitV13Op>;
 // clang-format on
 
 //===----------------------------------------------------------------------===//

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

@@ -59,14 +59,25 @@ LogicalResult ONNXCommonSplitOpShapeHelper<OP_TYPE>::customComputeShape(
   } else {
     // If split parameter is not specified, the dimension is split to
     // equal-sized parts.
+    bool hasNumOutputsAttr = std::is_same_v<OP_TYPE, ONNXSplitOp>;
+    // TODO figure out how to handle when numResults is determined by
+    // num_outputs attribute introduced for Split in opset 18
+    // Currently the whole graph depends on the number of outputs being
+    // determined at the ONNX to ONNX-MLIR ingestion stage.
     IndexExpr splitInputDim = createIE->getShapeAsDim(input, axisIndex);
     LiteralIndexExpr numOfPartitions(numOfResults);
     if (splitInputDim.isLiteral() &&
-        (splitInputDim.getLiteral() % numOfResults != 0))
+        (splitInputDim.getLiteral() % numOfResults != 0) && !hasNumOutputsAttr)
       return op->emitError("The dimension at the split axis is "
                            "expected to be divisible by the number of results");
+
+    unsigned numBiggerChunks = splitInputDim.isLiteral()
+                                   ? splitInputDim.getLiteral() % numOfResults
+                                   : numOfResults;
     for (unsigned int i = 0; i < numOfResults; ++i) {
-      IndexExpr splitDim = splitInputDim.ceilDiv(numOfPartitions);
+      IndexExpr splitDim = (i < numBiggerChunks)
+                               ? splitInputDim.ceilDiv(numOfPartitions)
+                               : splitInputDim.floorDiv(numOfPartitions);
       splitDims.emplace_back(splitDim);
     }
   }
@@ -103,6 +114,22 @@ LogicalResult ONNXSplitOpShapeHelper::computeShape() {
   return customComputeShape(indexExprArray);
 }
 
+// Code for SplitV13Op compute shape.
+template <>
+LogicalResult ONNXSplitV13OpShapeHelper::computeShape() {
+  ONNXSplitOpAdaptor operandAdaptor(operands, op->getAttrDictionary());
+  Value split = operandAdaptor.getSplit();
+  SmallVector<IndexExpr, 4> indexExprArray;
+  if (isNoneValue(split)) {
+    // None is fine, indexExprArray will be empty.
+  } else {
+    createIE->getIntFromArrayAsSymbols(split, indexExprArray);
+    assert(IndexExpr::isLiteral(indexExprArray) &&
+           "dynamic split not yet supported");
+  }
+  return customComputeShape(indexExprArray);
+}
+
 // Code for SplitV11Op compute shape.
 template <>
 LogicalResult ONNXSplitV11OpShapeHelper::computeShape() {
@@ -157,6 +184,19 @@ LogicalResult ONNXSplitOp::inferShapes(
   return shapeHelper.computeShapeAndUpdateType(elementType);
 }
 
+LogicalResult ONNXSplitV13Op::inferShapes(
+    std::function<void(Region &)> doShapeInference) {
+  // Cannot infer the output shape if the input shape isn't known yet.
+  if (!hasShapeAndRank(getInput()))
+    return success();
+
+  auto inputType = getInput().getType().cast<ShapedType>();
+  Type elementType = inputType.getElementType();
+  ONNXSplitV13OpShapeHelper shapeHelper(getOperation(), {});
+  // Same time for all results.
+  return shapeHelper.computeShapeAndUpdateType(elementType);
+}
+
 LogicalResult ONNXSplitV11Op::inferShapes(
     std::function<void(Region &)> doShapeInference) {
   // Cannot infer the output shape if the input shape isn't known yet.
@@ -180,5 +220,6 @@ LogicalResult ONNXSplitV11Op::inferShapes(
 
 namespace onnx_mlir {
 template struct ONNXCommonSplitOpShapeHelper<ONNXSplitOp>;
+template struct ONNXCommonSplitOpShapeHelper<ONNXSplitV13Op>;
 template struct ONNXCommonSplitOpShapeHelper<ONNXSplitV11Op>;
 } // namespace onnx_mlir

src/Transform/ONNX/Decompose.td

@@ -323,13 +323,13 @@ def ClipV12Pattern : Pat<
 
 def SplitV11PatternNoAttr : Pat<
   (ONNXSplitV11Op $x, $axis, $split),
-  (ONNXSplitOp $x, (CreateNoneValue), $axis),
+  (ONNXSplitOp $x, (CreateNoneValue), $axis, (GetNullIntegerAttr)),
   [(AttributeIsNull:$split)], (addBenefit 1)
 >;
 
 def SplitV11Pattern : Pat<
   (ONNXSplitV11Op $x, $axis, $split),
-  (ONNXSplitOp $x, (ONNXConstantOpFromDenseAttr(createDenseArrayAttr $split)), $axis),
+  (ONNXSplitOp $x, (ONNXConstantOpFromDenseAttr(createDenseArrayAttr $split)), $axis, (GetNullIntegerAttr)),
   [], (addBenefit 0)
 >;
 

test/mlir/onnx/onnx_shape_inference.mlir

@@ -1363,6 +1363,32 @@ func.func @test_split_5(%arg0 : tensor<16x?x64xf32>) -> tensor<*xf32> {
 
 // -----
 
+func.func @test_split_6(%arg0 : tensor<16x39x64xf32>) -> tensor<*xf32> {
+  %cst = "onnx.NoValue"() {value} : () -> none
+  %0, %1 = "onnx.Split"(%arg0, %cst) { axis = 1 : si64, num_outputs = 2 : si64} : (tensor<16x39x64xf32>, none) -> (tensor<*xf32>, tensor<*xf32>)
+  "func.return"(%0) : (tensor<*xf32>) -> ()
+
+  // CHECK-LABEL: test_split_6
+  // CHECK: [[CST:%.+]] = "onnx.NoValue"() {value} : () -> none
+  // CHECK-NEXT: [[RES:%.+]]:2 = "onnx.Split"(%arg0, [[CST]]) {axis = 1 : si64, num_outputs = 2 : si64} : (tensor<16x39x64xf32>, none) -> (tensor<16x20x64xf32>, tensor<16x19x64xf32>)
+  // CHECK: return [[RES]]#0 : tensor<16x20x64xf32>
+}
+
+// -----
+
+func.func @test_split_7(%arg0 : tensor<16x38x64xf32>) -> tensor<*xf32> {
+  %cst = "onnx.NoValue"() {value} : () -> none
+  %0, %1, %2 = "onnx.Split"(%arg0, %cst) { axis = 1 : si64, num_outputs = 3 : si64} : (tensor<16x38x64xf32>, none) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>)
+  "func.return"(%0) : (tensor<*xf32>) -> ()
+
+  // CHECK-LABEL: test_split_7
+  // CHECK: [[CST:%.+]] = "onnx.NoValue"() {value} : () -> none
+  // CHECK-NEXT: [[RES:%.+]]:3 = "onnx.Split"(%arg0, [[CST]]) {axis = 1 : si64, num_outputs = 3 : si64} : (tensor<16x38x64xf32>, none) -> (tensor<16x13x64xf32>, tensor<16x13x64xf32>, tensor<16x12x64xf32>)
+  // CHECK: return [[RES]]#0 : tensor<16x13x64xf32>
+}
+
+// -----
+
 func.func @test_splitv11_1(%arg0 : tensor<16x32x64xf32>) -> tensor<*xf32> {
   %0, %1 = "onnx.SplitV11"(%arg0) { axis = 1 : si64} : (tensor<16x32x64xf32>) -> (tensor<*xf32>, tensor<*xf32>)
   "func.return"(%0) : (tensor<*xf32>) -> ()

utils/gen_onnx_mlir.py

@@ -256,7 +256,7 @@
  'Softplus': [1],
  'Softsign': [1],
  'SpaceToDepth': [13],
- 'Split': [13, 11],
+ 'Split': [18, 13, 11],
  'SplitToSequence': [11],
  'Sqrt': [13],
  'Squeeze': [13, 11],
@@ -500,6 +500,7 @@
     'ReduceSumV11',
     'Softmax',
     'Split',
+    'SplitV13',
     'Sqrt',
     'Squeeze',
     'SqueezeV11',