mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Fix slice upstream - Incompatible dimensions (#16818)
### Fix slice upstream - (MatMul) [ShapeInferenceError] Incompatible
dimensions
```
2023-07-22 14:58:16.918478478 [I:onnxruntime:Default, constant_sharing.cc:256 ApplyImpl] Total shared scalar initializer count: 10
2023-07-22 14:58:16.919494252 [W:onnxruntime:Default, graph.cc:108 MergeShapeInfo] Error merging shape info for output. 'onnx::Cast_424' source:{-1,31,-1,-1} target:{-1,32,-1,-1}. Falling back to lenient merge.
2023-07-22 14:58:16.921014114 [W:onnxruntime:Default, graph.cc:108 MergeShapeInfo] Error merging shape info for output. 'onnx::MatMul_425' source:{-1,31,-1,-1} target:{-1,32,-1,-1}. Falling back to lenient merge.
Traceback (most recent call last):
File "examples/onnxruntime/training/language-modeling/run_clm.py", line 594, in <module>
main()
File "examples/onnxruntime/training/language-modeling/run_clm.py", line 542, in main
train_result = trainer.train(resume_from_checkpoint=checkpoint)
File "/bert_ort/pengwa/optimum/optimum/onnxruntime/trainer.py", line 454, in train
return inner_training_loop(
File "/bert_ort/pengwa/optimum/optimum/onnxruntime/trainer.py", line 755, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/transformers/trainer.py", line 2735, in training_step
loss = self.compute_loss(model, inputs)
File "/bert_ort/pengwa/optimum/optimum/onnxruntime/trainer.py", line 363, in compute_loss
return model_with_loss(dict_inputs, return_outputs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1724, in forward
loss = self.module(*inputs, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_utils.py", line 384, in _forward
return ortmodule._torch_module.forward(*inputs, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_utils.py", line 364, in _forward
return torch_module_ort._execution_manager(torch_module_ort.is_training()).forward(*inputs, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 345, in forward
self._fallback_manager.handle_exception(
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_fallback.py", line 157, in handle_exception
raise exception
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 280, in forward
self._build_graph(graph_transformer_config)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_logger.py", line 218, in wrapper
result = func(graph_execution_manager, *args, **kwargs)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 360, in _build_graph
super()._build_graph(graph_transformer_config)
File "/bert_ort/pengwa/py38/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_graph_execution_manager.py", line 186, in _build_graph
self._graph_builder.build(config)
RuntimeError: /bert_ort/pengwa/onnxruntime/orttraining/orttraining/python/orttraining_pybind_state.cc:823 onnxruntime::python::addObjectMethodsForTraining(pybind11::module&, onnxruntime::python::ExecutionProviderRegistrationFn)::<lambda(onnxruntime::training::OrtModuleGraphBuilder*, const onnxruntime::training::TrainingGraphTransformerConfiguration&)> [ONNXRuntimeError] : 1 : FAIL : Node (MatMul_403) Op (MatMul) [ShapeInferenceError] Incompatible dimensions
```
Missed using `axis` attribute for `Slice` op, so change to use `axes`
inputs instead.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
pengwa
GitHub
parent
b0279b14d8
commit
f2c0470436
4 changed files with 397 additions and 89 deletions
|
|
@ -184,7 +184,9 @@ NodeArg* CreateInitializerFromVector(Graph& graph,
|
|||
total_count *= dim;
|
||||
}
|
||||
|
||||
ORT_ENFORCE(total_count == static_cast<int64_t>(values.size()));
|
||||
ORT_ENFORCE(total_count == static_cast<int64_t>(values.size()),
|
||||
"The total count of dims does not match the size of values. ",
|
||||
"total_count: ", total_count, " values.size(): ", values.size());
|
||||
|
||||
const_tensor.set_raw_data(values.data(), values.size() * sizeof(int64_t));
|
||||
return &graph_utils::AddInitializer(graph, const_tensor);
|
||||
|
|
|
|||
|
|
@ -138,20 +138,65 @@ SliceInfo UpStreamGatherGraphTransformer::PropagateSlicingForInput(
|
|||
std::to_string(!info.is_scalar_slice));
|
||||
|
||||
InlinedVector<NodeArg*> input_args;
|
||||
input_args.reserve(slice_node.InputDefs().size());
|
||||
input_args.resize(slice_node.InputDefs().size());
|
||||
|
||||
int axis_input_index = -1; // -1 means axis is passed in attribute.
|
||||
if (std::holds_alternative<int>(info.axis_attr_name_or_input_index)) {
|
||||
axis_input_index = std::get<int>(info.axis_attr_name_or_input_index);
|
||||
}
|
||||
|
||||
auto create_axes_input = [&info, new_axis, &graph]() -> NodeArg* {
|
||||
InlinedVector<int64_t> dims;
|
||||
if (info.rank_of_axis_value == 1) {
|
||||
dims.push_back(1);
|
||||
}
|
||||
return CreateInitializerFromVector(graph, dims, {new_axis}, graph.GenerateNodeArgName("axes"));
|
||||
};
|
||||
|
||||
// The first slice op's data input should be current_node's current_node_input_index-th input.
|
||||
// For some cases when rank changes, slice op's slice input should also be adapted.
|
||||
input_args.push_back(current_node.MutableInputDefs()[current_node_input_index]);
|
||||
for (size_t i = 1; i < slice_node.InputDefs().size(); ++i) {
|
||||
input_args.push_back(slice_node.MutableInputDefs()[i]);
|
||||
int i = 0;
|
||||
for (; i < static_cast<int>(slice_node.InputDefs().size()); ++i) {
|
||||
if (i == info.GetDataInputIndex()) {
|
||||
input_args[i] = current_node.MutableInputDefs()[current_node_input_index];
|
||||
} else if (axis_input_index != -1 && i == axis_input_index) {
|
||||
if (info.non_negative_axis == new_axis) {
|
||||
input_args[i] = slice_node.MutableInputDefs()[i];
|
||||
} else {
|
||||
input_args[i] = create_axes_input();
|
||||
}
|
||||
} else {
|
||||
input_args[i] = slice_node.MutableInputDefs()[i];
|
||||
}
|
||||
}
|
||||
|
||||
// It is possible axes input is null.
|
||||
if (axis_input_index != -1 && info.non_negative_axis != new_axis) {
|
||||
for (; i <= axis_input_index; ++i) {
|
||||
if (i == axis_input_index) {
|
||||
input_args.push_back(create_axes_input());
|
||||
} else {
|
||||
NodeArg& empty_input = graph.GetOrCreateNodeArg("", nullptr);
|
||||
input_args.push_back(&empty_input);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update the axis attribute if new_axis is not the same as the original slicing axis (which happens when data
|
||||
// layout got changed by Transpose or Reshape ops)
|
||||
onnxruntime::NodeAttributes attributes = slice_node.GetAttributes();
|
||||
if (info.non_negative_axis != new_axis) {
|
||||
attributes[info.axis_attr_name] =
|
||||
ONNX_NAMESPACE::MakeAttribute(info.axis_attr_name, static_cast<int64_t>(new_axis));
|
||||
|
||||
if (axis_input_index == -1 && info.non_negative_axis != new_axis) {
|
||||
std::string attr_name = std::get<std::string>(info.axis_attr_name_or_input_index);
|
||||
if (info.rank_of_axis_value == 0) {
|
||||
attributes[attr_name] =
|
||||
ONNX_NAMESPACE::MakeAttribute(attr_name, static_cast<int64_t>(new_axis));
|
||||
} else if (info.rank_of_axis_value == 1) {
|
||||
attributes[attr_name] =
|
||||
ONNX_NAMESPACE::MakeAttribute(attr_name, std::vector<int64_t>{static_cast<int64_t>(new_axis)});
|
||||
} else {
|
||||
ORT_THROW("Unexpected rank of axis attribute value: " + std::to_string(info.rank_of_axis_value));
|
||||
}
|
||||
}
|
||||
|
||||
InlinedVector<NodeArg*> output_args;
|
||||
|
|
@ -183,7 +228,8 @@ SliceInfo UpStreamGatherGraphTransformer::PropagateSlicingForInput(
|
|||
auto new_slice_out_arg = new_slice_node->MutableOutputDefs()[new_slice_output_index_to_connect];
|
||||
UpdateSliceOutputShape(*new_slice_out_arg, new_axis, info.output_dim_on_axis);
|
||||
|
||||
auto new_slice_info = SliceInfo(graph, new_slice_node, info.is_scalar_slice, info.axis_attr_name, new_axis);
|
||||
auto new_slice_info = SliceInfo(graph, new_slice_node, info.is_scalar_slice, info.axis_attr_name_or_input_index,
|
||||
new_axis, info.rank_of_axis_value);
|
||||
new_slice_info.entry_node_name = info.entry_node_name;
|
||||
new_slice_info.entry_slice_arg_name = info.entry_slice_arg_name;
|
||||
return new_slice_info;
|
||||
|
|
@ -263,7 +309,8 @@ std::optional<SliceInfo> IsSupportedGatherND(Graph& graph, Node& node,
|
|||
return std::nullopt;
|
||||
}
|
||||
|
||||
return SliceInfo(graph, &node, false, "batch_dims", static_cast<int>(batch_dims), true);
|
||||
return SliceInfo(graph, &node, false, "batch_dims", static_cast<int>(batch_dims),
|
||||
0 /* rank of axis attribute value */, true);
|
||||
}
|
||||
|
||||
std::optional<SliceInfo> IsSupportedGather(Graph& graph, Node& node,
|
||||
|
|
@ -304,7 +351,7 @@ std::optional<SliceInfo> IsSupportedGather(Graph& graph, Node& node,
|
|||
}
|
||||
}
|
||||
|
||||
return SliceInfo(graph, &node, dim_size == 0, "axis", axis, true);
|
||||
return SliceInfo(graph, &node, dim_size == 0, "axis", axis, 0 /* rank of axis attribute value */, true);
|
||||
}
|
||||
|
||||
std::optional<SliceInfo> IsSupportedShrunkenGather(Graph& graph, Node& node,
|
||||
|
|
@ -342,7 +389,7 @@ std::optional<SliceInfo> IsSupportedShrunkenGather(Graph& graph, Node& node,
|
|||
return std::nullopt;
|
||||
}
|
||||
|
||||
return SliceInfo(graph, &node, false /*is_slice_scalar*/, "axis", axis, true);
|
||||
return SliceInfo(graph, &node, false /*is_slice_scalar*/, "axis", axis, 0 /* rank of axis attribute value */, true);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -366,34 +413,37 @@ std::optional<SliceInfo> IsSupportedSlice(Graph& graph, Node& node,
|
|||
const NodeArg* axes_input = node.InputDefs().size() > 3 ? node.InputDefs()[3] : nullptr;
|
||||
|
||||
if (data_input->Shape() == nullptr || starts_input->Shape() == nullptr || ends_input->Shape() == nullptr ||
|
||||
(axes_input && axes_input->Shape() == nullptr)) {
|
||||
(axes_input && axes_input->Exists() && axes_input->Shape() == nullptr)) {
|
||||
LOG_DEBUG_INFO(logger, "Skip Slice node " + node.Name() + " due to undefined shape.");
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
// Make sure starts/ends/axes/steps are all 1D tensors, since we only support single-dimension slicing.
|
||||
if (starts_input->Shape()->dim_size() != 1 || ends_input->Shape()->dim_size() != 1 ||
|
||||
(axes_input && axes_input->Shape()->dim_size() != 1)) {
|
||||
(axes_input && axes_input->Exists() && axes_input->Shape()->dim_size() != 1)) {
|
||||
LOG_DEBUG_INFO(logger, "Skip Slice node " + node.Name() + " due to unsupported dim size: " +
|
||||
std::to_string(starts_input->Shape()->dim_size()) + ", " +
|
||||
std::to_string(ends_input->Shape()->dim_size()) + ", " +
|
||||
std::to_string(axes_input ? axes_input->Shape()->dim_size() : 0));
|
||||
std::to_string(axes_input && axes_input->Exists() ? axes_input->Shape()->dim_size() : 0));
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
// Try to parse the 'axes' value.
|
||||
int axis = 0;
|
||||
if (axes_input) {
|
||||
if (axes_input && axes_input->Exists()) {
|
||||
InlinedVector<int64_t> axes_values;
|
||||
if (!graph_utils::IsConstantInitializer(graph, axes_input->Name()) ||
|
||||
!optimizer_utils::AppendTensorFromInitializer(graph, *axes_input, axes_values, true) ||
|
||||
axes_values.size() != 1) {
|
||||
LOG_DEBUG_INFO(logger, "Skip Slice node " + node.Name() + " due to unsupported axes value.");
|
||||
return std::nullopt;
|
||||
}
|
||||
axis = static_cast<int>(axes_values[0]);
|
||||
} else {
|
||||
// If 'axes' is not specified, then it is [0, .., r-1], so we force data rank to be 1.
|
||||
if (data_input->Shape()->dim_size() != 1) {
|
||||
LOG_DEBUG_INFO(logger, "Skip Slice node " + node.Name() + " due to unsupported data rank: " +
|
||||
std::to_string(data_input->Shape()->dim_size()));
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
|
|
@ -401,7 +451,8 @@ std::optional<SliceInfo> IsSupportedSlice(Graph& graph, Node& node,
|
|||
if (axis < 0)
|
||||
axis += data_input->Shape()->dim_size();
|
||||
|
||||
return SliceInfo(graph, &node, false /*is_slice_scalar*/, "axis", axis, true);
|
||||
return SliceInfo(graph, &node, false /*is_slice_scalar*/, 3 /* axis input index */, axis,
|
||||
1 /* rank of axes value */, true);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
|
|
|||
|
|
@ -25,11 +25,21 @@ struct SliceInfo : public UpstreamOperatorInfoBase {
|
|||
public:
|
||||
SliceInfo(const Graph& graph, Node* slice_node,
|
||||
bool is_slice_scalar,
|
||||
const std::string& slice_axis_attr_name,
|
||||
std::variant<std::string, int> axis_name_or_index,
|
||||
int slice_axis,
|
||||
int rank_of_axis,
|
||||
bool is_entry_node_ptr = false)
|
||||
: UpstreamOperatorInfoBase(slice_node, is_entry_node_ptr), is_scalar_slice(is_slice_scalar) {
|
||||
axis_attr_name = slice_axis_attr_name;
|
||||
axis_attr_name_or_input_index = axis_name_or_index;
|
||||
rank_of_axis_value = rank_of_axis;
|
||||
|
||||
if (std::holds_alternative<int>(axis_name_or_index)) {
|
||||
int axis_input_index = std::get<int>(axis_name_or_index);
|
||||
ORT_ENFORCE(axis_input_index >= 0, "Axis input index is invalid");
|
||||
}
|
||||
|
||||
ORT_ENFORCE(rank_of_axis_value == 0 || rank_of_axis_value == 1, "Rank of axis value is invalid: " +
|
||||
std::to_string(rank_of_axis_value));
|
||||
|
||||
const NodeArg* input = node_ptr->InputDefs()[kSliceDataInputIndex_];
|
||||
const NodeArg* output = node_ptr->OutputDefs()[kSliceOutputIndex_];
|
||||
|
|
@ -65,8 +75,16 @@ struct SliceInfo : public UpstreamOperatorInfoBase {
|
|||
}
|
||||
|
||||
bool is_scalar_slice; // whether the slice is a scalar, if it is after Gather, the rank will be reduced by 1.
|
||||
std::string axis_attr_name;
|
||||
|
||||
// The index of the input that contains the axis value. If it is a string, then axis will be treated as an attribute.
|
||||
std::variant<std::string, int> axis_attr_name_or_input_index;
|
||||
|
||||
int non_negative_axis; // The axis to slice on
|
||||
|
||||
// The rank of value for axis attribute. For example, for Gather, its axis attribute is a scalar, so the rank is 0.
|
||||
// For Slice, its axes attribute is a 1D tensor, so the rank is 1.
|
||||
int rank_of_axis_value;
|
||||
|
||||
std::string entry_slice_arg_name;
|
||||
|
||||
int input_rank; // rank of the Gather data input tensor
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@
|
|||
#include "core/common/span_utils.h"
|
||||
#include "core/framework/data_types.h"
|
||||
#include "core/framework/ort_value.h"
|
||||
#include "core/framework/tensorprotoutils.h"
|
||||
#include "core/graph/graph_utils.h"
|
||||
#include "core/graph/graph_viewer.h"
|
||||
#include "core/graph/model.h"
|
||||
|
|
@ -1597,93 +1598,329 @@ Test graph includes multiple equivalent subgraphs as below.
|
|||
Add an Identity node because currently we don't allow Slice generates graph output.
|
||||
*/
|
||||
TEST(ComputeOptimizerTests, SliceElementwiseOps_PropagationOnTwoBranches) {
|
||||
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
|
||||
InlinedVector<int64_t> starts_indices;
|
||||
auto pre_graph_checker = [&starts_indices](Graph& graph) -> Status {
|
||||
auto op_count_pre = CountOpsInGraph(graph);
|
||||
TEST_RETURN_IF_NOT(op_count_pre.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Add"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Slice"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
|
||||
// 0: no input, 1: has input, 2: empty input
|
||||
std::vector<std::tuple<std::optional<int>, std::vector<int64_t>, int, int, bool>> has_axes_and_has_steps_pairs{
|
||||
{std::nullopt, {4, 32, 256}, 0, 0, false}, // {axis, data_shape, has_axes, has_steps, expected to propagate}
|
||||
{1, {4, 32, 256}, 1, 0, true},
|
||||
{1, {4, 32, 256}, 1, 1, true},
|
||||
{1, {4, 32, 256}, 1, 2, true},
|
||||
{std::nullopt, {4, 32, 256}, 2, 0, false},
|
||||
{std::nullopt, {4, 32, 256}, 2, 1, false},
|
||||
{std::nullopt, {4, 32, 256}, 2, 2, false},
|
||||
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Slice") {
|
||||
TEST_RETURN_IF_NOT(starts_indices.empty());
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(node.InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, starts_indices,
|
||||
require_constant));
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
{std::nullopt, {256}, 0, 0, true},
|
||||
{0, {256}, 1, 0, true},
|
||||
{0, {256}, 1, 1, true},
|
||||
{0, {256}, 1, 2, true},
|
||||
{std::nullopt, {256}, 2, 0, true},
|
||||
{std::nullopt, {256}, 2, 1, true},
|
||||
{std::nullopt, {256}, 2, 2, true},
|
||||
};
|
||||
|
||||
auto post_graph_checker = [&starts_indices](Graph& graph) {
|
||||
auto op_count_post = CountOpsInGraph(graph);
|
||||
TEST_RETURN_IF_NOT(op_count_post.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Add"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Slice"] == 2);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Identity"] == 1);
|
||||
for (auto p : has_axes_and_has_steps_pairs) {
|
||||
std::optional<int> axis = std::get<0>(p);
|
||||
std::vector<int64_t> data_shape = std::get<1>(p);
|
||||
int has_axes = std::get<2>(p);
|
||||
int has_steps = std::get<3>(p);
|
||||
bool expected_to_propagate = std::get<4>(p);
|
||||
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Add") {
|
||||
const auto& input_defs = node.InputDefs();
|
||||
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
|
||||
InlinedVector<int64_t> starts_indices;
|
||||
auto pre_graph_checker = [&starts_indices](Graph& graph) -> Status {
|
||||
auto op_count_pre = CountOpsInGraph(graph);
|
||||
TEST_RETURN_IF_NOT(op_count_pre.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Add"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Slice"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
|
||||
|
||||
{
|
||||
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
|
||||
TEST_RETURN_IF_NOT(producer_node != nullptr);
|
||||
TEST_RETURN_IF_NOT(producer_node->OpType() == "Slice");
|
||||
|
||||
InlinedVector<int64_t> values;
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Slice") {
|
||||
TEST_RETURN_IF_NOT(starts_indices.empty());
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values,
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(node.InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, starts_indices,
|
||||
require_constant));
|
||||
for (size_t i = 0; i < values.size(); i++) {
|
||||
TEST_RETURN_IF_NOT(values[i] == starts_indices[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
};
|
||||
|
||||
{
|
||||
auto producer_node = graph.GetProducerNode(input_defs[1]->Name());
|
||||
TEST_RETURN_IF_NOT(producer_node != nullptr);
|
||||
TEST_RETURN_IF_NOT(producer_node->OpType() == "Slice");
|
||||
auto post_graph_checker = [&starts_indices, expected_to_propagate](Graph& graph) {
|
||||
auto op_count_post = CountOpsInGraph(graph);
|
||||
TEST_RETURN_IF_NOT(op_count_post.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Add"] == 1);
|
||||
if (expected_to_propagate) {
|
||||
TEST_RETURN_IF_NOT(op_count_post["Slice"] == 2);
|
||||
} else {
|
||||
TEST_RETURN_IF_NOT(op_count_post["Slice"] == 1);
|
||||
}
|
||||
TEST_RETURN_IF_NOT(op_count_post["Identity"] == 1);
|
||||
|
||||
InlinedVector<int64_t> values;
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values, require_constant));
|
||||
for (size_t i = 0; i < values.size(); i++) {
|
||||
TEST_RETURN_IF_NOT(values[i] == starts_indices[i]);
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Add") {
|
||||
const auto& input_defs = node.InputDefs();
|
||||
|
||||
{
|
||||
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
|
||||
|
||||
if (expected_to_propagate) {
|
||||
TEST_RETURN_IF_NOT(producer_node != nullptr);
|
||||
TEST_RETURN_IF_NOT(producer_node->OpType() == "Slice");
|
||||
|
||||
InlinedVector<int64_t> values;
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values,
|
||||
require_constant));
|
||||
for (size_t i = 0; i < values.size(); i++) {
|
||||
TEST_RETURN_IF_NOT(values[i] == starts_indices[i]);
|
||||
}
|
||||
} else {
|
||||
TEST_RETURN_IF_NOT(producer_node == nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
auto producer_node = graph.GetProducerNode(input_defs[1]->Name());
|
||||
|
||||
if (expected_to_propagate) {
|
||||
TEST_RETURN_IF_NOT(producer_node != nullptr);
|
||||
TEST_RETURN_IF_NOT(producer_node->OpType() == "Slice");
|
||||
|
||||
InlinedVector<int64_t> values;
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values, require_constant));
|
||||
for (size_t i = 0; i < values.size(); i++) {
|
||||
TEST_RETURN_IF_NOT(values[i] == starts_indices[i]);
|
||||
}
|
||||
} else {
|
||||
TEST_RETURN_IF_NOT(producer_node == nullptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
return Status::OK();
|
||||
};
|
||||
|
||||
auto build_test_case = [has_axes, has_steps, &data_shape, axis](ModelTestBuilder& builder) {
|
||||
auto* input1_arg = builder.MakeInput<int64_t>(data_shape);
|
||||
auto* input2_arg = builder.MakeInput<int64_t>(data_shape);
|
||||
auto* add_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Add", {input1_arg, input2_arg}, {add_out});
|
||||
|
||||
auto* starts_initializer = builder.MakeInitializer<int64_t>({1}, {0});
|
||||
auto* ends_initializer = builder.MakeInitializer<int64_t>({1}, {-1});
|
||||
|
||||
std::vector<NodeArg*> slice_inputs;
|
||||
slice_inputs = {add_out, starts_initializer, ends_initializer};
|
||||
|
||||
NodeArg* axes_initializer = nullptr;
|
||||
NodeArg* steps_initializer = nullptr;
|
||||
if (has_axes == 0 && has_steps == 0) {
|
||||
// nothing
|
||||
} else if (has_axes == 1 && has_steps == 0) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {axis.value()});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
} else if (has_axes == 1 && has_steps == 1) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {axis.value()});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 1 && has_steps == 2) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {axis.value()});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 0) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 1) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 2) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
}
|
||||
|
||||
auto* slice_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Slice", slice_inputs,
|
||||
{slice_out});
|
||||
|
||||
auto* identity_out = builder.MakeOutput();
|
||||
builder.AddNode("Identity", {slice_out}, {identity_out});
|
||||
};
|
||||
|
||||
std::unique_ptr<GraphTransformer> transformer = std::make_unique<UpStreamGatherGraphTransformer>();
|
||||
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger, std::move(transformer),
|
||||
TransformerLevel::Level1,
|
||||
1, pre_graph_checker, post_graph_checker));
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
Test graph includes multiple equivalent subgraphs as below.
|
||||
graph input [2, 4, 32, 256] (float)
|
||||
|
|
||||
Transpose[perms=[0, 2, 1, 3]]
|
||||
|
|
||||
[2, 32, 4, 256]
|
||||
| starts:(0) ends: (-1) axes: (1) steps: (1)
|
||||
\ \ | / /
|
||||
\ \ | / /
|
||||
\ \ | / /
|
||||
\ \ | / /
|
||||
\ \ | / /
|
||||
Slice
|
||||
|
|
||||
Identity
|
||||
|
|
||||
graph output [2, 31, 4, 256] (float)
|
||||
|
||||
Add an Identity node because currently, we don't allow Slice generates graph output.
|
||||
*/
|
||||
TEST(ComputeOptimizerTests, SliceTranspose_Propagation) {
|
||||
// 0: no input, 1: has input, 2: empty input
|
||||
std::vector<std::tuple<int, int, bool>> has_axes_and_has_steps_pairs{
|
||||
{0, 0, false}, // {has_axes, has_steps, expected to propagate}
|
||||
{1, 0, true},
|
||||
{1, 1, true},
|
||||
{1, 2, true},
|
||||
{2, 0, false},
|
||||
{2, 1, false},
|
||||
{2, 2, false},
|
||||
};
|
||||
|
||||
auto build_test_case = [](ModelTestBuilder& builder) {
|
||||
auto* input1_arg = builder.MakeInput<int64_t>({{4, 32, 256}});
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{4, 32, 256}});
|
||||
auto* add_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Add", {input1_arg, input2_arg}, {add_out});
|
||||
for (auto p : has_axes_and_has_steps_pairs) {
|
||||
int has_axes = std::get<0>(p);
|
||||
int has_steps = std::get<1>(p);
|
||||
bool expected_to_propagate = std::get<2>(p);
|
||||
|
||||
auto* starts_initializer = builder.MakeInitializer<int64_t>({1}, {0});
|
||||
auto* ends_initializer = builder.MakeInitializer<int64_t>({1}, {-1});
|
||||
auto* axes_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
auto* steps_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
auto* slice_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Slice", {add_out, starts_initializer, ends_initializer, axes_initializer, steps_initializer},
|
||||
{slice_out});
|
||||
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
|
||||
InlinedVector<int64_t> starts_indices;
|
||||
auto pre_graph_checker = [&starts_indices](Graph& graph) -> Status {
|
||||
auto op_count_pre = CountOpsInGraph(graph);
|
||||
TEST_RETURN_IF_NOT(op_count_pre.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Transpose"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Slice"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
|
||||
|
||||
auto* identity_out = builder.MakeOutput();
|
||||
builder.AddNode("Identity", {slice_out}, {identity_out});
|
||||
};
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Slice") {
|
||||
TEST_RETURN_IF_NOT(starts_indices.empty());
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(node.InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, starts_indices,
|
||||
require_constant));
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
};
|
||||
|
||||
std::unique_ptr<GraphTransformer> transformer = std::make_unique<UpStreamGatherGraphTransformer>();
|
||||
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger, std::move(transformer),
|
||||
TransformerLevel::Level1,
|
||||
1, pre_graph_checker, post_graph_checker));
|
||||
auto post_graph_checker = [&starts_indices, expected_to_propagate](Graph& graph) {
|
||||
auto op_count_post = CountOpsInGraph(graph);
|
||||
|
||||
TEST_RETURN_IF_NOT(op_count_post.size() == 3U);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Transpose"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Slice"] == 1);
|
||||
TEST_RETURN_IF_NOT(op_count_post["Identity"] == 1);
|
||||
|
||||
for (Node& node : graph.Nodes()) {
|
||||
if (node.OpType() == "Transpose") {
|
||||
const auto& input_defs = node.InputDefs();
|
||||
|
||||
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
|
||||
if (expected_to_propagate) {
|
||||
TEST_RETURN_IF_NOT(producer_node != nullptr);
|
||||
TEST_RETURN_IF_NOT(producer_node->OpType() == "Slice");
|
||||
|
||||
InlinedVector<int64_t> values;
|
||||
constexpr bool require_constant = true;
|
||||
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
|
||||
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values,
|
||||
require_constant));
|
||||
for (size_t i = 0; i < values.size(); i++) {
|
||||
TEST_RETURN_IF_NOT(values[i] == starts_indices[i]);
|
||||
}
|
||||
|
||||
const ONNX_NAMESPACE::TensorShapeProto* slice_out_shape = producer_node->OutputDefs()[0]->Shape();
|
||||
TEST_RETURN_IF_NOT(slice_out_shape != nullptr);
|
||||
TEST_RETURN_IF_NOT(slice_out_shape->dim_size() == 4);
|
||||
TEST_RETURN_IF_NOT(utils::HasDimValue(slice_out_shape->dim(0)) && slice_out_shape->dim(0).dim_value() == 2);
|
||||
TEST_RETURN_IF_NOT(utils::HasDimValue(slice_out_shape->dim(1)) && slice_out_shape->dim(1).dim_value() == 4);
|
||||
TEST_RETURN_IF_NOT(utils::HasDimValue(slice_out_shape->dim(2)) && slice_out_shape->dim(2).dim_value() == 31);
|
||||
TEST_RETURN_IF_NOT(utils::HasDimValue(slice_out_shape->dim(3)) && slice_out_shape->dim(3).dim_value() == 256);
|
||||
} else {
|
||||
TEST_RETURN_IF_NOT(producer_node == nullptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
};
|
||||
|
||||
auto build_test_case = [has_axes, has_steps](ModelTestBuilder& builder) {
|
||||
auto* input1_arg = builder.MakeInput<int64_t>({{2, 4, 32, 256}});
|
||||
auto* trans_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Transpose", {input1_arg}, {trans_out})
|
||||
.AddAttribute("perm", std::vector<int64_t>{0, 2, 1, 3});
|
||||
|
||||
std::vector<NodeArg*> slice_inputs;
|
||||
|
||||
auto* starts_initializer = builder.MakeInitializer<int64_t>({1}, {0});
|
||||
auto* ends_initializer = builder.MakeInitializer<int64_t>({1}, {-1});
|
||||
|
||||
slice_inputs = {trans_out, starts_initializer, ends_initializer};
|
||||
|
||||
NodeArg* axes_initializer = nullptr;
|
||||
NodeArg* steps_initializer = nullptr;
|
||||
if (has_axes == 0 && has_steps == 0) {
|
||||
// nothing
|
||||
} else if (has_axes == 1 && has_steps == 0) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
} else if (has_axes == 1 && has_steps == 1) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 1 && has_steps == 2) {
|
||||
axes_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 0) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 1) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeInitializer<int64_t>({1}, {1});
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
} else if (has_axes == 2 && has_steps == 2) {
|
||||
axes_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(axes_initializer);
|
||||
steps_initializer = builder.MakeEmptyInput();
|
||||
slice_inputs.push_back(steps_initializer);
|
||||
}
|
||||
|
||||
auto* slice_out = builder.MakeIntermediate();
|
||||
builder.AddNode("Slice", slice_inputs,
|
||||
{slice_out});
|
||||
|
||||
auto* identity_out = builder.MakeOutput();
|
||||
builder.AddNode("Identity", {slice_out}, {identity_out});
|
||||
};
|
||||
|
||||
std::unique_ptr<GraphTransformer> transformer = std::make_unique<UpStreamGatherGraphTransformer>();
|
||||
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger, std::move(transformer),
|
||||
TransformerLevel::Level1,
|
||||
1, pre_graph_checker, post_graph_checker));
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
|
|
|
|||
Loading…
Reference in a new issue