mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Refactor the cost check used by the transpose optimizer (#14690)
### Description <!-- Describe your changes. --> Refactor the cost check used by the transpose optimizer to separate out ORT specific logic. Change the post-layout transform optimization to only skip the cost check when moving the layout transform nodes around. Fall back to the normal cost check for all other transpose nodes. Cleanup some const correctness. Refactor usage of ResizeHandler slightly so the clang-formatting is nicer. ### 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. --> Address performance issue seen in SNPE model where a non-layout transpose node was moved. See https://github.com/microsoft/onnxruntime/pull/14547 for more details. Improve separation between generic transpose optimization code and any ORT specific code.
This commit is contained in:
Scott McKay
GitHub
parent
755161100a
commit
b234df3dd0
5 changed files with 196 additions and 87 deletions
|
|
@ -74,7 +74,8 @@ Status NhwcTransformer::ApplyImpl(Graph& graph, bool& modified, int graph_level,
|
|||
}
|
||||
|
||||
if (modified) {
|
||||
Optimize(*api_graph, /*allow_extended_ops*/ true, kCpuExecutionProvider);
|
||||
Optimize(*api_graph, /*allow_extended_ops*/ true, kCpuExecutionProvider, OptimizerMode::OPTIMIZE_TRANSPOSE,
|
||||
OrtEPCostCheck);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
|
|||
|
|
@ -3,13 +3,14 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
namespace onnx_layout_transformation {
|
||||
namespace api {
|
||||
|
|
@ -229,7 +230,7 @@ class NodeRef {
|
|||
/// not assigned to any EP.
|
||||
/// </summary>
|
||||
/// <returns>EP type or empty string</returns>
|
||||
virtual const std::string& GetExecutionProviderType() const = 0;
|
||||
virtual std::string_view GetExecutionProviderType() const = 0;
|
||||
|
||||
/// <summary>
|
||||
/// Returns the schema since version for the op_type of this node. Value of -1 means it is not set.
|
||||
|
|
@ -439,6 +440,30 @@ class GraphRef {
|
|||
constexpr int64_t kMinSupportedOpset = 7;
|
||||
constexpr int64_t kMaxSupportedOpset = 18;
|
||||
|
||||
// enum of results that a CostCheckFn can return.
|
||||
enum class CostCheckResult {
|
||||
kStop, // pushing Transpose is expected to negatively impact performance
|
||||
kPushTranspose, // pushing Transpose is expected to improve performance
|
||||
kFallThrough // fall through to default cost check
|
||||
};
|
||||
|
||||
/// <summary>
|
||||
/// Function to allow overriding the default cost check to determine whether it is worth pushing a Transpose through
|
||||
/// a node.
|
||||
/// </summary>
|
||||
/// <param name="graph">The graph being optimized</param>
|
||||
/// <param name="node">The node we're considering pushing a Transpose through</param>
|
||||
/// <param name="perm">The perm value of the Transpose</param>
|
||||
/// <param name="outputs_leading_to_transpose">The set of outputs that lead to another Transpose in the graph.
|
||||
/// If we can successfully push the Transpose until it meets another Transpose they can either cancel each other out,
|
||||
/// or be merged into a single Transpose.
|
||||
/// </param>
|
||||
using CostCheckFn =
|
||||
std::function<CostCheckResult(const api::GraphRef& graph,
|
||||
const api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm,
|
||||
const std::unordered_set<std::string>& outputs_leading_to_transpose)>;
|
||||
|
||||
enum class OptimizerMode {
|
||||
OPTIMIZE_TRANSPOSE, // simple transpose optimization
|
||||
OPTIMIZE_LAYOUT_TRANSFORM // transpose optimization post layout transformation
|
||||
|
|
@ -469,6 +494,8 @@ struct OptimizeResult {
|
|||
/// <param name="provider_type">Execution provider if applicable.</param>
|
||||
/// <param name="mode">Current mode. Optimizer can be called in the context of transpose optimizations or during
|
||||
/// layout transformations.</param>
|
||||
/// <param name="cost_check_fn">Optional cost checking function to determine whether it is worth pushing a Transpose
|
||||
/// through a node.</param>
|
||||
/// <param name="layout_sensitive_ops">List of ops which are treated as layout sensitive by the ONNX standard
|
||||
/// as well as any runtime specific ops. These ops should be provided when mode is set to OPTIMIZE_LAYOUT_TRANSFORM.
|
||||
/// If these ops are not provided, transpose optimizer may convert the layout for these ops </param>
|
||||
|
|
@ -477,6 +504,7 @@ struct OptimizeResult {
|
|||
OptimizeResult Optimize(api::GraphRef& graph, bool allow_extended_ops,
|
||||
const std::string& provider_type = "",
|
||||
OptimizerMode mode = OptimizerMode::OPTIMIZE_TRANSPOSE,
|
||||
CostCheckFn cost_check_fn = nullptr,
|
||||
const std::unordered_set<std::string_view>& layout_sensitive_ops = {});
|
||||
|
||||
/* Layout Transformation Tools
|
||||
|
|
|
|||
|
|
@ -90,7 +90,7 @@ class ApiNode final : public api::NodeRef {
|
|||
void CopyAttributes(const api::NodeRef& node) override;
|
||||
void ClearAttribute(std::string_view name) override;
|
||||
void SetInput(size_t i, std::string_view name) override;
|
||||
const std::string& GetExecutionProviderType() const override;
|
||||
std::string_view GetExecutionProviderType() const override;
|
||||
virtual int SinceVersion() const override;
|
||||
|
||||
private:
|
||||
|
|
@ -401,7 +401,7 @@ void ApiNode::SetInput(size_t i, std::string_view name) {
|
|||
}
|
||||
}
|
||||
|
||||
const std::string& ApiNode::GetExecutionProviderType() const {
|
||||
std::string_view ApiNode::GetExecutionProviderType() const {
|
||||
return node_.GetExecutionProviderType();
|
||||
}
|
||||
|
||||
|
|
@ -831,16 +831,57 @@ onnxruntime::Node& NodeFromApiNode(onnx_layout_transformation::api::NodeRef& nod
|
|||
return static_cast<ApiNode&>(node).Node();
|
||||
}
|
||||
|
||||
CostCheckResult OrtEPCostCheck(const api::GraphRef& graph, const api::NodeRef& node,
|
||||
const std::vector<int64_t>& /*perm*/,
|
||||
const std::unordered_set<std::string>& /*outputs_leading_to_transpose*/) {
|
||||
// special case some kernels based on the ORT implementation details
|
||||
if (node.GetExecutionProviderType() == kCpuExecutionProvider) {
|
||||
if (node.IsOp("MaxPool")) {
|
||||
// MaxPool has higher perf in the NHWC variant when supported. HandleMaxPool does the support checks.
|
||||
return CostCheckResult::kPushTranspose;
|
||||
}
|
||||
|
||||
if (node.IsOp("Resize")) {
|
||||
// Resize is included because it has higher perf in the NHWC variant when
|
||||
// the input X is 4D int8 tensor and the mode is linear
|
||||
auto X_value_info = graph.GetValueInfo(node.Inputs()[0]);
|
||||
auto X_shape = X_value_info->Shape();
|
||||
auto X_dtype = X_value_info->DType();
|
||||
auto mode = node.GetAttributeString("mode");
|
||||
if (X_shape && X_shape->size() == 4 &&
|
||||
(X_dtype == api::DataType::UINT8 || X_dtype == api::DataType::INT8) &&
|
||||
mode && *mode == "linear") {
|
||||
return CostCheckResult::kPushTranspose;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return CostCheckResult::kFallThrough;
|
||||
}
|
||||
|
||||
namespace layout_transformer {
|
||||
|
||||
// Layout sensitive NCHW ops. TransformLayoutForEP will wrap these with Transpose nodes to convert the input
|
||||
// data to NHWC and output data back to NCHW, and move the op to the internal NHWC domain (kMSInternalNHWCDomain).
|
||||
// The EP requesting these ops MUST be able to handle the node with the operator in the kMSInternalNHWCDomain.
|
||||
// Once all the layout sensitive ops requested by the EP are wrapped the transpose optimizer will attempt to remove
|
||||
// as many of the layout transposes as possible.
|
||||
const std::unordered_set<std::string_view>& GetORTLayoutSensitiveOps() {
|
||||
static std::unordered_set<std::string_view> ort_layout_sensitive_ops = []() {
|
||||
const auto& layout_sensitive_ops = onnx_layout_transformation::GetLayoutSensitiveOps();
|
||||
#if !defined(USE_CUDA) && !defined(USE_ROCM)
|
||||
std::unordered_set<std::string_view> ort_specific_ops = {"FusedConv", "QLinearAveragePool", "QLinearGlobalAveragePool"};
|
||||
#else
|
||||
std::unordered_set<std::string_view> ort_specific_ops = {"Resize", "FusedConv", "QLinearAveragePool", "QLinearGlobalAveragePool"};
|
||||
std::unordered_set<std::string_view> ort_specific_ops =
|
||||
{ "FusedConv",
|
||||
"QLinearAveragePool",
|
||||
"QLinearGlobalAveragePool"
|
||||
#if defined(USE_CUDA) || defined(USE_ROCM)
|
||||
// The CUDA/ROCM Resize kernel is layout sensitive as it only handles NCHW input.
|
||||
// The CPU kernel and ONNX spec are not limited to handling NCHW input so are not layout sensitive, and
|
||||
// onnx_layout_transformation::HandleResize is used.
|
||||
,
|
||||
"Resize"
|
||||
#endif
|
||||
};
|
||||
|
||||
ort_specific_ops.insert(layout_sensitive_ops.cbegin(), layout_sensitive_ops.cend());
|
||||
return ort_specific_ops;
|
||||
}();
|
||||
|
|
@ -848,6 +889,21 @@ const std::unordered_set<std::string_view>& GetORTLayoutSensitiveOps() {
|
|||
return ort_layout_sensitive_ops;
|
||||
}
|
||||
|
||||
// Cost check for aggressively pushing the Transpose nodes involved in the layout transformation further out.
|
||||
static CostCheckResult
|
||||
PostLayoutTransformCostCheck(const api::GraphRef& graph, const api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm,
|
||||
const std::unordered_set<std::string>& outputs_leading_to_transpose) {
|
||||
// we aggressively push the layout transpose nodes
|
||||
if (perm == ChannelFirstToLastPerm(perm.size()) ||
|
||||
perm == ChannelLastToFirstPerm(perm.size())) {
|
||||
return CostCheckResult::kPushTranspose;
|
||||
}
|
||||
|
||||
// for other nodes use the default ORT cost check
|
||||
return OrtEPCostCheck(graph, node, perm, outputs_leading_to_transpose);
|
||||
}
|
||||
|
||||
Status TransformLayoutForEP(Graph& graph, bool& modified, const IExecutionProvider& execution_provider) {
|
||||
// sub graph recurse will be added later
|
||||
auto api_graph = MakeApiGraph(graph, execution_provider.GetAllocator(OrtMemTypeDefault), nullptr);
|
||||
|
|
@ -924,6 +980,7 @@ Status TransformLayoutForEP(Graph& graph, bool& modified, const IExecutionProvid
|
|||
OptimizeResult result =
|
||||
onnx_layout_transformation::Optimize(*api_graph, /*allow_extended_ops*/ true, execution_provider.Type(),
|
||||
onnx_layout_transformation::OptimizerMode::OPTIMIZE_LAYOUT_TRANSFORM,
|
||||
PostLayoutTransformCostCheck,
|
||||
layout_sensitive_ops);
|
||||
if (result.error_msg) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Optimization after layout transformation failed: ",
|
||||
|
|
|
|||
|
|
@ -42,6 +42,24 @@ onnxruntime::Graph& GraphFromApiGraph(onnx_layout_transformation::api::GraphRef&
|
|||
/// <returns>ORT node</returns>
|
||||
onnxruntime::Node& NodeFromApiNode(onnx_layout_transformation::api::NodeRef& node);
|
||||
|
||||
/// <summary>
|
||||
/// Cost check function for transpose optimizer that takes into account implementation details of the
|
||||
/// ORT execution provider kernels.
|
||||
/// </summary>
|
||||
/// <param name="graph">The graph being optimized</param>
|
||||
/// <param name="node">The node we're considering pushing a Transpose through</param>
|
||||
/// <param name="perm">The perm value of the Transpose</param>
|
||||
/// <param name="outputs_leading_to_transpose">The set of outputs that lead to another Transpose in the graph.
|
||||
/// If we can successfully push the Transpose until it meets another Transpose they can either cancel each other out,
|
||||
/// or be merged into a single Transpose.
|
||||
/// </param>
|
||||
/// <returns>CostCheckResult indicating the action the transpose optimizer should perform.</returns>
|
||||
onnx_layout_transformation::CostCheckResult OrtEPCostCheck(
|
||||
const onnx_layout_transformation::api::GraphRef& graph,
|
||||
const onnx_layout_transformation::api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm,
|
||||
const std::unordered_set<std::string>& outputs_leading_to_transpose);
|
||||
|
||||
namespace layout_transformer {
|
||||
/// <summary>
|
||||
/// Gets a list of layout sensitive ops for ORT. This list contains onnx standard defined
|
||||
|
|
|
|||
|
|
@ -3,15 +3,14 @@
|
|||
|
||||
#include "optimizer_api.h"
|
||||
|
||||
#include "core/graph/constants.h"
|
||||
#include "core/common/make_string.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include "core/common/gsl.h"
|
||||
#include <iostream>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <cstring>
|
||||
|
||||
#include "core/common/gsl.h"
|
||||
#include "core/common/make_string.h"
|
||||
#include "core/graph/constants.h"
|
||||
|
||||
namespace onnx_layout_transformation {
|
||||
|
||||
|
|
@ -19,7 +18,7 @@ struct OptimizerCtx {
|
|||
int64_t opset;
|
||||
api::GraphRef& graph;
|
||||
bool allow_extended_ops;
|
||||
bool skip_cost_check;
|
||||
CostCheckFn cost_check_fn;
|
||||
const std::string provider_type;
|
||||
OptimizerMode mode;
|
||||
std::unordered_set<std::string_view> layout_sensitive_ops;
|
||||
|
|
@ -678,7 +677,7 @@ static void TransposeOutputs(OptimizerCtx& ctx, api::NodeRef& node, const std::v
|
|||
// and an unknown rank likely corresponds to a data-carrying (non-weight) tensor, which will be large.
|
||||
|
||||
// Given a value, returns the rank of the value excluding dimensions of value 1. Returns 5 if the rank is unknown.
|
||||
static int EstimateValueRank(api::GraphRef& graph, std::string_view input) {
|
||||
static int EstimateValueRank(const api::GraphRef& graph, std::string_view input) {
|
||||
auto value_info = graph.GetValueInfo(input);
|
||||
std::optional<std::vector<int64_t>> shape = value_info->Shape();
|
||||
if (shape == std::nullopt) {
|
||||
|
|
@ -696,7 +695,7 @@ static int EstimateValueRank(api::GraphRef& graph, std::string_view input) {
|
|||
static const HandlerInfo* GetHandler(api::NodeRef& node, bool allow_extended_ops);
|
||||
|
||||
// Returns true if the provided transpose node is only consumed by nodes we can likely push it through.
|
||||
static bool CanLikelyRemoveTranspose(api::GraphRef& graph, api::NodeRef& transpose) {
|
||||
static bool CanLikelyRemoveTranspose(const api::GraphRef& graph, api::NodeRef& transpose) {
|
||||
auto consumers = graph.GetValueConsumers(transpose.Outputs()[0]);
|
||||
if (!consumers->comprehensive) {
|
||||
return false;
|
||||
|
|
@ -711,7 +710,7 @@ static bool CanLikelyRemoveTranspose(api::GraphRef& graph, api::NodeRef& transpo
|
|||
|
||||
// Estimates the cost of transposing an input. Currently uses rank heuristic. Negative if transpose is removed.
|
||||
// Feel free to improve as needed.
|
||||
static int EstimateTransposeValueCost(api::GraphRef& graph, std::string_view input,
|
||||
static int EstimateTransposeValueCost(const api::GraphRef& graph, std::string_view input,
|
||||
const std::vector<int64_t>& perm_inv) {
|
||||
// Case 1: Transposing constants probably costs nothing.
|
||||
std::unique_ptr<api::TensorRef> constant = graph.GetConstant(input);
|
||||
|
|
@ -737,7 +736,8 @@ static int EstimateTransposeValueCost(api::GraphRef& graph, std::string_view inp
|
|||
}
|
||||
|
||||
// Estimates total cost of transposing a node's inputs. Negative if transposing is beneficial.
|
||||
static int EstimateTransposeInputsCost(api::GraphRef& graph, api::NodeRef& node, const std::vector<int64_t>& perm_inv,
|
||||
static int EstimateTransposeInputsCost(const api::GraphRef& graph, const api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm_inv,
|
||||
const std::vector<size_t>& input_indices) {
|
||||
auto inputs = node.Inputs();
|
||||
int cost = 0;
|
||||
|
|
@ -970,8 +970,16 @@ static void PermuteInput(api::GraphRef& graph, api::NodeRef& node, size_t i, con
|
|||
node.SetInput(i, gather_output);
|
||||
}
|
||||
|
||||
#if !defined(USE_CUDA) && !defined(USE_ROCM)
|
||||
static bool HandleResize(HandlerArgs& args) {
|
||||
static bool HandleResize([[maybe_unused]] HandlerArgs& args) {
|
||||
#if defined(USE_CUDA) || defined(USE_ROCM)
|
||||
// The CUDA Resize kernel requires that the input is NCHW, so we can't push a Transpose through a Resize
|
||||
// in ORT builds with CUDA enabled.
|
||||
// The ROCm EP is generated from the CUDA EP kernel so the same applies to builds with ROCm enabled.
|
||||
// TODO: Remove this special case once the CUDA Resize kernel is implemented "generically" (i.e.) aligning with the
|
||||
// generic nature of the ONNX spec.
|
||||
// See https://github.com/microsoft/onnxruntime/pull/10824 for a similar fix applied to the CPU Resize kernel.
|
||||
return false;
|
||||
#else
|
||||
auto inputs = args.node.Inputs();
|
||||
int64_t rank_int = gsl::narrow_cast<int64_t>(args.perm.size());
|
||||
|
||||
|
|
@ -997,10 +1005,10 @@ static bool HandleResize(HandlerArgs& args) {
|
|||
TransposeOutputs(args.ctx, args.node, args.perm);
|
||||
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
constexpr HandlerInfo resize_handler = {&FirstInput, &HandleResize};
|
||||
#endif
|
||||
|
||||
static bool HandlePad(HandlerArgs& args) {
|
||||
size_t rank = args.perm.size();
|
||||
|
|
@ -1623,7 +1631,6 @@ constexpr HandlerInfo max_pool_op_handler = {&FirstInput, &HandleMaxPool};
|
|||
|
||||
// TODO: check binary size of this and replace it with constexpr if large
|
||||
static const std::unordered_map<std::string_view, const HandlerInfo&> handler_map{
|
||||
|
||||
{"Cast", simple_node_handler},
|
||||
{"Exp", simple_node_handler},
|
||||
{"Identity", simple_node_handler},
|
||||
|
|
@ -1696,19 +1703,7 @@ static const std::unordered_map<std::string_view, const HandlerInfo&> handler_ma
|
|||
{"Split", split_handler},
|
||||
{"Shape", shape_handler},
|
||||
{"Pad", pad_handler},
|
||||
// The CUDA Resize kernel assumes that the input is NCHW and
|
||||
// Resize can't be supported in ORT builds with CUDA enabled.
|
||||
// TODO: Enable this once the CUDA Resize kernel is implemented
|
||||
// "generically" (i.e.) aligning with the generic nature of the
|
||||
// ONNX spec.
|
||||
// See https://github.com/microsoft/onnxruntime/pull/10824 for
|
||||
// a similar fix applied to the CPU Resize kernel.
|
||||
// Per tests included in #10824, the ROCM EP also generates
|
||||
// incorrect results when this handler is used, so the Resize
|
||||
// handler is not enabled even for those builds.
|
||||
#if !defined(USE_CUDA) && !defined(USE_ROCM)
|
||||
{"Resize", resize_handler},
|
||||
#endif
|
||||
{"ReduceSum", reduce_op_handler},
|
||||
|
||||
{"ReduceLogSum", reduce_op_handler},
|
||||
|
|
@ -1735,7 +1730,8 @@ static const std::unordered_map<std::string_view, const HandlerInfo&> handler_ma
|
|||
|
||||
{"QuantizeLinear", quantize_dequantize_linear_handler},
|
||||
{"DequantizeLinear", quantize_dequantize_linear_handler},
|
||||
{"Reshape", reshape_handler}};
|
||||
{"Reshape", reshape_handler},
|
||||
};
|
||||
|
||||
static const std::unordered_map<std::string_view, const HandlerInfo&> extended_handler_map{
|
||||
{"com.microsoft.QLinearReduceMean", reduce_op_handler},
|
||||
|
|
@ -1773,28 +1769,51 @@ static const HandlerInfo* GetHandler(api::NodeRef& node, bool allow_extended_ops
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
// Some op should be optimized any time there is a transpose as input and a handler is available.
|
||||
static bool CanNodeSkipCostCheck(const OptimizerCtx& ctx, const api::NodeRef& node) {
|
||||
static int CalculateCost(const api::GraphRef& graph, const api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm,
|
||||
const std::unordered_set<std::string>& outputs_leading_to_transpose,
|
||||
const HandlerInfo& info,
|
||||
const std::vector<size_t>& input_indices) {
|
||||
// We require the input cost (number of transposes before the op) and the total cost to strictly decrease.
|
||||
// Strict decrease of the input cost ensures the optimization is stable, since the total cost decrease is just an
|
||||
// estimate (the transpose after the op may or may not cancel with a subsequent transpose). We don't want
|
||||
// repeated runs of the optimizer to have a transpose toggle between two inputs of a binary op.
|
||||
int cost = EstimateTransposeInputsCost(graph, node, perm, input_indices);
|
||||
|
||||
if (cost < 0 && info.transposes_outputs) {
|
||||
// If the output will be transposed and won't ultimately cancel, factor in that cost.
|
||||
bool has_output_leading_to_transpose = false;
|
||||
auto outputs = node.Outputs();
|
||||
int out_cost = 0;
|
||||
// Having multiple outputs is rare. When it happens (Split), the total size of the outputs isn't much larger
|
||||
// than the largest input. Cost is rank currently, so just use the largest cost (rank) over all outputs.
|
||||
for (auto out : outputs) {
|
||||
out_cost = std::max(out_cost, EstimateValueRank(graph, out));
|
||||
if (outputs_leading_to_transpose.find(std::string(out)) != outputs_leading_to_transpose.end()) {
|
||||
has_output_leading_to_transpose = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!has_output_leading_to_transpose) {
|
||||
cost += out_cost;
|
||||
}
|
||||
}
|
||||
|
||||
return cost;
|
||||
}
|
||||
|
||||
// Default cost check. Returns `true` if pushing the Transpose through the node is considered to be beneficial.
|
||||
static bool ShouldPushTranspose(const api::GraphRef& graph, const api::NodeRef& node,
|
||||
const std::vector<int64_t>& perm,
|
||||
const std::unordered_set<std::string>& outputs_leading_to_transpose,
|
||||
const HandlerInfo& info,
|
||||
const std::vector<size_t> transposable_input_indices) {
|
||||
if (node.IsOp("Transpose")) {
|
||||
return true;
|
||||
}
|
||||
if (node.IsOp("MaxPool")) {
|
||||
// Inclusion of MaxPool is a hack because it has higher perf in the NHWC variant when supported.
|
||||
return true;
|
||||
}
|
||||
if (node.IsOp("Resize")) {
|
||||
// Resize is included because it has higher perf in the NHWC variant when
|
||||
// the input X is 4D int8 tensor and the mode is linear
|
||||
auto X_value_info = ctx.graph.GetValueInfo(node.Inputs()[0]);
|
||||
auto X_shape = X_value_info->Shape();
|
||||
auto X_dtype = X_value_info->DType();
|
||||
auto mode = node.GetAttributeString("mode");
|
||||
if (X_shape && X_shape->size() == 4 && (X_dtype == api::DataType::UINT8 || X_dtype == api::DataType::INT8) &&
|
||||
mode && *mode == "linear") {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
|
||||
int cost = CalculateCost(graph, node, perm, outputs_leading_to_transpose, info, transposable_input_indices);
|
||||
return cost < 0;
|
||||
}
|
||||
|
||||
// Finds a handler for the node and estimates the cost of pushing a transpose. Does so if deemed beneficial.
|
||||
|
|
@ -1812,35 +1831,20 @@ bool ProcessTranspose(OptimizerCtx& ctx, api::NodeRef& transpose, api::NodeRef&
|
|||
return false;
|
||||
}
|
||||
|
||||
if (!ctx.skip_cost_check && !CanNodeSkipCostCheck(ctx, node)) {
|
||||
// We require the input cost (number of transposes before the op) and the total cost to strictly decrease.
|
||||
// Strict decrease of the input cost ensures the optimization is stable, since the total cost decrease is just an
|
||||
// estimate (the transpose after the op may or may not cancel with a subsequent transpose). We don't want
|
||||
// repeated runs of the optimizer to have a transpose toggle between two inputs of a binary op.
|
||||
int cost = EstimateTransposeInputsCost(ctx.graph, node, perm, input_indices);
|
||||
CostCheckResult cost = CostCheckResult::kFallThrough;
|
||||
|
||||
if (cost < 0 && info->transposes_outputs) {
|
||||
// If the output will be transposed and won't ultimately cancel, factor in that cost.
|
||||
bool has_output_leading_to_transpose = false;
|
||||
auto outputs = node.Outputs();
|
||||
int out_cost = 0;
|
||||
// Having multiple outputs is rare. When it happens (Split), the total size of the outputs isn't much larger
|
||||
// than the largest input. Cost is rank currently, so just use the largest cost (rank) over all outputs.
|
||||
for (auto out : outputs) {
|
||||
out_cost = std::max(out_cost, EstimateValueRank(ctx.graph, out));
|
||||
if (outputs_leading_to_transpose.find(std::string(out)) != outputs_leading_to_transpose.end()) {
|
||||
has_output_leading_to_transpose = true;
|
||||
}
|
||||
}
|
||||
if (ctx.cost_check_fn) {
|
||||
cost = ctx.cost_check_fn(ctx.graph, node, perm, outputs_leading_to_transpose);
|
||||
}
|
||||
|
||||
if (!has_output_leading_to_transpose) {
|
||||
cost += out_cost;
|
||||
}
|
||||
}
|
||||
if (cost == CostCheckResult::kFallThrough) {
|
||||
cost = ShouldPushTranspose(ctx.graph, node, perm, outputs_leading_to_transpose, *info, input_indices)
|
||||
? CostCheckResult::kPushTranspose
|
||||
: CostCheckResult::kStop;
|
||||
}
|
||||
|
||||
if (cost >= 0) {
|
||||
return false;
|
||||
}
|
||||
if (cost == CostCheckResult::kStop) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<int64_t> perm_inv = InvertPerm(perm);
|
||||
|
|
@ -1850,7 +1854,9 @@ bool ProcessTranspose(OptimizerCtx& ctx, api::NodeRef& transpose, api::NodeRef&
|
|||
|
||||
// Returns nullopt if graph opset is unsupported.
|
||||
std::optional<OptimizerCtx> MakeOptimizerContext(api::GraphRef& graph, bool allow_extended_ops,
|
||||
const std::string& provider_type, OptimizerMode mode,
|
||||
const std::string& provider_type,
|
||||
OptimizerMode mode,
|
||||
CostCheckFn cost_check_fn,
|
||||
const std::unordered_set<std::string_view>& layout_sensitive_ops,
|
||||
std::string& error_msg) {
|
||||
auto opset = graph.Opset("");
|
||||
|
|
@ -1874,10 +1880,7 @@ std::optional<OptimizerCtx> MakeOptimizerContext(api::GraphRef& graph, bool allo
|
|||
}
|
||||
}
|
||||
|
||||
// during layout transformation we want to push the transposes as far out as possible.
|
||||
// it is important that the EP gets the entire graph in the layout it prefers.
|
||||
bool skip_cost_check = mode == OptimizerMode::OPTIMIZE_LAYOUT_TRANSFORM;
|
||||
OptimizerCtx ctx{*opset, graph, allow_extended_ops, skip_cost_check, provider_type, mode, layout_sensitive_ops};
|
||||
OptimizerCtx ctx{*opset, graph, allow_extended_ops, cost_check_fn, provider_type, mode, layout_sensitive_ops};
|
||||
return ctx;
|
||||
}
|
||||
|
||||
|
|
@ -2047,11 +2050,13 @@ const std::unordered_set<std::string_view>& GetLayoutSensitiveOps() {
|
|||
|
||||
OptimizeResult Optimize(api::GraphRef& graph, bool allow_extended_ops,
|
||||
const std::string& provider_type, OptimizerMode mode,
|
||||
CostCheckFn cost_check_fn,
|
||||
const std::unordered_set<std::string_view>& layout_sensitive_ops) {
|
||||
OptimizeResult result{};
|
||||
|
||||
std::string error_msg;
|
||||
auto ctx = MakeOptimizerContext(graph, allow_extended_ops, provider_type, mode, layout_sensitive_ops, error_msg);
|
||||
auto ctx = MakeOptimizerContext(graph, allow_extended_ops, provider_type, mode, cost_check_fn, layout_sensitive_ops,
|
||||
error_msg);
|
||||
if (ctx == std::nullopt) {
|
||||
if (!error_msg.empty()) {
|
||||
result.error_msg = error_msg;
|
||||
|
|
|
|||
Loading…
Reference in a new issue