Address review feedback

Author: qti-yuduo

onnxruntime/core/providers/qnn/builder/qnn_node_group/scale_softmax_fusion.cc

@@ -9,7 +9,8 @@
 #include <utility>
 #include <string>
 #include <array>
-#include <vector>
+#include <memory>
+#include <unordered_map>
 
 #include "core/providers/qnn/builder/qnn_utils.h"
 #include "core/providers/qnn/builder/op_builder_factory.h"
@@ -37,11 +38,12 @@ std::optional<size_t> GetMulScalarInputIndex(const NodeUnit* mul) {
   if (!is_x_scalar && !is_y_scalar) {
     return std::nullopt;
   }
-  return is_y_scalar ? 1 : 0;
+  return is_y_scalar ? 1U : 0U;
 }
 
 /// @brief Get the axis for softmax
-/// @param node_units The node units containing the softmax and mul nodes
+/// @param mul Multiply node unit
+/// @param softmax Softmax node unit
 /// @return The axis for softmax
 std::optional<uint32_t> GetPositiveSoftmaxAxis(const NodeUnit* mul, const NodeUnit* softmax) {
   NodeAttrHelper softmax_attr_helper(softmax->GetNode());
@@ -52,7 +54,7 @@ std::optional<uint32_t> GetPositiveSoftmaxAxis(const NodeUnit* mul, const NodeUn
   int64_t axis_value = param_axis.value();
   if (axis_value < 0) {
     size_t input_scale_index = GetMulScalarInputIndex(mul).value();
-    size_t input_other_index = 1 - input_scale_index;
+    size_t input_other_index = 1U - input_scale_index;
     int rank = mul->GetNode().InputDefs()[input_other_index]->Shape()->dim_size();
     axis_value += static_cast<int64_t>(rank);
   }
@@ -92,7 +94,7 @@ std::optional<float> ExtractScalarValueFromMul(const GraphViewer& graph_viewer,
 /// @return Status
 Status CreateOrValidateOnQnn(
     QnnModelWrapper* qnn_model_wrapper,
-    std::array<const NodeUnit*, 2> node_units,
+    gsl::span<const NodeUnit* const> node_units,
     bool validate) {
   const NodeUnit* mul = node_units[0];
   const NodeUnit* softmax = node_units[1];
@@ -101,7 +103,7 @@ Status CreateOrValidateOnQnn(
   ORT_RETURN_IF_NOT(softmax->OpType() == kOpSoftmax,
                     "Expected softmax node to be of type Softmax, got ", softmax->OpType());
   size_t input_scale_index = GetMulScalarInputIndex(mul).value();
-  size_t input_other_index = 1 - input_scale_index;
+  size_t input_other_index = 1U - input_scale_index;
   const NodeUnitIODef& mul_input_other = mul->Inputs()[input_other_index];
   const NodeUnitIODef& softmax_output = softmax->Outputs()[0];
 
@@ -188,22 +190,26 @@ std::unique_ptr<IQnnNodeGroup> ScaleSoftmaxFusion::TryFusion(
     return nullptr;
   }
 
-  std::array<const NodeUnit*, 2> node_units{&mul_node_unit, softmax};
+  std::array<const NodeUnit*, 2> node_unit_array{&mul_node_unit, softmax};
+  auto node_units = gsl::make_span<const NodeUnit*>(node_unit_array.data(), 2);
   if (CreateOrValidateOnQnn(&qnn_model_wrapper, node_units, /*validate=*/true) != Status::OK()) {
     return nullptr;
   }
-
   return std::make_unique<ScaleSoftmaxFusion>(node_units);
 }
 
+gsl::span<const NodeUnit* const> ScaleSoftmaxFusion::GetNodeUnits() const {
+  return gsl::span<const NodeUnit* const>{node_units_.data(), node_units_.size()};
+}
+
 Status ScaleSoftmaxFusion::IsSupported(
     QnnModelWrapper& qnn_model_wrapper, [[maybe_unused]] const logging::Logger& logger) const {
-  return CreateOrValidateOnQnn(&qnn_model_wrapper, node_units_, /*validate=*/true);
+  return CreateOrValidateOnQnn(&qnn_model_wrapper, GetNodeUnits(), /*validate=*/true);
 }
 
 Status ScaleSoftmaxFusion::AddToModelBuilder(
     QnnModelWrapper& qnn_model_wrapper, [[maybe_unused]] const logging::Logger& logger) const {
-  return CreateOrValidateOnQnn(&qnn_model_wrapper, node_units_, /*validate=*/false);
+  return CreateOrValidateOnQnn(&qnn_model_wrapper, GetNodeUnits(), /*validate=*/false);
 }
 
 }  // namespace qnn

onnxruntime/core/providers/qnn/builder/qnn_node_group/scale_softmax_fusion.h

@@ -22,12 +22,16 @@ class QnnModelWrapper;
 /// </summary>
 class ScaleSoftmaxFusion : public IQnnNodeGroup {
  public:
-  explicit ScaleSoftmaxFusion(const std::array<const NodeUnit*, 2>& node_units) : node_units_(node_units) {}
+  explicit ScaleSoftmaxFusion(gsl::span<const NodeUnit* const> node_units) {
+    ORT_ENFORCE(node_units.size() == 2, "Pattern expect exactly 2 NodeUnits.");
+    node_units_[0] = node_units[0];
+    node_units_[1] = node_units[1];
+  }
   ORT_DISALLOW_COPY_AND_ASSIGNMENT(ScaleSoftmaxFusion);
 
   Status IsSupported(QnnModelWrapper& qnn_model_wrapper, const logging::Logger& logger) const override;
   Status AddToModelBuilder(QnnModelWrapper& qnn_model_wrapper, const logging::Logger& logger) const override;
-  gsl::span<const NodeUnit* const> GetNodeUnits() const override { return node_units_; }
+  gsl::span<const NodeUnit* const> GetNodeUnits() const override;
   const NodeUnit* GetTargetNodeUnit() const override { return node_units_[1]; }
   std::string_view Type() const override { return "ScaleSoftmaxFusion"; }
 

onnxruntime/test/providers/qnn/qnn_node_group/scale_softmax_fusion_test.cc

@@ -71,8 +71,6 @@ TEST_F(QnnHTPBackendTests, ScaleSoftmaxFusionScalarInitializer) {
                   /*opset_version=*/13,
                   /*expected_ep_assignment=*/ExpectedEPNodeAssignment::All,
                   /*fp32_abs_err=*/1e-2f);
-
-
 }
 
 TEST_F(QnnHTPBackendTests, ScaleSoftmaxFusionScalarConstant) {