[executorch][native] Parse packed quant metadata and validate PTN method structure (pytorch#23333)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#23228 by
@digantdesai
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/108/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/108/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/107/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/108/orig
Differential Revision:
[D121704317](https://our.internmc.facebook.com/intern/diff/D121704317/)
@diff-train-skip-merge

Co-authored-by: Digant Desai <digantdesai@meta.com>

Author: pytorchbot

backends/native/extension/module/test/TestData.h

@@ -18,6 +18,61 @@
 namespace executorch::extension::native_module::testing {
 
 // cppcheck-suppress-begin useStlAlgorithm
+// Placeholders for input0..N-1, one opaque op producing output0..M-1 from them,
+// and the output node.
+inline flatbuffers::Offset<
+    flatbuffers::Vector<flatbuffers::Offset<native_backend::Node>>>
+make_signature_nodes(
+    flatbuffers::FlatBufferBuilder& builder,
+    size_t num_inputs,
+    size_t num_outputs) {
+  const auto tensor_arg = [&builder](const std::string& name) {
+    return native_backend::CreateNamedArgument(
+        builder,
+        /*name=*/0,
+        native_backend::CreateArgument(
+            builder,
+            native_backend::ArgumentValue::TensorArg,
+            native_backend::CreateTensorArgDirect(builder, name.c_str())
+                .Union()));
+  };
+  std::vector<flatbuffers::Offset<native_backend::Node>> nodes;
+  std::vector<flatbuffers::Offset<native_backend::NamedArgument>> op_inputs;
+  for (size_t i = 0; i < num_inputs; ++i) {
+    const std::string name = "input" + std::to_string(i);
+    const std::vector<flatbuffers::Offset<native_backend::Output>> outputs = {
+        native_backend::CreateOutputDirect(builder, name.c_str())};
+    nodes.push_back(native_backend::CreateNodeDirect(
+        builder,
+        name.c_str(),
+        native_backend::OpKind::PLACEHOLDER,
+        "",
+        nullptr,
+        &outputs));
+    op_inputs.push_back(tensor_arg(name));
+  }
+  std::vector<flatbuffers::Offset<native_backend::Output>> op_outputs;
+  std::vector<flatbuffers::Offset<native_backend::NamedArgument>> results;
+  for (size_t i = 0; i < num_outputs; ++i) {
+    const std::string name = "output" + std::to_string(i);
+    op_outputs.push_back(
+        native_backend::CreateOutputDirect(builder, name.c_str()));
+    results.push_back(tensor_arg(name));
+  }
+  if (num_outputs > 0) {
+    nodes.push_back(native_backend::CreateNodeDirect(
+        builder,
+        "op",
+        native_backend::OpKind::CALL_FUNCTION,
+        "test.op",
+        &op_inputs,
+        &op_outputs));
+  }
+  nodes.push_back(native_backend::CreateNodeDirect(
+      builder, "output", native_backend::OpKind::OUTPUT, "", &results));
+  return builder.CreateVector(nodes);
+}
+
 inline std::vector<uint8_t> make_tensor_package(
     const std::vector<std::string>& method_names,
     size_t num_inputs = 1,
@@ -95,8 +150,7 @@ inline std::vector<uint8_t> make_tensor_package(
     }
     const auto graph = native_backend::CreateGraph(
         builder,
-        builder.CreateVector(
-            std::vector<flatbuffers::Offset<native_backend::Node>>{}),
+        make_signature_nodes(builder, num_inputs, num_outputs),
         builder.CreateVector(input_names),
         builder.CreateVector(output_names),
         builder.CreateVector(tensor_values));

backends/native/runtime/Deserialize.cpp

@@ -179,6 +179,44 @@ TensorMeta build_tensor_meta(
       out.dim_order_hint.push_back(static_cast<int32_t>(dord->Get(i)));
     }
   }
+  if (const auto* quant = m->quant()) {
+    switch (quant->scheme_type()) {
+      case fbs::QuantScheme::AffineGroup: {
+        const fbs::AffineGroup* affine = quant->scheme_as_AffineGroup();
+        if (affine == nullptr) {
+          throw std::runtime_error(
+              "build_tensor_meta: " + name +
+              " has a malformed affine quant scheme");
+        }
+        out.quant = AffineGroupQuant{
+            .scale_data_key = str_of(affine->scale_data_key()),
+            .scale_dtype = map_scalar_type(affine->scale_dtype()),
+            .quant_min = affine->quant_min(),
+            .quant_max = affine->quant_max(),
+            .group_size = affine->group_size(),
+            .zero_point_data_key = str_of(affine->zero_point_data_key()),
+            .zero_point_dtype = map_scalar_type(affine->zero_point_dtype()),
+        };
+        break;
+      }
+      case fbs::QuantScheme::PackedQuant: {
+        const fbs::PackedQuant* packed = quant->scheme_as_PackedQuant();
+        if (packed == nullptr) {
+          throw std::runtime_error(
+              "build_tensor_meta: " + name +
+              " has a malformed packed quant scheme");
+        }
+        out.quant = PackedQuant{.codec = str_of(packed->codec())};
+        break;
+      }
+      case fbs::QuantScheme::NONE:
+        throw std::runtime_error(
+            "build_tensor_meta: " + name + " has an empty quant scheme");
+      default:
+        throw std::runtime_error(
+            "build_tensor_meta: " + name + " has an unsupported quant scheme");
+    }
+  }
   return out;
 }
 

backends/native/runtime/Method.h

@@ -6,7 +6,9 @@
 
 #pragma once
 
+#include <algorithm>
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include <executorch/backends/native/runtime/graph/Graph.h>
@@ -59,4 +61,18 @@ struct Method {
   std::vector<OutputSpec> output_specs; // aligned to graph.output_ids by index
 };
 
+// Returns the binding that ships data under `key` (e.g. a quantized tensor's
+// scale_data_key), or nullptr.
+inline const DataBinding* find_data_binding(
+    const Method& method,
+    std::string_view key) {
+  const auto it = std::find_if(
+      method.data_bindings.begin(),
+      method.data_bindings.end(),
+      [key](const DataBinding& binding) {
+        return binding.has_data && binding.key == key;
+      });
+  return it == method.data_bindings.end() ? nullptr : &*it;
+}
+
 } // namespace ptn