Skip to content

allow tensor alias for Loop's carried params #497

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fs-eire wants to merge 3 commits into from
Closed

allow tensor alias for Loop's carried params #497

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1d60776
allow tensor alias for Loop's carried params
fs-eire Feb 15, 2019
2360bae
Merge remote-tracking branch 'ms/master' into allow-loop-subgraph-ten…
fs-eire Mar 6, 2019
0914ff4
generalize the allocation conditions
fs-eire Mar 8, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 57 additions & 9 deletions onnxruntime/core/framework/allocation_planner.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,8 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP

class PlannerImpl {
public:
PlannerImpl(const onnxruntime::GraphViewer& graph_viewer,
PlannerImpl(const Node* parent_node,
const onnxruntime::GraphViewer& graph_viewer,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
Expand All @@ -106,6 +107,7 @@ class PlannerImpl {
SequentialExecutionPlan& plan)
: context_{context},
plan_{plan},
parent_node_{parent_node},
graph_viewer_{graph_viewer},
outer_scope_node_args_{outer_scope_node_args},
execution_providers_{providers},
Expand All @@ -119,6 +121,7 @@ class PlannerImpl {
const ISequentialPlannerContext& context_;
SequentialExecutionPlan& plan_;

const Node* parent_node_;
const onnxruntime::GraphViewer& graph_viewer_;
const std::vector<const NodeArg*>& outer_scope_node_args_;
const ExecutionProviders& execution_providers_;
Expand Down Expand Up @@ -193,8 +196,7 @@ class PlannerImpl {
}

// Find if there exists some input tensor that we can use in-place for output_arg
bool FindReusableInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* reusable_input) {
auto p_output_arg = node.OutputDefs()[output_arg_num];
bool FindAliasInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* alias_input) {
const KernelCreateInfo* ci;
Status st = kernel_registry_.SearchKernelRegistry(node, &ci);
if (!st.IsOK() || ci == nullptr || ci->kernel_def == nullptr) {
Expand All @@ -209,14 +211,26 @@ class PlannerImpl {
if ((0 <= pair.first) && (static_cast<size_t>(pair.first) < input_args.size())) {
auto p_input_arg = input_args[pair.first];
if (p_input_arg->Exists()) {
*reusable_input = Index(p_input_arg->Name());
*alias_input = Index(p_input_arg->Name());
return true;
}
}
}
}

return false;
}

bool FindInplaceInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* inplace_input) {
auto p_output_arg = node.OutputDefs()[output_arg_num];
const KernelCreateInfo* ci;
Status st = kernel_registry_.SearchKernelRegistry(node, &ci);
if (!st.IsOK() || ci == nullptr || ci->kernel_def == nullptr) {
return false;
}

const std::vector<std::pair<int, int>>& inplace_map = ci->kernel_def->MayInplace();
auto& input_args = node.InputDefs();
for (auto pair : inplace_map) {
if (pair.second == output_arg_num) {
if ((0 <= pair.first) && (static_cast<size_t>(pair.first) < input_args.size())) {
Expand All @@ -227,14 +241,45 @@ class PlannerImpl {
if (1 == UseCount(original)) {
if (SameSize(*p_input_arg, *p_output_arg)) {
// we can reuse this input since it is its last use and permitted for in-place update
*reusable_input = input_arg_index; // or original; both should be okay
*inplace_input = input_arg_index; // or original; both should be okay
return true;
}
}
}
}
}
}

return false;
}

bool FindReusableGraphInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* reusable_input) {
// skip if this is main graph
if (parent_node_) {
// only check a single layer. ie. an Identity/Dropout node connecting graph's input and output
auto& graph_inputs = graph_viewer_.GetInputs();

@skottmckay skottmckay Feb 20, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you clarify this comment? Not quite sure what you mean by 'only check a single layer'. FindReusableInput looks at the alias and 'may inplace' maps and I'm not quite translating that into limiting the check to a 'single layer'.

@fs-eire fs-eire Feb 21, 2019 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic only checks one node instead of triversal the whole graph.

Triversal the whole graph will make this logic very complicated; for most of the cases, it's one Identity node connected the graph's input and output

@skottmckay skottmckay Feb 21, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A would expect the comment to be it just checks this graph rather than 'layer'. A graph could contain multiple layers of say RNNs, which I thought was a completely different concept.

Also I don't quite buy that 'most of the cases' will be one Identity node. If there's some state being carried between iterations, I would expect it is changed across each iteration and therefore there would be other nodes involved. I can understand unit test or very simple example models using an Identity node, but I wouldn't expect that in a real model.

In reply to: 258747489 [](ancestors = 258747489)

auto& graph_outputs = graph_viewer_.GetOutputs();

if (FindAliasInput(node, output_arg_num, reusable_input)) {
auto& arg = ml_value_info_.at(Buffer(*reusable_input)).p_def_site;
auto it = std::find(graph_inputs.cbegin(), graph_inputs.cend(), arg);
if (it != graph_inputs.end()) {
auto graph_input_index = std::distance(graph_inputs.cbegin(), it);
ORT_ENFORCE(graph_input_index >= 0 && static_cast<size_t>(graph_input_index) < graph_inputs.size());
auto graph_output_index = std::distance(graph_outputs.cbegin(),
std::find(graph_outputs.cbegin(), graph_outputs.cend(), node.OutputDefs()[output_arg_num]));

@skottmckay skottmckay Feb 20, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

node.OutputDefs()[output_arg_num]) [](start = 106, length = 34)

This should probably be outside the iteration so OutputDefs() and operator[] aren't called every time. Will probably get optimized out, but no need to rely on that.

ORT_ENFORCE(graph_output_index >= 0 && static_cast<size_t>(graph_output_index) < graph_outputs.size());

if (parent_node_->OpType() == "Loop") {
// the matching carried parameters
return graph_input_index >= 2;
} else {
// TODO: other operators like Scan
}

@skottmckay skottmckay Feb 20, 2019 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be better if this information is provided in a more generic manner so that we're not hardcoding special casing for a list of operators here.

e.g. something like the kernel def could be expanded and we lookup that information based on OpType.

Otherwise changes to operator specs (e.g. say Loop in opset 10 changes the order of things) are easily missed and the code becomes fragile.

@fs-eire fs-eire Feb 21, 2019

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am fully agree with the point. However I think we need a design agreement before I make a change since this change may apply to the core type KernelOp.

}
}
}

return false;
}

Expand Down Expand Up @@ -497,13 +542,15 @@ class PlannerImpl {
auto current = Index(node_output->Name());
AllocPlan(current).value_type = utils::GetMLDataType(*node_output);
MLValueIndex reused;
if (std::find(graph_outputs.begin(), graph_outputs.end(), node_output) != graph_outputs.end()) {
if (std::find(graph_outputs.begin(), graph_outputs.end(), node_output) != graph_outputs.end() &&
!FindReusableGraphInput(*pnode, output_arg_num, &reused)) {

@gramalingam gramalingam Feb 21, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand this fallthrough logic, because the lines below will invoke "FindReusableInput" again. (In this case, correctness depends on FindReusableGraph & FindReusableInput "finding" the same "reused".) I think it will be better to use a nested "if FindReusableGraphInput(…)" and do the right thing inside this if-statement than cascading to subsequent else-branches.

@gramalingam gramalingam Mar 8, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, How about this comment? Can we change this to: if (this is a graph output) { if (FindReusableGraphInput(…)) Reuse(reused, current) else ...kAllocateOutput" ?

// node_output is graph's output, so we can't reuse intermedia buffer
AllocPlan(current).alloc_kind = AllocKind::kAllocateOutput;
} else if (IsNonTensor(*node_output)) {
// we do not try sharing-optimization for non-tensors
AllocPlan(current).alloc_kind = AllocKind::kAllocate;
} else if (FindReusableInput(*pnode, output_arg_num, &reused)) {
} else if (FindAliasInput(*pnode, output_arg_num, &reused) ||
FindInplaceInput(*pnode, output_arg_num, &reused)) {
// Reuse one of this node's input buffers as the output buffer (for in-place update)
Reuse(reused, current);
} else if (!context_.EnableParallelExecution() && FindReusableTensor(*node_output, &reused)) {
Expand Down Expand Up @@ -610,7 +657,8 @@ Status PlannerImpl::CreatePlan() {
return Status::OK();
}

Status SequentialPlanner::CreatePlan(const onnxruntime::GraphViewer& graph_viewer,
Status SequentialPlanner::CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph_viewer,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
Expand All @@ -620,7 +668,7 @@ Status SequentialPlanner::CreatePlan(const onnxruntime::GraphViewer& graph_viewe
// allocate/reset here so we know it's clean
plan = std::make_unique<SequentialExecutionPlan>();

PlannerImpl planner(graph_viewer, outer_scope_node_args,
PlannerImpl planner(parent_node, graph_viewer, outer_scope_node_args,
providers, kernel_registry, mlvalue_name_idx_map, context, *plan);

return planner.CreatePlan();
Expand Down
8 changes: 5 additions & 3 deletions onnxruntime/core/framework/allocation_planner.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ class SequentialPlannerContext : public ISequentialPlannerContext {
class SequentialPlanner {
public:
// This API allows user to provide a custom planner context.
static Status CreatePlan(const onnxruntime::GraphViewer& graph,
static Status CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
Expand All @@ -60,14 +61,15 @@ class SequentialPlanner {

// This uses a standard planner context and is meant to be the primary API for creating a plan
// as the context is primarily used in test scenarios.
static Status CreatePlan(const onnxruntime::GraphViewer& graph,
static Status CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
const MLValueNameIdxMap& mlvalue_name_idx_map,
std::unique_ptr<SequentialExecutionPlan>& plan) {
SequentialPlannerContext context;
return CreatePlan(graph, outer_scope_node_args, providers, kernel_registry, mlvalue_name_idx_map, context, plan);
return CreatePlan(parent_node, graph, outer_scope_node_args, providers, kernel_registry, mlvalue_name_idx_map, context, plan);
}
};

Expand Down
7 changes: 4 additions & 3 deletions onnxruntime/core/framework/session_state_initializer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,8 @@ SessionStateInitializer::SessionStateInitializer(const std::basic_string<PATH_CH
kernel_registry_manager_{kernel_registry_manager},
logger_{session_state.Logger()} {}

common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>& outer_scope_node_args,
common::Status SessionStateInitializer::CreatePlan(const Node* parent_node,
const std::vector<NodeArg*>& outer_scope_node_args,
bool enable_sequential_execution) {
auto graph_viewer = std::make_unique<onnxruntime::GraphViewer>(graph_);

Expand All @@ -83,15 +84,15 @@ common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>&
// CreatePlan will create a new SequentialExecutionPlan instance that we will
// save into the session state.
ORT_RETURN_IF_ERROR(
SequentialPlanner::CreatePlan(*graph_viewer, valid_outer_scope_node_args, execution_providers_,
SequentialPlanner::CreatePlan(parent_node, *graph_viewer, valid_outer_scope_node_args, execution_providers_,
kernel_registry_manager_, mlvalue_name_idx_map, exec_plan));

session_state_.SetExecutionPlan(std::move(exec_plan));
} else {
// Parallel execution still uses same allocation plan, but has limitation of memory buffer reuse.
SequentialPlannerContext context(true /* enable parallel execution */);
ORT_RETURN_IF_ERROR(
SequentialPlanner::CreatePlan(*graph_viewer, valid_outer_scope_node_args, execution_providers_,
SequentialPlanner::CreatePlan(parent_node, *graph_viewer, valid_outer_scope_node_args, execution_providers_,
kernel_registry_manager_, mlvalue_name_idx_map, context, exec_plan));

session_state_.SetExecutionPlan(std::move(exec_plan));
Expand Down
4 changes: 3 additions & 1 deletion onnxruntime/core/framework/session_state_initializer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ class Graph;
class GraphTransformerManager;
class InsertCastTransformer;
class KernelRegistryManager;
class Node;
class NodeArg;
class SessionState;

Expand All @@ -33,7 +34,8 @@ class SessionStateInitializer {
KernelRegistryManager& kernel_registry_manager);

// First perform any transformations and create the execution plan
common::Status CreatePlan(const std::vector<NodeArg*>& outer_scope_node_args,
common::Status CreatePlan(const Node* parent_node,
const std::vector<NodeArg*>& outer_scope_node_args,
bool enable_sequential_execution);

// initialize tensors, and save. save kernels and input/output node mappings
Expand Down
5 changes: 3 additions & 2 deletions onnxruntime/core/providers/cpu/controlflow/loop.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,8 +328,9 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca
CreateInitialFeeds(feeds);

auto& iter_num_value = *iter_num_mlvalue_.GetMutable<Tensor>()->MutableData<int64_t>();
auto current_condition_mlvalue = condition_mlvalue_;

while (iter_num_value < max_trip_count_ && *condition_mlvalue_.GetMutable<Tensor>()->MutableData<bool>()) {
while (iter_num_value < max_trip_count_ && *current_condition_mlvalue.GetMutable<Tensor>()->MutableData<bool>()) {
if (iter_num_value != 0) {
UpdateFeeds(fetches, feeds);
fetches.clear();
Expand All @@ -353,7 +354,7 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca

ORT_RETURN_IF_ERROR(status);

condition_mlvalue_ = fetches[0];
current_condition_mlvalue = fetches[0];

++iter_num_value;
}
Expand Down
6 changes: 3 additions & 3 deletions onnxruntime/core/session/inference_session.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -344,8 +344,8 @@ class InferenceSession::Impl {
SessionStateInitializer initializer{model_location_, subgraph, *subgraph_session_state, execution_providers_,
kernel_registry_manager_};

ORT_RETURN_IF_ERROR(initializer.CreatePlan(node.ImplicitInputDefs(),
session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(initializer.CreatePlan(&node, node.ImplicitInputDefs(),
session_options_.enable_sequential_execution));

ORT_RETURN_IF_ERROR(initializer.InitializeAndSave(&node.ImplicitInputDefs()));

Expand Down Expand Up @@ -412,7 +412,7 @@ class InferenceSession::Impl {
// now that all the transforms are done, call Resolve on the main graph. this will recurse into the subgraphs.
ORT_RETURN_IF_ERROR(graph.Resolve());

ORT_RETURN_IF_ERROR(session_initializer.CreatePlan({}, session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(session_initializer.CreatePlan(nullptr, {}, session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(session_initializer.InitializeAndSave(nullptr));

// handle any subgraphs
Expand Down
2 changes: 1 addition & 1 deletion onnxruntime/test/framework/allocation_planner_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,7 +238,7 @@ class PlannerTest : public ::testing::Test {
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();

SequentialPlannerTestContext test_context(&shape_map_);
status = SequentialPlanner::CreatePlan(GraphViewer(graph_), outer_scope_node_args, execution_providers,
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph_), outer_scope_node_args, execution_providers,
kernel_registry_manager, mlvalue_name_idx_map, test_context, plan_);

EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
Expand Down
4 changes: 2 additions & 2 deletions onnxruntime/test/framework/execution_frame_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ TEST(ExecutionFrameTest, TensorAllocationTest) {

std::unique_ptr<SequentialExecutionPlan> p_seq_exec_plan;
// TODO below line is for testing only. In production use SequentialPlanner::CreatePlan()
status = SequentialPlanner::CreatePlan(GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
p_seq_exec_plan);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
state.SetExecutionPlan(std::move(p_seq_exec_plan));
Expand Down Expand Up @@ -214,7 +214,7 @@ TEST(ExecutionFrameTest, MemPatternTest) {
std::vector<float>(6, 1.0f), &v3);

std::unique_ptr<SequentialExecutionPlan> p_seq_exec_plan = std::make_unique<SequentialExecutionPlan>();
status = SequentialPlanner::CreatePlan(GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
p_seq_exec_plan);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();

Expand Down
2 changes: 1 addition & 1 deletion onnxruntime/test/providers/memcpy_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ TEST(MemcpyTest, copy1) {
PutAllNodesOnOneProvider(model.MainGraph(), onnxruntime::kCpuExecutionProvider);
SessionStateInitializer session_initializer{ORT_TSTR(""), model.MainGraph(), s, execution_providers,
kernel_registry_manager};
st = session_initializer.CreatePlan({}, true);
st = session_initializer.CreatePlan(nullptr, {}, true);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();
st = session_initializer.InitializeAndSave(nullptr);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();
Expand Down