From 5f300c4abee606b3e8684eac06bb939f2b6919b3 Mon Sep 17 00:00:00 2001 From: Lagos Date: Wed, 12 Aug 2026 15:52:54 -0400 Subject: [PATCH] Fix Tier-0 fusion failure on graphs with QLinearConcat Two bugs in the multi-op-translation path blocked QLinearConcat from fusing, forcing per-op fallback: 1. RebuildTensorDescPointers / RebuildSubNodePointers set op_desc.Desc = desc_storage.get(), assuming the operator desc is the first member of the storage object. The concat's primary dequant and Join sub-node point desc_storage at the shared QLinConcatStorage (first member is a vector), so op_desc.Desc was left pointing at the wrong bytes. Restore the correct pointer at the end of both fixups. 2. Sub-node graph_inputs wiring handled constants and partition-boundary inputs but not partition-internal producers, leaving such inputs with no graph edge. Add a value_producer branch in both Compile and TryCompilePartition. Also switch the concat's dequant/quantize decomposition to the modern DML_OPERATOR_DEQUANTIZE / DML_OPERATOR_QUANTIZE, which take per-tensor scalar scale/zero-point natively via QuantizationTensors. --- src/directx/dml_op_translators.cc | 67 +++++++++++++++++++++---------- src/directx/full_graph_fusion.cc | 45 +++++++++++++++++++++ 2 files changed, 91 insertions(+), 21 deletions(-) diff --git a/src/directx/dml_op_translators.cc b/src/directx/dml_op_translators.cc index 123e0ce..0964d84 100644 --- a/src/directx/dml_op_translators.cc +++ b/src/directx/dml_op_translators.cc @@ -9529,7 +9529,6 @@ static std::optional TranslateQLinearConcat( struct QLinConcatStorage { std::vector> sub_storages; DML_JOIN_OPERATOR_DESC join_desc{}; - DML_ELEMENT_WISE_QUANTIZE_LINEAR_OPERATOR_DESC quant_desc{}; std::vector deq_out_infos; std::vector deq_out_bufs; std::vector deq_out_tds; @@ -9543,13 +9542,16 @@ static std::optional TranslateQLinearConcat( auto storage = std::make_shared(); storage->axis = dml_axis; - // Build all input tensors (y_scale, y_zp, then all tuples). + // Build all input tensors (y_scale, y_zp, then all tuples). Uses the modern + // DML_OPERATOR_DEQUANTIZE / DML_OPERATOR_QUANTIZE, which accept per-tensor + // scalar scale/zero-point natively (via QuantizationTensors) — so scale/zp + // keep their raw sizes, no broadcast needed. TranslatedOp result; for (size_t i = 0; i < inputs.size(); ++i) { if (inputs[i].empty() || !value_shapes.count(inputs[i])) return std::nullopt; auto* info = LookupShape(value_shapes, inputs[i]); if (!info) return std::nullopt; - auto t = MakeTensorInfo(info->sizes, info->data_type); + DmlTensorInfo t = MakeTensorInfo(info->sizes, info->data_type); result.input_tensors.push_back(t); result.input_buffer_descs.push_back(t.ToBufferDesc()); result.input_name_reorder.push_back(i); @@ -9611,22 +9613,31 @@ static std::optional TranslateQLinearConcat( storage->deq_out_bufs.push_back(deq0_out.ToBufferDesc()); storage->deq_out_tds.push_back({ DML_TENSOR_TYPE_BUFFER, &storage->deq_out_bufs.back() }); - struct DeqStorage { DML_ELEMENT_WISE_DEQUANTIZE_LINEAR_OPERATOR_DESC desc{}; }; + struct DeqStorage { + DML_DEQUANTIZE_OPERATOR_DESC desc{}; + std::vector quant_tensor_descs; + }; auto deq0_store = std::make_shared(); storage->sub_storages.push_back(deq0_store); result.desc_storage = storage; - result.op_desc = { DML_OPERATOR_ELEMENT_WISE_DEQUANTIZE_LINEAR, &deq0_store->desc }; + deq0_store->desc.QuantizationType = DML_QUANTIZATION_TYPE_SCALE_ZERO_POINT; + result.op_desc = { DML_OPERATOR_DEQUANTIZE, &deq0_store->desc }; result.fixup = [storage, deq0_store](TranslatedOp& self) { RebuildTensorDescPointers(self); storage->deq_out_bufs[0].Sizes = storage->deq_out_infos[0].sizes.data(); storage->deq_out_bufs[0].Strides = storage->deq_out_infos[0].strides.empty() ? nullptr : storage->deq_out_infos[0].strides.data(); storage->deq_out_tds[0] = { DML_TENSOR_TYPE_BUFFER, &storage->deq_out_bufs[0] }; - deq0_store->desc.InputTensor = &self.input_tensor_descs[0]; - deq0_store->desc.ScaleTensor = &self.input_tensor_descs[1]; - deq0_store->desc.ZeroPointTensor= &self.input_tensor_descs[2]; - deq0_store->desc.OutputTensor = &storage->deq_out_tds[0]; + deq0_store->quant_tensor_descs = { self.input_tensor_descs[1], self.input_tensor_descs[2] }; + deq0_store->desc.InputTensor = &self.input_tensor_descs[0]; + deq0_store->desc.QuantizationTensorCount = static_cast(deq0_store->quant_tensor_descs.size()); + deq0_store->desc.QuantizationTensors = deq0_store->quant_tensor_descs.data(); + deq0_store->desc.OutputTensor = &storage->deq_out_tds[0]; + // RebuildTensorDescPointers set op_desc.Desc = desc_storage.get() (the + // QLinConcatStorage). The primary's desc lives in the separate deq0_store, + // so restore the correct pointer. + self.op_desc.Desc = &deq0_store->desc; }; result.FixupPointers(); @@ -9649,7 +9660,8 @@ static std::optional TranslateQLinearConcat( deq_node.input_from = { {-2, 0}, {-2, 0}, {-2, 0} }; deq_node.graph_inputs = { {onnx_base, 0}, {onnx_base + 1, 1}, {onnx_base + 2, 2} }; - // Fill input tensor info for desc pointers. + // Fill input tensor info for desc pointers. Modern DML_OPERATOR_DEQUANTIZE + // takes per-tensor scalar scale/zp natively — all inputs keep raw sizes. for (size_t j = 0; j < 3; ++j) { auto* ti = LookupShape(value_shapes, inputs[onnx_base + j]); if (!ti) return std::nullopt; @@ -9661,13 +9673,15 @@ static std::optional TranslateQLinearConcat( deq_node.output_buffer_descs = { deq_out.ToBufferDesc() }; deq_node.output_tensor_descs.resize(1); deq_node.desc_storage = deq_store; - deq_node.op_desc = { DML_OPERATOR_ELEMENT_WISE_DEQUANTIZE_LINEAR, &deq_store->desc }; + deq_store->desc.QuantizationType = DML_QUANTIZATION_TYPE_SCALE_ZERO_POINT; + deq_node.op_desc = { DML_OPERATOR_DEQUANTIZE, &deq_store->desc }; deq_node.fixup = [deq_store](SubNode& self) { RebuildSubNodePointers(self); - deq_store->desc.InputTensor = &self.input_tensor_descs[0]; - deq_store->desc.ScaleTensor = &self.input_tensor_descs[1]; - deq_store->desc.ZeroPointTensor= &self.input_tensor_descs[2]; - deq_store->desc.OutputTensor = &self.output_tensor_descs[0]; + deq_store->quant_tensor_descs = { self.input_tensor_descs[1], self.input_tensor_descs[2] }; + deq_store->desc.InputTensor = &self.input_tensor_descs[0]; + deq_store->desc.QuantizationTensorCount = static_cast(deq_store->quant_tensor_descs.size()); + deq_store->desc.QuantizationTensors = deq_store->quant_tensor_descs.data(); + deq_store->desc.OutputTensor = &self.output_tensor_descs[0]; }; deq_node.FixupPointers(); result.sub_nodes.push_back(std::move(deq_node)); @@ -9704,12 +9718,21 @@ static std::optional TranslateQLinearConcat( storage->join_desc.InputTensors = storage->join_in_tds.data(); storage->join_desc.OutputTensor = &self.output_tensor_descs[0]; storage->join_desc.Axis = local_axis; + // RebuildSubNodePointers set op_desc.Desc = desc_storage.get() (the + // QLinConcatStorage, first member is a vector). join_desc lives mid-struct, + // so restore the correct pointer. + self.op_desc.Desc = &storage->join_desc; }; join_node.FixupPointers(); result.sub_nodes.push_back(std::move(join_node)); // QuantizeLinear sub_node: quantize joined output with y_scale, y_zp. - auto quant_store = std::make_shared(); + // Modern DML_OPERATOR_QUANTIZE takes per-tensor scalar scale/zp natively. + struct QuantStorage { + DML_QUANTIZE_OPERATOR_DESC desc{}; + std::vector quant_tensor_descs; + }; + auto quant_store = std::make_shared(); storage->sub_storages.push_back(quant_store); auto ys_tensor = MakeTensorInfo(ys_info->sizes, ys_info->data_type); @@ -9728,13 +9751,15 @@ static std::optional TranslateQLinearConcat( quant_node.output_buffer_descs = { out_tensor.ToBufferDesc() }; quant_node.output_tensor_descs.resize(1); quant_node.desc_storage = quant_store; - quant_node.op_desc = { DML_OPERATOR_ELEMENT_WISE_QUANTIZE_LINEAR, quant_store.get() }; + quant_store->desc.QuantizationType = DML_QUANTIZATION_TYPE_SCALE_ZERO_POINT; + quant_node.op_desc = { DML_OPERATOR_QUANTIZE, &quant_store->desc }; quant_node.fixup = [quant_store](SubNode& self) { RebuildSubNodePointers(self); - quant_store->InputTensor = &self.input_tensor_descs[0]; - quant_store->ScaleTensor = &self.input_tensor_descs[1]; - quant_store->ZeroPointTensor= &self.input_tensor_descs[2]; - quant_store->OutputTensor = &self.output_tensor_descs[0]; + quant_store->quant_tensor_descs = { self.input_tensor_descs[1], self.input_tensor_descs[2] }; + quant_store->desc.InputTensor = &self.input_tensor_descs[0]; + quant_store->desc.QuantizationTensorCount = static_cast(quant_store->quant_tensor_descs.size()); + quant_store->desc.QuantizationTensors = quant_store->quant_tensor_descs.data(); + quant_store->desc.OutputTensor = &self.output_tensor_descs[0]; }; quant_node.FixupPointers(); result.sub_nodes.push_back(std::move(quant_node)); diff --git a/src/directx/full_graph_fusion.cc b/src/directx/full_graph_fusion.cc index ed56b82..8e36744 100644 --- a/src/directx/full_graph_fusion.cc +++ b/src/directx/full_graph_fusion.cc @@ -1401,6 +1401,34 @@ OrtNodeComputeInfo* FullGraphFusion::Compile( edge.ToNodeIndex = static_cast(sn_dml_idx); edge.ToNodeInputIndex = static_cast(to_input); input_edge_storage.push_back(edge); + continue; + } + // Partition-internal producer (e.g. a Resize feeding a concat + // sub_node dequant). Mirror the primary-input producer logic. + auto prod_it = value_producer.find(gi_name); + if (prod_it != value_producer.end()) { + size_t prod_compiled_idx = prod_it->second.first; + size_t prod_dml_idx = dml_node_offset[prod_compiled_idx]; + size_t prod_output_slot = prod_it->second.second; + UINT from_output_index; + const auto& osrc = compiled_nodes[prod_compiled_idx].translated->output_source; + if (!osrc.empty() && prod_output_slot < osrc.size()) { + auto [src_sub, src_slot] = osrc[prod_output_slot]; + if (src_sub >= 0) + prod_dml_idx += 1 + static_cast(src_sub); + from_output_index = static_cast(src_slot); + } else { + size_t num_subs = compiled_nodes[prod_compiled_idx].translated->sub_nodes.size(); + if (num_subs > 0) + prod_dml_idx += num_subs; + from_output_index = static_cast(prod_output_slot); + } + DML_INTERMEDIATE_GRAPH_EDGE_DESC edge{}; + edge.FromNodeIndex = static_cast(prod_dml_idx); + edge.FromNodeOutputIndex = from_output_index; + edge.ToNodeIndex = static_cast(sn_dml_idx); + edge.ToNodeInputIndex = static_cast(to_input); + intermediate_edge_storage.push_back(edge); } } } @@ -2089,6 +2117,23 @@ bool FullGraphFusion::TryCompilePartition( } else if (input_map.dml_input_map.count(gi_name)) { ie.push_back({(UINT)input_map.dml_input_map[gi_name], (UINT)sn_dml_idx, (UINT)to_input}); + } else if (value_producer.count(gi_name)) { + // Partition-internal producer (e.g. a Resize feeding a concat + // sub_node dequant). Mirror the primary-input producer logic. + auto [pci, pos] = value_producer[gi_name]; + size_t pdi = dml_node_offset[pci]; + UINT foi; + const auto& osrc = compiled_nodes[pci].translated->output_source; + if (!osrc.empty() && pos < osrc.size()) { + auto [ss, sl] = osrc[pos]; + if (ss >= 0) pdi += 1+ss; + foi = (UINT)sl; + } else { + if (!compiled_nodes[pci].translated->sub_nodes.empty()) + pdi += compiled_nodes[pci].translated->sub_nodes.size(); + foi = (UINT)pos; + } + me.push_back({(UINT)pdi, foi, (UINT)sn_dml_idx, (UINT)to_input}); } } }