Fix one input case

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

Author: srcarroll

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -992,11 +992,13 @@ struct SumToAddPattern : public OpRewritePattern<ONNXSumOp> {
   LogicalResult matchAndRewrite(
       ONNXSumOp sumOp, PatternRewriter &rewriter) const final {
     SmallVector<Value> inputs(sumOp.getData_0());
-    assert(inputs.size() >= 2 && "expected at least two inputs");
+    assert(inputs.size() > 0 && "expected at least one input");
     Value result = inputs[0];
-    inputs.erase(inputs.begin());
-    for (auto input : inputs) {
-      result = rewriter.create<ONNXAddOp>(sumOp.getLoc(), result, input);
+    if (inputs.size() > 1) {
+      inputs.erase(inputs.begin());
+      for (auto input : inputs) {
+        result = rewriter.create<ONNXAddOp>(sumOp.getLoc(), result, input);
+      }
     }
     rewriter.replaceOp(sumOp, result);
     return success();

test/mlir/onnx/onnx_decompose.mlir

@@ -710,3 +710,13 @@ func.func @test_sum(%arg0: tensor<64x128x10xf32>, %arg1: tensor<128x10xf32>, %ar
   // CHECK-NEXT:      %[[SUM2:.*]] = "onnx.Add"(%[[SUM1]], %[[ARG3]])
   // CHECK-NEXT:      onnx.Return %[[SUM2]]
 }
+
+// -----
+
+func.func @test_sum_single_input(%arg0: tensor<64x128x10xf32>) -> tensor<64x128x10xf32> {
+  %0 = "onnx.Sum"(%arg0) : (tensor<64x128x10xf32>) -> tensor<64x128x10xf32>
+  onnx.Return %0 : tensor<64x128x10xf32> 
+  // CHECK-LABEL:       func @test_sum_single_input
+  // CHECK-SAME:     (%[[ARG0:.*]]: {{.*}})
+  // CHECK-NEXT:      onnx.Return %[[ARG0]]
+}