Add a new fusion pass with pattern of constant and reshape

onnxoptimizer/pass_registry.h

@@ -52,6 +52,7 @@
 #include "onnxoptimizer/passes/fuse_pad_into_conv.h"
 #include "onnxoptimizer/passes/fuse_pad_into_pool.h"
 #include "onnxoptimizer/passes/fuse_transpose_into_gemm.h"
+#include "onnxoptimizer/passes/fuse_constant_reshape.h"
 #include "onnxoptimizer/passes/lift_lexical_references.h"
 #include "onnxoptimizer/passes/nop.h"
 #include "onnxoptimizer/passes/rename_input_output.h"
@@ -94,6 +95,7 @@ struct GlobalPassRegistry {
     registerPass<EliminateShapeGather>();
     registerPass<EliminateSliceAfterShape>();
     registerPass<EliminateNopTranspose>();
+    registerPass<FuseConstantReshape>();
     registerPass<FuseAddBiasIntoConv>();
     registerPass<FuseBNIntoConv>();
     registerPass<FuseConsecutiveConcats>();

onnxoptimizer/passes/fuse_constant_reshape.h

@@ -0,0 +1,144 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+#pragma once
+
+// Before:
+//   B = Reshape(Constant)
+// After:
+//   B = Constant (Constant with new shape)
+
+#include <limits>
+
+#include "onnx/defs/tensor_util.h"
+#include "onnxoptimizer/pass.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+struct FuseConstantReshape final : public PredicateBasedPass {
+  explicit FuseConstantReshape()
+      : PredicateBasedPass(PassType::Fuse, PassEfficiency::Complete,
+                           PassOptimizationType::Compute) {}
+  std::string getPassName() const override {
+    return "fuse_constant_reshape";
+  }
+
+  bool patternMatchPredicate(Node* node) override {
+    return node->kind() == kReshape && !node->inputs().empty() &&
+           node->inputs()[0]->node()->kind() == kConstant;
+  }
+  bool runTransform(Node* n, Graph& graph,
+                    NodeDestroyType& destroy_current) override {
+    destroy_current = NodeDestroyType::DestroyZero;
+
+    // check if Constant is only used by Reshape
+    if (n->inputs()[0]->uses().size() > 1) {
+      return false;
+    }
+
+    Node* reshape = n;
+    Node* constant = n->inputs()[0]->node();
+    if (!constant->hasAttribute(kvalue)) {
+      return false;
+    }
+
+    // Process 'reshape' data
+    std::vector<int64_t> shape;
+    if (reshape->hasAttribute(kshape)) {
+      // Before opset 5.
+      shape = reshape->is(kshape);
+    } else {
+      if (reshape->inputs().size() < 2 ||
+          reshape->inputs()[1]->node()->kind() != kConstant) {
+        return false;
+      }
+      if (reshape->inputs()[1]->uses().size() > 1) {
+        return false;
+      }
+      Node* shape_const = reshape->inputs()[1]->node();
+      if (!shape_const->hasAttribute(kvalue)) {
+        return false;
+      }
+      Tensor t = shape_const->t(kvalue);
+      if (t.elem_type() != TensorProto_DataType_INT64 ||
+          t.sizes().size() != 1) {
+        return false;
+      }
+      shape = ParseData<int64_t>(&t);
+    }
+
+    int64_t allow_zero = 0;
+    Symbol sym = Symbol("allowzero");
+    if (reshape->hasAttribute(sym)) {
+      allow_zero = reshape->i(sym);
+    }
+    if (allow_zero != 0 && allow_zero != 1) {
+      return false;
+    }
+
+    Tensor t = constant->t(kvalue);
+    const auto& ori_size = t.sizes();
+
+    const auto multiply = [](int64_t& product, int64_t dimension) {
+      if (dimension < 0 ||
+          (dimension != 0 &&
+           product > std::numeric_limits<int64_t>::max() / dimension)) {
+        return false;
+      }
+      product *= dimension;
+      return true;
+    };
+    int64_t numel = 1;
+    for (int64_t dimension : ori_size) {
+      if (!multiply(numel, dimension)) {
+        return false;
+      }
+    }
+
+    size_t unknown_index = shape.size();
+    int64_t known_product = 1;
+    for (size_t i = 0; i < shape.size(); ++i) {
+      if (shape[i] == 0 && allow_zero == 0) {
+        if (i >= ori_size.size()) {
+          return false;
+        }
+        shape[i] = ori_size[i];
+      }
+      if (shape[i] == -1) {
+        if (unknown_index != shape.size()) {
+          return false;
+        }
+        unknown_index = i;
+      } else if (!multiply(known_product, shape[i])) {
+        return false;
+      }
+    }
+    if (unknown_index != shape.size()) {
+      if (known_product == 0 || numel % known_product != 0) {
+        return false;
+      }
+      shape[unknown_index] = numel / known_product;
+    } else if (known_product != numel) {
+      return false;
+    }
+
+    const bool replacing_success = tryReplacingAllUsesWith(reshape->output(), reshape->inputs()[0]);
+    if (!replacing_success) {
+      return false;
+    }
+    t.sizes() = shape;
+    constant->output()->setSizes(
+        std::vector<Dimension>(shape.begin(), shape.end()));
+    constant->output()->setElemType(t.elem_type());
+    constant->t_(kvalue, std::move(t));
+    destroy_current = NodeDestroyType::DestroyOne;
+    return true;
+
+  }
+};
+
+}  // namespace optimization
+}  // namespace ONNX_NAMESPACE

onnxoptimizer/test/constant_reshape_test.py

@@ -0,0 +1,95 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from onnx import TensorProto, checker, helper, numpy_helper
+from onnx.reference import ReferenceEvaluator
+
+import onnxoptimizer
+
+
+def make_model(input_shape: list[int], target_shape: list[int], allowzero: int | None):
+    values = np.arange(np.prod(input_shape, dtype=np.int64), dtype=np.float32).reshape(
+        input_shape
+    )
+    attributes = {} if allowzero is None else {"allowzero": allowzero}
+    graph = helper.make_graph(
+        [
+            helper.make_node(
+                "Constant", [], ["data"], value=numpy_helper.from_array(values)
+            ),
+            helper.make_node(
+                "Constant",
+                [],
+                ["shape"],
+                value=numpy_helper.from_array(np.array(target_shape, dtype=np.int64)),
+            ),
+            helper.make_node("Reshape", ["data", "shape"], ["Y"], **attributes),
+        ],
+        "constant_reshape",
+        [],
+        [
+            helper.make_tensor_value_info(
+                "Y", TensorProto.FLOAT, [None] * len(target_shape)
+            )
+        ],
+    )
+    return helper.make_model(
+        graph, ir_version=7, opset_imports=[helper.make_opsetid("", 14)]
+    )
+
+
+@pytest.mark.parametrize(
+    "input_shape,target_shape,allowzero,expected_shape",
+    [
+        ([2, 3], [0, 3], None, [2, 3]),
+        ([2, 3], [0, -1], 0, [2, 3]),
+        ([2, 3], [3, -1], None, [3, 2]),
+        ([2, 3], [-1], None, [6]),
+        ([0, 3], [0, 3], 1, [0, 3]),
+        ([0, 3], [-1, 3], None, [0, 3]),
+        ([], [], None, []),
+        ([1], [], None, []),
+    ],
+)
+def test_constant_reshape_preserves_values_and_shape(
+    input_shape, target_shape, allowzero, expected_shape
+) -> None:
+    model = make_model(input_shape, target_shape, allowzero)
+    checker.check_model(model)
+    expected = ReferenceEvaluator(model).run(None, {})[0]
+    optimized = onnxoptimizer.optimize(model, ["fuse_constant_reshape"])
+    checker.check_model(optimized)
+    assert all(node.op_type != "Reshape" for node in optimized.graph.node)
+    actual = ReferenceEvaluator(optimized).run(None, {})[0]
+    assert list(actual.shape) == expected_shape
+    assert [
+        dimension.dim_value
+        for dimension in optimized.graph.output[0].type.tensor_type.shape.dim
+    ] == expected_shape
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "input_shape,target_shape,allowzero",
+    [
+        ([0, 3], [0, -1], 0),
+        ([0, 3], [0, -1], 1),
+        ([2, 3], [-1, -1], None),
+        ([2, 3], [-2, 3], None),
+        ([2, 3], [4, -1], None),
+        ([2, 3], [5], None),
+        ([2, 3], [1, 1, 0], None),
+        ([2, 3], [2**32, 2**32], None),
+    ],
+)
+def test_constant_reshape_invalid_shape_is_not_folded(
+    input_shape, target_shape, allowzero
+) -> None:
+    model = make_model(input_shape, target_shape, allowzero)
+    optimized = onnxoptimizer.optimize(model, ["fuse_constant_reshape"])
+    assert any(node.op_type == "Reshape" for node in optimized.graph.node)