diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index e9d9c7de45..29c8de161a 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -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(); diff --git a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api.h b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api.h index bb0974c2b9..7a66e17e60 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api.h +++ b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api.h @@ -3,13 +3,14 @@ #pragma once +#include #include #include #include #include #include -#include #include +#include namespace onnx_layout_transformation { namespace api { @@ -229,7 +230,7 @@ class NodeRef { /// not assigned to any EP. /// /// EP type or empty string - virtual const std::string& GetExecutionProviderType() const = 0; + virtual std::string_view GetExecutionProviderType() const = 0; /// /// 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 +}; + +/// +/// Function to allow overriding the default cost check to determine whether it is worth pushing a Transpose through +/// a node. +/// +/// The graph being optimized +/// The node we're considering pushing a Transpose through +/// The perm value of the 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. +/// +using CostCheckFn = + std::function& perm, + const std::unordered_set& 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 { /// Execution provider if applicable. /// Current mode. Optimizer can be called in the context of transpose optimizations or during /// layout transformations. +/// Optional cost checking function to determine whether it is worth pushing a Transpose +/// through a node. /// 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 @@ -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& layout_sensitive_ops = {}); /* Layout Transformation Tools diff --git a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api_impl.cc b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api_impl.cc index ec46893233..0f941596c9 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api_impl.cc +++ b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_api_impl.cc @@ -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(node).Node(); } +CostCheckResult OrtEPCostCheck(const api::GraphRef& graph, const api::NodeRef& node, + const std::vector& /*perm*/, + const std::unordered_set& /*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& GetORTLayoutSensitiveOps() { static std::unordered_set ort_layout_sensitive_ops = []() { const auto& layout_sensitive_ops = onnx_layout_transformation::GetLayoutSensitiveOps(); -#if !defined(USE_CUDA) && !defined(USE_ROCM) - std::unordered_set ort_specific_ops = {"FusedConv", "QLinearAveragePool", "QLinearGlobalAveragePool"}; -#else - std::unordered_set ort_specific_ops = {"Resize", "FusedConv", "QLinearAveragePool", "QLinearGlobalAveragePool"}; + std::unordered_set 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& 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& perm, + const std::unordered_set& 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: ", diff --git a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_utils.h b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_utils.h index 88861af56d..584b61f4bd 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/optimizer_utils.h +++ b/onnxruntime/core/optimizer/transpose_optimizer/optimizer_utils.h @@ -42,6 +42,24 @@ onnxruntime::Graph& GraphFromApiGraph(onnx_layout_transformation::api::GraphRef& /// ORT node onnxruntime::Node& NodeFromApiNode(onnx_layout_transformation::api::NodeRef& node); +/// +/// Cost check function for transpose optimizer that takes into account implementation details of the +/// ORT execution provider kernels. +/// +/// The graph being optimized +/// The node we're considering pushing a Transpose through +/// The perm value of the 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. +/// +/// CostCheckResult indicating the action the transpose optimizer should perform. +onnx_layout_transformation::CostCheckResult OrtEPCostCheck( + const onnx_layout_transformation::api::GraphRef& graph, + const onnx_layout_transformation::api::NodeRef& node, + const std::vector& perm, + const std::unordered_set& outputs_leading_to_transpose); + namespace layout_transformer { /// /// Gets a list of layout sensitive ops for ORT. This list contains onnx standard defined diff --git a/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc b/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc index 700c91ab85..4d8c881663 100644 --- a/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc +++ b/onnxruntime/core/optimizer/transpose_optimizer/transpose_optimizer.cc @@ -3,15 +3,14 @@ #include "optimizer_api.h" -#include "core/graph/constants.h" -#include "core/common/make_string.h" - #include -#include "core/common/gsl.h" #include #include #include -#include + +#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 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> 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& perm_inv) { // Case 1: Transposing constants probably costs nothing. std::unique_ptr 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& perm_inv, +static int EstimateTransposeInputsCost(const api::GraphRef& graph, const api::NodeRef& node, + const std::vector& perm_inv, const std::vector& 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(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 handler_map{ - {"Cast", simple_node_handler}, {"Exp", simple_node_handler}, {"Identity", simple_node_handler}, @@ -1696,19 +1703,7 @@ static const std::unordered_map 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 handler_ma {"QuantizeLinear", quantize_dequantize_linear_handler}, {"DequantizeLinear", quantize_dequantize_linear_handler}, - {"Reshape", reshape_handler}}; + {"Reshape", reshape_handler}, +}; static const std::unordered_map 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& perm, + const std::unordered_set& outputs_leading_to_transpose, + const HandlerInfo& info, + const std::vector& 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& perm, + const std::unordered_set& outputs_leading_to_transpose, + const HandlerInfo& info, + const std::vector 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 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 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& layout_sensitive_ops, std::string& error_msg) { auto opset = graph.Opset(""); @@ -1874,10 +1880,7 @@ std::optional 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& 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& 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;