Skip to content

Preserve double constants in graph fusions #32983

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

Open
Yuhx141 wants to merge 6 commits into
base: main
Choose a base branch
from
+504 −99
Open

Preserve double constants in graph fusions #32983

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d5020e0
Preserve double constants in graph fusions
Yuhx141 Sep 30, 2026
ba7bbf5
Fix fusion precision CI regressions
Yuhx141 Sep 30, 2026
02e42fa
Preserve double scale precision in fusions
Yuhx141 Oct 1, 2026
aadacfc
Fix unused LayerNorm epsilon constant
Yuhx141 Oct 1, 2026
bf0854c
Preserve exact LayerNorm epsilon conversions
Yuhx141 Oct 2, 2026
d6b3bbf
Reject lossy double scales before combining MatMul alpha
Yuhx141 Oct 11, 2026
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
96 changes: 56 additions & 40 deletions onnxruntime/core/optimizer/layer_norm_fusion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@ namespace onnxruntime {

// LayerNorm supports limited data types.
static constexpr std::array<std::string_view, 4> supported_data_types{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"};
// Default epsilon
static constexpr float DEFAULT_LAYERNORM_EPSILON = 1e-5f;

static bool IsSupportedDataType(const Node& node, int first_n_inputs = -1) {
int input_index = 0;
Expand Down Expand Up @@ -79,7 +77,8 @@ static std::vector<int64_t> GetAxesFromReduceMeanNode(Node& reduce_mean_node, co
return axes_values;
};

static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) {
template <typename T>
static bool TryGetScalarInitializer(const Graph& graph, const NodeArg& node_arg, T& value) {
const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, node_arg.Name());
if (tensor_proto == nullptr) {
return false;
Expand All @@ -92,54 +91,78 @@ static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& n

switch (tensor_proto->data_type()) {
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
value = static_cast<double>(initializer.data<float>()[0]);
value = static_cast<T>(initializer.data<float>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
value = static_cast<double>(initializer.data<MLFloat16>()[0]);
value = static_cast<T>(initializer.data<MLFloat16>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE:
value = initializer.data<double>()[0];
value = static_cast<T>(initializer.data<double>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16:
value = static_cast<double>(initializer.data<BFloat16>()[0]);
value = static_cast<T>(initializer.data<BFloat16>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_INT8:
value = static_cast<double>(initializer.data<int8_t>()[0]);
value = static_cast<T>(initializer.data<int8_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_INT16:
value = static_cast<double>(initializer.data<int16_t>()[0]);
value = static_cast<T>(initializer.data<int16_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_INT32:
value = static_cast<double>(initializer.data<int32_t>()[0]);
value = static_cast<T>(initializer.data<int32_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_INT64:
value = static_cast<double>(initializer.data<int64_t>()[0]);
value = static_cast<T>(initializer.data<int64_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_UINT8:
value = static_cast<double>(initializer.data<uint8_t>()[0]);
value = static_cast<T>(initializer.data<uint8_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_UINT16:
value = static_cast<double>(initializer.data<uint16_t>()[0]);
value = static_cast<T>(initializer.data<uint16_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_UINT32:
value = static_cast<double>(initializer.data<uint32_t>()[0]);
value = static_cast<T>(initializer.data<uint32_t>()[0]);
return true;
case ONNX_NAMESPACE::TensorProto_DataType_UINT64:
value = static_cast<double>(initializer.data<uint64_t>()[0]);
value = static_cast<T>(initializer.data<uint64_t>()[0]);
return true;
default:
return false;
}
}

static bool TryGetScalarInitializerOrFloatCastAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) {
if (TryGetScalarInitializer(graph, node_arg, value)) {
return true;
}

const Node* cast_node = graph.GetProducerNode(node_arg.Name());
if (cast_node == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) ||
cast_node->InputDefs().empty() || cast_node->InputDefs()[0] == nullptr) {
return false;
}

const auto to_it = cast_node->GetAttributes().find("to");
float cast_value = 0.0f;
if (to_it == cast_node->GetAttributes().end() ||
to_it->second.i() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT ||
!TryGetScalarInitializer(graph, *cast_node->InputDefs()[0], cast_value)) {
return false;
}

value = static_cast<double>(cast_value);
return true;
}

static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) {
const auto& pow_inputs = pow_node.InputDefs();
if (pow_inputs.size() < 2 || pow_inputs[1] == nullptr) {
return false;
}

double exponent_value = 0.0;
if (TryGetScalarInitializerAsDouble(graph, *pow_inputs[1], exponent_value)) {
if (TryGetScalarInitializer(graph, *pow_inputs[1], exponent_value)) {
return exponent_value == 2.0;
}

Expand All @@ -150,7 +173,7 @@ static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) {
return false;
}

return TryGetScalarInitializerAsDouble(graph, *exponent_input_node->InputDefs()[0], exponent_value) &&
return TryGetScalarInitializer(graph, *exponent_input_node->InputDefs()[0], exponent_value) &&
exponent_value == 2.0;
}

Expand Down Expand Up @@ -562,20 +585,21 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
continue;
}

const NodeArg* epsilon_input = GetOtherAddInput(add2_node, *reduce_mean2_node.MutableOutputDefs()[0]);
double epsilon = 0.0;
if (epsilon_input == nullptr ||
!TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) ||
static_cast<double>(static_cast<float>(epsilon)) != epsilon) {
continue;
}

InlinedVector<NodeArg*> layer_norm_input_defs{x_input, scale, bias};
Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/LayerNormFusion/"),
"LayerNormalization",
"fused LayerNorm subgraphs ",
layer_norm_input_defs,
{}, mul_node, nullptr, kOnnxDomain);

// Get constant "epsilon" from "Add2" node if available. Else, default value will be used.
double epsilon = 0.0;
if (TryGetScalarInitializerAsDouble(graph, *add2_node.MutableInputDefs()[1], epsilon)) {
layer_norm_node.AddAttribute("epsilon", static_cast<float>(epsilon));
} else {
layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON);
}
layer_norm_node.AddAttribute("epsilon", static_cast<float>(epsilon));

// The axis definition of layer_norm is ranging from axis to the last dim
layer_norm_node.AddAttribute("axis", static_cast<int64_t>(axes_values[0]));
Expand Down Expand Up @@ -814,25 +838,17 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr
continue;
}

double epsilon = 0.0;
if (!TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) ||
static_cast<double>(static_cast<float>(epsilon)) != epsilon) {
continue;
}

InlinedVector<NodeArg*> layer_norm_input_defs{x_input, scale};
Node& layer_norm_node =
graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/SimplifiedLayerNormFusion/"), "SimplifiedLayerNormalization",
"fused LayerNorm subgraphs ", layer_norm_input_defs, {}, mul_node, nullptr, kOnnxDomain);

// Get constant "epsilon" from "Add" node if available. Else, default value will be used.
const ONNX_NAMESPACE::TensorProto* tensor_proto =
graph_utils::GetConstantInitializer(graph, epsilon_input->Name());
if (tensor_proto != nullptr && tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
Initializer initializer{graph, *tensor_proto, graph.ModelPath()};
// epsilon must be a scalar/1-element tensor; fall back to default otherwise.
if (initializer.size() == 1) {
layer_norm_node.AddAttribute("epsilon", initializer.data<float>()[0]);
} else {
layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON);
}
} else {
layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON);
}
layer_norm_node.AddAttribute("epsilon", static_cast<float>(epsilon));

// Set stash_type to double if any input is double, default value if float.
if (x_input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE ||
Expand Down
106 changes: 87 additions & 19 deletions onnxruntime/core/optimizer/matmul_scale_fusion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@

#include "core/optimizer/matmul_scale_fusion.h"

#include <limits>
#include <type_traits>

#include "onnx/defs/attr_proto_util.h"

#include "core/common/optional.h"
Expand All @@ -16,12 +19,52 @@ namespace onnxruntime {

namespace {
template <typename T>
struct ExtractScalarAsFloatDispatchTarget {
bool IsIntegralExactlyRepresentableAsFloat(T value) {
static_assert(std::is_integral_v<T>);
using U = std::make_unsigned_t<T>;
const U unsigned_value = static_cast<U>(value);
const U magnitude = [&]() {
if constexpr (std::is_signed_v<T>) {
return value < 0 ? U{0} - unsigned_value : unsigned_value;
} else {
return unsigned_value;
}
}();

size_t significant_bits = 0;
for (U remaining = magnitude; remaining != 0; remaining >>= 1) {
++significant_bits;
}

constexpr size_t float_digits = std::numeric_limits<float>::digits;
if (significant_bits <= float_digits) {
return true;
}

const size_t discarded_bits = significant_bits - float_digits;
return (magnitude & ((U{1} << discarded_bits) - 1)) == 0;
}

struct ScalarScale {
double value;
bool is_double;
};

template <typename T>
struct ExtractScalarDispatchTarget {
Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path,
float& scalar_float) {
ScalarScale& scale, bool& lossless) {
T scalar;
ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1));
scalar_float = static_cast<float>(scalar);
const float scalar_float = static_cast<float>(scalar);
if constexpr (std::is_integral_v<T>) {
lossless = IsIntegralExactlyRepresentableAsFloat(scalar);
} else {
lossless = true;
}
Comment thread
Yuhx141 marked this conversation as resolved.
Outdated
scale = {std::is_same_v<T, double> ? static_cast<double>(scalar)
: static_cast<double>(scalar_float),
std::is_same_v<T, double>};
return Status::OK();
}
};
Comment thread
Copilot marked this conversation as resolved.
Expand All @@ -41,7 +84,7 @@ std::optional<TensorShape> GetTensorShape(const NodeArg& node_arg) {
// if it does not have any broadcasting effect on the Mul or Div output shape.
// Because the dimension lengths can only be 1, we only need to consider additional leading dimensions being prepended.
// `max_rank` should be set to the rank of the other Mul or Div input to avoid that.
std::optional<float> GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) {
std::optional<ScalarScale> GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) {
const auto* initializer = graph_utils::GetConstantInitializer(graph, node_arg.Name());

if (!initializer) {
Expand All @@ -59,18 +102,20 @@ std::optional<float> GetScalarConstantInitializer(const Graph& graph, const Node
return {};
}

float scalar{};
ScalarScale scale{};
bool lossless = false;
utils::MLTypeCallDispatcher<
uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double, BFloat16>
dispatcher{initializer->data_type()};
ORT_THROW_IF_ERROR(
(dispatcher.InvokeRet<Status, ExtractScalarAsFloatDispatchTarget>(*initializer, graph.ModelPath(), scalar)));
(dispatcher.InvokeRet<Status, ExtractScalarDispatchTarget>(
*initializer, graph.ModelPath(), scale, lossless)));

return {scalar};
return lossless ? std::optional<ScalarScale>{scale} : std::nullopt;
}

// gets the scale value and its input index if node is a fusable scale (Mul or Div by scalar constant)
std::optional<std::pair<float, int>> GetScaleFromNode(
std::optional<std::pair<ScalarScale, int>> GetScaleFromNode(
const Graph& graph, const Node& scale_node,
const InlinedHashSet<std::string>& excluded_initializer_names) {
const auto is_excluded_initializer =
Expand All @@ -95,7 +140,11 @@ std::optional<std::pair<float, int>> GetScaleFromNode(

if (!divisor.has_value()) return std::nullopt;

return {std::make_pair(1.0f / *divisor, scale_reciprocal_arg_index)};
ScalarScale scale = *divisor;
scale.value = scale.is_double
? 1.0 / scale.value
: static_cast<double>(1.0f / static_cast<float>(scale.value));
return {std::make_pair(scale, scale_reciprocal_arg_index)};
}

if (graph_utils::IsSupportedOptypeVersionAndDomain(scale_node, "Mul", {7, 13, 14})) {
Expand Down Expand Up @@ -128,7 +177,7 @@ struct ScaleMergeInfo {
// the edge from the base node to the original node
Node::EdgeConstIterator node_to_merge_edge;
// the scale of the original node
float scale;
ScalarScale scale;
// the index of the input or output def on the original node
// this def is moved to the fused node
// for a leading scale (scale -> MatMul), it will be the unscaled input
Expand Down Expand Up @@ -236,15 +285,34 @@ Status ProcessNode(

{
ONNX_NAMESPACE::AttributeProto& alpha_attr = fused_node_attrs["alpha"];
float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f;
auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) {
total_scale *= fusion.scale;
};

std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale);
std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale);

alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale);
const bool uses_double_scale =
std::any_of(input_node_merges.begin(), input_node_merges.end(),
[](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; }) ||
std::any_of(output_node_merges.begin(), output_node_merges.end(),
[](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; });

if (uses_double_scale) {
double total_scale = utils::HasFloat(alpha_attr) ? static_cast<double>(alpha_attr.f()) : 1.0;
auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) {
total_scale *= fusion.scale.value;
};
std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale);
std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale);

const float fused_scale = static_cast<float>(total_scale);
if (static_cast<double>(fused_scale) != total_scale) {
return Status::OK();
}
alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", fused_scale);
} else {
float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f;
auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) {
total_scale *= static_cast<float>(fusion.scale.value);
};
std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale);
std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale);
alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale);
}
}

auto get_mutable_node_to_merge = [&graph](const ScaleMergeInfo& merge) -> Node& {
Expand Down
Loading
Loading