Refactor QDQ node group selection infrastructure (#10195)

* Separate out the QDQ node group selection from the SAT specific NodeSelector to make re-use in NNAPI etc. cleaner.

* Make MatMulIntegerToFloat matching optional.
Add move ctor to BaseSelector. Required now that it has a unique_ptr member.

* Avoid Guardian warning by using rvalue unique_ptr created with make_unique

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc

@@ -31,7 +31,7 @@ void DropQDQNodesRules(SelectorActionRegistry& qdq_selector_action_registry) {
   std::unique_ptr<Action> action = std::make_unique<MergeIntoTarget>(std::move(moves));
 
 #if !defined(ORT_MINIMAL_BUILD)
-  std::unique_ptr<NodeSelector> selector = std::make_unique<QDQ::DropDQDNodesSelector>();
+  std::unique_ptr<NodeSelector> selector = std::make_unique<QDQ::DropQDQNodesSelector>();
   qdq_selector_action_registry.RegisterSelectorAndAction(action_name,
                                                          {{"Gather", {}},
                                                           {"Reshape", {}},

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc

@@ -37,10 +37,10 @@ static std::vector<const Node*> FindQDQNodes(const GraphViewer& graph_viewer, co
   return nodes;
 }
 
-bool BaseSelector::CheckQDQNodes(const GraphViewer& graph_viewer, const Node& node,
-                                 const std::vector<const Node*>& dq_nodes,
-                                 const std::vector<const Node*>& q_nodes,
-                                 int num_dq_inputs) const {
+bool NodeGroupSelector::CheckQDQNodes(const GraphViewer& graph_viewer, const Node& node,
+                                      const std::vector<const Node*>& dq_nodes,
+                                      const std::vector<const Node*>& q_nodes,
+                                      int num_dq_inputs) const {
   if (num_dq_inputs == -1) {
     num_dq_inputs = NumActualValues(node, true);
   }
@@ -54,7 +54,7 @@ bool BaseSelector::CheckQDQNodes(const GraphViewer& graph_viewer, const Node& no
          !graph_viewer.NodeProducesGraphOutput(node);
 }
 
-std::optional<NodeGroup> BaseSelector::GetQDQSelection(const GraphViewer& graph_viewer, const Node& node) const {
+std::optional<NodeGroup> NodeGroupSelector::GetQDQSelection(const GraphViewer& graph_viewer, const Node& node) const {
   std::vector<const Node*> dq_nodes = FindQDQNodes(graph_viewer, node, true);
   std::vector<const Node*> q_nodes = FindQDQNodes(graph_viewer, node, false);
   if (!Check(graph_viewer, node, dq_nodes, q_nodes)) {
@@ -72,7 +72,7 @@ std::optional<NodeGroup> BaseSelector::GetQDQSelection(const GraphViewer& graph_
 }
 
 std::optional<NodesToOptimizeIndices> BaseSelector::Select(const GraphViewer& graph_viewer, const Node& node) const {
-  const auto qdq_group = GetQDQSelection(graph_viewer, node);
+  const auto qdq_group = node_group_selector_->GetQDQSelection(graph_viewer, node);
   if (!qdq_group.has_value()) {
     return std::nullopt;
   }
@@ -86,10 +86,10 @@ std::optional<NodesToOptimizeIndices> BaseSelector::Select(const GraphViewer& gr
   return builder.Build();
 }
 
-bool DropDQDNodesSelector::Check(const GraphViewer& graph_viewer,
-                                 const Node& node,
-                                 const std::vector<const Node*>& dq_nodes,
-                                 const std::vector<const Node*>& q_nodes) const {
+bool DropQDQNodeGroupSelector::Check(const GraphViewer& graph_viewer,
+                                     const Node& node,
+                                     const std::vector<const Node*>& dq_nodes,
+                                     const std::vector<const Node*>& q_nodes) const {
   if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes, 1)) {
     return false;
   }
@@ -104,9 +104,9 @@ bool DropDQDNodesSelector::Check(const GraphViewer& graph_viewer,
   return IsQDQPairSupported(q_node, dq_node, get_const_initializer, graph_viewer.ModelPath());
 }
 
-bool UnarySelector::Check(const GraphViewer& graph_viewer, const Node& node,
-                          const std::vector<const Node*>& dq_nodes,
-                          const std::vector<const Node*>& q_nodes) const {
+bool UnaryNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& node,
+                                   const std::vector<const Node*>& dq_nodes,
+                                   const std::vector<const Node*>& q_nodes) const {
   if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes, 1)) {
     return false;
   }
@@ -117,10 +117,10 @@ bool UnarySelector::Check(const GraphViewer& graph_viewer, const Node& node,
   return dt_input == dt_output;
 }
 
-bool BinarySelector::Check(const GraphViewer& graph_viewer,
-                           const Node& node,
-                           const std::vector<const Node*>& dq_nodes,
-                           const std::vector<const Node*>& q_nodes) const {
+bool BinaryNodeGroupSelector::Check(const GraphViewer& graph_viewer,
+                                    const Node& node,
+                                    const std::vector<const Node*>& dq_nodes,
+                                    const std::vector<const Node*>& q_nodes) const {
   if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes)) {
     return false;
   }
@@ -132,10 +132,10 @@ bool BinarySelector::Check(const GraphViewer& graph_viewer,
          dt_input_1 == dt_output;
 }
 
-bool VariadicSelector::Check(const GraphViewer& graph_viewer,
-                             const Node& node,
-                             const std::vector<const Node*>& dq_nodes,
-                             const std::vector<const Node*>& q_nodes) const {
+bool VariadicNodeGroupSelector::Check(const GraphViewer& graph_viewer,
+                                      const Node& node,
+                                      const std::vector<const Node*>& dq_nodes,
+                                      const std::vector<const Node*>& q_nodes) const {
   if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes)) {
     return false;
   }
@@ -156,10 +156,10 @@ void VariadicSelector::UpdateBuilder(NodesToOptimizeIndicesBuilder& builder) con
   builder.num_input_defs = 1;  // set to 1 as the first input is variadic
 }
 
-bool ConvSelector::Check(const GraphViewer& graph_viewer,
-                         const Node& node,
-                         const std::vector<const Node*>& dq_nodes,
-                         const std::vector<const Node*>& q_nodes) const {
+bool ConvNodeGroupSelector::Check(const GraphViewer& graph_viewer,
+                                  const Node& node,
+                                  const std::vector<const Node*>& dq_nodes,
+                                  const std::vector<const Node*>& q_nodes) const {
   if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes)) {
     return false;
   }
@@ -190,10 +190,10 @@ void ConvSelector::UpdateBuilder(NodesToOptimizeIndicesBuilder& builder) const {
   builder.input_nodes.resize(3, NodesToOptimizeIndices::kEmptyNodeIndex);
 }
 
-bool MatMulSelector::Check(const GraphViewer& graph_viewer,
-                           const Node& node,
-                           const std::vector<const Node*>& dq_nodes,
-                           const std::vector<const Node*>& q_nodes) const {
+bool MatMulNodeGroupSelector::Check(const GraphViewer& graph_viewer,
+                                    const Node& node,
+                                    const std::vector<const Node*>& dq_nodes,
+                                    const std::vector<const Node*>& q_nodes) const {
   if (dq_nodes.size() != 2) {
     return false;
   }
@@ -218,9 +218,10 @@ bool MatMulSelector::Check(const GraphViewer& graph_viewer,
 
     int32_t dt_output = q_nodes[0]->OutputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
     return dt_input == dt_output;
+  } else {
+    // can be converted to MatMulIntegerToFloat if EP supports that.
+    return matmulintegertofloat_allowed_;
   }
-
-  return true;
 }
 
 }  // namespace QDQ

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h

@@ -20,19 +20,15 @@ struct NodeGroup {
   NodeIndex target_node;
 };
 
-// Base QDQ checker. Finds and provides the DQ and Q nodes to the operator specific checkers, as the QDQ optimizations
-// always involve those nodes.
-class BaseSelector : public NodeSelector {
+class NodeGroupSelector {
  public:
-  std::optional<NodesToOptimizeIndices> Select(const GraphViewer& graph_viewer, const Node& node) const override;
-
   // This is a QDQ Selectors only function, will return QDQ::NodeGroup instead of NodesToOptimizeIndices
   // Can be used in QDQ handling in EPs such as NNAPI
   std::optional<NodeGroup> GetQDQSelection(const GraphViewer& graph_viewer, const Node& node) const;
 
- protected:
-  BaseSelector() = default;
+  virtual ~NodeGroupSelector() = default;
 
+ protected:
   // base check that we have the expected number of QDQ inputs/outputs, and `node` isn't producing a graph output.
   // num_dq_inputs defaults to the number of inputs `node` has if not explicitly specified
   bool CheckQDQNodes(const GraphViewer& graph_viewer, const Node& node,
@@ -45,42 +41,37 @@ class BaseSelector : public NodeSelector {
   bool virtual Check(const GraphViewer& graph_viewer, const Node& node,
                      const std::vector<const Node*>& dq_nodes,
                      const std::vector<const Node*>& q_nodes) const = 0;
-
-  // override if you need to adjust the values in NodesToOptimize.
-  // e.g. add entries for missing optional DQ inputs or set num_inputs to handle variadic inputs
-  // Called post-Check, if Check returned `true`
-  virtual void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const {}
 };
 
+/*
+ * NodeGroup selectors. These are general purpose and used in both the QDQ SelectorActionTransformer setup that the
+ * CPU EP has, and directly in compiling EPs such as NNAPI and CoreML.
+ */
+
 // Single DQ -> node that does not change data -> Q.
 // Zero point and scale are constant scalars and must match
-class DropDQDNodesSelector : public BaseSelector {
- private:
+class DropQDQNodeGroupSelector : public NodeGroupSelector {
   bool Check(const GraphViewer& graph_viewer, const Node& node,
              const std::vector<const Node*>& dq_nodes,
              const std::vector<const Node*>& q_nodes) const override;
 };
 
 // single input. default is to only support uint8.
-class UnarySelector : public BaseSelector {
+class UnaryNodeGroupSelector : public NodeGroupSelector {
   bool Check(const GraphViewer& graph_viewer, const Node& node,
              const std::vector<const Node*>& dq_nodes,
              const std::vector<const Node*>& q_nodes) const override;
-
 };
 
 // 2 DQ nodes providing input -> node -> Q
-class BinarySelector : public BaseSelector {
+class BinaryNodeGroupSelector : public NodeGroupSelector {
   bool Check(const GraphViewer& graph_viewer, const Node& node,
              const std::vector<const Node*>& dq_nodes,
              const std::vector<const Node*>& q_nodes) const override;
 };
 
 // Variadic DQ nodes -> node -> Q
-class VariadicSelector : public BaseSelector {
- public:
-  void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const override;
-
+class VariadicNodeGroupSelector : public NodeGroupSelector {
  private:
   bool Check(const GraphViewer& graph_viewer, const Node& node,
              const std::vector<const Node*>& dq_nodes,
@@ -88,11 +79,10 @@ class VariadicSelector : public BaseSelector {
 };
 
 // DQ nodes for X, W and optionally B -> node -> Q
-class ConvSelector : public BaseSelector {
+class ConvNodeGroupSelector : public NodeGroupSelector {
  public:
-  ConvSelector(bool int8_allowed = false) : int8_allowed_(int8_allowed) {}
-
-  void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const override;
+  // default to 'true'
+  ConvNodeGroupSelector(bool int8_allowed = true) : int8_allowed_(int8_allowed) {}
 
  private:
   bool Check(const GraphViewer& graph_viewer, const Node& node,
@@ -103,16 +93,88 @@ class ConvSelector : public BaseSelector {
 };
 
 // 2 DQ nodes for input -> node -> optional Q if QLinearMatMul, MatMulIntegerToFloat if not
-class MatMulSelector : public BaseSelector {
+// The lack of a trailing Q isn't really a QDQ node group, so we default support for that to off.
+class MatMulNodeGroupSelector : public NodeGroupSelector {
  public:
-  MatMulSelector(bool int8_allowed = false) : int8_allowed_(int8_allowed) {}
+  MatMulNodeGroupSelector(bool int8_allowed = true,
+                          bool matmulintegertofloat_allowed = false)
+      : int8_allowed_(int8_allowed),
+        matmulintegertofloat_allowed_(matmulintegertofloat_allowed) {
+  }
 
  private:
   bool Check(const GraphViewer& graph_viewer, const Node& node,
              const std::vector<const Node*>& dq_nodes,
              const std::vector<const Node*>& q_nodes) const override;
   bool int8_allowed_;
+  bool matmulintegertofloat_allowed_;
 };
+
+/*
+ * NodeSelector instances for use in the QDQ::SelectorActionTransformer.
+ */
+// Base QDQ checker. Finds and provides the DQ and Q nodes to the operator specific checkers, as the QDQ optimizations
+// always involve those nodes.
+class BaseSelector : public NodeSelector {
+ public:
+  std::optional<NodesToOptimizeIndices> Select(const GraphViewer& graph_viewer, const Node& node) const override;
+
+  // We std::move SelectorActionRegistry into the SelectorActionTransformer so this class needs to have a move ctor
+  BaseSelector(BaseSelector&& rhs) noexcept
+      : node_group_selector_{std::move(rhs.node_group_selector_)} {
+  }
+
+ protected:
+  BaseSelector(std::unique_ptr<NodeGroupSelector> node_group_selector)
+      : node_group_selector_{std::move(node_group_selector)} {}
+
+  // override if you need to adjust the values in NodesToOptimize.
+  // e.g. add entries for missing optional DQ inputs or set num_inputs to handle variadic inputs
+  // Called post-Check, if Check returned `true`
+  virtual void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const {}
+
+ private:
+  std::unique_ptr<NodeGroupSelector> node_group_selector_;
+};
+
+class DropQDQNodesSelector : public BaseSelector {
+ public:
+  DropQDQNodesSelector() : BaseSelector(std::make_unique<DropQDQNodeGroupSelector>()) {}
+};
+
+class UnarySelector : public BaseSelector {
+ public:
+  UnarySelector() : BaseSelector(std::make_unique<UnaryNodeGroupSelector>()) {}
+};
+
+class BinarySelector : public BaseSelector {
+ public:
+  BinarySelector() : BaseSelector(std::make_unique<BinaryNodeGroupSelector>()) {}
+};
+
+// Variadic DQ nodes -> node -> Q
+class VariadicSelector : public BaseSelector {
+ public:
+  VariadicSelector() : BaseSelector(std::make_unique<VariadicNodeGroupSelector>()) {}
+
+  void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const override;
+};
+
+// DQ nodes for X, W and optionally B -> node -> Q
+class ConvSelector : public BaseSelector {
+ public:
+  ConvSelector(bool int8_allowed = false) : BaseSelector(std::make_unique<ConvNodeGroupSelector>(int8_allowed)) {}
+
+  void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const override;
+};
+
+// 2 DQ nodes for input -> node -> optional Q if QLinearMatMul, MatMulIntegerToFloat if not
+class MatMulSelector : public BaseSelector {
+ public:
+  MatMulSelector(bool int8_allowed)
+      : BaseSelector(std::make_unique<MatMulNodeGroupSelector>(int8_allowed, /*matmulintegertofloat_allowed*/ true)) {}
+};
+
 }  // namespace QDQ
 }  // namespace onnxruntime
 

onnxruntime/test/optimizer/qdq_transformer_test.cc

@@ -1789,7 +1789,7 @@ TEST(QDQTransformerTests, QDQ_Selector_Test) {
   ASSERT_TRUE(nullptr != conv_node);
   ASSERT_EQ("Conv", conv_node->OpType());
 
-  onnxruntime::QDQ::ConvSelector conv_selector;
+  onnxruntime::QDQ::ConvNodeGroupSelector conv_selector;
 
   // Create a GraphViewer covers the whole graph
   const GraphViewer whole_graph_viewer(graph);