From e0b0c85e8da61c78806786c775d24aa719e9116d Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 1 Jul 2026 16:23:32 -0700 Subject: [PATCH 1/6] Add native backend serialization: partitioner, preprocess (ET-within-ET), and graph IR interface Cherry-picked from metascroy/gnn and adapted: - backends/native/partitioner/ - NativePartitioner marks all ops as supported - backends/native/preprocess.py - NativeBackend preprocess pipeline (ET-within-ET serialization) - backends/native/ir/ - GraphTypes.h, Plan.h, Step.h graph interface headers reinplace_pass: uses main's existing exir/passes/reinplace.py interface (ops_to_inplace as frozenset, auto-derived in-place forms via schema matching) --- backends/native/ir/GraphTypes.h | 1074 +++++++++++++++++++++++++++++++ backends/native/ir/Plan.h | 106 +++ backends/native/ir/Step.h | 164 +++++ backends/native/partitioner.py | 229 +++++++ backends/native/preprocess.py | 390 +++++++++++ 5 files changed, 1963 insertions(+) create mode 100644 backends/native/ir/GraphTypes.h create mode 100644 backends/native/ir/Plan.h create mode 100644 backends/native/ir/Step.h create mode 100644 backends/native/partitioner.py create mode 100644 backends/native/preprocess.py diff --git a/backends/native/ir/GraphTypes.h b/backends/native/ir/GraphTypes.h new file mode 100644 index 00000000000..7cce34d9122 --- /dev/null +++ b/backends/native/ir/GraphTypes.h @@ -0,0 +1,1074 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +/** + * Graph: backend-facing IR adapter. + * + * These types provide our backends' view of the program IR. The current + * implementation adapts ExecuTorch's flatbuffer + * (executorch_flatbuffer::ExecutionPlan / KernelCall) but the API is + * intentionally backing-agnostic — a different serialization could + * replace the underlying storage without changing this header or any + * backend that uses it. + * + * Some methods are documented IR concepts that the current adapter does + * not yet back (e.g., mutable_buffer_ids, version): they are placeholders + * pending serialization support. + * + * ---------------------------------------------------------------------- + * Fictional IR schema (what Graph effectively presents) + * ---------------------------------------------------------------------- + * If the IR were serialized in its own format (independent of the + * underlying ExecuTorch flatbuffer), it would look like this: + * + * // Top-level + * table Graph { + * version: string; // schema/program version + * values: [Value]; // dense pool indexed by value_id (uint) + * inputs: [uint]; // graph input value_ids + * outputs: [uint]; // graph output value_ids + * mutable_buffers: [uint]; // values that persist across executes + * operators: [OperatorDef]; // op-name registry (deduped) + * chains: [Chain]; // chains[0] is main; others used by + * // future control flow (if/while/call) + * } + * + * table OperatorDef { + * name: string; // e.g. "aten.add.Tensor" + * } + * + * // Op chains + * table Chain { + * instructions: [Instruction]; + * } + * + * table Instruction { + * // Today only one variant; future variants for control flow + * // (if/while/call) would extend this union. + * body: KernelCall; + * } + * + * table KernelCall { + * op_index: uint; // → Graph.operators[op_index].name + * args: [uint]; // value_ids: args[0..n-2] = inputs, + * // args[n-1] = output. Single-output + * // assumed today; multi-output is a + * // future extension. + * } + * + * // Values + * union Value { + * None, + * Int { v: int64; }, + * Double { v: float64; }, + * Bool { v: bool; }, + * String { v: string; }, + * Tensor, + * IntList, + * DoubleList, + * BoolList, + * OptionalTensor, + * // ... + * } + * + * table Tensor { + * scalar_type: ScalarType; // dtype + * sizes: [int32]; // for DYNAMIC_BOUND, this is max-shape + * dim_order: [uint8]; // permutation defining memory layout + * shape_dynamism: ShapeDynamism; // STATIC | DYNAMIC_BOUND | + * DYNAMIC_UNBOUND allocation: AllocationInfo?;// null = no AOT plan data: + * TensorData; + * } + * + * union TensorData { + * None, + * Inline { buffer_idx: uint; }, // bytes embedded in program + * External { ndm_key: string;}, // bytes in NamedDataMap, FQN-keyed + * } + * + * table AllocationInfo { + * pool_id: int32; + * offset: uint64; // raw byte offset within pool_id + * } + * + * table IntList { + * member_ids: [int64]; // value_ids of list elements + * } + * + * ---------------------------------------------------------------------- + * Derived views computed by the adapter (not in the schema itself) + * ---------------------------------------------------------------------- + * mem_obj_id(vid) sort-and-index over (pool_id, offset) + * → dense small int identifying shared + * storage slots. Two values with the + * same id were memory-planned to share + * storage (used by router for + * AllocRequest grouping; backends MAY + * honor it as actual storage aliasing). + * value_kind(vid) from membership in inputs/outputs + + * the data field + * → INPUT / OUTPUT / CONSTANT / + * INTERMEDIATE / MUTABLE_BUFFER. + * tensor_constant_data_key(vid) convenience accessor for + * TensorData.External.ndm_key. + * tensor_nbytes_max(vid) dtype_size × prod(sizes); upper bound + * for DYNAMIC_BOUND tensors. + * + * ---------------------------------------------------------------------- + * Adapter cost + * ---------------------------------------------------------------------- + * All accessors are inline and compile down to essentially the same + * machine code as direct flatbuffer access. The Graph constructor pays + * a one-time O(N log N) precompute (over tensor values) for mem_obj_id, + * and O(num_inputs + num_outputs) for the value_kind sets. Per-call + * overhead in the runtime hot path is a few inline indirections plus + * predictable ET_CHECK branches; release builds compile away these + * checks entirely under -DNDEBUG. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +namespace executorch { +namespace backends { +namespace portable { + +// Forward declare +class Graph; + +/** + * Value kind — matches design doc's TensorKind + */ +enum class ValueKind : uint8_t { + INPUT = 0, // Graph input (user provides) + OUTPUT, // Graph output (user reads) + CONSTANT, // Immutable weight + MUTABLE_BUFFER, // Mutable state (e.g., KV cache) + INTERMEDIATE, // Temporary (produced/consumed internally) +}; + +/** + * Type of an EValue stored at a value_id. Mirrors the runtime EValue + * sum-type but in adapter-level form (no flatbuffer types in the API). + */ +enum class ValueType : uint8_t { + None = 0, + Int, + Double, + Bool, + Tensor, + IntList, + TensorList, + OptionalTensorList, + Other, // String, OptionalTensor, BoolList, DoubleList, ... + // Adapter doesn't surface these yet; executor falls back to + // default-constructed EValue. +}; + +/** + * Kind of an Instruction in a Chain. Mirrors ET's InstructionArguments + * union but as an adapter-level enum (no flatbuffer types in the API). + * + * See CONTROL_FLOW_DESIGN.md §4. + * + * - Kernel: the only kind a CompiledSegment may contain. + * - JumpFalse: control-flow boundary; segment break. + * - Move: copies one EValue into another; segment break (host-host + * transfer). + * - Free: informational hint to release a value's storage; v2 ignores + * (memory plan handles deallocation at delegate destroy). + * - Delegate: nested executorch_call_delegate inside our chain. + * REJECTED at init; partitioner contract forbids nested delegates. + */ +enum class InstructionKind : uint8_t { + Kernel = 0, + JumpFalse, + Move, + Free, + Delegate, +}; + +/** + * Decoded JumpFalseCall payload. Backing-agnostic. + */ +struct JumpFalseInfo { + uint32_t cond_value_id; // EValue index whose value is the predicate + uint32_t destination_pc; // Source-PC of the jump target +}; + +/** + * Decoded MoveCall payload. Backing-agnostic. + */ +struct MoveInfo { + uint32_t src_value_id; + uint32_t dst_value_id; +}; + +/** + * Thin wrapper around ExecuTorch's flatbuffer KernelCall providing + * convenient access to op name and input/output value_ids. + * + * ExecuTorch packs all op args into a single `args` array per the + * convention "the last arg is the (single) output." This wrapper + * exposes inputs/outputs accordingly without copying — `inputs()` / + * `output()` return Spans / values that reference the underlying + * flatbuffer storage directly. + * + * Per-call cost: just stores two pointers + an int. No allocations. + * Safe to construct repeatedly in hot dispatch loops. + * + * NOTE: Multi-output ops (e.g., `aten.split`, `aten.max.dim`) are not + * supported by this wrapper — `num_outputs()` always returns 0 or 1. + * Adding multi-output support requires per-op schema knowledge that + * isn't in the flatbuffer. ET_CHECK guards the single-output access + * path. + */ +class OperatorCall { + public: + explicit OperatorCall( + const executorch_flatbuffer::KernelCall* call, + const Graph* graph) + : call_(call), graph_(graph) {} + + // node_id for error messages/profiling + uint32_t node_id() const { + return node_id_; + } + void set_node_id(uint32_t id) { + node_id_ = id; + } + + // Op base name (e.g., "aten::add"). Does NOT include the overload + // suffix; for that use full_name(). + const char* name() const; + + // Op overload (e.g., "Tensor", "Scalar", "out"). May be empty string + // for ops with no overload disambiguation. Use full_name() to get + // the combined "name.overload" form most callers want. + const char* overload() const; + + // Combined unique key "." (e.g., "aten::add.Tensor"), + // or just "" if the overload is empty. This is the canonical + // identity used by the CPU op registry to dispatch — registering + // by base name alone collapses every overload of an op into one + // handler, which is wrong for ops like aten::pow_ that have multiple + // overloads with different schemas. + std::string full_name() const; + + // All op args (flat). Last entry is the output; the rest are inputs. + // Optional args are NOT indicated by -1 — instead the index points to + // a value with isNone() == true. (ExecuTorch convention.) + runtime::Span args() const { + auto* a = call_->args(); + return a ? runtime::Span(a->data(), a->size()) + : runtime::Span{}; + } + + // Inputs = args[0..n-2]. Returns empty span if op has no args. + runtime::Span inputs() const { + auto a = args(); + return a.empty() ? a : runtime::Span(a.data(), a.size() - 1); + } + + size_t num_inputs() const { + auto a = args(); + return a.empty() ? 0 : a.size() - 1; + } + uint32_t input(size_t i) const { + ET_CHECK_MSG( + i < num_inputs(), + "OperatorCall::input: index %zu >= num_inputs()=%zu", + i, + num_inputs()); + return static_cast(args()[i]); + } + + // Single-output assumption (see class doc). Multi-output ops will + // break here. + size_t num_outputs() const { + return args().empty() ? 0 : 1; + } + uint32_t output(size_t i) const { + ET_CHECK_MSG( + i < num_outputs(), + "OperatorCall::output: index %zu >= num_outputs()=%zu", + i, + num_outputs()); + auto a = args(); + return static_cast(a[a.size() - 1]); + } + + private: + const executorch_flatbuffer::KernelCall* call_; + const Graph* graph_; + uint32_t node_id_ = 0; +}; + +/** + * IR view of a program. Adapts executorch_flatbuffer::ExecutionPlan; + * exposes value metadata, input/output IDs, the operator table, and + * chains of operator calls. See the file-header comment for the + * adapter-pattern rationale. + */ +class Graph { + public: + explicit Graph( + const executorch_flatbuffer::ExecutionPlan* plan, + const executorch_flatbuffer::Program* program = nullptr) + : plan_(plan), program_(program) { + // Precompute input/output value_id sets for O(1) value_kind lookup. + if (auto* in = plan_->inputs()) { + input_ids_.reserve(in->size()); + for (size_t i = 0; i < in->size(); ++i) { + input_ids_.insert(static_cast(in->Get(i))); + } + } + if (auto* out = plan_->outputs()) { + output_ids_.reserve(out->size()); + for (size_t i = 0; i < out->size(); ++i) { + output_ids_.insert(static_cast(out->Get(i))); + } + } + + // Precompute mem_obj_id for every tensor value. + // + // Algorithm: collect (pool_id, offset) keys for all aliasable tensor + // values; sort the unique keys; assign mem_obj_id = sort rank. Two + // values with the same (pool_id, offset) get the same id (they share + // storage). Sort-and-index is deterministic across runs and depends + // only on the AOT memory plan. + size_t n_vals = num_values(); + mem_obj_ids_.assign(n_vals, -1); + if (n_vals == 0) + return; + + // 1. Collect (key, value_id) entries for tensor values with + // allocation_info. + struct Entry { + uint64_t key; // (pool_id << 32) | offset + uint32_t value_id; + }; + std::vector entries; + entries.reserve(n_vals); + for (uint32_t i = 0; i < n_vals; ++i) { + auto* val = value_meta(i); + if (!val || + val->val_type() != executorch_flatbuffer::KernelTypes::Tensor) { + continue; + } + auto* t = val->val_as_Tensor(); + if (!t) + continue; + auto* alloc = t->allocation_info(); + if (!alloc) + continue; + uint64_t pool = static_cast(alloc->memory_id()); + uint64_t off = alloc->memory_offset_low(); + entries.push_back({(pool << 32) | off, i}); + } + if (entries.empty()) + return; + + // 2. Sort by key (lex order on (pool_id, offset)). + std::sort( + entries.begin(), entries.end(), [](const Entry& a, const Entry& b) { + return a.key < b.key; + }); + + // 3. Assign mem_obj_id = sort rank (same key → same id). + int32_t next_id = -1; + uint64_t prev_key = ~0ULL; + for (const auto& e : entries) { + if (e.key != prev_key) { + ++next_id; + prev_key = e.key; + } + mem_obj_ids_[e.value_id] = next_id; + } + + // Precompute mutable_buffer_ids_: tensor values with allocation_info, + // not graph IO, not constants, and NOT produced by any op (i.e. + // placeholders that aren't graph inputs). These are mutable buffer + // placeholders pulled into the delegate by tag_mutated_buffer; their + // state persists across execute() calls. + // + // Used by the router to distinguish semantic alias groups (buffer + // mutation: AOT spec-shared the buffer placeholder with its mutation + // source) from lifetime-reuse aliasing (the planner happened to put + // two values at the same offset because their lifetimes don't + // overlap). Only semantic groups need the "all touching ops on same + // runtime else home=host" coordination. + { + // Collect all op-output value_ids across all chains. Skip non-Kernel + // instructions (JumpFalseCall, MoveCall, FreeCall, DelegateCall) — + // see CONTROL_FLOW_DESIGN.md §4.3. MoveCall's dst is a special case: + // it appears here under produced_vids only if MoveCall handling is + // wired through the router (currently MoveCall is treated as a + // host->host TransferStep, which does not classify the dst as + // "produced by an op"). Tolerable today; revisit if MoveCall-only + // mutable-buffer patterns appear. + std::unordered_set produced_vids; + for (size_t ci = 0; ci < num_chains(); ++ci) { + auto chains = plan_->chains(); + auto instrs = chains->Get(ci)->instructions(); + size_t n_instr = instrs ? instrs->size() : 0; + for (size_t oi = 0; oi < n_instr; ++oi) { + if (instruction_kind(ci, oi) != InstructionKind::Kernel) + continue; + OperatorCall op = get_op(ci, oi); + for (size_t k = 0; k < op.num_outputs(); ++k) { + produced_vids.insert(op.output(k)); + } + } + } + for (uint32_t i = 0; i < n_vals; ++i) { + if (mem_obj_ids_[i] < 0) + continue; // not an allocated tensor + if (input_ids_.count(i) > 0) + continue; // graph input + if (output_ids_.count(i) > 0) + continue; // graph output + if (tensor_constant_data_key(i) != nullptr) + continue; // constant + if (produced_vids.count(i) > 0) + continue; // produced by an op + // Tensor with alloc, not IO, not constant, not produced → it's a + // mutable buffer placeholder. + mutable_buffer_ids_.push_back(i); + } + } + } + + //===------------------------------------------------------------------===// + // Version + //===------------------------------------------------------------------===// + + const char* version() const { + // TODO: Return actual version when available + return "1.0"; + } + + //===------------------------------------------------------------------===// + // Values + //===------------------------------------------------------------------===// + + size_t num_values() const { + auto v = plan_->values(); + return v ? v->size() : 0; + } + + // Access serialized value metadata. + // NOTE: returns the raw flatbuffer EValue. This is the construction- + // seam escape hatch — backends and routers should prefer the typed + // accessors below (value_type, int_value, tensor_*, etc) so they + // don't couple to the underlying serialization. + const executorch_flatbuffer::EValue* value_meta(uint32_t value_id) const { + auto values = plan_->values(); + if (!values || value_id >= values->size()) + return nullptr; + return values->Get(value_id); + } + + // Value metadata helpers + ValueKind value_kind(uint32_t value_id) const; + int32_t mem_obj_id(uint32_t value_id) const; + + //===------------------------------------------------------------------===// + // Typed value accessors (adapter-level — no flatbuffer types leak) + //===------------------------------------------------------------------===// + + // Returns the kind of the EValue stored at value_id. + ValueType value_type(uint32_t value_id) const; + + // Scalar accessors — ET_CHECK if the value isn't of the expected kind. + int64_t int_value(uint32_t value_id) const; + double double_value(uint32_t value_id) const; + bool bool_value(uint32_t value_id) const; + + // Tensor accessors — ET_CHECK if the value isn't a tensor. + ::executorch::aten::ScalarType tensor_dtype(uint32_t value_id) const; + ::executorch::runtime::Span tensor_sizes( + uint32_t value_id) const; + ::executorch::runtime::Span tensor_dim_order( + uint32_t value_id) const; + ::executorch::aten::TensorShapeDynamism tensor_shape_dynamism( + uint32_t value_id) const; + // Returns NDM key (FQN) for an external constant tensor, or nullptr + // if the tensor isn't an NDM-stored constant. + const char* tensor_constant_data_key(uint32_t value_id) const; + + // Returns the raw bytes of an inline constant (stored in the + // program's constant_buffer field), or an empty span if the tensor + // isn't an inline constant. Inline constants are constants the AOT + // didn't promote to NDM (e.g., literals lifted into _lifted_tensor_* + // placeholders). Mutually exclusive with tensor_constant_data_key: + // a constant is either NDM-stored (key != nullptr) or inline + // (this returns non-empty), never both. + ::executorch::runtime::Span tensor_inline_data( + uint32_t value_id) const; + + // True if the tensor is a constant of either flavor (NDM-stored or + // inline). Use this for "is this an immutable constant?" filtering + // checks in the router; the source matters only at upload time. + bool is_constant(uint32_t value_id) const { + if (tensor_constant_data_key(value_id) != nullptr) + return true; + return !tensor_inline_data(value_id).empty(); + } + + // dtype-size × prod(sizes); 0 if not a tensor or sizes empty. + size_t tensor_nbytes_max(uint32_t value_id) const; + + // IntList accessors — ET_CHECK if the value isn't an IntList. + // Returns the EValue indices that the list elements reference (stored + // as int64 in the serialization); the caller resolves them through + // the values array. + ::executorch::runtime::Span int_list_member_ids( + uint32_t value_id) const; + + // For a TensorList or OptionalTensorList value, returns the EValue + // indices that the list contains. For OptionalTensorList, indices may + // point at None values (representing nullopt). + ::executorch::runtime::Span tensor_list_member_ids( + uint32_t value_id) const; + + //===------------------------------------------------------------------===// + // Input/Output IDs + //===------------------------------------------------------------------===// + + size_t num_input_ids() const { + auto in = plan_->inputs(); + return in ? in->size() : 0; + } + + uint32_t input_id(size_t i) const { + auto in = plan_->inputs(); + ET_CHECK_MSG( + in && i < in->size(), + "Graph::input_id(%zu) out of range (have %zu inputs)", + i, + in ? in->size() : 0); + return static_cast(in->Get(i)); + } + + size_t num_output_ids() const { + auto out = plan_->outputs(); + return out ? out->size() : 0; + } + + uint32_t output_id(size_t i) const { + auto out = plan_->outputs(); + ET_CHECK_MSG( + out && i < out->size(), + "Graph::output_id(%zu) out of range (have %zu outputs)", + i, + out ? out->size() : 0); + return static_cast(out->Get(i)); + } + + //===------------------------------------------------------------------===// + // Mutable Buffer IDs (values that persist across execute() calls) + //===------------------------------------------------------------------===// + + size_t num_mutable_buffer_ids() const { + return mutable_buffer_ids_.size(); + } + + uint32_t mutable_buffer_id(size_t i) const { + ET_CHECK_MSG( + i < mutable_buffer_ids_.size(), + "Graph::mutable_buffer_id: index %zu out of range " + "(have %zu mutable buffers)", + i, + mutable_buffer_ids_.size()); + return mutable_buffer_ids_[i]; + } + + //===------------------------------------------------------------------===// + // Operators (for op name lookup) + //===------------------------------------------------------------------===// + + size_t num_operators() const { + auto ops = plan_->operators(); + return ops ? ops->size() : 0; + } + + const char* operator_name(size_t idx) const { + auto ops = plan_->operators(); + if (!ops || idx >= ops->size()) + return nullptr; + auto op = ops->Get(idx); + return op && op->name() ? op->name()->c_str() : nullptr; + } + + const char* operator_overload(size_t idx) const { + auto ops = plan_->operators(); + if (!ops || idx >= ops->size()) + return nullptr; + auto op = ops->Get(idx); + return op && op->overload() ? op->overload()->c_str() : nullptr; + } + + //===------------------------------------------------------------------===// + // Chains + //===------------------------------------------------------------------===// + + size_t num_chains() const { + auto chains = plan_->chains(); + return chains ? chains->size() : 0; + } + + int32_t main_chain_idx() const { + return 0; // Default: first chain is main + } + + // Get number of ops in a chain + size_t num_ops_in_chain(size_t chain_idx) const { + auto chains = plan_->chains(); + ET_CHECK_MSG( + chains && chain_idx < chains->size(), + "Graph::num_ops_in_chain(%zu) out of range (have %zu chains)", + chain_idx, + chains ? chains->size() : 0); + auto instrs = chains->Get(chain_idx)->instructions(); + return instrs ? instrs->size() : 0; + } + + // Get OperatorCall for op in chain + // PRECONDITION: instruction_kind(chain_idx, op_idx) == + // InstructionKind::Kernel. Use instruction_kind() first to dispatch on kind. + OperatorCall get_op(size_t chain_idx, size_t op_idx) const { + auto chains = plan_->chains(); + ET_CHECK_MSG( + chains && chain_idx < chains->size(), + "Graph::get_op: chain_idx=%zu out of range (have %zu chains)", + chain_idx, + chains ? chains->size() : 0); + auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG( + instrs && op_idx < instrs->size(), + "Graph::get_op: op_idx=%zu out of range in chain %zu " + "(have %zu ops)", + op_idx, + chain_idx, + instrs ? instrs->size() : 0); + auto instr = instrs->Get(op_idx); + ET_CHECK_MSG( + instr->instr_args_type() == + executorch_flatbuffer::InstructionArguments::KernelCall, + "Graph::get_op: instruction at chain=%zu op_idx=%zu is not a KernelCall " + "(type=%u). Use instruction_kind() to dispatch.", + chain_idx, + op_idx, + static_cast(instr->instr_args_type())); + auto kernel = static_cast( + instr->instr_args()); + return OperatorCall(kernel, this); + } + + //===------------------------------------------------------------------===// + // Typed instruction accessors (kind-aware; see CONTROL_FLOW_DESIGN.md §4) + //===------------------------------------------------------------------===// + + // Returns the InstructionKind at (chain_idx, op_idx). + InstructionKind instruction_kind(size_t chain_idx, size_t op_idx) const { + auto chains = plan_->chains(); + ET_CHECK_MSG( + chains && chain_idx < chains->size(), + "Graph::instruction_kind: chain_idx=%zu out of range (have %zu chains)", + chain_idx, + chains ? chains->size() : 0); + auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG( + instrs && op_idx < instrs->size(), + "Graph::instruction_kind: op_idx=%zu out of range in chain %zu", + op_idx, + chain_idx); + auto instr = instrs->Get(op_idx); + using IA = executorch_flatbuffer::InstructionArguments; + switch (instr->instr_args_type()) { + case IA::KernelCall: + return InstructionKind::Kernel; + case IA::JumpFalseCall: + return InstructionKind::JumpFalse; + case IA::MoveCall: + return InstructionKind::Move; + case IA::FreeCall: + return InstructionKind::Free; + case IA::DelegateCall: + return InstructionKind::Delegate; + default: + ET_CHECK_MSG( + false, + "Graph::instruction_kind: unsupported InstructionArguments type %u", + static_cast(instr->instr_args_type())); + } + } + + // Convenience: main-chain shortcut. + InstructionKind instruction_kind(size_t op_idx) const { + return instruction_kind(main_chain_idx(), op_idx); + } + + // Typed accessors per kind. ET_CHECK if the kind doesn't match. + OperatorCall get_kernel_call(size_t chain_idx, size_t op_idx) const { + return get_op(chain_idx, op_idx); + } + + JumpFalseInfo get_jump_false(size_t chain_idx, size_t op_idx) const { + auto chains = plan_->chains(); + ET_CHECK_MSG( + chains && chain_idx < chains->size(), + "Graph::get_jump_false: chain_idx=%zu out of range", + chain_idx); + auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG( + instrs && op_idx < instrs->size(), + "Graph::get_jump_false: op_idx=%zu out of range in chain %zu", + op_idx, + chain_idx); + auto instr = instrs->Get(op_idx); + auto* jf = instr->instr_args_as_JumpFalseCall(); + ET_CHECK_MSG( + jf != nullptr, + "Graph::get_jump_false: instruction at chain=%zu op_idx=%zu is not a " + "JumpFalseCall", + chain_idx, + op_idx); + JumpFalseInfo info; + info.cond_value_id = static_cast(jf->cond_value_index()); + info.destination_pc = static_cast(jf->destination_instruction()); + return info; + } + + MoveInfo get_move(size_t chain_idx, size_t op_idx) const { + auto chains = plan_->chains(); + ET_CHECK_MSG( + chains && chain_idx < chains->size(), + "Graph::get_move: chain_idx=%zu out of range", + chain_idx); + auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG( + instrs && op_idx < instrs->size(), + "Graph::get_move: op_idx=%zu out of range in chain %zu", + op_idx, + chain_idx); + auto instr = instrs->Get(op_idx); + auto* mv = instr->instr_args_as_MoveCall(); + ET_CHECK_MSG( + mv != nullptr, + "Graph::get_move: instruction at chain=%zu op_idx=%zu is not a MoveCall", + chain_idx, + op_idx); + MoveInfo info; + info.src_value_id = static_cast(mv->move_from()); + info.dst_value_id = static_cast(mv->move_to()); + return info; + } + + //===------------------------------------------------------------------===// + // Convenience: main chain accessors + //===------------------------------------------------------------------===// + + size_t num_instructions() const { + return num_ops_in_chain(main_chain_idx()); + } + + OperatorCall get_instruction(size_t idx) const { + return get_op(main_chain_idx(), idx); + } + + private: + const executorch_flatbuffer::ExecutionPlan* plan_; + // Optional reference to the parent Program, needed only for + // tensor_inline_data() (which dereferences program_->constant_buffer). + // Pre-existing constructions that pass only the plan get nullptr and + // tensor_inline_data() returns empty for them. + const executorch_flatbuffer::Program* program_; + // Precomputed at construction for O(1) value_kind lookup. + std::unordered_set input_ids_; + std::unordered_set output_ids_; + // mem_obj_ids_[value_id] = dense small int identifying the storage slot + // (sort rank of (pool_id, offset) pairs across all aliasable tensor + // values). -1 for non-tensor / non-allocated values. Same id ⇒ same + // storage. Computed once at construction; O(1) lookup at use sites. + std::vector mem_obj_ids_; + + // Mutable buffer placeholder value_ids: tensor values with allocation + // info that aren't graph IO, aren't constants, and aren't produced by + // any op. These persist across execute() calls (their storage is + // preserved between invocations). Identified by tag_mutated_buffer at + // AOT time. + std::vector mutable_buffer_ids_; +}; + +// Implement OperatorCall::name() after Graph is defined +inline const char* OperatorCall::name() const { + // In ExecuTorch, op names are in the operators table, indexed by op_index + return graph_->operator_name(call_->op_index()); +} + +inline const char* OperatorCall::overload() const { + return graph_->operator_overload(call_->op_index()); +} + +inline std::string OperatorCall::full_name() const { + const char* base = name(); + const char* ovl = overload(); + if (!base) + return {}; + if (!ovl || *ovl == '\0') + return std::string(base); + std::string s; + s.reserve(std::strlen(base) + 1 + std::strlen(ovl)); + s.append(base); + s.push_back('.'); + s.append(ovl); + return s; +} + +// Implement value metadata accessors +inline ValueKind Graph::value_kind(uint32_t value_id) const { + if (input_ids_.count(value_id)) + return ValueKind::INPUT; + if (output_ids_.count(value_id)) + return ValueKind::OUTPUT; + + // Constant if the tensor has a baked data buffer. + auto val = value_meta(value_id); + if (val && val->val_type() == executorch_flatbuffer::KernelTypes::Tensor) { + auto* tensor = val->val_as_Tensor(); + if (tensor && tensor->data_buffer_idx() > 0) { + return ValueKind::CONSTANT; + } + } + return ValueKind::INTERMEDIATE; +} + +inline int32_t Graph::mem_obj_id(uint32_t value_id) const { + return value_id < mem_obj_ids_.size() ? mem_obj_ids_[value_id] : -1; +} + +//===----------------------------------------------------------------------===// +// Typed value accessors +//===----------------------------------------------------------------------===// + +inline ValueType Graph::value_type(uint32_t value_id) const { + auto* val = value_meta(value_id); + if (!val) + return ValueType::None; + using KT = executorch_flatbuffer::KernelTypes; + switch (val->val_type()) { + case KT::Null: + return ValueType::None; + case KT::Int: + return ValueType::Int; + case KT::Double: + return ValueType::Double; + case KT::Bool: + return ValueType::Bool; + case KT::Tensor: + return ValueType::Tensor; + case KT::IntList: + return ValueType::IntList; + case KT::TensorList: + return ValueType::TensorList; + case KT::OptionalTensorList: + return ValueType::OptionalTensorList; + default: + return ValueType::Other; + } +} + +inline int64_t Graph::int_value(uint32_t value_id) const { + auto* val = value_meta(value_id); + ET_CHECK_MSG( + val && val->val_type() == executorch_flatbuffer::KernelTypes::Int, + "Graph::int_value(%u): value is not an Int", + value_id); + return static_cast(val->val())->int_val(); +} + +inline double Graph::double_value(uint32_t value_id) const { + auto* val = value_meta(value_id); + ET_CHECK_MSG( + val && val->val_type() == executorch_flatbuffer::KernelTypes::Double, + "Graph::double_value(%u): value is not a Double", + value_id); + return static_cast(val->val()) + ->double_val(); +} + +inline bool Graph::bool_value(uint32_t value_id) const { + auto* val = value_meta(value_id); + ET_CHECK_MSG( + val && val->val_type() == executorch_flatbuffer::KernelTypes::Bool, + "Graph::bool_value(%u): value is not a Bool", + value_id); + return static_cast(val->val()) + ->bool_val(); +} + +namespace detail { +inline const executorch_flatbuffer::Tensor* tensor_or_null( + const executorch_flatbuffer::EValue* val) { + if (!val || val->val_type() != executorch_flatbuffer::KernelTypes::Tensor) { + return nullptr; + } + return val->val_as_Tensor(); +} +} // namespace detail + +inline ::executorch::aten::ScalarType Graph::tensor_dtype( + uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + ET_CHECK_MSG(t, "Graph::tensor_dtype(%u): value is not a Tensor", value_id); + return static_cast<::executorch::aten::ScalarType>(t->scalar_type()); +} + +inline ::executorch::runtime::Span Graph::tensor_sizes( + uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + ET_CHECK_MSG(t, "Graph::tensor_sizes(%u): value is not a Tensor", value_id); + auto* s = t->sizes(); + return s ? ::executorch::runtime::Span(s->data(), s->size()) + : ::executorch::runtime::Span{}; +} + +inline ::executorch::runtime::Span Graph::tensor_dim_order( + uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + ET_CHECK_MSG( + t, "Graph::tensor_dim_order(%u): value is not a Tensor", value_id); + auto* d = t->dim_order(); + return d ? ::executorch::runtime::Span(d->data(), d->size()) + : ::executorch::runtime::Span{}; +} + +inline ::executorch::aten::TensorShapeDynamism Graph::tensor_shape_dynamism( + uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + ET_CHECK_MSG( + t, "Graph::tensor_shape_dynamism(%u): value is not a Tensor", value_id); + return static_cast<::executorch::aten::TensorShapeDynamism>( + t->shape_dynamism()); +} + +inline const char* Graph::tensor_constant_data_key(uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + if (!t) + return nullptr; + auto* eti = t->extra_tensor_info(); + if (!eti) + return nullptr; + if (eti->location() != executorch_flatbuffer::TensorDataLocation::EXTERNAL) { + return nullptr; + } + auto* fqn = eti->fully_qualified_name(); + return (fqn && fqn->size() > 0) ? fqn->c_str() : nullptr; +} + +inline ::executorch::runtime::Span Graph::tensor_inline_data( + uint32_t value_id) const { + if (!program_) + return {}; + auto* t = detail::tensor_or_null(value_meta(value_id)); + if (!t) + return {}; + uint32_t idx = static_cast(t->data_buffer_idx()); + // Index 0 is reserved (placeholder for "no inline data"). External + // constants also have idx == 0; they're handled by + // tensor_constant_data_key. + if (idx == 0) + return {}; + auto* buffers = program_->constant_buffer(); + if (!buffers || idx >= buffers->size()) + return {}; + auto* buf = buffers->Get(idx); + if (!buf) + return {}; + auto* storage = buf->storage(); + if (!storage || storage->size() == 0) + return {}; + return ::executorch::runtime::Span( + storage->data(), storage->size()); +} + +inline size_t Graph::tensor_nbytes_max(uint32_t value_id) const { + auto* t = detail::tensor_or_null(value_meta(value_id)); + if (!t || !t->sizes()) + return 0; + size_t numel = 1; + for (size_t i = 0; i < t->sizes()->size(); ++i) { + int dim = t->sizes()->Get(i); + if (dim < 0) + return 0; + numel *= static_cast(dim); + } + auto stype = static_cast<::executorch::aten::ScalarType>(t->scalar_type()); + return numel * ::executorch::runtime::elementSize(stype); +} + +inline ::executorch::runtime::Span Graph::int_list_member_ids( + uint32_t value_id) const { + auto* val = value_meta(value_id); + ET_CHECK_MSG( + val && val->val_type() == executorch_flatbuffer::KernelTypes::IntList, + "Graph::int_list_member_ids(%u): value is not an IntList", + value_id); + auto* items = + static_cast(val->val())->items(); + return items + ? ::executorch::runtime::Span(items->data(), items->size()) + : ::executorch::runtime::Span{}; +} + +inline ::executorch::runtime::Span Graph::tensor_list_member_ids( + uint32_t value_id) const { + auto* val = value_meta(value_id); + ET_CHECK_MSG( + val && + (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList || + val->val_type() == + executorch_flatbuffer::KernelTypes::OptionalTensorList), + "Graph::tensor_list_member_ids(%u): value is not a TensorList " + "or OptionalTensorList", + value_id); + // Both TensorList and OptionalTensorList have the same shape: + // table { items: [int]; }. Cast to either to access items(). + const flatbuffers::Vector* items = nullptr; + if (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList) { + items = static_cast(val->val()) + ->items(); + } else { + items = static_cast( + val->val()) + ->items(); + } + return items + ? ::executorch::runtime::Span(items->data(), items->size()) + : ::executorch::runtime::Span{}; +} + +} // namespace portable +} // namespace backends +} // namespace executorch diff --git a/backends/native/ir/Plan.h b/backends/native/ir/Plan.h new file mode 100644 index 00000000000..f3fd9eecebf --- /dev/null +++ b/backends/native/ir/Plan.h @@ -0,0 +1,106 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace native { + +/** + * Each event slot remembers which Engine created it. Required so the + * executor can call the right Engine::wait at the host boundary. + */ +struct EventSlot { + std::unique_ptr event; + Engine* owner; +}; + +/** + * Frozen output of Router::route. Holds: + * - Runtime/Engine arrays indexed by RuntimeIndex, + * - the issue-ordered Step list, + * - per-input/output bindings, + * - pre-allocated event slots, + * - per-provider allocation request lists. + * + * Buffer ownership lives in the Engines themselves (each engine releases + * everything it allocated in its destructor). NativeBackend tracks no + * per-vid storage state; engines are self-sufficient via their own + * value->Buffer tables and via set_io_bindings at init. + * + * See §4.9 of the design doc. + */ +struct Plan { + // Parallel arrays indexed by RuntimeIndex. By convention, index 0 is + // the HostPool (canonical home for boundary values; not a compute + // runtime — its can_run() always returns false). Compute providers + // (CpuRuntime, MetalRuntime, ...) occupy slots 1+. + std::vector providers; + std::vector instances; // non-owning; lifetime is DelegateInstance + + std::vector steps; // ordered by *issue* (not completion) + + std::vector inputs; + std::vector outputs; + + // Maximum control-flow hops the executor walks before declaring a + // malformed back-edge. Stamped from RouterOptions::max_hops at + // route() time so different programs can carry different limits. + size_t max_hops = 10'000'000; + + // Pre-allocated event slots. Each event is reset lazily by its + // producing Engine immediately before signaling. + std::vector events; + + // EventIds whose signal must complete before execute() returns. + // Replaces a blanket per-Engine drain. Populated by the router as + // the signals of the last step on each Engine. + std::vector terminal_events; + + // Mirror identities are conveyed entirely through AllocRequest: + // device-side mirror requests carry host_mirror_value_id, and + // engines materialize their own TensorImpls for mirror_ids when they + // claim the AllocRequest. NativeBackend has no separate mirror table. + + // Per-provider list of allocation requests emitted by the router. + // Allocation itself is performed by a post-route step in NativeBackend + // (the bid auction). Engines claim requests they want; HostPool is the + // fallback claimant for any unclaimed HostMirror. + // + // alloc_plans[runtime_idx] is the request list for that provider. + std::vector> alloc_plans; + + // Per-provider list of constant upload requests emitted by the router, + // partitioned by which engines actually consume each constant. The + // router fills this purely as planning output — engine I/O is driven + // by NativeBackend's post-route upload_constants pass (symmetric to + // alloc_plans / materialize_buffers). + // + // const_plans[runtime_idx] is the constant-request list for that + // provider. Each engine independently materializes its constants + // (zero-copy NDM alias on CPU / Apple-Silicon Metal; device-side + // load on discrete GPU). No cross-engine coordination is required; + // multiple engines may hold independent Buffer wrappers for the + // same NDM key without duplicating bytes (UMA platforms). + std::vector> const_plans; +}; + +} // namespace native +} // namespace backends +} // namespace executorch diff --git a/backends/native/ir/Step.h b/backends/native/ir/Step.h new file mode 100644 index 00000000000..4a90aa2e653 --- /dev/null +++ b/backends/native/ir/Step.h @@ -0,0 +1,164 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace executorch { +namespace backends { +namespace native { + +/** + * Sentinel for an unresolved step index. JumpFalseStep::dst_step_idx + * carries this between router-side emit and PCResolver-side patch. An + * unresolved value at execute time is a router bug (asserted in debug + * builds in NativeBackend.cpp::execute). + * + * See CONTROL_FLOW_DESIGN.md §5, §8. + */ +inline constexpr size_t kUnresolvedStep = std::numeric_limits::max(); + +/** + * Source-PC sentinel meaning "no source PC" (for synthesized steps the + * router emitted with no underlying instruction, e.g. a host-pool prep). + */ +inline constexpr uint32_t kNoSourcePc = std::numeric_limits::max(); + +/** + * One unit of issued work. Carries dense RuntimeIndex (not opaque + * RuntimeId). See §4.9 of PORTABLE_BACKEND_API_PROPOSAL.md. + * + * The Step variant is intentionally extensible. Today's variants: + * - ComputeStep — dispatch a compiled segment on one Engine. + * - TransferStep — cross-runtime byte movement (host↔device). + * - JumpFalseStep — host-side conditional PC mutation + * (CONTROL_FLOW_DESIGN.md §5). + * - MoveStep — host-side EValue assignment + * (CONTROL_FLOW_DESIGN.md §16). + * + * Anticipated future variants (see CONTROL_FLOW_DESIGN.md, "Region IR" + * follow-up): + * - RegionStep — delegate a whole control-flow region + * (cond/loop/while) to a single Engine that opted + * in via region_capabilities(). Acts as one logical + * step from the executor's POV; the Engine + * decides internally whether to do GPU-side + * branching, host-driven sync, etc. + * + * When a new variant is added, three places must be extended: + * 1. NativeBackend.cpp::execute_step (std::visit branches). + * 2. NativeBackend.cpp::execute (PC walker tracking + * observed_signals; the source_pc helper). + * 3. routers/GreedyRouter.cpp Phase 6 (host_visible_producer, + * step_input_values, step_signal lambdas) — see comments at + * those sites for extension guidance. + * + * std::visit's exhaustive-dispatch rule means the compiler will flag + * every site that needs updating. Do not add a default-fallthrough + * case — let the compiler enforce. + */ +struct ComputeStep { + RuntimeIndex runtime_idx; // dense index into Plan::instances + CompiledSegment* segment; + std::vector wait_for; + EventId signal = kNoEvent; + + // Source PC (instruction index in the chain) of the FIRST kernel in + // this segment. Used by PCResolver to map source-PC jump destinations + // to step indices. kNoSourcePc for synthesized steps. + // See CONTROL_FLOW_DESIGN.md §8. + uint32_t source_pc = kNoSourcePc; +}; + +struct TransferStep { + // Both ends are looked up via bindings at execute time. + uint32_t src_value_id; + uint32_t dst_value_id; + RuntimeIndex src_idx; // hot-path; kHostIdx if host + RuntimeIndex dst_idx; + std::vector wait_for; + EventId signal = kNoEvent; + + // Source PC of the originating instruction (the producer or consumer + // KernelCall this transfer serves; for MoveCall, the MoveCall PC; for + // PredicateLocator-emitted predicate downloads, the JumpFalseCall PC). + uint32_t source_pc = kNoSourcePc; +}; + +/** + * A control-flow step. Reads `pred_value_id` on host, mutates the + * executor's PC. See CONTROL_FLOW_DESIGN.md §5. + * + * Always runs on host (Invariant CF-1). PredicateLocator (§6) guarantees + * `pred_value_id` is host-resident at execute time. + * + * No `signal` field: a JumpFalseStep produces no value. Steps that + * follow it but consume values the predicate depends on already wait on + * those values' producing-step signals via their own wait_for; the jump + * is data-flow transparent. + */ +struct JumpFalseStep { + // EValue to inspect (Bool scalar or Bool tensor). Always host-resident + // at execute time per PredicateLocator (CONTROL_FLOW_DESIGN.md §6). + uint32_t pred_value_id; + + // Resolved destination as a step index into Plan::steps. Filled by + // PCResolver in a second pass; emitted with kUnresolvedStep. + size_t dst_step_idx = kUnresolvedStep; + + // Source-PC of the JumpFalseCall's destination (held alongside + // dst_step_idx for PCResolver and for diagnostics). + uint32_t unresolved_dst_pc = kNoSourcePc; + + // Precise transitive dependency closure of pred_value_id, as the + // signals that produce those dependencies. Filled by DepClosure + // (CONTROL_FLOW_DESIGN.md §7). + std::vector wait_for; + + // Source PC of the originating JumpFalseCall. + uint32_t source_pc = kNoSourcePc; +}; + +/** + * EValue-level move: `values[dst] = values[src]`. Mirrors ET's + * MoveCall semantics exactly (runtime/executor/method.cpp). + * + * NOT a byte copy. After this step, the destination value SHARES + * storage with the source value (for tensors, the dst's TensorImpl + * points at src's data). The router emits a MoveStep in two cases: + * 1. ET's chain contains a MoveCall (HOP output plumbing). + * 2. (Future) PredicateLocator host-mirror aliasing. + * + * Always runs on host. No signal (synchronous EValue assignment). + * + * Distinct from TransferStep (which is for cross-runtime byte movement + * via Engine::upload_from_host / download_to_host) — see + * CONTROL_FLOW_DESIGN.md §16. + */ +struct MoveStep { + uint32_t src_value_id; + uint32_t dst_value_id; + uint32_t source_pc = kNoSourcePc; +}; + +using Step = std::variant; + +} // namespace native +} // namespace backends +} // namespace executorch diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py new file mode 100644 index 00000000000..b728f69fc79 --- /dev/null +++ b/backends/native/partitioner.py @@ -0,0 +1,229 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Native Backend Partitioner + +The native partitioner marks ALL nodes as supported. The NativeBackend +runtime has a CPU fallback that can execute any portable op; runtime +partitioning across accelerators (Metal, etc.) happens in C++ based on +has_op() queries against the per-runtime op registries. +""" + +from typing import Any, Callable, Dict, final, List, Mapping, Optional, Tuple + +import torch + +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer + +from torch.export.exported_program import ExportedProgram +from torch.fx import Node +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner +from torch.fx.passes.operator_support import OperatorSupportBase + + +# Canonical list of ops the native backend preserves (does NOT decompose) +# during edge lowering when used with `to_edge_transform_and_lower`. Part +# of the universal-IR specification of NativeBackend; not user-configurable +# per call. +# +# Add an op here when you ship a dedicated C++ handler for it on the Metal +# (or other accelerator) side; otherwise the default decomposition pass at +# edge lowering will break it apart and we'll just dispatch the pieces. +_DEFAULT_PRESERVED_OPS = [ + torch.ops.aten.matmul.default, + torch.ops.aten.linear.default, + torch.ops.aten.addmm.default, + torch.ops.aten.baddbmm.default, + torch.ops.aten.scaled_dot_product_attention.default, +] + + +class NativeSupportedOperators(OperatorSupportBase): + """ + Operator support checker for the Native backend. + + NativeBackend has a CPU fallback, so it supports ALL operators. The + actual runtime partitioning across accelerators happens in C++. + """ + + def is_node_supported( + self, submodules: Mapping[str, torch.nn.Module], node: Node + ) -> bool: + # Skip placeholder and output nodes - they shouldn't be partitioned + if node.op in ("placeholder", "output", "get_attr"): + return False + if node.op != "call_function": + return False + # TODO(native): Claim HOPs (cond, while_loop, scan, ...) so the + # delegate's inner program is byte-equivalent to ET stock emit + # (modulo preserved ops + reinplace), including JumpFalseCall + + # branch instructions. Today we skip them so they stay at the + # outer level and run through ET's standard control-flow + # executor; ops INSIDE the branches still get partitioned via + # to_backend's recursive is_submodule=True path. + # + # To claim HOPs we'd need: + # 1. Pre-populate `val` on the branch-graph `get_attr` nodes + # so the wrap step's get_attr→placeholder conversion in + # fuse_as_graphmodule produces a valid placeholder. ARM + # does this in + # backends/arm/operator_support/control_flow_support.py + # (_submodules_fully_partitioned). + # 2. Make NativeBackend.preprocess HOP-aware: SpecPropPass, + # MemoryPlanningPass, ConstraintBasedSymShapeEvalPass, + # and emit_program all need to recursively walk branch + # graphs. Today they crash with `'ProxyValue' has no + # attribute 'graph'` when they hit a cond. + # 3. Make our runtime Graph adapter resolve the nested branch + # programs that emit_program inlines under JumpFalseCall. + # ARM's stack of control-flow passes (control_flow_const_inline.py + # etc.) is the reference implementation. + if isinstance(node.target, torch._ops.HigherOrderOperator): + return False + return True + + +@final +class NativePartitioner(Partitioner): + """ + Partitioner for the Native Backend. + + Unlike accelerator-specific partitioners that only claim ops they can + accelerate, the native partitioner claims ALL ops since: + 1. CPU fallback can execute any portable op. + 2. Runtime partitioning in C++ handles dispatch to accelerators. + + Two usage paths: + + A) Classic to_edge + to_backend (default decomposition runs first): + edge_program = to_edge(exported_program) + edge_program = edge_program.to_backend(NativePartitioner()) + + B) to_edge_transform_and_lower (preserves listed ops from decomposition): + edge_program = to_edge_transform_and_lower( + exported_program, + partitioner=[NativePartitioner()], + ) + # The ops in _DEFAULT_PRESERVED_OPS are kept intact (no decomposition) + # so accelerator kernels see them whole. The list is maintained + # by this partitioner — extend `_DEFAULT_PRESERVED_OPS` (in code) when + # you ship a new dedicated handler. Per-call overrides are NOT + # exposed: NativeBackend is a universal IR and the preserve list + # is part of its specification. + """ + + def __init__( + self, + compile_options: Optional[Dict[str, Any]] = None, + ) -> None: + self.options = compile_options or {} + compile_spec = self._parse_compile_options(self.options) + # Import here to avoid circular dependency + from executorch.backends.native.preprocess import NativeBackend + + self.delegation_spec = DelegationSpec(NativeBackend.__name__, compile_spec) + + def _parse_compile_options(self, options: Dict[str, Any]) -> List[CompileSpec]: + """Convert compile options dict to CompileSpec list.""" + compile_specs = [] + + for key, value in options.items(): + if isinstance(value, bool): + value_bytes = value.to_bytes(1, byteorder="little") + compile_specs.append(CompileSpec(key, value_bytes)) + elif isinstance(value, int): + value_bytes = value.to_bytes(4, byteorder="little") + compile_specs.append(CompileSpec(key, value_bytes)) + + return compile_specs + + def ops_to_not_decompose( + self, ep: ExportedProgram + ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[Node], bool]]]: + """ + Return ops that should NOT be decomposed during edge lowering. + + Called by `to_edge_transform_and_lower` BEFORE partitioning. The ops + returned here are kept whole (skipped by the default decomposition + pass), so backend kernels see them in their original form. + + The preserve list (`_DEFAULT_PRESERVED_OPS`) is the canonical list + maintained by NativeBackend — part of the universal-IR specification + and not user-configurable. Extend the constant in this file when + you ship a new dedicated handler. + + Behavior: + - If the graph is already partitioned (contains lowered_module + get_attr nodes), return empty — partitioning has run, decomposition + decisions are settled. + - Otherwise return the preserve list, intersected with ops that + actually appear in `ep` (no point listing ops the graph doesn't + have). + + The second tuple element (filter callable) is None: the rule applies + uniformly to every node whose target is in the preserve list. + """ + # Already-partitioned graph -> nothing to preserve. + for node in ep.graph.nodes: + if node.op == "get_attr" and "lowered_module" in node.name: + return ([], None) + + # Intersect preserve list with ops actually present in the graph. + present: List[torch._ops.OpOverload] = [] + seen = set() + for node in ep.graph.nodes: + if node.op != "call_function": + continue + if not isinstance(node.target, torch._ops.OpOverload): + continue + if node.target in _DEFAULT_PRESERVED_OPS and node.target not in seen: + present.append(node.target) + seen.add(node.target) + + return (present, None) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + """ + Partition the exported program for the native backend. + + Since native supports everything, this partitions the entire graph + into a single delegation block. + """ + partition_tags = {} + + # Use CapabilityBasedPartitioner with our "support everything" checker + capability_partitioner = CapabilityBasedPartitioner( + exported_program.graph_module, + NativeSupportedOperators(), + allows_single_node_partition=True, + ) + + partition_list = capability_partitioner.propose_partitions() + + for partition in partition_list: + for node in partition.nodes: + tag = f"tag{partition.id}" + node.meta["delegation_tag"] = tag + partition_tags[tag] = self.delegation_spec + + # Tag constant data for proper handling + tag_constant_data(exported_program) + # Tag mutated buffer placeholders so they're owned by our delegate + # (KV-cache style state flows through the delegated subgraph; the + # writeback copy_ collapses out of the top-level chain). + tag_mutated_buffer(exported_program) + + return PartitionResult( + tagged_exported_program=exported_program, + partition_tags=partition_tags, + ) diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py new file mode 100644 index 00000000000..4c8fdf01a44 --- /dev/null +++ b/backends/native/preprocess.py @@ -0,0 +1,390 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +NativeBackend — AOT BackendDetails for the v2 portable runtime. + +The class name `NativeBackend` is the backend_id used at runtime +to find the matching BackendInterface. C++ side registers it via +`register_backend({"NativeBackend", ...})` in +native/NativeBackend.cpp. + +The preprocess() pipeline: +1. SpecPropPass to populate tensor specs. +2. NEW: reinplace_pass with the backend's BACKEND_INPLACE_OPS registry + (~85 edge-functional ops mapped to aten in-place targets). Rewrites + `op(self, ...)` -> `op_(self, ...)` when ET's safety check passes. + Re-runs SpecPropPass so new in-place nodes get spec metadata. +3. (If buffer mutations exist) insert_write_back_for_buffers_pass to + add an explicit aten::copy_(buf, mut_src) at end of subgraph. When + step 2 already rewrote the buffer mutation in-place, + `_inplace_lineage` detects this and skips inserting the writeback. +4. Spec-sharing fallback for any remaining writebacks: alias the + mutation source's TensorSpec to the buffer's. Used only when step 2 + didn't catch the buffer mutation (e.g., custom op outside the + registry). +5. ExternalConstantsPass to tag constants for NDM storage. +6. Memory planning (greedy, allow_overlapping_allocations). The + upstream `_alias_inplace_result_specs` aliases in-place op result + specs to mutated input specs so emit's `_emit_spec` dedup gives + them one value_id. +7. Emit and serialize. +""" + +from functools import partial +from typing import Any, Dict, final, List + +import torch +from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer + +from executorch.exir.backend.backend_details import ( + BackendDetails, + CompileSpec, + ExportedProgram, + PreprocessResult, +) +from executorch.exir.emit import emit_program +from executorch.exir.memory_planning import greedy, MemoryPlanningAlgorithmSuite +from executorch.exir.passes import MemoryPlanningPass, SpecPropPass +from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + insert_write_back_for_buffers_pass, +) +from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass +from executorch.exir.program._program import _transform + +from torch._export.verifier import Verifier + + +# --------------------------------------------------------------------------- +# Backend-owned in-place op registry (consumed by ET's reinplace_pass). +# +# Maps edge-dialect functional ops -> aten in-place targets. The pass +# runs on the edge-dialect graph; targets use the aten in-place form +# because the edge dialect doesn't register in-place variants for most +# ops. The lowered IR ends up carrying `aten::_` instructions +# which the executor accepts. +# +# Built once at module load by intersecting `ops.edge.aten.` with +# `torch.ops.aten._`. This table lives in the backend (not +# upstream) because it's a per-runtime kernel-availability statement. +# --------------------------------------------------------------------------- + + +def _build_backend_inplace_ops() -> frozenset: + """Construct the v2 backend's in-place op set by listing the + edge-functional ops we support. ET's reinplace_pass auto-derives + the in-place counterpart for each via schema matching. + """ + from executorch.exir.dialects._ops import ops as _edge_ops + from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS + + edge = _edge_ops.edge.aten + + pairs = [ + # ------- pointwise unary ------- + ("relu", ["default"]), + ("relu6", ["default"]), + ("sigmoid", ["default"]), + ("tanh", ["default"]), + ("exp", ["default"]), + ("expm1", ["default"]), + ("log", ["default"]), + ("log1p", ["default"]), + ("log2", ["default"]), + ("log10", ["default"]), + ("neg", ["default"]), + ("abs", ["default"]), + ("sqrt", ["default"]), + ("rsqrt", ["default"]), + ("reciprocal", ["default"]), + ("square", ["default"]), + ("cos", ["default"]), + ("sin", ["default"]), + ("tan", ["default"]), + ("cosh", ["default"]), + ("sinh", ["default"]), + ("asin", ["default"]), + ("acos", ["default"]), + ("atan", ["default"]), + ("asinh", ["default"]), + ("acosh", ["default"]), + ("atanh", ["default"]), + ("erf", ["default"]), + ("erfc", ["default"]), + ("sign", ["default"]), + ("ceil", ["default"]), + ("floor", ["default"]), + ("round", ["default"]), + ("trunc", ["default"]), + ("frac", ["default"]), + ("silu", ["default"]), + ("gelu", ["default"]), + ("elu", ["default"]), + ("leaky_relu", ["default"]), + ("hardtanh", ["default"]), + ("hardsigmoid", ["default"]), + ("hardswish", ["default"]), + ("logical_not", ["default"]), + ("bitwise_not", ["default"]), + # ------- binary ------- + ("add", ["Tensor", "Scalar"]), + ("sub", ["Tensor", "Scalar"]), + ("mul", ["Tensor", "Scalar"]), + ("div", ["Tensor", "Scalar"]), + ("pow", ["Scalar", "Tensor_Scalar"]), + ("remainder", ["Tensor", "Scalar"]), + ("fmod", ["Tensor", "Scalar"]), + ("atan2", ["default"]), + ("logical_and", ["default"]), + ("logical_or", ["default"]), + ("logical_xor", ["default"]), + ("bitwise_and", ["Tensor", "Scalar"]), + ("bitwise_or", ["Tensor", "Scalar"]), + ("bitwise_xor", ["Tensor", "Scalar"]), + # ------- scatter / index ------- + ("index_copy", ["default"]), + ("index_fill", ["int_Scalar"]), + ("index_add", ["default"]), + ("scatter", ["src", "value"]), + ("scatter_add", ["default"]), + ("masked_fill", ["Scalar", "Tensor"]), + ("masked_scatter", ["default"]), + # ------- misc ------- + ("fill", ["Scalar", "Tensor"]), + ("clamp", ["default"]), + ("clamp_min", ["default"]), + ("clamp_max", ["default"]), + ("addcmul", ["default"]), + ("addcdiv", ["default"]), + ("lerp", ["Scalar", "Tensor"]), + ] + + functional_ops = set() + for name, overloads in pairs: + edge_pkg = getattr(edge, name, None) + if edge_pkg is None: + continue + for ovld in overloads: + op = getattr(edge_pkg, ovld, None) + if op is not None: + functional_ops.add(op) + + return DEFAULT_INPLACEABLE_OPS | functional_ops + + +BACKEND_INPLACE_OPS: frozenset = _build_backend_inplace_ops() + + +class _AnyOp(Verifier): + """Permissive verifier that allows any op (skip functional check).""" + + dialect = "TRAINING" + + def allowed_op_types(self): + from typing import Callable + + return (Callable,) + + +def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram: + """Apply a sequence of passes to an ExportedProgram.""" + from executorch.exir.pass_base import ExportPass, PassBase + + for p in passes: + if isinstance(p, MemoryPlanningPass) and hasattr(p, "run"): + p.run(program.graph_module) + elif issubclass(type(p), (ExportPass, PassBase)): + if hasattr(p, "_exported_program"): + p._exported_program = program + program = _transform(program, p, override_verifiers=[_AnyOp]) + if isinstance(p, SpecPropPass): + p.update_placeholder_tensor_specs(program, program.graph_module) + else: + program = p(program) + + return program + + +def _parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: + """Parse compile specs into options dict.""" + options: Dict[str, Any] = { + "skip_memory_planning": False, + # Default ON: the v2 portable runtime's CPU provider has + # dispatches for in-place ops (`aten::add_` etc.) that route to + # the existing `.out` variant kernels. The router falls back to + # CPU when no other provider claims the op. + "enable_reinplace": True, + } + for spec in compile_specs: + if spec.key in options and isinstance(options[spec.key], bool): + options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + return options + + +@final +class NativeBackend(BackendDetails): + """ + BackendDetails for the v2 portable backend. + + Class name `NativeBackend` matches the runtime backend_id + registered in native/NativeBackend.cpp. + """ + + @classmethod + def preprocess( + cls, + program: ExportedProgram, + module_compile_spec: List[CompileSpec], + ) -> PreprocessResult: + """ + Preprocess the partitioned subgraph for v2 portable backend execution. + """ + compile_options = _parse_compile_spec(module_compile_spec) + skip_memory_planning = compile_options["skip_memory_planning"] + enable_reinplace = compile_options["enable_reinplace"] + + # Step 1: SpecPropPass to propagate tensor specs. + program = _apply_passes(program, [SpecPropPass()]) + + # Step 2: NEW — reinplace_pass with the backend's in-place op + # registry. Rewrites functional ops to their `*_` form when ET's + # safety check passes (sole consumer; mutable inputs OK; no + # later use). For buffer mutations (KV-cache pattern), this + # turns `index_put(buf, ...)` into `index_put_(buf, ...)`, + # which then makes `insert_write_back_for_buffers_pass` (step 3) + # skip inserting a redundant `copy_` writeback via its + # `_inplace_lineage` check. + if enable_reinplace: + from executorch.exir.passes.reinplace import ( + reinplace_pass as _et_reinplace_pass, + ) + + program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS) + # ET's reinplace_pass copies meta["val"] to the new in-place + # node but not meta["spec"]. Re-run SpecPropPass so the new + # nodes have specs that downstream passes can read. + program = _apply_passes(program, [SpecPropPass()]) + + # Step 3: Insert writeback copy_ ops for any mutable buffers + # that weren't already mutated in place by step 2. The pass + # uses `_inplace_lineage` to detect in-place chains and skip + # writeback insertion for them. + from torch.export.graph_signature import InputKind, OutputKind + + has_buffer_mutation = any( + ospec.kind == OutputKind.BUFFER_MUTATION + for ospec in program.graph_signature.output_specs + ) + if has_buffer_mutation: + gm, new_sig = insert_write_back_for_buffers_pass(program) + program._graph_module = gm + program._graph_signature = new_sig + # Re-propagate specs onto the newly inserted copy_ nodes. + program = _apply_passes(program, [SpecPropPass()]) + + # Spec-sharing trick: for each (buffer_placeholder, mutation + # source) pair, make the mutation source's TensorSpec be the + # SAME object as the buffer's TensorSpec. The greedy memory + # planner walks specs by identity, so a shared spec yields a + # single allocation. No wasted slot, no runtime override + # needed (the .pte naturally reports both at the same offset). + + sig = program.graph_signature + nodes_by_name = {n.name: n for n in program.graph_module.graph.nodes} + buf_target_to_node = { + ispec.target: nodes_by_name.get(ispec.arg.name) + for ispec in sig.input_specs + if ispec.kind == InputKind.BUFFER and ispec.target + } + for ospec in sig.output_specs: + if ospec.kind != OutputKind.BUFFER_MUTATION: + continue + buf_node = buf_target_to_node.get(ospec.target) + wb_node = nodes_by_name.get(ospec.arg.name) + if ( + buf_node is None + or wb_node is None + or wb_node.op != "call_function" + or wb_node.target != torch.ops.aten.copy_.default + or len(wb_node.args) < 2 + ): + continue + src_node = wb_node.args[1] + if not hasattr(src_node, "meta"): + continue + buf_spec = buf_node.meta.get("spec") + if buf_spec is None or "spec" not in src_node.meta: + continue + # Alias: src now shares buf's spec object. + src_node.meta["spec"] = buf_spec + + # Step 4: External constants pass — tag named parameters/buffers + # for NDM storage. The stock pass deliberately skips lifted + # tensor constants ("they are closer to code than data"). We + # used to override that to force lifted constants to NDM, but + # the runtime now supports inline constants directly via + # Engine::ConstRequest.inline_data + Graph::tensor_inline_data, + # so we keep the stock semantics. + from executorch.exir.passes.external_constants_pass import ( + external_constants_pass, + ) + + external_constants_pass(program.graph_module) + + # Step 5: Memory planning (greedy, allows overlapping allocs). + if not skip_memory_planning: + greedy_memory_planning = partial(greedy, allow_overlapping_allocations=True) + mem_planning_suite = MemoryPlanningAlgorithmSuite( + algo_list=[greedy_memory_planning] + ) + + # Workaround for memory planning without ToOutVarPass + program.graph_module.encounter_to_out_var_failure = True + + program = _apply_passes( + program, + [ + ConstraintBasedSymShapeEvalPass(), + MemoryPlanningPass(memory_planning_algo=mem_planning_suite), + ], + ) + + # Step 6: Emit the program. + emitter_output = emit_program(program) + + # Step 7: Build named data store from external constants. + from executorch.exir._serialize._named_data_store import NamedDataStore + + named_data_store = NamedDataStore() + if emitter_output.external_constant_buffer: + for _tag, fqn_to_idx in emitter_output.external_constant_map.items(): + for fqn, idx in fqn_to_idx.items(): + data = emitter_output.external_constant_buffer[idx] + named_data_store.add_named_data(fqn, data) + + # Step 8: Serialize the inner program to bytes. + # + # We deliberately do NOT use serialize_pte_binary here, because + # it extracts constants from program.constant_buffer into a + # separate constant_segment of the OUTER PTE file — which makes + # the inner program's constant_buffer empty when the runtime + # parses it. The constant data would end up in segments that + # the delegate has no view into. + # + # _program_to_flatbuffer is the lower-level serializer that + # keeps program.constant_buffer in place. The runtime's + # Graph::tensor_inline_data() reads directly from + # program->constant_buffer()->Get(idx), so inline lifted + # constants reach Engine::upload_constants via the inline_data + # path on ConstRequest. + fb_result = _program_to_flatbuffer(emitter_output.program) + serialized_bytes = bytes(fb_result.data) + + return PreprocessResult( + processed_bytes=serialized_bytes, + debug_handle_map=emitter_output.debug_handle_map, + data_store_output=named_data_store.get_named_data_store_output(), + ) From 456b365488c71d8d328b29cfe8bd3daa5e793038 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 1 Jul 2026 16:56:34 -0700 Subject: [PATCH 2/6] up --- backends/native/ir/GraphTypes.h | 7 +- backends/native/ir/Plan.h | 106 ------------------- backends/native/ir/Step.h | 164 ----------------------------- backends/native/partitioner.py | 22 +++- backends/native/preprocess.py | 180 ++++++++++++++++++-------------- 5 files changed, 128 insertions(+), 351 deletions(-) delete mode 100644 backends/native/ir/Plan.h delete mode 100644 backends/native/ir/Step.h diff --git a/backends/native/ir/GraphTypes.h b/backends/native/ir/GraphTypes.h index 7cce34d9122..98b942bf572 100644 --- a/backends/native/ir/GraphTypes.h +++ b/backends/native/ir/GraphTypes.h @@ -462,8 +462,7 @@ class Graph { // Version //===------------------------------------------------------------------===// - const char* version() const { - // TODO: Return actual version when available + static const char* version() { return "1.0"; } @@ -637,8 +636,8 @@ class Graph { return chains ? chains->size() : 0; } - int32_t main_chain_idx() const { - return 0; // Default: first chain is main + static int32_t main_chain_idx() { + return 0; } // Get number of ops in a chain diff --git a/backends/native/ir/Plan.h b/backends/native/ir/Plan.h deleted file mode 100644 index f3fd9eecebf..00000000000 --- a/backends/native/ir/Plan.h +++ /dev/null @@ -1,106 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include -#include -#include -#include - -#include -#include -#include - -namespace executorch { -namespace backends { -namespace native { - -/** - * Each event slot remembers which Engine created it. Required so the - * executor can call the right Engine::wait at the host boundary. - */ -struct EventSlot { - std::unique_ptr event; - Engine* owner; -}; - -/** - * Frozen output of Router::route. Holds: - * - Runtime/Engine arrays indexed by RuntimeIndex, - * - the issue-ordered Step list, - * - per-input/output bindings, - * - pre-allocated event slots, - * - per-provider allocation request lists. - * - * Buffer ownership lives in the Engines themselves (each engine releases - * everything it allocated in its destructor). NativeBackend tracks no - * per-vid storage state; engines are self-sufficient via their own - * value->Buffer tables and via set_io_bindings at init. - * - * See §4.9 of the design doc. - */ -struct Plan { - // Parallel arrays indexed by RuntimeIndex. By convention, index 0 is - // the HostPool (canonical home for boundary values; not a compute - // runtime — its can_run() always returns false). Compute providers - // (CpuRuntime, MetalRuntime, ...) occupy slots 1+. - std::vector providers; - std::vector instances; // non-owning; lifetime is DelegateInstance - - std::vector steps; // ordered by *issue* (not completion) - - std::vector inputs; - std::vector outputs; - - // Maximum control-flow hops the executor walks before declaring a - // malformed back-edge. Stamped from RouterOptions::max_hops at - // route() time so different programs can carry different limits. - size_t max_hops = 10'000'000; - - // Pre-allocated event slots. Each event is reset lazily by its - // producing Engine immediately before signaling. - std::vector events; - - // EventIds whose signal must complete before execute() returns. - // Replaces a blanket per-Engine drain. Populated by the router as - // the signals of the last step on each Engine. - std::vector terminal_events; - - // Mirror identities are conveyed entirely through AllocRequest: - // device-side mirror requests carry host_mirror_value_id, and - // engines materialize their own TensorImpls for mirror_ids when they - // claim the AllocRequest. NativeBackend has no separate mirror table. - - // Per-provider list of allocation requests emitted by the router. - // Allocation itself is performed by a post-route step in NativeBackend - // (the bid auction). Engines claim requests they want; HostPool is the - // fallback claimant for any unclaimed HostMirror. - // - // alloc_plans[runtime_idx] is the request list for that provider. - std::vector> alloc_plans; - - // Per-provider list of constant upload requests emitted by the router, - // partitioned by which engines actually consume each constant. The - // router fills this purely as planning output — engine I/O is driven - // by NativeBackend's post-route upload_constants pass (symmetric to - // alloc_plans / materialize_buffers). - // - // const_plans[runtime_idx] is the constant-request list for that - // provider. Each engine independently materializes its constants - // (zero-copy NDM alias on CPU / Apple-Silicon Metal; device-side - // load on discrete GPU). No cross-engine coordination is required; - // multiple engines may hold independent Buffer wrappers for the - // same NDM key without duplicating bytes (UMA platforms). - std::vector> const_plans; -}; - -} // namespace native -} // namespace backends -} // namespace executorch diff --git a/backends/native/ir/Step.h b/backends/native/ir/Step.h deleted file mode 100644 index 4a90aa2e653..00000000000 --- a/backends/native/ir/Step.h +++ /dev/null @@ -1,164 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include -#include - -#include - -#include -#include -#include -#include -#include - -namespace executorch { -namespace backends { -namespace native { - -/** - * Sentinel for an unresolved step index. JumpFalseStep::dst_step_idx - * carries this between router-side emit and PCResolver-side patch. An - * unresolved value at execute time is a router bug (asserted in debug - * builds in NativeBackend.cpp::execute). - * - * See CONTROL_FLOW_DESIGN.md §5, §8. - */ -inline constexpr size_t kUnresolvedStep = std::numeric_limits::max(); - -/** - * Source-PC sentinel meaning "no source PC" (for synthesized steps the - * router emitted with no underlying instruction, e.g. a host-pool prep). - */ -inline constexpr uint32_t kNoSourcePc = std::numeric_limits::max(); - -/** - * One unit of issued work. Carries dense RuntimeIndex (not opaque - * RuntimeId). See §4.9 of PORTABLE_BACKEND_API_PROPOSAL.md. - * - * The Step variant is intentionally extensible. Today's variants: - * - ComputeStep — dispatch a compiled segment on one Engine. - * - TransferStep — cross-runtime byte movement (host↔device). - * - JumpFalseStep — host-side conditional PC mutation - * (CONTROL_FLOW_DESIGN.md §5). - * - MoveStep — host-side EValue assignment - * (CONTROL_FLOW_DESIGN.md §16). - * - * Anticipated future variants (see CONTROL_FLOW_DESIGN.md, "Region IR" - * follow-up): - * - RegionStep — delegate a whole control-flow region - * (cond/loop/while) to a single Engine that opted - * in via region_capabilities(). Acts as one logical - * step from the executor's POV; the Engine - * decides internally whether to do GPU-side - * branching, host-driven sync, etc. - * - * When a new variant is added, three places must be extended: - * 1. NativeBackend.cpp::execute_step (std::visit branches). - * 2. NativeBackend.cpp::execute (PC walker tracking - * observed_signals; the source_pc helper). - * 3. routers/GreedyRouter.cpp Phase 6 (host_visible_producer, - * step_input_values, step_signal lambdas) — see comments at - * those sites for extension guidance. - * - * std::visit's exhaustive-dispatch rule means the compiler will flag - * every site that needs updating. Do not add a default-fallthrough - * case — let the compiler enforce. - */ -struct ComputeStep { - RuntimeIndex runtime_idx; // dense index into Plan::instances - CompiledSegment* segment; - std::vector wait_for; - EventId signal = kNoEvent; - - // Source PC (instruction index in the chain) of the FIRST kernel in - // this segment. Used by PCResolver to map source-PC jump destinations - // to step indices. kNoSourcePc for synthesized steps. - // See CONTROL_FLOW_DESIGN.md §8. - uint32_t source_pc = kNoSourcePc; -}; - -struct TransferStep { - // Both ends are looked up via bindings at execute time. - uint32_t src_value_id; - uint32_t dst_value_id; - RuntimeIndex src_idx; // hot-path; kHostIdx if host - RuntimeIndex dst_idx; - std::vector wait_for; - EventId signal = kNoEvent; - - // Source PC of the originating instruction (the producer or consumer - // KernelCall this transfer serves; for MoveCall, the MoveCall PC; for - // PredicateLocator-emitted predicate downloads, the JumpFalseCall PC). - uint32_t source_pc = kNoSourcePc; -}; - -/** - * A control-flow step. Reads `pred_value_id` on host, mutates the - * executor's PC. See CONTROL_FLOW_DESIGN.md §5. - * - * Always runs on host (Invariant CF-1). PredicateLocator (§6) guarantees - * `pred_value_id` is host-resident at execute time. - * - * No `signal` field: a JumpFalseStep produces no value. Steps that - * follow it but consume values the predicate depends on already wait on - * those values' producing-step signals via their own wait_for; the jump - * is data-flow transparent. - */ -struct JumpFalseStep { - // EValue to inspect (Bool scalar or Bool tensor). Always host-resident - // at execute time per PredicateLocator (CONTROL_FLOW_DESIGN.md §6). - uint32_t pred_value_id; - - // Resolved destination as a step index into Plan::steps. Filled by - // PCResolver in a second pass; emitted with kUnresolvedStep. - size_t dst_step_idx = kUnresolvedStep; - - // Source-PC of the JumpFalseCall's destination (held alongside - // dst_step_idx for PCResolver and for diagnostics). - uint32_t unresolved_dst_pc = kNoSourcePc; - - // Precise transitive dependency closure of pred_value_id, as the - // signals that produce those dependencies. Filled by DepClosure - // (CONTROL_FLOW_DESIGN.md §7). - std::vector wait_for; - - // Source PC of the originating JumpFalseCall. - uint32_t source_pc = kNoSourcePc; -}; - -/** - * EValue-level move: `values[dst] = values[src]`. Mirrors ET's - * MoveCall semantics exactly (runtime/executor/method.cpp). - * - * NOT a byte copy. After this step, the destination value SHARES - * storage with the source value (for tensors, the dst's TensorImpl - * points at src's data). The router emits a MoveStep in two cases: - * 1. ET's chain contains a MoveCall (HOP output plumbing). - * 2. (Future) PredicateLocator host-mirror aliasing. - * - * Always runs on host. No signal (synchronous EValue assignment). - * - * Distinct from TransferStep (which is for cross-runtime byte movement - * via Engine::upload_from_host / download_to_host) — see - * CONTROL_FLOW_DESIGN.md §16. - */ -struct MoveStep { - uint32_t src_value_id; - uint32_t dst_value_id; - uint32_t source_pc = kNoSourcePc; -}; - -using Step = std::variant; - -} // namespace native -} // namespace backends -} // namespace executorch diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py index b728f69fc79..d0bddf7f01e 100644 --- a/backends/native/partitioner.py +++ b/backends/native/partitioner.py @@ -13,7 +13,17 @@ has_op() queries against the per-runtime op registries. """ -from typing import Any, Callable, Dict, final, List, Mapping, Optional, Tuple +from typing import ( + Any, + Callable, + Dict, + final, + List, + Mapping, + Optional, + Tuple, + TYPE_CHECKING, +) import torch @@ -30,6 +40,9 @@ from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase +if TYPE_CHECKING: + from executorch.backends.native.preprocess import Specialization + # Canonical list of ops the native backend preserves (does NOT decompose) # during edge lowering when used with `to_edge_transform_and_lower`. Part @@ -125,8 +138,10 @@ class NativePartitioner(Partitioner): def __init__( self, compile_options: Optional[Dict[str, Any]] = None, + specializations: Optional[List["Specialization"]] = None, ) -> None: self.options = compile_options or {} + self.specializations = specializations compile_spec = self._parse_compile_options(self.options) # Import here to avoid circular dependency from executorch.backends.native.preprocess import NativeBackend @@ -199,6 +214,11 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: Since native supports everything, this partitions the entire graph into a single delegation block. """ + # Wire specializations to the backend class before partitioning. + from executorch.backends.native.preprocess import NativeBackend + + NativeBackend._specializations = self.specializations + partition_tags = {} # Use CapabilityBasedPartitioner with our "support everything" checker diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py index 4c8fdf01a44..b1834850c0d 100644 --- a/backends/native/preprocess.py +++ b/backends/native/preprocess.py @@ -34,8 +34,10 @@ 7. Emit and serialize. """ +import copy +import struct from functools import partial -from typing import Any, Dict, final, List +from typing import Any, Dict, final, List, Optional, Sequence, Tuple, Type import torch from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer @@ -57,6 +59,9 @@ from torch._export.verifier import Verifier +# Type alias for specialization config: (BackendDetails subclass, compile_specs) +Specialization = Tuple[Type[BackendDetails], List[CompileSpec]] + # --------------------------------------------------------------------------- # Backend-owned in-place op registry (consumed by ET's reinplace_pass). @@ -79,87 +84,15 @@ def _build_backend_inplace_ops() -> frozenset: the in-place counterpart for each via schema matching. """ from executorch.exir.dialects._ops import ops as _edge_ops - from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS edge = _edge_ops.edge.aten pairs = [ - # ------- pointwise unary ------- ("relu", ["default"]), - ("relu6", ["default"]), - ("sigmoid", ["default"]), - ("tanh", ["default"]), - ("exp", ["default"]), - ("expm1", ["default"]), - ("log", ["default"]), - ("log1p", ["default"]), - ("log2", ["default"]), - ("log10", ["default"]), - ("neg", ["default"]), - ("abs", ["default"]), - ("sqrt", ["default"]), - ("rsqrt", ["default"]), - ("reciprocal", ["default"]), - ("square", ["default"]), - ("cos", ["default"]), - ("sin", ["default"]), - ("tan", ["default"]), - ("cosh", ["default"]), - ("sinh", ["default"]), - ("asin", ["default"]), - ("acos", ["default"]), - ("atan", ["default"]), - ("asinh", ["default"]), - ("acosh", ["default"]), - ("atanh", ["default"]), - ("erf", ["default"]), - ("erfc", ["default"]), - ("sign", ["default"]), - ("ceil", ["default"]), - ("floor", ["default"]), - ("round", ["default"]), - ("trunc", ["default"]), - ("frac", ["default"]), - ("silu", ["default"]), ("gelu", ["default"]), - ("elu", ["default"]), - ("leaky_relu", ["default"]), - ("hardtanh", ["default"]), - ("hardsigmoid", ["default"]), - ("hardswish", ["default"]), - ("logical_not", ["default"]), - ("bitwise_not", ["default"]), - # ------- binary ------- - ("add", ["Tensor", "Scalar"]), - ("sub", ["Tensor", "Scalar"]), - ("mul", ["Tensor", "Scalar"]), - ("div", ["Tensor", "Scalar"]), - ("pow", ["Scalar", "Tensor_Scalar"]), - ("remainder", ["Tensor", "Scalar"]), - ("fmod", ["Tensor", "Scalar"]), - ("atan2", ["default"]), - ("logical_and", ["default"]), - ("logical_or", ["default"]), - ("logical_xor", ["default"]), - ("bitwise_and", ["Tensor", "Scalar"]), - ("bitwise_or", ["Tensor", "Scalar"]), - ("bitwise_xor", ["Tensor", "Scalar"]), - # ------- scatter / index ------- + ("sigmoid", ["default"]), + ("index_put", ["default"]), ("index_copy", ["default"]), - ("index_fill", ["int_Scalar"]), - ("index_add", ["default"]), - ("scatter", ["src", "value"]), - ("scatter_add", ["default"]), - ("masked_fill", ["Scalar", "Tensor"]), - ("masked_scatter", ["default"]), - # ------- misc ------- - ("fill", ["Scalar", "Tensor"]), - ("clamp", ["default"]), - ("clamp_min", ["default"]), - ("clamp_max", ["default"]), - ("addcmul", ["default"]), - ("addcdiv", ["default"]), - ("lerp", ["Scalar", "Tensor"]), ] functional_ops = set() @@ -172,12 +105,53 @@ def _build_backend_inplace_ops() -> frozenset: if op is not None: functional_ops.add(op) - return DEFAULT_INPLACEABLE_OPS | functional_ops + return frozenset(functional_ops) BACKEND_INPLACE_OPS: frozenset = _build_backend_inplace_ops() +# --------------------------------------------------------------------------- +# Fat PTE container format +# +# A fat blob packs multiple backend specializations into one delegate payload. +# The runtime parses the header and picks the best available specialization. +# +# Layout: +# [4 bytes] magic "NFAT" +# [4 bytes] version (1) +# [4 bytes] num_specializations +# For each specialization: +# [32 bytes] backend_id (utf-8, null-padded) +# [8 bytes] offset into data section +# [8 bytes] length +# [payload bytes for specialization 0] +# [payload bytes for specialization 1] +# ... +# --------------------------------------------------------------------------- + +_FAT_MAGIC = b"NFAT" +_FAT_VERSION = 1 +_FAT_ENTRY_FMT = "32sQQ" # backend_id(32) + offset(u64) + size(u64) +_FAT_ENTRY_SIZE = struct.calcsize(_FAT_ENTRY_FMT) + + +def _pack_fat_blob( + specializations: List[Tuple[str, bytes]], +) -> bytes: + """Pack multiple (backend_id, payload) pairs into a fat blob.""" + header = struct.pack("<4sII", _FAT_MAGIC, _FAT_VERSION, len(specializations)) + + entries = [] + offset = 0 + for backend_id, payload in specializations: + bid = backend_id.encode("utf-8")[:32].ljust(32, b"\x00") + entries.append(struct.pack("<" + _FAT_ENTRY_FMT, bid, offset, len(payload))) + offset += len(payload) + + return header + b"".join(entries) + b"".join(p for _, p in specializations) + + class _AnyOp(Verifier): """Permissive verifier that allows any op (skip functional check).""" @@ -231,13 +205,67 @@ class NativeBackend(BackendDetails): Class name `NativeBackend` matches the runtime backend_id registered in native/NativeBackend.cpp. + + Supports "fat PTE" mode: when specializations are configured, + preprocess runs each backend's pipeline on the same subgraph and + packs all results into a single fat blob. The runtime picks the + best available specialization at load time. All specializations + share the same PTD (named data store). """ + # Set by NativePartitioner before preprocess is called. + _specializations: Optional[Sequence[Specialization]] = None + @classmethod def preprocess( cls, program: ExportedProgram, module_compile_spec: List[CompileSpec], + ) -> PreprocessResult: + native_result = cls._preprocess_native(program, module_compile_spec) + + specializations = cls._specializations + if not specializations: + return native_result + + # Fat PTE: run each specialization backend and pack results. + from executorch.exir._serialize._named_data_store import NamedDataStore + + fat_entries: List[Tuple[str, bytes]] = [ + ("NativeBackend", native_result.processed_bytes), + ] + + merged_data_store = NamedDataStore() + # Seed with native's named data. + if native_result.data_store_output: + for key, entry in native_result.data_store_output.pte_data.items(): + buf = native_result.data_store_output.buffers[entry.buffer_index] + merged_data_store.add_named_data(key, buf, alignment=entry.alignment) + + for backend_cls, spec_compile_specs in specializations: + spec_program = copy.deepcopy(program) + spec_result = backend_cls.preprocess(spec_program, spec_compile_specs) + + fat_entries.append((backend_cls.__name__, spec_result.processed_bytes)) + + if spec_result.data_store_output: + for key, entry in spec_result.data_store_output.pte_data.items(): + buf = spec_result.data_store_output.buffers[entry.buffer_index] + merged_data_store.add_named_data( + key, buf, alignment=entry.alignment + ) + + return PreprocessResult( + processed_bytes=_pack_fat_blob(fat_entries), + debug_handle_map=native_result.debug_handle_map, + data_store_output=merged_data_store.get_named_data_store_output(), + ) + + @classmethod + def _preprocess_native( + cls, + program: ExportedProgram, + module_compile_spec: List[CompileSpec], ) -> PreprocessResult: """ Preprocess the partitioned subgraph for v2 portable backend execution. From cd4548c6af1588d85cc8605c8f0120b10795f4ac Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 2 Jul 2026 12:23:34 -0700 Subject: [PATCH 3/6] up --- backends/native/fat_pte.py | 72 +++ backends/native/ir/GraphTypes.h | 545 +++-------------------- backends/native/partitioner.py | 171 +------ backends/native/preprocess.py | 335 ++++---------- backends/native/specializations.py | 47 ++ backends/native/test/__init__.py | 0 backends/native/test/test_fat_pte.py | 105 +++++ backends/native/test/test_partitioner.py | 125 ++++++ backends/native/test/test_preprocess.py | 163 +++++++ pytest.ini | 1 + 10 files changed, 689 insertions(+), 875 deletions(-) create mode 100644 backends/native/fat_pte.py create mode 100644 backends/native/specializations.py create mode 100644 backends/native/test/__init__.py create mode 100644 backends/native/test/test_fat_pte.py create mode 100644 backends/native/test/test_partitioner.py create mode 100644 backends/native/test/test_preprocess.py diff --git a/backends/native/fat_pte.py b/backends/native/fat_pte.py new file mode 100644 index 00000000000..93f240a3551 --- /dev/null +++ b/backends/native/fat_pte.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Fat PTE: packs multiple backend specializations into one delegate payload. +Runtime selects the best specialization at load time. All specializations +share a single PTD (named data store). + +Binary layout: + [4B] magic "NFAT" + [4B] version (1) + [4B] num_specializations + Per specialization: + [32B] backend_id (utf-8, null-padded) + [8B] offset into data section + [8B] length + [payload bytes ...] +""" + +import struct +from typing import Dict, List, Optional, Tuple + +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import PreprocessResult + +FAT_MAGIC = b"NFAT" +FAT_VERSION = 1 +_ENTRY_FMT = "32sQQ" +_ENTRY_SIZE = struct.calcsize(_ENTRY_FMT) + + +def pack_fat_blob(specializations: List[Tuple[str, bytes]]) -> bytes: + """Pack (backend_id, payload) pairs into a fat blob.""" + header = struct.pack("<4sII", FAT_MAGIC, FAT_VERSION, len(specializations)) + + entries = [] + offset = 0 + for backend_id, payload in specializations: + bid = backend_id.encode("utf-8")[:32].ljust(32, b"\x00") + entries.append(struct.pack("<" + _ENTRY_FMT, bid, offset, len(payload))) + offset += len(payload) + + return header + b"".join(entries) + b"".join(p for _, p in specializations) + + +def build_fat_result( + results: List[Tuple[str, PreprocessResult]], + debug_handle_map: Optional[Dict] = None, +) -> PreprocessResult: + """Merge (backend_id, PreprocessResult) pairs into a single fat result.""" + fat_entries: List[Tuple[str, bytes]] = [] + merged_data_store = NamedDataStore() + + for backend_id, result in results: + fat_entries.append((backend_id, result.processed_bytes)) + + if result.data_store_output: + for key, entry in result.data_store_output.pte_data.items(): + buf = result.data_store_output.buffers[entry.buffer_index] + merged_data_store.add_named_data(key, buf, alignment=entry.alignment) + + if debug_handle_map is None and results: + debug_handle_map = results[0][1].debug_handle_map + + return PreprocessResult( + processed_bytes=pack_fat_blob(fat_entries), + debug_handle_map=debug_handle_map, + data_store_output=merged_data_store.get_named_data_store_output(), + ) diff --git a/backends/native/ir/GraphTypes.h b/backends/native/ir/GraphTypes.h index 98b942bf572..ecec59f64bc 100644 --- a/backends/native/ir/GraphTypes.h +++ b/backends/native/ir/GraphTypes.h @@ -11,126 +11,31 @@ /** * Graph: backend-facing IR adapter. * - * These types provide our backends' view of the program IR. The current - * implementation adapts ExecuTorch's flatbuffer - * (executorch_flatbuffer::ExecutionPlan / KernelCall) but the API is - * intentionally backing-agnostic — a different serialization could - * replace the underlying storage without changing this header or any - * backend that uses it. + * Wraps ExecuTorch's flatbuffer (ExecutionPlan / KernelCall) behind a + * backing-agnostic API. A different serialization could replace the + * underlying storage without changing this header. * - * Some methods are documented IR concepts that the current adapter does - * not yet back (e.g., mutable_buffer_ids, version): they are placeholders - * pending serialization support. + * IR schema (what Graph presents): * - * ---------------------------------------------------------------------- - * Fictional IR schema (what Graph effectively presents) - * ---------------------------------------------------------------------- - * If the IR were serialized in its own format (independent of the - * underlying ExecuTorch flatbuffer), it would look like this: - * - * // Top-level * table Graph { - * version: string; // schema/program version - * values: [Value]; // dense pool indexed by value_id (uint) - * inputs: [uint]; // graph input value_ids - * outputs: [uint]; // graph output value_ids - * mutable_buffers: [uint]; // values that persist across executes - * operators: [OperatorDef]; // op-name registry (deduped) - * chains: [Chain]; // chains[0] is main; others used by - * // future control flow (if/while/call) - * } - * - * table OperatorDef { - * name: string; // e.g. "aten.add.Tensor" - * } - * - * // Op chains - * table Chain { - * instructions: [Instruction]; - * } - * - * table Instruction { - * // Today only one variant; future variants for control flow - * // (if/while/call) would extend this union. - * body: KernelCall; + * version: string; + * values: [Value]; // dense pool, indexed by value_id + * inputs: [uint]; + * outputs: [uint]; + * mutable_buffers: [uint]; // values that persist across executes + * operators: [OperatorDef]; // deduped op-name registry + * instructions: [KernelCall]; // flat list of operator calls * } * * table KernelCall { - * op_index: uint; // → Graph.operators[op_index].name - * args: [uint]; // value_ids: args[0..n-2] = inputs, - * // args[n-1] = output. Single-output - * // assumed today; multi-output is a - * // future extension. - * } - * - * // Values - * union Value { - * None, - * Int { v: int64; }, - * Double { v: float64; }, - * Bool { v: bool; }, - * String { v: string; }, - * Tensor, - * IntList, - * DoubleList, - * BoolList, - * OptionalTensor, - * // ... - * } - * - * table Tensor { - * scalar_type: ScalarType; // dtype - * sizes: [int32]; // for DYNAMIC_BOUND, this is max-shape - * dim_order: [uint8]; // permutation defining memory layout - * shape_dynamism: ShapeDynamism; // STATIC | DYNAMIC_BOUND | - * DYNAMIC_UNBOUND allocation: AllocationInfo?;// null = no AOT plan data: - * TensorData; - * } - * - * union TensorData { - * None, - * Inline { buffer_idx: uint; }, // bytes embedded in program - * External { ndm_key: string;}, // bytes in NamedDataMap, FQN-keyed + * op_index: uint; // -> operators[op_index] + * args: [uint]; // value_ids; last = output * } * - * table AllocationInfo { - * pool_id: int32; - * offset: uint64; // raw byte offset within pool_id - * } - * - * table IntList { - * member_ids: [int64]; // value_ids of list elements - * } - * - * ---------------------------------------------------------------------- - * Derived views computed by the adapter (not in the schema itself) - * ---------------------------------------------------------------------- - * mem_obj_id(vid) sort-and-index over (pool_id, offset) - * → dense small int identifying shared - * storage slots. Two values with the - * same id were memory-planned to share - * storage (used by router for - * AllocRequest grouping; backends MAY - * honor it as actual storage aliasing). - * value_kind(vid) from membership in inputs/outputs + - * the data field - * → INPUT / OUTPUT / CONSTANT / - * INTERMEDIATE / MUTABLE_BUFFER. - * tensor_constant_data_key(vid) convenience accessor for - * TensorData.External.ndm_key. - * tensor_nbytes_max(vid) dtype_size × prod(sizes); upper bound - * for DYNAMIC_BOUND tensors. - * - * ---------------------------------------------------------------------- - * Adapter cost - * ---------------------------------------------------------------------- - * All accessors are inline and compile down to essentially the same - * machine code as direct flatbuffer access. The Graph constructor pays - * a one-time O(N log N) precompute (over tensor values) for mem_obj_id, - * and O(num_inputs + num_outputs) for the value_kind sets. Per-call - * overhead in the runtime hot path is a few inline indirections plus - * predictable ET_CHECK branches; release builds compile away these - * checks entirely under -DNDEBUG. + * Derived views (computed at construction): + * mem_obj_id(vid) — dense int from (pool_id, offset) sort-rank + * value_kind(vid) — INPUT / OUTPUT / CONSTANT / INTERMEDIATE / + * MUTABLE_BUFFER */ #include @@ -152,24 +57,16 @@ namespace executorch { namespace backends { namespace portable { -// Forward declare class Graph; -/** - * Value kind — matches design doc's TensorKind - */ enum class ValueKind : uint8_t { - INPUT = 0, // Graph input (user provides) - OUTPUT, // Graph output (user reads) - CONSTANT, // Immutable weight - MUTABLE_BUFFER, // Mutable state (e.g., KV cache) - INTERMEDIATE, // Temporary (produced/consumed internally) + INPUT = 0, + OUTPUT, + CONSTANT, + MUTABLE_BUFFER, + INTERMEDIATE, }; -/** - * Type of an EValue stored at a value_id. Mirrors the runtime EValue - * sum-type but in adapter-level form (no flatbuffer types in the API). - */ enum class ValueType : uint8_t { None = 0, Int, @@ -179,68 +76,12 @@ enum class ValueType : uint8_t { IntList, TensorList, OptionalTensorList, - Other, // String, OptionalTensor, BoolList, DoubleList, ... - // Adapter doesn't surface these yet; executor falls back to - // default-constructed EValue. + Other, }; /** - * Kind of an Instruction in a Chain. Mirrors ET's InstructionArguments - * union but as an adapter-level enum (no flatbuffer types in the API). - * - * See CONTROL_FLOW_DESIGN.md §4. - * - * - Kernel: the only kind a CompiledSegment may contain. - * - JumpFalse: control-flow boundary; segment break. - * - Move: copies one EValue into another; segment break (host-host - * transfer). - * - Free: informational hint to release a value's storage; v2 ignores - * (memory plan handles deallocation at delegate destroy). - * - Delegate: nested executorch_call_delegate inside our chain. - * REJECTED at init; partitioner contract forbids nested delegates. - */ -enum class InstructionKind : uint8_t { - Kernel = 0, - JumpFalse, - Move, - Free, - Delegate, -}; - -/** - * Decoded JumpFalseCall payload. Backing-agnostic. - */ -struct JumpFalseInfo { - uint32_t cond_value_id; // EValue index whose value is the predicate - uint32_t destination_pc; // Source-PC of the jump target -}; - -/** - * Decoded MoveCall payload. Backing-agnostic. - */ -struct MoveInfo { - uint32_t src_value_id; - uint32_t dst_value_id; -}; - -/** - * Thin wrapper around ExecuTorch's flatbuffer KernelCall providing - * convenient access to op name and input/output value_ids. - * - * ExecuTorch packs all op args into a single `args` array per the - * convention "the last arg is the (single) output." This wrapper - * exposes inputs/outputs accordingly without copying — `inputs()` / - * `output()` return Spans / values that reference the underlying - * flatbuffer storage directly. - * - * Per-call cost: just stores two pointers + an int. No allocations. - * Safe to construct repeatedly in hot dispatch loops. - * - * NOTE: Multi-output ops (e.g., `aten.split`, `aten.max.dim`) are not - * supported by this wrapper — `num_outputs()` always returns 0 or 1. - * Adding multi-output support requires per-op schema knowledge that - * isn't in the flatbuffer. ET_CHECK guards the single-output access - * path. + * Wraps a flatbuffer KernelCall. Provides access to op name and + * input/output value_ids. Last arg is the (single) output. */ class OperatorCall { public: @@ -249,7 +90,6 @@ class OperatorCall { const Graph* graph) : call_(call), graph_(graph) {} - // node_id for error messages/profiling uint32_t node_id() const { return node_id_; } @@ -257,33 +97,16 @@ class OperatorCall { node_id_ = id; } - // Op base name (e.g., "aten::add"). Does NOT include the overload - // suffix; for that use full_name(). const char* name() const; - - // Op overload (e.g., "Tensor", "Scalar", "out"). May be empty string - // for ops with no overload disambiguation. Use full_name() to get - // the combined "name.overload" form most callers want. const char* overload() const; - - // Combined unique key "." (e.g., "aten::add.Tensor"), - // or just "" if the overload is empty. This is the canonical - // identity used by the CPU op registry to dispatch — registering - // by base name alone collapses every overload of an op into one - // handler, which is wrong for ops like aten::pow_ that have multiple - // overloads with different schemas. std::string full_name() const; - // All op args (flat). Last entry is the output; the rest are inputs. - // Optional args are NOT indicated by -1 — instead the index points to - // a value with isNone() == true. (ExecuTorch convention.) runtime::Span args() const { auto* a = call_->args(); return a ? runtime::Span(a->data(), a->size()) : runtime::Span{}; } - // Inputs = args[0..n-2]. Returns empty span if op has no args. runtime::Span inputs() const { auto a = args(); return a.empty() ? a : runtime::Span(a.data(), a.size() - 1); @@ -293,6 +116,7 @@ class OperatorCall { auto a = args(); return a.empty() ? 0 : a.size() - 1; } + uint32_t input(size_t i) const { ET_CHECK_MSG( i < num_inputs(), @@ -302,11 +126,10 @@ class OperatorCall { return static_cast(args()[i]); } - // Single-output assumption (see class doc). Multi-output ops will - // break here. size_t num_outputs() const { return args().empty() ? 0 : 1; } + uint32_t output(size_t i) const { ET_CHECK_MSG( i < num_outputs(), @@ -324,10 +147,7 @@ class OperatorCall { }; /** - * IR view of a program. Adapts executorch_flatbuffer::ExecutionPlan; - * exposes value metadata, input/output IDs, the operator table, and - * chains of operator calls. See the file-header comment for the - * adapter-pattern rationale. + * IR view of a program. Adapts executorch_flatbuffer::ExecutionPlan. */ class Graph { public: @@ -335,7 +155,6 @@ class Graph { const executorch_flatbuffer::ExecutionPlan* plan, const executorch_flatbuffer::Program* program = nullptr) : plan_(plan), program_(program) { - // Precompute input/output value_id sets for O(1) value_kind lookup. if (auto* in = plan_->inputs()) { input_ids_.reserve(in->size()); for (size_t i = 0; i < in->size(); ++i) { @@ -349,20 +168,13 @@ class Graph { } } - // Precompute mem_obj_id for every tensor value. - // - // Algorithm: collect (pool_id, offset) keys for all aliasable tensor - // values; sort the unique keys; assign mem_obj_id = sort rank. Two - // values with the same (pool_id, offset) get the same id (they share - // storage). Sort-and-index is deterministic across runs and depends - // only on the AOT memory plan. + // Compute mem_obj_id: sort (pool_id, offset) pairs, assign dense rank. + // Same (pool_id, offset) → same id (shared storage). size_t n_vals = num_values(); mem_obj_ids_.assign(n_vals, -1); if (n_vals == 0) return; - // 1. Collect (key, value_id) entries for tensor values with - // allocation_info. struct Entry { uint64_t key; // (pool_id << 32) | offset uint32_t value_id; @@ -388,13 +200,11 @@ class Graph { if (entries.empty()) return; - // 2. Sort by key (lex order on (pool_id, offset)). std::sort( entries.begin(), entries.end(), [](const Entry& a, const Entry& b) { return a.key < b.key; }); - // 3. Assign mem_obj_id = sort rank (same key → same id). int32_t next_id = -1; uint64_t prev_key = ~0ULL; for (const auto& e : entries) { @@ -405,54 +215,31 @@ class Graph { mem_obj_ids_[e.value_id] = next_id; } - // Precompute mutable_buffer_ids_: tensor values with allocation_info, - // not graph IO, not constants, and NOT produced by any op (i.e. - // placeholders that aren't graph inputs). These are mutable buffer - // placeholders pulled into the delegate by tag_mutated_buffer; their - // state persists across execute() calls. - // - // Used by the router to distinguish semantic alias groups (buffer - // mutation: AOT spec-shared the buffer placeholder with its mutation - // source) from lifetime-reuse aliasing (the planner happened to put - // two values at the same offset because their lifetimes don't - // overlap). Only semantic groups need the "all touching ops on same - // runtime else home=host" coordination. + // Mutable buffers: allocated tensors that aren't IO, constants, or + // produced by any op. These are buffer placeholders (e.g., KV cache) + // that persist across execute() calls. { - // Collect all op-output value_ids across all chains. Skip non-Kernel - // instructions (JumpFalseCall, MoveCall, FreeCall, DelegateCall) — - // see CONTROL_FLOW_DESIGN.md §4.3. MoveCall's dst is a special case: - // it appears here under produced_vids only if MoveCall handling is - // wired through the router (currently MoveCall is treated as a - // host->host TransferStep, which does not classify the dst as - // "produced by an op"). Tolerable today; revisit if MoveCall-only - // mutable-buffer patterns appear. std::unordered_set produced_vids; - for (size_t ci = 0; ci < num_chains(); ++ci) { - auto chains = plan_->chains(); - auto instrs = chains->Get(ci)->instructions(); - size_t n_instr = instrs ? instrs->size() : 0; - for (size_t oi = 0; oi < n_instr; ++oi) { - if (instruction_kind(ci, oi) != InstructionKind::Kernel) - continue; - OperatorCall op = get_op(ci, oi); - for (size_t k = 0; k < op.num_outputs(); ++k) { - produced_vids.insert(op.output(k)); - } + auto chains = plan_->chains(); + auto instrs = chains->Get(0)->instructions(); + size_t n_instr = instrs ? instrs->size() : 0; + for (size_t oi = 0; oi < n_instr; ++oi) { + OperatorCall op = get_op(oi); + for (size_t k = 0; k < op.num_outputs(); ++k) { + produced_vids.insert(op.output(k)); } } for (uint32_t i = 0; i < n_vals; ++i) { if (mem_obj_ids_[i] < 0) - continue; // not an allocated tensor + continue; if (input_ids_.count(i) > 0) - continue; // graph input + continue; if (output_ids_.count(i) > 0) - continue; // graph output + continue; if (tensor_constant_data_key(i) != nullptr) - continue; // constant + continue; if (produced_vids.count(i) > 0) - continue; // produced by an op - // Tensor with alloc, not IO, not constant, not produced → it's a - // mutable buffer placeholder. + continue; mutable_buffer_ids_.push_back(i); } } @@ -475,11 +262,6 @@ class Graph { return v ? v->size() : 0; } - // Access serialized value metadata. - // NOTE: returns the raw flatbuffer EValue. This is the construction- - // seam escape hatch — backends and routers should prefer the typed - // accessors below (value_type, int_value, tensor_*, etc) so they - // don't couple to the underlying serialization. const executorch_flatbuffer::EValue* value_meta(uint32_t value_id) const { auto values = plan_->values(); if (!values || value_id >= values->size()) @@ -487,23 +269,19 @@ class Graph { return values->Get(value_id); } - // Value metadata helpers ValueKind value_kind(uint32_t value_id) const; int32_t mem_obj_id(uint32_t value_id) const; //===------------------------------------------------------------------===// - // Typed value accessors (adapter-level — no flatbuffer types leak) + // Typed value accessors //===------------------------------------------------------------------===// - // Returns the kind of the EValue stored at value_id. ValueType value_type(uint32_t value_id) const; - // Scalar accessors — ET_CHECK if the value isn't of the expected kind. int64_t int_value(uint32_t value_id) const; double double_value(uint32_t value_id) const; bool bool_value(uint32_t value_id) const; - // Tensor accessors — ET_CHECK if the value isn't a tensor. ::executorch::aten::ScalarType tensor_dtype(uint32_t value_id) const; ::executorch::runtime::Span tensor_sizes( uint32_t value_id) const; @@ -511,42 +289,25 @@ class Graph { uint32_t value_id) const; ::executorch::aten::TensorShapeDynamism tensor_shape_dynamism( uint32_t value_id) const; - // Returns NDM key (FQN) for an external constant tensor, or nullptr - // if the tensor isn't an NDM-stored constant. + + // NDM key (FQN) for external constants, or nullptr. const char* tensor_constant_data_key(uint32_t value_id) const; - // Returns the raw bytes of an inline constant (stored in the - // program's constant_buffer field), or an empty span if the tensor - // isn't an inline constant. Inline constants are constants the AOT - // didn't promote to NDM (e.g., literals lifted into _lifted_tensor_* - // placeholders). Mutually exclusive with tensor_constant_data_key: - // a constant is either NDM-stored (key != nullptr) or inline - // (this returns non-empty), never both. + // Raw bytes for inline constants, or empty span. ::executorch::runtime::Span tensor_inline_data( uint32_t value_id) const; - // True if the tensor is a constant of either flavor (NDM-stored or - // inline). Use this for "is this an immutable constant?" filtering - // checks in the router; the source matters only at upload time. bool is_constant(uint32_t value_id) const { if (tensor_constant_data_key(value_id) != nullptr) return true; return !tensor_inline_data(value_id).empty(); } - // dtype-size × prod(sizes); 0 if not a tensor or sizes empty. size_t tensor_nbytes_max(uint32_t value_id) const; - // IntList accessors — ET_CHECK if the value isn't an IntList. - // Returns the EValue indices that the list elements reference (stored - // as int64 in the serialization); the caller resolves them through - // the values array. ::executorch::runtime::Span int_list_member_ids( uint32_t value_id) const; - // For a TensorList or OptionalTensorList value, returns the EValue - // indices that the list contains. For OptionalTensorList, indices may - // point at None values (representing nullopt). ::executorch::runtime::Span tensor_list_member_ids( uint32_t value_id) const; @@ -585,7 +346,7 @@ class Graph { } //===------------------------------------------------------------------===// - // Mutable Buffer IDs (values that persist across execute() calls) + // Mutable Buffers //===------------------------------------------------------------------===// size_t num_mutable_buffer_ids() const { @@ -603,7 +364,7 @@ class Graph { } //===------------------------------------------------------------------===// - // Operators (for op name lookup) + // Operators //===------------------------------------------------------------------===// size_t num_operators() const { @@ -628,55 +389,31 @@ class Graph { } //===------------------------------------------------------------------===// - // Chains + // Instructions //===------------------------------------------------------------------===// - size_t num_chains() const { - auto chains = plan_->chains(); - return chains ? chains->size() : 0; - } - - static int32_t main_chain_idx() { - return 0; - } - - // Get number of ops in a chain - size_t num_ops_in_chain(size_t chain_idx) const { + size_t num_instructions() const { auto chains = plan_->chains(); - ET_CHECK_MSG( - chains && chain_idx < chains->size(), - "Graph::num_ops_in_chain(%zu) out of range (have %zu chains)", - chain_idx, - chains ? chains->size() : 0); - auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG(chains && chains->size() > 0, "Graph has no chains"); + auto instrs = chains->Get(0)->instructions(); return instrs ? instrs->size() : 0; } - // Get OperatorCall for op in chain - // PRECONDITION: instruction_kind(chain_idx, op_idx) == - // InstructionKind::Kernel. Use instruction_kind() first to dispatch on kind. - OperatorCall get_op(size_t chain_idx, size_t op_idx) const { + OperatorCall get_op(size_t op_idx) const { auto chains = plan_->chains(); - ET_CHECK_MSG( - chains && chain_idx < chains->size(), - "Graph::get_op: chain_idx=%zu out of range (have %zu chains)", - chain_idx, - chains ? chains->size() : 0); - auto instrs = chains->Get(chain_idx)->instructions(); + ET_CHECK_MSG(chains && chains->size() > 0, "Graph has no chains"); + auto instrs = chains->Get(0)->instructions(); ET_CHECK_MSG( instrs && op_idx < instrs->size(), - "Graph::get_op: op_idx=%zu out of range in chain %zu " - "(have %zu ops)", + "Graph::get_op: op_idx=%zu out of range (have %zu instructions)", op_idx, - chain_idx, instrs ? instrs->size() : 0); auto instr = instrs->Get(op_idx); ET_CHECK_MSG( instr->instr_args_type() == executorch_flatbuffer::InstructionArguments::KernelCall, - "Graph::get_op: instruction at chain=%zu op_idx=%zu is not a KernelCall " - "(type=%u). Use instruction_kind() to dispatch.", - chain_idx, + "Graph::get_op: instruction at op_idx=%zu is not a KernelCall " + "(type=%u)", op_idx, static_cast(instr->instr_args_type())); auto kernel = static_cast( @@ -684,145 +421,20 @@ class Graph { return OperatorCall(kernel, this); } - //===------------------------------------------------------------------===// - // Typed instruction accessors (kind-aware; see CONTROL_FLOW_DESIGN.md §4) - //===------------------------------------------------------------------===// - - // Returns the InstructionKind at (chain_idx, op_idx). - InstructionKind instruction_kind(size_t chain_idx, size_t op_idx) const { - auto chains = plan_->chains(); - ET_CHECK_MSG( - chains && chain_idx < chains->size(), - "Graph::instruction_kind: chain_idx=%zu out of range (have %zu chains)", - chain_idx, - chains ? chains->size() : 0); - auto instrs = chains->Get(chain_idx)->instructions(); - ET_CHECK_MSG( - instrs && op_idx < instrs->size(), - "Graph::instruction_kind: op_idx=%zu out of range in chain %zu", - op_idx, - chain_idx); - auto instr = instrs->Get(op_idx); - using IA = executorch_flatbuffer::InstructionArguments; - switch (instr->instr_args_type()) { - case IA::KernelCall: - return InstructionKind::Kernel; - case IA::JumpFalseCall: - return InstructionKind::JumpFalse; - case IA::MoveCall: - return InstructionKind::Move; - case IA::FreeCall: - return InstructionKind::Free; - case IA::DelegateCall: - return InstructionKind::Delegate; - default: - ET_CHECK_MSG( - false, - "Graph::instruction_kind: unsupported InstructionArguments type %u", - static_cast(instr->instr_args_type())); - } - } - - // Convenience: main-chain shortcut. - InstructionKind instruction_kind(size_t op_idx) const { - return instruction_kind(main_chain_idx(), op_idx); - } - - // Typed accessors per kind. ET_CHECK if the kind doesn't match. - OperatorCall get_kernel_call(size_t chain_idx, size_t op_idx) const { - return get_op(chain_idx, op_idx); - } - - JumpFalseInfo get_jump_false(size_t chain_idx, size_t op_idx) const { - auto chains = plan_->chains(); - ET_CHECK_MSG( - chains && chain_idx < chains->size(), - "Graph::get_jump_false: chain_idx=%zu out of range", - chain_idx); - auto instrs = chains->Get(chain_idx)->instructions(); - ET_CHECK_MSG( - instrs && op_idx < instrs->size(), - "Graph::get_jump_false: op_idx=%zu out of range in chain %zu", - op_idx, - chain_idx); - auto instr = instrs->Get(op_idx); - auto* jf = instr->instr_args_as_JumpFalseCall(); - ET_CHECK_MSG( - jf != nullptr, - "Graph::get_jump_false: instruction at chain=%zu op_idx=%zu is not a " - "JumpFalseCall", - chain_idx, - op_idx); - JumpFalseInfo info; - info.cond_value_id = static_cast(jf->cond_value_index()); - info.destination_pc = static_cast(jf->destination_instruction()); - return info; - } - - MoveInfo get_move(size_t chain_idx, size_t op_idx) const { - auto chains = plan_->chains(); - ET_CHECK_MSG( - chains && chain_idx < chains->size(), - "Graph::get_move: chain_idx=%zu out of range", - chain_idx); - auto instrs = chains->Get(chain_idx)->instructions(); - ET_CHECK_MSG( - instrs && op_idx < instrs->size(), - "Graph::get_move: op_idx=%zu out of range in chain %zu", - op_idx, - chain_idx); - auto instr = instrs->Get(op_idx); - auto* mv = instr->instr_args_as_MoveCall(); - ET_CHECK_MSG( - mv != nullptr, - "Graph::get_move: instruction at chain=%zu op_idx=%zu is not a MoveCall", - chain_idx, - op_idx); - MoveInfo info; - info.src_value_id = static_cast(mv->move_from()); - info.dst_value_id = static_cast(mv->move_to()); - return info; - } - - //===------------------------------------------------------------------===// - // Convenience: main chain accessors - //===------------------------------------------------------------------===// - - size_t num_instructions() const { - return num_ops_in_chain(main_chain_idx()); - } - OperatorCall get_instruction(size_t idx) const { - return get_op(main_chain_idx(), idx); + return get_op(idx); } private: const executorch_flatbuffer::ExecutionPlan* plan_; - // Optional reference to the parent Program, needed only for - // tensor_inline_data() (which dereferences program_->constant_buffer). - // Pre-existing constructions that pass only the plan get nullptr and - // tensor_inline_data() returns empty for them. const executorch_flatbuffer::Program* program_; - // Precomputed at construction for O(1) value_kind lookup. std::unordered_set input_ids_; std::unordered_set output_ids_; - // mem_obj_ids_[value_id] = dense small int identifying the storage slot - // (sort rank of (pool_id, offset) pairs across all aliasable tensor - // values). -1 for non-tensor / non-allocated values. Same id ⇒ same - // storage. Computed once at construction; O(1) lookup at use sites. std::vector mem_obj_ids_; - - // Mutable buffer placeholder value_ids: tensor values with allocation - // info that aren't graph IO, aren't constants, and aren't produced by - // any op. These persist across execute() calls (their storage is - // preserved between invocations). Identified by tag_mutated_buffer at - // AOT time. std::vector mutable_buffer_ids_; }; -// Implement OperatorCall::name() after Graph is defined inline const char* OperatorCall::name() const { - // In ExecuTorch, op names are in the operators table, indexed by op_index return graph_->operator_name(call_->op_index()); } @@ -845,14 +457,12 @@ inline std::string OperatorCall::full_name() const { return s; } -// Implement value metadata accessors inline ValueKind Graph::value_kind(uint32_t value_id) const { if (input_ids_.count(value_id)) return ValueKind::INPUT; if (output_ids_.count(value_id)) return ValueKind::OUTPUT; - // Constant if the tensor has a baked data buffer. auto val = value_meta(value_id); if (val && val->val_type() == executorch_flatbuffer::KernelTypes::Tensor) { auto* tensor = val->val_as_Tensor(); @@ -867,10 +477,6 @@ inline int32_t Graph::mem_obj_id(uint32_t value_id) const { return value_id < mem_obj_ids_.size() ? mem_obj_ids_[value_id] : -1; } -//===----------------------------------------------------------------------===// -// Typed value accessors -//===----------------------------------------------------------------------===// - inline ValueType Graph::value_type(uint32_t value_id) const { auto* val = value_meta(value_id); if (!val) @@ -902,7 +508,7 @@ inline int64_t Graph::int_value(uint32_t value_id) const { auto* val = value_meta(value_id); ET_CHECK_MSG( val && val->val_type() == executorch_flatbuffer::KernelTypes::Int, - "Graph::int_value(%u): value is not an Int", + "Graph::int_value(%u): not an Int", value_id); return static_cast(val->val())->int_val(); } @@ -911,7 +517,7 @@ inline double Graph::double_value(uint32_t value_id) const { auto* val = value_meta(value_id); ET_CHECK_MSG( val && val->val_type() == executorch_flatbuffer::KernelTypes::Double, - "Graph::double_value(%u): value is not a Double", + "Graph::double_value(%u): not a Double", value_id); return static_cast(val->val()) ->double_val(); @@ -921,7 +527,7 @@ inline bool Graph::bool_value(uint32_t value_id) const { auto* val = value_meta(value_id); ET_CHECK_MSG( val && val->val_type() == executorch_flatbuffer::KernelTypes::Bool, - "Graph::bool_value(%u): value is not a Bool", + "Graph::bool_value(%u): not a Bool", value_id); return static_cast(val->val()) ->bool_val(); @@ -940,14 +546,14 @@ inline const executorch_flatbuffer::Tensor* tensor_or_null( inline ::executorch::aten::ScalarType Graph::tensor_dtype( uint32_t value_id) const { auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_dtype(%u): value is not a Tensor", value_id); + ET_CHECK_MSG(t, "Graph::tensor_dtype(%u): not a Tensor", value_id); return static_cast<::executorch::aten::ScalarType>(t->scalar_type()); } inline ::executorch::runtime::Span Graph::tensor_sizes( uint32_t value_id) const { auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_sizes(%u): value is not a Tensor", value_id); + ET_CHECK_MSG(t, "Graph::tensor_sizes(%u): not a Tensor", value_id); auto* s = t->sizes(); return s ? ::executorch::runtime::Span(s->data(), s->size()) : ::executorch::runtime::Span{}; @@ -956,8 +562,7 @@ inline ::executorch::runtime::Span Graph::tensor_sizes( inline ::executorch::runtime::Span Graph::tensor_dim_order( uint32_t value_id) const { auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG( - t, "Graph::tensor_dim_order(%u): value is not a Tensor", value_id); + ET_CHECK_MSG(t, "Graph::tensor_dim_order(%u): not a Tensor", value_id); auto* d = t->dim_order(); return d ? ::executorch::runtime::Span(d->data(), d->size()) : ::executorch::runtime::Span{}; @@ -966,8 +571,7 @@ inline ::executorch::runtime::Span Graph::tensor_dim_order( inline ::executorch::aten::TensorShapeDynamism Graph::tensor_shape_dynamism( uint32_t value_id) const { auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG( - t, "Graph::tensor_shape_dynamism(%u): value is not a Tensor", value_id); + ET_CHECK_MSG(t, "Graph::tensor_shape_dynamism(%u): not a Tensor", value_id); return static_cast<::executorch::aten::TensorShapeDynamism>( t->shape_dynamism()); } @@ -994,9 +598,7 @@ inline ::executorch::runtime::Span Graph::tensor_inline_data( if (!t) return {}; uint32_t idx = static_cast(t->data_buffer_idx()); - // Index 0 is reserved (placeholder for "no inline data"). External - // constants also have idx == 0; they're handled by - // tensor_constant_data_key. + // Index 0 is reserved (no inline data). if (idx == 0) return {}; auto* buffers = program_->constant_buffer(); @@ -1032,7 +634,7 @@ inline ::executorch::runtime::Span Graph::int_list_member_ids( auto* val = value_meta(value_id); ET_CHECK_MSG( val && val->val_type() == executorch_flatbuffer::KernelTypes::IntList, - "Graph::int_list_member_ids(%u): value is not an IntList", + "Graph::int_list_member_ids(%u): not an IntList", value_id); auto* items = static_cast(val->val())->items(); @@ -1049,11 +651,8 @@ inline ::executorch::runtime::Span Graph::tensor_list_member_ids( (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList || val->val_type() == executorch_flatbuffer::KernelTypes::OptionalTensorList), - "Graph::tensor_list_member_ids(%u): value is not a TensorList " - "or OptionalTensorList", + "Graph::tensor_list_member_ids(%u): not a TensorList or OptionalTensorList", value_id); - // Both TensorList and OptionalTensorList have the same shape: - // table { items: [int]; }. Cast to either to access items(). const flatbuffers::Vector* items = nullptr; if (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList) { items = static_cast(val->val()) diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py index d0bddf7f01e..cc5099dfe2d 100644 --- a/backends/native/partitioner.py +++ b/backends/native/partitioner.py @@ -5,29 +5,17 @@ # LICENSE file in the root directory of this source tree. """ -Native Backend Partitioner +Native backend partitioner. -The native partitioner marks ALL nodes as supported. The NativeBackend -runtime has a CPU fallback that can execute any portable op; runtime -partitioning across accelerators (Metal, etc.) happens in C++ based on -has_op() queries against the per-runtime op registries. +Claims core ATen ops (torch.Tag.core) plus an explicit opt-in set. """ -from typing import ( - Any, - Callable, - Dict, - final, - List, - Mapping, - Optional, - Tuple, - TYPE_CHECKING, -) +from typing import Callable, final, List, Mapping, Optional, Tuple + +import executorch.backends.native.specializations # noqa: F401 — register recipes import torch -from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import ( DelegationSpec, Partitioner, @@ -40,160 +28,59 @@ from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase -if TYPE_CHECKING: - from executorch.backends.native.preprocess import Specialization - - -# Canonical list of ops the native backend preserves (does NOT decompose) -# during edge lowering when used with `to_edge_transform_and_lower`. Part -# of the universal-IR specification of NativeBackend; not user-configurable -# per call. -# -# Add an op here when you ship a dedicated C++ handler for it on the Metal -# (or other accelerator) side; otherwise the default decomposition pass at -# edge lowering will break it apart and we'll just dispatch the pieces. -_DEFAULT_PRESERVED_OPS = [ +# Non-core ops the native backend supports. Also preserved (not decomposed). +_SUPPORTED_NON_CORE_OPS = [ torch.ops.aten.matmul.default, torch.ops.aten.linear.default, torch.ops.aten.addmm.default, - torch.ops.aten.baddbmm.default, torch.ops.aten.scaled_dot_product_attention.default, ] class NativeSupportedOperators(OperatorSupportBase): - """ - Operator support checker for the Native backend. - - NativeBackend has a CPU fallback, so it supports ALL operators. The - actual runtime partitioning across accelerators happens in C++. - """ + _NON_CORE = set(_SUPPORTED_NON_CORE_OPS) def is_node_supported( self, submodules: Mapping[str, torch.nn.Module], node: Node ) -> bool: - # Skip placeholder and output nodes - they shouldn't be partitioned if node.op in ("placeholder", "output", "get_attr"): return False if node.op != "call_function": return False - # TODO(native): Claim HOPs (cond, while_loop, scan, ...) so the - # delegate's inner program is byte-equivalent to ET stock emit - # (modulo preserved ops + reinplace), including JumpFalseCall + - # branch instructions. Today we skip them so they stay at the - # outer level and run through ET's standard control-flow - # executor; ops INSIDE the branches still get partitioned via - # to_backend's recursive is_submodule=True path. - # - # To claim HOPs we'd need: - # 1. Pre-populate `val` on the branch-graph `get_attr` nodes - # so the wrap step's get_attr→placeholder conversion in - # fuse_as_graphmodule produces a valid placeholder. ARM - # does this in - # backends/arm/operator_support/control_flow_support.py - # (_submodules_fully_partitioned). - # 2. Make NativeBackend.preprocess HOP-aware: SpecPropPass, - # MemoryPlanningPass, ConstraintBasedSymShapeEvalPass, - # and emit_program all need to recursively walk branch - # graphs. Today they crash with `'ProxyValue' has no - # attribute 'graph'` when they hit a cond. - # 3. Make our runtime Graph adapter resolve the nested branch - # programs that emit_program inlines under JumpFalseCall. - # ARM's stack of control-flow passes (control_flow_const_inline.py - # etc.) is the reference implementation. if isinstance(node.target, torch._ops.HigherOrderOperator): return False - return True - - -@final -class NativePartitioner(Partitioner): - """ - Partitioner for the Native Backend. - - Unlike accelerator-specific partitioners that only claim ops they can - accelerate, the native partitioner claims ALL ops since: - 1. CPU fallback can execute any portable op. - 2. Runtime partitioning in C++ handles dispatch to accelerators. - Two usage paths: + from executorch.exir.dialects.edge._ops import EdgeOpOverload - A) Classic to_edge + to_backend (default decomposition runs first): - edge_program = to_edge(exported_program) - edge_program = edge_program.to_backend(NativePartitioner()) + target = node.target + if isinstance(target, EdgeOpOverload): + target = target._op + if isinstance(target, torch._ops.OpOverload): + if target in self._NON_CORE: + return True + return torch.Tag.core in target.tags or torch.Tag.view_copy in target.tags + return False - B) to_edge_transform_and_lower (preserves listed ops from decomposition): - edge_program = to_edge_transform_and_lower( - exported_program, - partitioner=[NativePartitioner()], - ) - # The ops in _DEFAULT_PRESERVED_OPS are kept intact (no decomposition) - # so accelerator kernels see them whole. The list is maintained - # by this partitioner — extend `_DEFAULT_PRESERVED_OPS` (in code) when - # you ship a new dedicated handler. Per-call overrides are NOT - # exposed: NativeBackend is a universal IR and the preserve list - # is part of its specification. - """ +@final +class NativePartitioner(Partitioner): def __init__( self, - compile_options: Optional[Dict[str, Any]] = None, - specializations: Optional[List["Specialization"]] = None, + specializations: Optional[List[str]] = None, ) -> None: - self.options = compile_options or {} self.specializations = specializations - compile_spec = self._parse_compile_options(self.options) - # Import here to avoid circular dependency from executorch.backends.native.preprocess import NativeBackend - self.delegation_spec = DelegationSpec(NativeBackend.__name__, compile_spec) - - def _parse_compile_options(self, options: Dict[str, Any]) -> List[CompileSpec]: - """Convert compile options dict to CompileSpec list.""" - compile_specs = [] - - for key, value in options.items(): - if isinstance(value, bool): - value_bytes = value.to_bytes(1, byteorder="little") - compile_specs.append(CompileSpec(key, value_bytes)) - elif isinstance(value, int): - value_bytes = value.to_bytes(4, byteorder="little") - compile_specs.append(CompileSpec(key, value_bytes)) - - return compile_specs + self.delegation_spec = DelegationSpec(NativeBackend.__name__, []) def ops_to_not_decompose( self, ep: ExportedProgram ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[Node], bool]]]: - """ - Return ops that should NOT be decomposed during edge lowering. - - Called by `to_edge_transform_and_lower` BEFORE partitioning. The ops - returned here are kept whole (skipped by the default decomposition - pass), so backend kernels see them in their original form. - - The preserve list (`_DEFAULT_PRESERVED_OPS`) is the canonical list - maintained by NativeBackend — part of the universal-IR specification - and not user-configurable. Extend the constant in this file when - you ship a new dedicated handler. - - Behavior: - - If the graph is already partitioned (contains lowered_module - get_attr nodes), return empty — partitioning has run, decomposition - decisions are settled. - - Otherwise return the preserve list, intersected with ops that - actually appear in `ep` (no point listing ops the graph doesn't - have). - - The second tuple element (filter callable) is None: the rule applies - uniformly to every node whose target is in the preserve list. - """ # Already-partitioned graph -> nothing to preserve. for node in ep.graph.nodes: if node.op == "get_attr" and "lowered_module" in node.name: return ([], None) - # Intersect preserve list with ops actually present in the graph. present: List[torch._ops.OpOverload] = [] seen = set() for node in ep.graph.nodes: @@ -201,27 +88,19 @@ def ops_to_not_decompose( continue if not isinstance(node.target, torch._ops.OpOverload): continue - if node.target in _DEFAULT_PRESERVED_OPS and node.target not in seen: + if node.target in _SUPPORTED_NON_CORE_OPS and node.target not in seen: present.append(node.target) seen.add(node.target) return (present, None) def partition(self, exported_program: ExportedProgram) -> PartitionResult: - """ - Partition the exported program for the native backend. - - Since native supports everything, this partitions the entire graph - into a single delegation block. - """ - # Wire specializations to the backend class before partitioning. from executorch.backends.native.preprocess import NativeBackend - NativeBackend._specializations = self.specializations + NativeBackend._specialization_names = self.specializations partition_tags = {} - # Use CapabilityBasedPartitioner with our "support everything" checker capability_partitioner = CapabilityBasedPartitioner( exported_program.graph_module, NativeSupportedOperators(), @@ -236,11 +115,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: node.meta["delegation_tag"] = tag partition_tags[tag] = self.delegation_spec - # Tag constant data for proper handling tag_constant_data(exported_program) - # Tag mutated buffer placeholders so they're owned by our delegate - # (KV-cache style state flows through the delegated subgraph; the - # writeback copy_ collapses out of the top-level chain). tag_mutated_buffer(exported_program) return PartitionResult( diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py index b1834850c0d..f668a406084 100644 --- a/backends/native/preprocess.py +++ b/backends/native/preprocess.py @@ -5,41 +5,19 @@ # LICENSE file in the root directory of this source tree. """ -NativeBackend — AOT BackendDetails for the v2 portable runtime. - -The class name `NativeBackend` is the backend_id used at runtime -to find the matching BackendInterface. C++ side registers it via -`register_backend({"NativeBackend", ...})` in -native/NativeBackend.cpp. - -The preprocess() pipeline: -1. SpecPropPass to populate tensor specs. -2. NEW: reinplace_pass with the backend's BACKEND_INPLACE_OPS registry - (~85 edge-functional ops mapped to aten in-place targets). Rewrites - `op(self, ...)` -> `op_(self, ...)` when ET's safety check passes. - Re-runs SpecPropPass so new in-place nodes get spec metadata. -3. (If buffer mutations exist) insert_write_back_for_buffers_pass to - add an explicit aten::copy_(buf, mut_src) at end of subgraph. When - step 2 already rewrote the buffer mutation in-place, - `_inplace_lineage` detects this and skips inserting the writeback. -4. Spec-sharing fallback for any remaining writebacks: alias the - mutation source's TensorSpec to the buffer's. Used only when step 2 - didn't catch the buffer mutation (e.g., custom op outside the - registry). -5. ExternalConstantsPass to tag constants for NDM storage. -6. Memory planning (greedy, allow_overlapping_allocations). The - upstream `_alias_inplace_result_specs` aliases in-place op result - specs to mutated input specs so emit's `_emit_spec` dedup gives - them one value_id. -7. Emit and serialize. +NativeBackend — AOT preprocess for the native portable runtime. + +Pipeline: SpecProp → reinplace → writeback → external constants → +memory planning → emit → serialize via _program_to_flatbuffer. """ -import copy -import struct from functools import partial -from typing import Any, Dict, final, List, Optional, Sequence, Tuple, Type +from typing import final, List, Optional, Sequence, Tuple import torch + +from executorch.backends.native.fat_pte import build_fat_result +from executorch.backends.native.specializations import _SPECIALIZATION_REGISTRY from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer from executorch.exir.backend.backend_details import ( @@ -59,102 +37,34 @@ from torch._export.verifier import Verifier -# Type alias for specialization config: (BackendDetails subclass, compile_specs) -Specialization = Tuple[Type[BackendDetails], List[CompileSpec]] - - -# --------------------------------------------------------------------------- -# Backend-owned in-place op registry (consumed by ET's reinplace_pass). -# -# Maps edge-dialect functional ops -> aten in-place targets. The pass -# runs on the edge-dialect graph; targets use the aten in-place form -# because the edge dialect doesn't register in-place variants for most -# ops. The lowered IR ends up carrying `aten::_` instructions -# which the executor accepts. -# -# Built once at module load by intersecting `ops.edge.aten.` with -# `torch.ops.aten._`. This table lives in the backend (not -# upstream) because it's a per-runtime kernel-availability statement. -# --------------------------------------------------------------------------- +# A specialization recipe: callable that takes an ExportedProgram and returns +# serialized bytes. Each recipe owns the full +# to_edge_transform_and_lower → to_executorch → serialize flow. def _build_backend_inplace_ops() -> frozenset: - """Construct the v2 backend's in-place op set by listing the - edge-functional ops we support. ET's reinplace_pass auto-derives - the in-place counterpart for each via schema matching. + """Edge-dialect ops the native runtime supports in-place. + + reinplace_pass auto-derives the in-place counterpart for each. """ from executorch.exir.dialects._ops import ops as _edge_ops edge = _edge_ops.edge.aten - - pairs = [ - ("relu", ["default"]), - ("gelu", ["default"]), - ("sigmoid", ["default"]), - ("index_put", ["default"]), - ("index_copy", ["default"]), - ] - - functional_ops = set() - for name, overloads in pairs: - edge_pkg = getattr(edge, name, None) - if edge_pkg is None: - continue - for ovld in overloads: - op = getattr(edge_pkg, ovld, None) - if op is not None: - functional_ops.add(op) - - return frozenset(functional_ops) + return frozenset( + [ + edge.relu.default, + edge.gelu.default, + edge.sigmoid.default, + edge.index_put.default, + edge.index_copy.default, + ] + ) BACKEND_INPLACE_OPS: frozenset = _build_backend_inplace_ops() -# --------------------------------------------------------------------------- -# Fat PTE container format -# -# A fat blob packs multiple backend specializations into one delegate payload. -# The runtime parses the header and picks the best available specialization. -# -# Layout: -# [4 bytes] magic "NFAT" -# [4 bytes] version (1) -# [4 bytes] num_specializations -# For each specialization: -# [32 bytes] backend_id (utf-8, null-padded) -# [8 bytes] offset into data section -# [8 bytes] length -# [payload bytes for specialization 0] -# [payload bytes for specialization 1] -# ... -# --------------------------------------------------------------------------- - -_FAT_MAGIC = b"NFAT" -_FAT_VERSION = 1 -_FAT_ENTRY_FMT = "32sQQ" # backend_id(32) + offset(u64) + size(u64) -_FAT_ENTRY_SIZE = struct.calcsize(_FAT_ENTRY_FMT) - - -def _pack_fat_blob( - specializations: List[Tuple[str, bytes]], -) -> bytes: - """Pack multiple (backend_id, payload) pairs into a fat blob.""" - header = struct.pack("<4sII", _FAT_MAGIC, _FAT_VERSION, len(specializations)) - - entries = [] - offset = 0 - for backend_id, payload in specializations: - bid = backend_id.encode("utf-8")[:32].ljust(32, b"\x00") - entries.append(struct.pack("<" + _FAT_ENTRY_FMT, bid, offset, len(payload))) - offset += len(payload) - - return header + b"".join(entries) + b"".join(p for _, p in specializations) - - class _AnyOp(Verifier): - """Permissive verifier that allows any op (skip functional check).""" - dialect = "TRAINING" def allowed_op_types(self): @@ -164,7 +74,6 @@ def allowed_op_types(self): def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram: - """Apply a sequence of passes to an ExportedProgram.""" from executorch.exir.pass_base import ExportPass, PassBase for p in passes: @@ -182,39 +91,10 @@ def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram: return program -def _parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: - """Parse compile specs into options dict.""" - options: Dict[str, Any] = { - "skip_memory_planning": False, - # Default ON: the v2 portable runtime's CPU provider has - # dispatches for in-place ops (`aten::add_` etc.) that route to - # the existing `.out` variant kernels. The router falls back to - # CPU when no other provider claims the op. - "enable_reinplace": True, - } - for spec in compile_specs: - if spec.key in options and isinstance(options[spec.key], bool): - options[spec.key] = bool.from_bytes(spec.value, byteorder="little") - return options - - @final class NativeBackend(BackendDetails): - """ - BackendDetails for the v2 portable backend. - - Class name `NativeBackend` matches the runtime backend_id - registered in native/NativeBackend.cpp. - - Supports "fat PTE" mode: when specializations are configured, - preprocess runs each backend's pipeline on the same subgraph and - packs all results into a single fat blob. The runtime picks the - best available specialization at load time. All specializations - share the same PTD (named data store). - """ - # Set by NativePartitioner before preprocess is called. - _specializations: Optional[Sequence[Specialization]] = None + _specialization_names: Optional[Sequence[str]] = None @classmethod def preprocess( @@ -222,84 +102,63 @@ def preprocess( program: ExportedProgram, module_compile_spec: List[CompileSpec], ) -> PreprocessResult: - native_result = cls._preprocess_native(program, module_compile_spec) - - specializations = cls._specializations - if not specializations: - return native_result - - # Fat PTE: run each specialization backend and pack results. - from executorch.exir._serialize._named_data_store import NamedDataStore - - fat_entries: List[Tuple[str, bytes]] = [ - ("NativeBackend", native_result.processed_bytes), - ] + names = cls._specialization_names + if not names: + return cls._preprocess_native(program) - merged_data_store = NamedDataStore() - # Seed with native's named data. - if native_result.data_store_output: - for key, entry in native_result.data_store_output.pte_data.items(): - buf = native_result.data_store_output.buffers[entry.buffer_index] - merged_data_store.add_named_data(key, buf, alignment=entry.alignment) + for name in names: + if name not in _SPECIALIZATION_REGISTRY: + raise ValueError( + f"Specialization '{name}' is not registered. " + f"Registered: {sorted(_SPECIALIZATION_REGISTRY)}" + ) - for backend_cls, spec_compile_specs in specializations: - spec_program = copy.deepcopy(program) - spec_result = backend_cls.preprocess(spec_program, spec_compile_specs) + import copy - fat_entries.append((backend_cls.__name__, spec_result.processed_bytes)) + spec_programs = [copy.deepcopy(program) for _ in names] + native_result = cls._preprocess_native(program) - if spec_result.data_store_output: - for key, entry in spec_result.data_store_output.pte_data.items(): - buf = spec_result.data_store_output.buffers[entry.buffer_index] - merged_data_store.add_named_data( - key, buf, alignment=entry.alignment - ) + results: List[Tuple[str, PreprocessResult]] = [ + ("NativeBackend", native_result), + ] + for name, spec_program in zip(names, spec_programs): + payload = _SPECIALIZATION_REGISTRY[name](spec_program) + results.append((name, PreprocessResult(processed_bytes=payload))) - return PreprocessResult( - processed_bytes=_pack_fat_blob(fat_entries), - debug_handle_map=native_result.debug_handle_map, - data_store_output=merged_data_store.get_named_data_store_output(), - ) + return build_fat_result(results) @classmethod def _preprocess_native( cls, program: ExportedProgram, - module_compile_spec: List[CompileSpec], ) -> PreprocessResult: - """ - Preprocess the partitioned subgraph for v2 portable backend execution. - """ - compile_options = _parse_compile_spec(module_compile_spec) - skip_memory_planning = compile_options["skip_memory_planning"] - enable_reinplace = compile_options["enable_reinplace"] - - # Step 1: SpecPropPass to propagate tensor specs. program = _apply_passes(program, [SpecPropPass()]) - # Step 2: NEW — reinplace_pass with the backend's in-place op - # registry. Rewrites functional ops to their `*_` form when ET's - # safety check passes (sole consumer; mutable inputs OK; no - # later use). For buffer mutations (KV-cache pattern), this - # turns `index_put(buf, ...)` into `index_put_(buf, ...)`, - # which then makes `insert_write_back_for_buffers_pass` (step 3) - # skip inserting a redundant `copy_` writeback via its - # `_inplace_lineage` check. - if enable_reinplace: - from executorch.exir.passes.reinplace import ( - reinplace_pass as _et_reinplace_pass, + from executorch.exir.passes.reinplace import ( + reinplace_pass as _et_reinplace_pass, + ) + + program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS) + + # reinplace may rename output nodes; fix the graph signature to match. + output_node = next( + n for n in program.graph_module.graph.nodes if n.op == "output" + ) + actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")] + new_output_specs = [] + for spec, name in zip(program.graph_signature.output_specs, actual_names): + new_output_specs.append( + type(spec)( + kind=spec.kind, + arg=type(spec.arg)(name=name), + target=spec.target, + ) ) + program._graph_signature.output_specs = new_output_specs - program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS) - # ET's reinplace_pass copies meta["val"] to the new in-place - # node but not meta["spec"]. Re-run SpecPropPass so the new - # nodes have specs that downstream passes can read. - program = _apply_passes(program, [SpecPropPass()]) + # Re-run SpecPropPass: reinplace copies meta["val"] but not meta["spec"]. + program = _apply_passes(program, [SpecPropPass()]) - # Step 3: Insert writeback copy_ ops for any mutable buffers - # that weren't already mutated in place by step 2. The pass - # uses `_inplace_lineage` to detect in-place chains and skip - # writeback insertion for them. from torch.export.graph_signature import InputKind, OutputKind has_buffer_mutation = any( @@ -310,16 +169,10 @@ def _preprocess_native( gm, new_sig = insert_write_back_for_buffers_pass(program) program._graph_module = gm program._graph_signature = new_sig - # Re-propagate specs onto the newly inserted copy_ nodes. program = _apply_passes(program, [SpecPropPass()]) - # Spec-sharing trick: for each (buffer_placeholder, mutation - # source) pair, make the mutation source's TensorSpec be the - # SAME object as the buffer's TensorSpec. The greedy memory - # planner walks specs by identity, so a shared spec yields a - # single allocation. No wasted slot, no runtime override - # needed (the .pte naturally reports both at the same offset). - + # Spec-share buffer placeholders with their mutation sources so the + # memory planner assigns them the same storage slot. sig = program.graph_signature nodes_by_name = {n.name: n for n in program.graph_module.graph.nodes} buf_target_to_node = { @@ -346,44 +199,31 @@ def _preprocess_native( buf_spec = buf_node.meta.get("spec") if buf_spec is None or "spec" not in src_node.meta: continue - # Alias: src now shares buf's spec object. src_node.meta["spec"] = buf_spec - # Step 4: External constants pass — tag named parameters/buffers - # for NDM storage. The stock pass deliberately skips lifted - # tensor constants ("they are closer to code than data"). We - # used to override that to force lifted constants to NDM, but - # the runtime now supports inline constants directly via - # Engine::ConstRequest.inline_data + Graph::tensor_inline_data, - # so we keep the stock semantics. from executorch.exir.passes.external_constants_pass import ( external_constants_pass, ) external_constants_pass(program.graph_module) - # Step 5: Memory planning (greedy, allows overlapping allocs). - if not skip_memory_planning: - greedy_memory_planning = partial(greedy, allow_overlapping_allocations=True) - mem_planning_suite = MemoryPlanningAlgorithmSuite( - algo_list=[greedy_memory_planning] - ) + greedy_memory_planning = partial(greedy, allow_overlapping_allocations=True) + mem_planning_suite = MemoryPlanningAlgorithmSuite( + algo_list=[greedy_memory_planning] + ) - # Workaround for memory planning without ToOutVarPass - program.graph_module.encounter_to_out_var_failure = True + program.graph_module.encounter_to_out_var_failure = True - program = _apply_passes( - program, - [ - ConstraintBasedSymShapeEvalPass(), - MemoryPlanningPass(memory_planning_algo=mem_planning_suite), - ], - ) + program = _apply_passes( + program, + [ + ConstraintBasedSymShapeEvalPass(), + MemoryPlanningPass(memory_planning_algo=mem_planning_suite), + ], + ) - # Step 6: Emit the program. emitter_output = emit_program(program) - # Step 7: Build named data store from external constants. from executorch.exir._serialize._named_data_store import NamedDataStore named_data_store = NamedDataStore() @@ -393,21 +233,8 @@ def _preprocess_native( data = emitter_output.external_constant_buffer[idx] named_data_store.add_named_data(fqn, data) - # Step 8: Serialize the inner program to bytes. - # - # We deliberately do NOT use serialize_pte_binary here, because - # it extracts constants from program.constant_buffer into a - # separate constant_segment of the OUTER PTE file — which makes - # the inner program's constant_buffer empty when the runtime - # parses it. The constant data would end up in segments that - # the delegate has no view into. - # - # _program_to_flatbuffer is the lower-level serializer that - # keeps program.constant_buffer in place. The runtime's - # Graph::tensor_inline_data() reads directly from - # program->constant_buffer()->Get(idx), so inline lifted - # constants reach Engine::upload_constants via the inline_data - # path on ConstRequest. + # Use _program_to_flatbuffer (not serialize_pte_binary) to keep + # constants inline — the delegate can't see outer PTE segments. fb_result = _program_to_flatbuffer(emitter_output.program) serialized_bytes = bytes(fb_result.data) diff --git a/backends/native/specializations.py b/backends/native/specializations.py new file mode 100644 index 00000000000..232346eda75 --- /dev/null +++ b/backends/native/specializations.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Specialization registry and recipes for the native backend fat PTE. + +Each recipe is a callable that takes an ExportedProgram and returns +serialized bytes. Automatically imported by NativePartitioner to +populate the registry. +""" + +from typing import Callable, Dict + +from torch.export import ExportedProgram + +SpecializationRecipe = Callable[[ExportedProgram], bytes] + +_SPECIALIZATION_REGISTRY: Dict[str, SpecializationRecipe] = {} + + +def register_specialization(name: str, recipe: SpecializationRecipe) -> None: + _SPECIALIZATION_REGISTRY[name] = recipe + + +# --------------------------------------------------------------------------- +# Built-in recipes +# --------------------------------------------------------------------------- + + +def _mlx_recipe(ep: ExportedProgram) -> bytes: + from executorch.backends.mlx.partitioner import MLXPartitioner + from executorch.backends.mlx.passes import get_default_passes + from executorch.exir import to_edge_transform_and_lower + + lowered = to_edge_transform_and_lower( + ep, + transform_passes=get_default_passes(), + partitioner=[MLXPartitioner()], + ) + et_program = lowered.to_executorch() + return et_program.buffer + + +register_specialization("MLXBackend", _mlx_recipe) diff --git a/backends/native/test/__init__.py b/backends/native/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/native/test/test_fat_pte.py b/backends/native/test/test_fat_pte.py new file mode 100644 index 00000000000..2509ea36107 --- /dev/null +++ b/backends/native/test/test_fat_pte.py @@ -0,0 +1,105 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import struct +import unittest + +from executorch.backends.native.fat_pte import ( + _ENTRY_FMT, + _ENTRY_SIZE, + build_fat_result, + FAT_MAGIC, + FAT_VERSION, + pack_fat_blob, +) +from executorch.exir.backend.backend_details import PreprocessResult + + +class TestPackFatBlob(unittest.TestCase): + def test_single_specialization(self): + payload = b"hello" + blob = pack_fat_blob([("TestBackend", payload)]) + + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(version, FAT_VERSION) + self.assertEqual(count, 1) + + header_size = 12 + bid, offset, size = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) + self.assertEqual(bid.rstrip(b"\x00"), b"TestBackend") + self.assertEqual(offset, 0) + self.assertEqual(size, len(payload)) + + data_start = header_size + _ENTRY_SIZE + self.assertEqual(blob[data_start:], payload) + + def test_multiple_specializations(self): + entries = [("A", b"aaa"), ("B", b"bbbb"), ("C", b"cc")] + blob = pack_fat_blob(entries) + + _, _, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 3) + + header_size = 12 + offsets_sizes = [] + for i in range(3): + bid, off, sz = struct.unpack_from( + "<" + _ENTRY_FMT, blob, header_size + i * _ENTRY_SIZE + ) + offsets_sizes.append((bid.rstrip(b"\x00"), off, sz)) + + self.assertEqual(offsets_sizes[0], (b"A", 0, 3)) + self.assertEqual(offsets_sizes[1], (b"B", 3, 4)) + self.assertEqual(offsets_sizes[2], (b"C", 7, 2)) + + data_start = header_size + 3 * _ENTRY_SIZE + self.assertEqual(blob[data_start:], b"aaabbbbcc") + + def test_empty_specializations(self): + blob = pack_fat_blob([]) + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 0) + self.assertEqual(len(blob), 12) + + def test_backend_id_truncated_to_32_bytes(self): + long_name = "A" * 100 + blob = pack_fat_blob([(long_name, b"x")]) + header_size = 12 + bid, _, _ = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) + self.assertEqual(len(bid), 32) + self.assertEqual(bid, b"A" * 32) + + +class TestBuildFatResult(unittest.TestCase): + def test_merges_results(self): + r1 = PreprocessResult(processed_bytes=b"native", debug_handle_map={1: [1]}) + r2 = PreprocessResult(processed_bytes=b"accel", debug_handle_map={2: [2]}) + + result = build_fat_result([("NativeBackend", r1), ("AccelBackend", r2)]) + + magic, version, count = struct.unpack_from("<4sII", result.processed_bytes, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(count, 2) + + # debug_handle_map defaults to first result's + self.assertEqual(result.debug_handle_map, {1: [1]}) + + def test_explicit_debug_handle_map(self): + r1 = PreprocessResult(processed_bytes=b"a", debug_handle_map={1: [1]}) + custom_map = {99: [99]} + result = build_fat_result([("X", r1)], debug_handle_map=custom_map) + self.assertEqual(result.debug_handle_map, custom_map) + + def test_single_result(self): + r = PreprocessResult(processed_bytes=b"only") + result = build_fat_result([("Solo", r)]) + _, _, count = struct.unpack_from("<4sII", result.processed_bytes, 0) + self.assertEqual(count, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/native/test/test_partitioner.py b/backends/native/test/test_partitioner.py new file mode 100644 index 00000000000..b045455d09f --- /dev/null +++ b/backends/native/test/test_partitioner.py @@ -0,0 +1,125 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from unittest.mock import MagicMock + +import torch +import torch.nn as nn + +from executorch.backends.native.partitioner import ( + _SUPPORTED_NON_CORE_OPS, + NativePartitioner, + NativeSupportedOperators, +) + + +class TestNativeSupportedOperators(unittest.TestCase): + def _make_node(self, op, target): + node = MagicMock() + node.op = op + node.target = target + return node + + def test_rejects_placeholder(self): + sup = NativeSupportedOperators() + node = self._make_node("placeholder", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_output(self): + sup = NativeSupportedOperators() + node = self._make_node("output", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_get_attr(self): + sup = NativeSupportedOperators() + node = self._make_node("get_attr", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_accepts_core_aten_op(self): + sup = NativeSupportedOperators() + node = self._make_node("call_function", torch.ops.aten.add.Tensor) + self.assertTrue(sup.is_node_supported({}, node)) + + def test_accepts_non_core_supported_op(self): + sup = NativeSupportedOperators() + for op in _SUPPORTED_NON_CORE_OPS: + node = self._make_node("call_function", op) + self.assertTrue( + sup.is_node_supported({}, node), + f"{op} should be supported as a non-core op", + ) + + def test_rejects_non_core_unsupported_op(self): + sup = NativeSupportedOperators() + op = torch.ops.aten.linalg_solve_triangular.default + if torch.Tag.core not in op.tags: + node = self._make_node("call_function", op) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_higher_order_operator(self): + sup = NativeSupportedOperators() + hop = torch.ops.higher_order.cond + node = self._make_node("call_function", hop) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_non_opoverload_callable(self): + sup = NativeSupportedOperators() + node = self._make_node("call_function", lambda x: x) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_accepts_edge_op_overlay(self): + """EdgeOpOverload wraps OpOverload; partitioner should unwrap and accept.""" + from executorch.exir.dialects._ops import ops as _edge_ops + + sup = NativeSupportedOperators() + edge_add = _edge_ops.edge.aten.add.Tensor + node = self._make_node("call_function", edge_add) + self.assertTrue(sup.is_node_supported({}, node)) + + def test_rejects_non_core_edge_op(self): + from executorch.exir.dialects._ops import ops as _edge_ops + + sup = NativeSupportedOperators() + edge_op = _edge_ops.edge.aten.linalg_solve_triangular.default + if torch.Tag.core not in edge_op._op.tags: + node = self._make_node("call_function", edge_op) + self.assertFalse(sup.is_node_supported({}, node)) + + +class TestNativePartitionerE2E(unittest.TestCase): + def test_linear_delegates_all_ops(self): + from executorch.exir import to_edge_transform_and_lower + + model = nn.Linear(4, 4) + ep = torch.export.export(model, (torch.randn(1, 4),)) + lowered = to_edge_transform_and_lower(ep, partitioner=[NativePartitioner()]) + graph = lowered._edge_programs["forward"].graph + delegate_calls = [ + n + for n in graph.nodes + if n.op == "call_function" and "executorch_call_delegate" in str(n.target) + ] + non_delegate_ops = [ + n + for n in graph.nodes + if n.op == "call_function" + and "executorch_call_delegate" not in str(n.target) + and "getitem" not in str(n.target) + ] + self.assertGreater( + len(delegate_calls), 0, "Expected at least one delegate call" + ) + self.assertEqual( + len(non_delegate_ops), + 0, + f"All ops should be delegated, but found: " + f"{[str(n.target) for n in non_delegate_ops]}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/native/test/test_preprocess.py b/backends/native/test/test_preprocess.py new file mode 100644 index 00000000000..11463ff7365 --- /dev/null +++ b/backends/native/test/test_preprocess.py @@ -0,0 +1,163 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import struct +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.native.fat_pte import FAT_MAGIC +from executorch.backends.native.partitioner import NativePartitioner +from executorch.backends.native.preprocess import NativeBackend +from executorch.backends.native.specializations import ( + _SPECIALIZATION_REGISTRY, + register_specialization, +) +from executorch.exir import to_edge_transform_and_lower +from executorch.exir._serialize._flatbuffer_program import _flatbuffer_to_program + + +def _lower(model, example_inputs, specializations=None): + ep = torch.export.export(model, example_inputs) + edge = to_edge_transform_and_lower( + ep, + partitioner=[NativePartitioner(specializations=specializations)], + ) + return edge + + +def _get_delegate_blob(edge): + """Extract the single delegate's processed bytes from the lowered program.""" + et = edge.to_executorch() + delegates = et.executorch_program.backend_delegate_data + assert len(delegates) == 1, f"Expected 1 delegate blob, got {len(delegates)}" + return bytes(delegates[0].data) + + +class TestSpecializationRegistry(unittest.TestCase): + def setUp(self): + self._saved = dict(_SPECIALIZATION_REGISTRY) + + def tearDown(self): + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(self._saved) + NativeBackend._specialization_names = None + + def test_register_and_lookup(self): + register_specialization("TestBackend", lambda ep: b"test") + self.assertIn("TestBackend", _SPECIALIZATION_REGISTRY) + + def test_unregistered_backend_rejected(self): + with self.assertRaises(ValueError) as ctx: + _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["UnregisteredBackend"], + ) + self.assertIn("UnregisteredBackend", str(ctx.exception)) + self.assertIn("not registered", str(ctx.exception)) + + def test_no_specializations_produces_plain_flatbuffer(self): + edge = _lower(nn.Linear(2, 2), (torch.randn(1, 2),)) + blob = _get_delegate_blob(edge) + self.assertNotEqual(blob[:4], FAT_MAGIC) + program = _flatbuffer_to_program(blob) + self.assertGreater(len(program.execution_plan[0].operators), 0) + + def test_registered_backend_produces_fat_pte(self): + register_specialization("TestAccel", lambda ep: b"accel_payload") + edge = _lower( + nn.Linear(2, 2), (torch.randn(1, 2),), specializations=["TestAccel"] + ) + blob = _get_delegate_blob(edge) + self.assertEqual(blob[:4], FAT_MAGIC) + + +class TestPreprocessSerialization(unittest.TestCase): + def _lower_and_deserialize(self, model, example_inputs): + edge = _lower(model, example_inputs) + blob = _get_delegate_blob(edge) + return _flatbuffer_to_program(blob) + + def test_linear_op_names(self): + program = self._lower_and_deserialize(nn.Linear(4, 4), (torch.randn(1, 4),)) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + + def test_add_op_name(self): + class AddModel(nn.Module): + def forward(self, x, y): + return x + y + + program = self._lower_and_deserialize( + AddModel(), (torch.randn(2, 3), torch.randn(2, 3)) + ) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::add", op_names) + + def test_reinplace_converts_relu(self): + class ReluModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + + def forward(self, x): + return torch.relu(self.linear(x)) + + program = self._lower_and_deserialize(ReluModel(), (torch.randn(1, 8),)) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + self.assertIn("aten::relu_", op_names) + + def test_fat_pte_contains_both_payloads(self): + saved = dict(_SPECIALIZATION_REGISTRY) + register_specialization("FakeAccel", lambda ep: b"FAKE_ACCEL_DATA") + try: + edge = _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["FakeAccel"], + ) + blob = _get_delegate_blob(edge) + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(count, 2) + self.assertIn(b"FAKE_ACCEL_DATA", blob) + finally: + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(saved) + + def test_fat_pte_native_payload_is_valid(self): + """The native slice inside a fat PTE is a valid flatbuffer program.""" + saved = dict(_SPECIALIZATION_REGISTRY) + register_specialization("FakeAccel", lambda ep: b"FAKE") + try: + edge = _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["FakeAccel"], + ) + blob = _get_delegate_blob(edge) + _, _, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 2) + + from executorch.backends.native.fat_pte import _ENTRY_FMT, _ENTRY_SIZE + + header_size = 12 + bid, offset, size = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) + data_start = header_size + count * _ENTRY_SIZE + native_blob = blob[data_start + offset : data_start + offset + size] + program = _flatbuffer_to_program(native_blob) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + finally: + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(saved) + + +if __name__ == "__main__": + unittest.main() diff --git a/pytest.ini b/pytest.ini index 05aea9d4da6..949d918a963 100644 --- a/pytest.ini +++ b/pytest.ini @@ -76,6 +76,7 @@ testpaths = # backends backends/apple/coreml/test + backends/native/test backends/test/harness/tests backends/test/suite/tests backends/transforms From 1754efdca3d507c3eaf5fb73b759a6fbbb598115 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 2 Jul 2026 14:58:01 -0700 Subject: [PATCH 4/6] up --- backends/mlx/passes.py | 41 ------------------------- backends/native/fat_pte.py | 7 ++++- backends/native/partitioner.py | 23 ++++++++++---- backends/native/preprocess.py | 31 ++++++------------- backends/native/test/test_fat_pte.py | 11 +++---- backends/native/test/test_preprocess.py | 2 -- exir/passes/reinplace.py | 22 +++++++++++++ exir/tests/test_reinplace_pass.py | 19 ++++++++++++ 8 files changed, 78 insertions(+), 78 deletions(-) diff --git a/backends/mlx/passes.py b/backends/mlx/passes.py index e88cd83f0bc..c22b8c0b15e 100644 --- a/backends/mlx/passes.py +++ b/backends/mlx/passes.py @@ -99,49 +99,8 @@ def call(self, exported_program) -> ExportedProgramPassResult: } if ops_to_inplace: reinplace_pass(exported_program, ops_to_inplace=ops_to_inplace) - self._resync_output_specs(exported_program) return ExportedProgramPassResult(exported_program, True) - @staticmethod - def _resync_output_specs(exported_program) -> None: - """Re-sync graph-signature output names after reinplace. - - ``reinplace_pass`` rewrites an output-producing node (e.g. the final - ``exp`` -> ``exp_``) via ``replace_all_uses_with`` + erase, but does not - update ``graph_signature.output_specs``. Output order is preserved, so we - positionally re-sync each spec's argument name to the current output node - arg; otherwise ``ExportedProgram.validate()`` (run by the pass manager) - raises a SpecViolationError. - - This positional pairing is only valid because reinplace does a 1:1 - ``replace_all_uses_with`` + erase and never drops, adds, or reorders - outputs. We assert ``len(output_specs) == len(out_args)`` so that a - future change violating that invariant fails loudly here instead of - silently mis-pairing names (``zip`` would otherwise truncate). Specs - whose ``arg`` is not a named tensor (e.g. ``ConstantArgument``) carry no - ``name`` and are skipped by the ``getattr`` guard below. - """ - out_node = next( - n for n in reversed(exported_program.graph.nodes) if n.op == "output" - ) - out_args = out_node.args[0] - if not isinstance(out_args, (tuple, list)): - out_args = (out_args,) - output_specs = exported_program.graph_signature.output_specs - assert len(output_specs) == len(out_args), ( - "reinplace changed graph output count: " - f"{len(output_specs)} output_specs vs {len(out_args)} output args. " - "Positional output-spec re-sync assumes a 1:1, order-preserving " - "rewrite." - ) - for spec, arg in zip(output_specs, out_args): - if ( - isinstance(arg, torch.fx.Node) - and getattr(spec.arg, "name", None) is not None - and spec.arg.name != arg.name - ): - spec.arg.name = arg.name - @dataclass class RMSNormMatch(PatternMatch): diff --git a/backends/native/fat_pte.py b/backends/native/fat_pte.py index 93f240a3551..d2700fcd019 100644 --- a/backends/native/fat_pte.py +++ b/backends/native/fat_pte.py @@ -39,7 +39,12 @@ def pack_fat_blob(specializations: List[Tuple[str, bytes]]) -> bytes: entries = [] offset = 0 for backend_id, payload in specializations: - bid = backend_id.encode("utf-8")[:32].ljust(32, b"\x00") + encoded = backend_id.encode("ascii") + if len(encoded) > 32: + raise ValueError( + f"Backend ID '{backend_id}' is {len(encoded)} bytes; max is 32" + ) + bid = encoded.ljust(32, b"\x00") entries.append(struct.pack("<" + _ENTRY_FMT, bid, offset, len(payload))) offset += len(payload) diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py index cc5099dfe2d..65577dee3b1 100644 --- a/backends/native/partitioner.py +++ b/backends/native/partitioner.py @@ -10,12 +10,15 @@ Claims core ATen ops (torch.Tag.core) plus an explicit opt-in set. """ +import json + from typing import Callable, final, List, Mapping, Optional, Tuple import executorch.backends.native.specializations # noqa: F401 — register recipes import torch +from executorch.exir.backend.backend_details import CompileSpec from executorch.exir.backend.partitioner import ( DelegationSpec, Partitioner, @@ -64,6 +67,8 @@ def is_node_supported( @final class NativePartitioner(Partitioner): + _SPECIALIZATIONS_KEY = "native_specializations" + def __init__( self, specializations: Optional[List[str]] = None, @@ -71,14 +76,24 @@ def __init__( self.specializations = specializations from executorch.backends.native.preprocess import NativeBackend - self.delegation_spec = DelegationSpec(NativeBackend.__name__, []) + compile_specs = [] + if specializations: + compile_specs.append( + CompileSpec( + self._SPECIALIZATIONS_KEY, + json.dumps(specializations).encode("utf-8"), + ) + ) + self.delegation_spec = DelegationSpec(NativeBackend.__name__, compile_specs) def ops_to_not_decompose( self, ep: ExportedProgram ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[Node], bool]]]: # Already-partitioned graph -> nothing to preserve. + from executorch.exir.lowered_backend_module import executorch_call_delegate + for node in ep.graph.nodes: - if node.op == "get_attr" and "lowered_module" in node.name: + if node.op == "call_function" and node.target is executorch_call_delegate: return ([], None) present: List[torch._ops.OpOverload] = [] @@ -95,10 +110,6 @@ def ops_to_not_decompose( return (present, None) def partition(self, exported_program: ExportedProgram) -> PartitionResult: - from executorch.backends.native.preprocess import NativeBackend - - NativeBackend._specialization_names = self.specializations - partition_tags = {} capability_partitioner = CapabilityBasedPartitioner( diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py index f668a406084..d886fd9da15 100644 --- a/backends/native/preprocess.py +++ b/backends/native/preprocess.py @@ -11,8 +11,9 @@ memory planning → emit → serialize via _program_to_flatbuffer. """ +import json from functools import partial -from typing import final, List, Optional, Sequence, Tuple +from typing import final, List, Tuple import torch @@ -65,6 +66,8 @@ def _build_backend_inplace_ops() -> frozenset: class _AnyOp(Verifier): + # "TRAINING" is the only dialect that allows non-functional (mutating) ops. + # After reinplace converts e.g. relu → relu_, the verifier must accept them. dialect = "TRAINING" def allowed_op_types(self): @@ -93,16 +96,17 @@ def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram: @final class NativeBackend(BackendDetails): - # Set by NativePartitioner before preprocess is called. - _specialization_names: Optional[Sequence[str]] = None - @classmethod def preprocess( cls, program: ExportedProgram, module_compile_spec: List[CompileSpec], ) -> PreprocessResult: - names = cls._specialization_names + names = None + for spec in module_compile_spec: + if spec.key == "native_specializations": + names = json.loads(spec.value.decode("utf-8")) + if not names: return cls._preprocess_native(program) @@ -140,22 +144,6 @@ def _preprocess_native( program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS) - # reinplace may rename output nodes; fix the graph signature to match. - output_node = next( - n for n in program.graph_module.graph.nodes if n.op == "output" - ) - actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")] - new_output_specs = [] - for spec, name in zip(program.graph_signature.output_specs, actual_names): - new_output_specs.append( - type(spec)( - kind=spec.kind, - arg=type(spec.arg)(name=name), - target=spec.target, - ) - ) - program._graph_signature.output_specs = new_output_specs - # Re-run SpecPropPass: reinplace copies meta["val"] but not meta["spec"]. program = _apply_passes(program, [SpecPropPass()]) @@ -212,6 +200,7 @@ def _preprocess_native( algo_list=[greedy_memory_planning] ) + # Tells the emitter to use out-variant kernels for ops that support them. program.graph_module.encounter_to_out_var_failure = True program = _apply_passes( diff --git a/backends/native/test/test_fat_pte.py b/backends/native/test/test_fat_pte.py index 2509ea36107..9ecec149ee2 100644 --- a/backends/native/test/test_fat_pte.py +++ b/backends/native/test/test_fat_pte.py @@ -65,13 +65,10 @@ def test_empty_specializations(self): self.assertEqual(count, 0) self.assertEqual(len(blob), 12) - def test_backend_id_truncated_to_32_bytes(self): - long_name = "A" * 100 - blob = pack_fat_blob([(long_name, b"x")]) - header_size = 12 - bid, _, _ = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) - self.assertEqual(len(bid), 32) - self.assertEqual(bid, b"A" * 32) + def test_backend_id_over_32_bytes_raises(self): + long_name = "A" * 33 + with self.assertRaises(ValueError): + pack_fat_blob([(long_name, b"x")]) class TestBuildFatResult(unittest.TestCase): diff --git a/backends/native/test/test_preprocess.py b/backends/native/test/test_preprocess.py index 11463ff7365..00077e116e9 100644 --- a/backends/native/test/test_preprocess.py +++ b/backends/native/test/test_preprocess.py @@ -12,7 +12,6 @@ from executorch.backends.native.fat_pte import FAT_MAGIC from executorch.backends.native.partitioner import NativePartitioner -from executorch.backends.native.preprocess import NativeBackend from executorch.backends.native.specializations import ( _SPECIALIZATION_REGISTRY, register_specialization, @@ -45,7 +44,6 @@ def setUp(self): def tearDown(self): _SPECIALIZATION_REGISTRY.clear() _SPECIALIZATION_REGISTRY.update(self._saved) - NativeBackend._specialization_names = None def test_register_and_lookup(self): register_specialization("TestBackend", lambda ep: b"test") diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index 332878b5190..1d14a29b081 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -507,4 +507,26 @@ def reinplace_pass( # noqa: C901 continue seen_nodes.update(node.all_input_nodes) + + # Reinplace rewrites may rename output nodes (e.g. aten_relu_default → + # aten_relu__default) but replace_all_uses_with does not update the + # graph signature. Fix output_specs to match the actual output node args. + output_node = next((n for n in ep.graph.nodes if n.op == "output"), None) + if output_node is not None: + out_args = output_node.args[0] + if not isinstance(out_args, (tuple, list)): + out_args = (out_args,) + output_specs = ep.graph_signature.output_specs + assert len(output_specs) == len(out_args), ( + f"reinplace: output spec count changed: " + f"{len(output_specs)} specs vs {len(out_args)} output args" + ) + for spec, arg in zip(output_specs, out_args): + if ( + isinstance(arg, torch.fx.Node) + and getattr(spec.arg, "name", None) is not None + and spec.arg.name != arg.name + ): + spec.arg.name = arg.name + return ep diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 6f1e07b7f32..d4244d1e4e8 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -326,6 +326,25 @@ def forward( # kernels in the other tests in this file. edge.to_executorch() + def test_output_specs_updated_after_reinplace(self) -> None: + """reinplace_pass must update graph_signature.output_specs to match + the renamed output nodes (e.g. relu -> relu_).""" + + class M(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x + 1.0) + + ep = export(M(), (torch.randn(4),), strict=True) + edge_program = to_edge(ep).exported_program() + + custom_set = {edge_ops.edge.aten.relu.default} + ep = reinplace_pass(edge_program, ops_to_inplace=custom_set) + + output_node = next(n for n in ep.graph.nodes if n.op == "output") + actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")] + spec_names = [s.arg.name for s in ep.graph_signature.output_specs] + self.assertEqual(actual_names, spec_names) + def test_broadcasting_self_not_reinplaced(self) -> None: """An op whose mutated arg (self) broadcasts up to a larger output must NOT be reinplaced: the in-place form cannot grow self From d3688864b602257068aeb46074f1160fb2d52fd4 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 2 Jul 2026 15:19:49 -0700 Subject: [PATCH 5/6] up --- backends/native/ir/GraphTypes.h | 672 -------------------------------- backends/native/preprocess.py | 48 +-- 2 files changed, 20 insertions(+), 700 deletions(-) delete mode 100644 backends/native/ir/GraphTypes.h diff --git a/backends/native/ir/GraphTypes.h b/backends/native/ir/GraphTypes.h deleted file mode 100644 index ecec59f64bc..00000000000 --- a/backends/native/ir/GraphTypes.h +++ /dev/null @@ -1,672 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -/** - * Graph: backend-facing IR adapter. - * - * Wraps ExecuTorch's flatbuffer (ExecutionPlan / KernelCall) behind a - * backing-agnostic API. A different serialization could replace the - * underlying storage without changing this header. - * - * IR schema (what Graph presents): - * - * table Graph { - * version: string; - * values: [Value]; // dense pool, indexed by value_id - * inputs: [uint]; - * outputs: [uint]; - * mutable_buffers: [uint]; // values that persist across executes - * operators: [OperatorDef]; // deduped op-name registry - * instructions: [KernelCall]; // flat list of operator calls - * } - * - * table KernelCall { - * op_index: uint; // -> operators[op_index] - * args: [uint]; // value_ids; last = output - * } - * - * Derived views (computed at construction): - * mem_obj_id(vid) — dense int from (pool_id, offset) sort-rank - * value_kind(vid) — INPUT / OUTPUT / CONSTANT / INTERMEDIATE / - * MUTABLE_BUFFER - */ - -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -#include - -namespace executorch { -namespace backends { -namespace portable { - -class Graph; - -enum class ValueKind : uint8_t { - INPUT = 0, - OUTPUT, - CONSTANT, - MUTABLE_BUFFER, - INTERMEDIATE, -}; - -enum class ValueType : uint8_t { - None = 0, - Int, - Double, - Bool, - Tensor, - IntList, - TensorList, - OptionalTensorList, - Other, -}; - -/** - * Wraps a flatbuffer KernelCall. Provides access to op name and - * input/output value_ids. Last arg is the (single) output. - */ -class OperatorCall { - public: - explicit OperatorCall( - const executorch_flatbuffer::KernelCall* call, - const Graph* graph) - : call_(call), graph_(graph) {} - - uint32_t node_id() const { - return node_id_; - } - void set_node_id(uint32_t id) { - node_id_ = id; - } - - const char* name() const; - const char* overload() const; - std::string full_name() const; - - runtime::Span args() const { - auto* a = call_->args(); - return a ? runtime::Span(a->data(), a->size()) - : runtime::Span{}; - } - - runtime::Span inputs() const { - auto a = args(); - return a.empty() ? a : runtime::Span(a.data(), a.size() - 1); - } - - size_t num_inputs() const { - auto a = args(); - return a.empty() ? 0 : a.size() - 1; - } - - uint32_t input(size_t i) const { - ET_CHECK_MSG( - i < num_inputs(), - "OperatorCall::input: index %zu >= num_inputs()=%zu", - i, - num_inputs()); - return static_cast(args()[i]); - } - - size_t num_outputs() const { - return args().empty() ? 0 : 1; - } - - uint32_t output(size_t i) const { - ET_CHECK_MSG( - i < num_outputs(), - "OperatorCall::output: index %zu >= num_outputs()=%zu", - i, - num_outputs()); - auto a = args(); - return static_cast(a[a.size() - 1]); - } - - private: - const executorch_flatbuffer::KernelCall* call_; - const Graph* graph_; - uint32_t node_id_ = 0; -}; - -/** - * IR view of a program. Adapts executorch_flatbuffer::ExecutionPlan. - */ -class Graph { - public: - explicit Graph( - const executorch_flatbuffer::ExecutionPlan* plan, - const executorch_flatbuffer::Program* program = nullptr) - : plan_(plan), program_(program) { - if (auto* in = plan_->inputs()) { - input_ids_.reserve(in->size()); - for (size_t i = 0; i < in->size(); ++i) { - input_ids_.insert(static_cast(in->Get(i))); - } - } - if (auto* out = plan_->outputs()) { - output_ids_.reserve(out->size()); - for (size_t i = 0; i < out->size(); ++i) { - output_ids_.insert(static_cast(out->Get(i))); - } - } - - // Compute mem_obj_id: sort (pool_id, offset) pairs, assign dense rank. - // Same (pool_id, offset) → same id (shared storage). - size_t n_vals = num_values(); - mem_obj_ids_.assign(n_vals, -1); - if (n_vals == 0) - return; - - struct Entry { - uint64_t key; // (pool_id << 32) | offset - uint32_t value_id; - }; - std::vector entries; - entries.reserve(n_vals); - for (uint32_t i = 0; i < n_vals; ++i) { - auto* val = value_meta(i); - if (!val || - val->val_type() != executorch_flatbuffer::KernelTypes::Tensor) { - continue; - } - auto* t = val->val_as_Tensor(); - if (!t) - continue; - auto* alloc = t->allocation_info(); - if (!alloc) - continue; - uint64_t pool = static_cast(alloc->memory_id()); - uint64_t off = alloc->memory_offset_low(); - entries.push_back({(pool << 32) | off, i}); - } - if (entries.empty()) - return; - - std::sort( - entries.begin(), entries.end(), [](const Entry& a, const Entry& b) { - return a.key < b.key; - }); - - int32_t next_id = -1; - uint64_t prev_key = ~0ULL; - for (const auto& e : entries) { - if (e.key != prev_key) { - ++next_id; - prev_key = e.key; - } - mem_obj_ids_[e.value_id] = next_id; - } - - // Mutable buffers: allocated tensors that aren't IO, constants, or - // produced by any op. These are buffer placeholders (e.g., KV cache) - // that persist across execute() calls. - { - std::unordered_set produced_vids; - auto chains = plan_->chains(); - auto instrs = chains->Get(0)->instructions(); - size_t n_instr = instrs ? instrs->size() : 0; - for (size_t oi = 0; oi < n_instr; ++oi) { - OperatorCall op = get_op(oi); - for (size_t k = 0; k < op.num_outputs(); ++k) { - produced_vids.insert(op.output(k)); - } - } - for (uint32_t i = 0; i < n_vals; ++i) { - if (mem_obj_ids_[i] < 0) - continue; - if (input_ids_.count(i) > 0) - continue; - if (output_ids_.count(i) > 0) - continue; - if (tensor_constant_data_key(i) != nullptr) - continue; - if (produced_vids.count(i) > 0) - continue; - mutable_buffer_ids_.push_back(i); - } - } - } - - //===------------------------------------------------------------------===// - // Version - //===------------------------------------------------------------------===// - - static const char* version() { - return "1.0"; - } - - //===------------------------------------------------------------------===// - // Values - //===------------------------------------------------------------------===// - - size_t num_values() const { - auto v = plan_->values(); - return v ? v->size() : 0; - } - - const executorch_flatbuffer::EValue* value_meta(uint32_t value_id) const { - auto values = plan_->values(); - if (!values || value_id >= values->size()) - return nullptr; - return values->Get(value_id); - } - - ValueKind value_kind(uint32_t value_id) const; - int32_t mem_obj_id(uint32_t value_id) const; - - //===------------------------------------------------------------------===// - // Typed value accessors - //===------------------------------------------------------------------===// - - ValueType value_type(uint32_t value_id) const; - - int64_t int_value(uint32_t value_id) const; - double double_value(uint32_t value_id) const; - bool bool_value(uint32_t value_id) const; - - ::executorch::aten::ScalarType tensor_dtype(uint32_t value_id) const; - ::executorch::runtime::Span tensor_sizes( - uint32_t value_id) const; - ::executorch::runtime::Span tensor_dim_order( - uint32_t value_id) const; - ::executorch::aten::TensorShapeDynamism tensor_shape_dynamism( - uint32_t value_id) const; - - // NDM key (FQN) for external constants, or nullptr. - const char* tensor_constant_data_key(uint32_t value_id) const; - - // Raw bytes for inline constants, or empty span. - ::executorch::runtime::Span tensor_inline_data( - uint32_t value_id) const; - - bool is_constant(uint32_t value_id) const { - if (tensor_constant_data_key(value_id) != nullptr) - return true; - return !tensor_inline_data(value_id).empty(); - } - - size_t tensor_nbytes_max(uint32_t value_id) const; - - ::executorch::runtime::Span int_list_member_ids( - uint32_t value_id) const; - - ::executorch::runtime::Span tensor_list_member_ids( - uint32_t value_id) const; - - //===------------------------------------------------------------------===// - // Input/Output IDs - //===------------------------------------------------------------------===// - - size_t num_input_ids() const { - auto in = plan_->inputs(); - return in ? in->size() : 0; - } - - uint32_t input_id(size_t i) const { - auto in = plan_->inputs(); - ET_CHECK_MSG( - in && i < in->size(), - "Graph::input_id(%zu) out of range (have %zu inputs)", - i, - in ? in->size() : 0); - return static_cast(in->Get(i)); - } - - size_t num_output_ids() const { - auto out = plan_->outputs(); - return out ? out->size() : 0; - } - - uint32_t output_id(size_t i) const { - auto out = plan_->outputs(); - ET_CHECK_MSG( - out && i < out->size(), - "Graph::output_id(%zu) out of range (have %zu outputs)", - i, - out ? out->size() : 0); - return static_cast(out->Get(i)); - } - - //===------------------------------------------------------------------===// - // Mutable Buffers - //===------------------------------------------------------------------===// - - size_t num_mutable_buffer_ids() const { - return mutable_buffer_ids_.size(); - } - - uint32_t mutable_buffer_id(size_t i) const { - ET_CHECK_MSG( - i < mutable_buffer_ids_.size(), - "Graph::mutable_buffer_id: index %zu out of range " - "(have %zu mutable buffers)", - i, - mutable_buffer_ids_.size()); - return mutable_buffer_ids_[i]; - } - - //===------------------------------------------------------------------===// - // Operators - //===------------------------------------------------------------------===// - - size_t num_operators() const { - auto ops = plan_->operators(); - return ops ? ops->size() : 0; - } - - const char* operator_name(size_t idx) const { - auto ops = plan_->operators(); - if (!ops || idx >= ops->size()) - return nullptr; - auto op = ops->Get(idx); - return op && op->name() ? op->name()->c_str() : nullptr; - } - - const char* operator_overload(size_t idx) const { - auto ops = plan_->operators(); - if (!ops || idx >= ops->size()) - return nullptr; - auto op = ops->Get(idx); - return op && op->overload() ? op->overload()->c_str() : nullptr; - } - - //===------------------------------------------------------------------===// - // Instructions - //===------------------------------------------------------------------===// - - size_t num_instructions() const { - auto chains = plan_->chains(); - ET_CHECK_MSG(chains && chains->size() > 0, "Graph has no chains"); - auto instrs = chains->Get(0)->instructions(); - return instrs ? instrs->size() : 0; - } - - OperatorCall get_op(size_t op_idx) const { - auto chains = plan_->chains(); - ET_CHECK_MSG(chains && chains->size() > 0, "Graph has no chains"); - auto instrs = chains->Get(0)->instructions(); - ET_CHECK_MSG( - instrs && op_idx < instrs->size(), - "Graph::get_op: op_idx=%zu out of range (have %zu instructions)", - op_idx, - instrs ? instrs->size() : 0); - auto instr = instrs->Get(op_idx); - ET_CHECK_MSG( - instr->instr_args_type() == - executorch_flatbuffer::InstructionArguments::KernelCall, - "Graph::get_op: instruction at op_idx=%zu is not a KernelCall " - "(type=%u)", - op_idx, - static_cast(instr->instr_args_type())); - auto kernel = static_cast( - instr->instr_args()); - return OperatorCall(kernel, this); - } - - OperatorCall get_instruction(size_t idx) const { - return get_op(idx); - } - - private: - const executorch_flatbuffer::ExecutionPlan* plan_; - const executorch_flatbuffer::Program* program_; - std::unordered_set input_ids_; - std::unordered_set output_ids_; - std::vector mem_obj_ids_; - std::vector mutable_buffer_ids_; -}; - -inline const char* OperatorCall::name() const { - return graph_->operator_name(call_->op_index()); -} - -inline const char* OperatorCall::overload() const { - return graph_->operator_overload(call_->op_index()); -} - -inline std::string OperatorCall::full_name() const { - const char* base = name(); - const char* ovl = overload(); - if (!base) - return {}; - if (!ovl || *ovl == '\0') - return std::string(base); - std::string s; - s.reserve(std::strlen(base) + 1 + std::strlen(ovl)); - s.append(base); - s.push_back('.'); - s.append(ovl); - return s; -} - -inline ValueKind Graph::value_kind(uint32_t value_id) const { - if (input_ids_.count(value_id)) - return ValueKind::INPUT; - if (output_ids_.count(value_id)) - return ValueKind::OUTPUT; - - auto val = value_meta(value_id); - if (val && val->val_type() == executorch_flatbuffer::KernelTypes::Tensor) { - auto* tensor = val->val_as_Tensor(); - if (tensor && tensor->data_buffer_idx() > 0) { - return ValueKind::CONSTANT; - } - } - return ValueKind::INTERMEDIATE; -} - -inline int32_t Graph::mem_obj_id(uint32_t value_id) const { - return value_id < mem_obj_ids_.size() ? mem_obj_ids_[value_id] : -1; -} - -inline ValueType Graph::value_type(uint32_t value_id) const { - auto* val = value_meta(value_id); - if (!val) - return ValueType::None; - using KT = executorch_flatbuffer::KernelTypes; - switch (val->val_type()) { - case KT::Null: - return ValueType::None; - case KT::Int: - return ValueType::Int; - case KT::Double: - return ValueType::Double; - case KT::Bool: - return ValueType::Bool; - case KT::Tensor: - return ValueType::Tensor; - case KT::IntList: - return ValueType::IntList; - case KT::TensorList: - return ValueType::TensorList; - case KT::OptionalTensorList: - return ValueType::OptionalTensorList; - default: - return ValueType::Other; - } -} - -inline int64_t Graph::int_value(uint32_t value_id) const { - auto* val = value_meta(value_id); - ET_CHECK_MSG( - val && val->val_type() == executorch_flatbuffer::KernelTypes::Int, - "Graph::int_value(%u): not an Int", - value_id); - return static_cast(val->val())->int_val(); -} - -inline double Graph::double_value(uint32_t value_id) const { - auto* val = value_meta(value_id); - ET_CHECK_MSG( - val && val->val_type() == executorch_flatbuffer::KernelTypes::Double, - "Graph::double_value(%u): not a Double", - value_id); - return static_cast(val->val()) - ->double_val(); -} - -inline bool Graph::bool_value(uint32_t value_id) const { - auto* val = value_meta(value_id); - ET_CHECK_MSG( - val && val->val_type() == executorch_flatbuffer::KernelTypes::Bool, - "Graph::bool_value(%u): not a Bool", - value_id); - return static_cast(val->val()) - ->bool_val(); -} - -namespace detail { -inline const executorch_flatbuffer::Tensor* tensor_or_null( - const executorch_flatbuffer::EValue* val) { - if (!val || val->val_type() != executorch_flatbuffer::KernelTypes::Tensor) { - return nullptr; - } - return val->val_as_Tensor(); -} -} // namespace detail - -inline ::executorch::aten::ScalarType Graph::tensor_dtype( - uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_dtype(%u): not a Tensor", value_id); - return static_cast<::executorch::aten::ScalarType>(t->scalar_type()); -} - -inline ::executorch::runtime::Span Graph::tensor_sizes( - uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_sizes(%u): not a Tensor", value_id); - auto* s = t->sizes(); - return s ? ::executorch::runtime::Span(s->data(), s->size()) - : ::executorch::runtime::Span{}; -} - -inline ::executorch::runtime::Span Graph::tensor_dim_order( - uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_dim_order(%u): not a Tensor", value_id); - auto* d = t->dim_order(); - return d ? ::executorch::runtime::Span(d->data(), d->size()) - : ::executorch::runtime::Span{}; -} - -inline ::executorch::aten::TensorShapeDynamism Graph::tensor_shape_dynamism( - uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - ET_CHECK_MSG(t, "Graph::tensor_shape_dynamism(%u): not a Tensor", value_id); - return static_cast<::executorch::aten::TensorShapeDynamism>( - t->shape_dynamism()); -} - -inline const char* Graph::tensor_constant_data_key(uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - if (!t) - return nullptr; - auto* eti = t->extra_tensor_info(); - if (!eti) - return nullptr; - if (eti->location() != executorch_flatbuffer::TensorDataLocation::EXTERNAL) { - return nullptr; - } - auto* fqn = eti->fully_qualified_name(); - return (fqn && fqn->size() > 0) ? fqn->c_str() : nullptr; -} - -inline ::executorch::runtime::Span Graph::tensor_inline_data( - uint32_t value_id) const { - if (!program_) - return {}; - auto* t = detail::tensor_or_null(value_meta(value_id)); - if (!t) - return {}; - uint32_t idx = static_cast(t->data_buffer_idx()); - // Index 0 is reserved (no inline data). - if (idx == 0) - return {}; - auto* buffers = program_->constant_buffer(); - if (!buffers || idx >= buffers->size()) - return {}; - auto* buf = buffers->Get(idx); - if (!buf) - return {}; - auto* storage = buf->storage(); - if (!storage || storage->size() == 0) - return {}; - return ::executorch::runtime::Span( - storage->data(), storage->size()); -} - -inline size_t Graph::tensor_nbytes_max(uint32_t value_id) const { - auto* t = detail::tensor_or_null(value_meta(value_id)); - if (!t || !t->sizes()) - return 0; - size_t numel = 1; - for (size_t i = 0; i < t->sizes()->size(); ++i) { - int dim = t->sizes()->Get(i); - if (dim < 0) - return 0; - numel *= static_cast(dim); - } - auto stype = static_cast<::executorch::aten::ScalarType>(t->scalar_type()); - return numel * ::executorch::runtime::elementSize(stype); -} - -inline ::executorch::runtime::Span Graph::int_list_member_ids( - uint32_t value_id) const { - auto* val = value_meta(value_id); - ET_CHECK_MSG( - val && val->val_type() == executorch_flatbuffer::KernelTypes::IntList, - "Graph::int_list_member_ids(%u): not an IntList", - value_id); - auto* items = - static_cast(val->val())->items(); - return items - ? ::executorch::runtime::Span(items->data(), items->size()) - : ::executorch::runtime::Span{}; -} - -inline ::executorch::runtime::Span Graph::tensor_list_member_ids( - uint32_t value_id) const { - auto* val = value_meta(value_id); - ET_CHECK_MSG( - val && - (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList || - val->val_type() == - executorch_flatbuffer::KernelTypes::OptionalTensorList), - "Graph::tensor_list_member_ids(%u): not a TensorList or OptionalTensorList", - value_id); - const flatbuffers::Vector* items = nullptr; - if (val->val_type() == executorch_flatbuffer::KernelTypes::TensorList) { - items = static_cast(val->val()) - ->items(); - } else { - items = static_cast( - val->val()) - ->items(); - } - return items - ? ::executorch::runtime::Span(items->data(), items->size()) - : ::executorch::runtime::Span{}; -} - -} // namespace portable -} // namespace backends -} // namespace executorch diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py index d886fd9da15..e9d853498b0 100644 --- a/backends/native/preprocess.py +++ b/backends/native/preprocess.py @@ -27,6 +27,13 @@ ExportedProgram, PreprocessResult, ) + +# A specialization recipe: callable that takes an ExportedProgram and returns +# serialized bytes. Each recipe owns the full +# to_edge_transform_and_lower → to_executorch → serialize flow. + + +from executorch.exir.dialects._ops import ops as _edge_ops from executorch.exir.emit import emit_program from executorch.exir.memory_planning import greedy, MemoryPlanningAlgorithmSuite from executorch.exir.passes import MemoryPlanningPass, SpecPropPass @@ -38,31 +45,16 @@ from torch._export.verifier import Verifier -# A specialization recipe: callable that takes an ExportedProgram and returns -# serialized bytes. Each recipe owns the full -# to_edge_transform_and_lower → to_executorch → serialize flow. - - -def _build_backend_inplace_ops() -> frozenset: - """Edge-dialect ops the native runtime supports in-place. - - reinplace_pass auto-derives the in-place counterpart for each. - """ - from executorch.exir.dialects._ops import ops as _edge_ops - - edge = _edge_ops.edge.aten - return frozenset( - [ - edge.relu.default, - edge.gelu.default, - edge.sigmoid.default, - edge.index_put.default, - edge.index_copy.default, - ] - ) - - -BACKEND_INPLACE_OPS: frozenset = _build_backend_inplace_ops() +_edge = _edge_ops.edge.aten +BACKEND_INPLACE_OPS: frozenset = frozenset( + [ + _edge.relu.default, + _edge.gelu.default, + _edge.sigmoid.default, + _edge.index_put.default, + _edge.index_copy.default, + ] +) class _AnyOp(Verifier): @@ -119,13 +111,13 @@ def preprocess( import copy - spec_programs = [copy.deepcopy(program) for _ in names] - native_result = cls._preprocess_native(program) + native_result = cls._preprocess_native(copy.deepcopy(program)) results: List[Tuple[str, PreprocessResult]] = [ ("NativeBackend", native_result), ] - for name, spec_program in zip(names, spec_programs): + for name in names: + spec_program = copy.deepcopy(program) payload = _SPECIALIZATION_REGISTRY[name](spec_program) results.append((name, PreprocessResult(processed_bytes=payload))) From 4e31e68e6f515d560611270a3d0480358467eab1 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 6 Jul 2026 10:42:31 -0700 Subject: [PATCH 6/6] up --- backends/native/preprocess.py | 4 ++-- backends/native/specializations.py | 25 +++++++++++++++----- backends/native/test/test_fat_pte.py | 31 +++++++++++++++++++++++++ backends/native/test/test_preprocess.py | 17 ++++++++++---- 4 files changed, 65 insertions(+), 12 deletions(-) diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py index e9d853498b0..627af8f0e2d 100644 --- a/backends/native/preprocess.py +++ b/backends/native/preprocess.py @@ -118,8 +118,8 @@ def preprocess( ] for name in names: spec_program = copy.deepcopy(program) - payload = _SPECIALIZATION_REGISTRY[name](spec_program) - results.append((name, PreprocessResult(processed_bytes=payload))) + result = _SPECIALIZATION_REGISTRY[name](spec_program) + results.append((name, result)) return build_fat_result(results) diff --git a/backends/native/specializations.py b/backends/native/specializations.py index 232346eda75..fa4499e5147 100644 --- a/backends/native/specializations.py +++ b/backends/native/specializations.py @@ -8,15 +8,20 @@ Specialization registry and recipes for the native backend fat PTE. Each recipe is a callable that takes an ExportedProgram and returns -serialized bytes. Automatically imported by NativePartitioner to -populate the registry. +a PreprocessResult. The processed_bytes contain the delegate payload +and data_store_output (if any) carries externalized constants that +build_fat_result merges into the shared named-data store, enabling +cross-specialization dedup of identical weight data. + +Automatically imported by NativePartitioner to populate the registry. """ from typing import Callable, Dict +from executorch.exir.backend.backend_details import PreprocessResult from torch.export import ExportedProgram -SpecializationRecipe = Callable[[ExportedProgram], bytes] +SpecializationRecipe = Callable[[ExportedProgram], PreprocessResult] _SPECIALIZATION_REGISTRY: Dict[str, SpecializationRecipe] = {} @@ -30,18 +35,26 @@ def register_specialization(name: str, recipe: SpecializationRecipe) -> None: # --------------------------------------------------------------------------- -def _mlx_recipe(ep: ExportedProgram) -> bytes: +def _mlx_recipe(ep: ExportedProgram) -> PreprocessResult: from executorch.backends.mlx.partitioner import MLXPartitioner from executorch.backends.mlx.passes import get_default_passes from executorch.exir import to_edge_transform_and_lower + from executorch.exir.lowered_backend_module import get_lowered_backend_modules lowered = to_edge_transform_and_lower( ep, transform_passes=get_default_passes(), partitioner=[MLXPartitioner()], ) - et_program = lowered.to_executorch() - return et_program.buffer + lbms = get_lowered_backend_modules( + lowered.exported_program().graph_module, + ) + assert len(lbms) == 1, f"Expected 1 MLX delegate, got {len(lbms)}" + lbm = lbms[0] + return PreprocessResult( + processed_bytes=lbm.processed_bytes, + data_store_output=lbm.named_data_store_output, + ) register_specialization("MLXBackend", _mlx_recipe) diff --git a/backends/native/test/test_fat_pte.py b/backends/native/test/test_fat_pte.py index 9ecec149ee2..283556115ba 100644 --- a/backends/native/test/test_fat_pte.py +++ b/backends/native/test/test_fat_pte.py @@ -15,6 +15,7 @@ FAT_VERSION, pack_fat_blob, ) +from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import PreprocessResult @@ -97,6 +98,36 @@ def test_single_result(self): _, _, count = struct.unpack_from("<4sII", result.processed_bytes, 0) self.assertEqual(count, 1) + def test_shared_constants_deduped(self): + """Two results with identical constant data share one buffer.""" + weight_data = b"\x01\x02\x03\x04" * 64 + + store_a = NamedDataStore() + store_a.add_named_data("model.weight", weight_data, alignment=16) + r_native = PreprocessResult( + processed_bytes=b"native", + data_store_output=store_a.get_named_data_store_output(), + ) + + store_b = NamedDataStore() + store_b.add_named_data("ab12/model.weight", weight_data, alignment=16) + r_accel = PreprocessResult( + processed_bytes=b"accel", + data_store_output=store_b.get_named_data_store_output(), + ) + + result = build_fat_result([("Native", r_native), ("Accel", r_accel)]) + merged = result.data_store_output + + self.assertIn("model.weight", merged.pte_data) + self.assertIn("ab12/model.weight", merged.pte_data) + # Different keys, but same buffer index (deduped by content). + self.assertEqual( + merged.pte_data["model.weight"].buffer_index, + merged.pte_data["ab12/model.weight"].buffer_index, + ) + self.assertEqual(len(merged.buffers), 1) + if __name__ == "__main__": unittest.main() diff --git a/backends/native/test/test_preprocess.py b/backends/native/test/test_preprocess.py index 00077e116e9..5995a02c792 100644 --- a/backends/native/test/test_preprocess.py +++ b/backends/native/test/test_preprocess.py @@ -18,6 +18,7 @@ ) from executorch.exir import to_edge_transform_and_lower from executorch.exir._serialize._flatbuffer_program import _flatbuffer_to_program +from executorch.exir.backend.backend_details import PreprocessResult def _lower(model, example_inputs, specializations=None): @@ -46,7 +47,9 @@ def tearDown(self): _SPECIALIZATION_REGISTRY.update(self._saved) def test_register_and_lookup(self): - register_specialization("TestBackend", lambda ep: b"test") + register_specialization( + "TestBackend", lambda ep: PreprocessResult(processed_bytes=b"test") + ) self.assertIn("TestBackend", _SPECIALIZATION_REGISTRY) def test_unregistered_backend_rejected(self): @@ -67,7 +70,9 @@ def test_no_specializations_produces_plain_flatbuffer(self): self.assertGreater(len(program.execution_plan[0].operators), 0) def test_registered_backend_produces_fat_pte(self): - register_specialization("TestAccel", lambda ep: b"accel_payload") + register_specialization( + "TestAccel", lambda ep: PreprocessResult(processed_bytes=b"accel_payload") + ) edge = _lower( nn.Linear(2, 2), (torch.randn(1, 2),), specializations=["TestAccel"] ) @@ -113,7 +118,9 @@ def forward(self, x): def test_fat_pte_contains_both_payloads(self): saved = dict(_SPECIALIZATION_REGISTRY) - register_specialization("FakeAccel", lambda ep: b"FAKE_ACCEL_DATA") + register_specialization( + "FakeAccel", lambda ep: PreprocessResult(processed_bytes=b"FAKE_ACCEL_DATA") + ) try: edge = _lower( nn.Linear(2, 2), @@ -132,7 +139,9 @@ def test_fat_pte_contains_both_payloads(self): def test_fat_pte_native_payload_is_valid(self): """The native slice inside a fat PTE is a valid flatbuffer program.""" saved = dict(_SPECIALIZATION_REGISTRY) - register_specialization("FakeAccel", lambda ep: b"FAKE") + register_specialization( + "FakeAccel", lambda ep: PreprocessResult(processed_bytes=b"FAKE") + ) try: edge = _lower( nn.Linear(2, 2),