diff --git a/Test.cmake b/Test.cmake index 6bd7a86e70b..b2fadd47579 100644 --- a/Test.cmake +++ b/Test.cmake @@ -10,6 +10,7 @@ if(BUILD_TESTING) # This contains the list of tests which are always built + add_subdirectory(backends/native/ir/test) add_subdirectory(extension/evalue_util/test) add_subdirectory(extension/kernel_util/test) add_subdirectory(extension/memory_allocator/test) diff --git a/backends/native/ir/GraphTypes.h b/backends/native/ir/GraphTypes.h new file mode 100644 index 00000000000..b37cecca7a7 --- /dev/null +++ b/backends/native/ir/GraphTypes.h @@ -0,0 +1,1107 @@ +/* + * 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 + * } + * + * table OperatorDef { + * name: string; // e.g. "aten.add.Tensor" + * } + * + * // Op chains + * table Chain { + * instructions: [Instruction]; + * } + * + * table Instruction { + * body: KernelCall; // Only Kernel instructions are supported. + * } + * + * 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. + * producer(vid) which instruction produces this value + * (nullptr for inputs/constants). + * users(vid) all instructions that consume this value. + * num_users(vid) shortcut for users(vid).size(). + * find_ops(name) all KernelCall instructions matching + * the given operator base name. + * + * ---------------------------------------------------------------------- + * 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, + * O(num_inputs + num_outputs) for the value_kind sets, and O(total_args) + * for the use-def indices (producer, users, op name index). 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 + +#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. + * + * Only Kernel instructions are supported; the native backend + * partitioner guarantees that no control-flow, move, free, or + * delegate instructions appear in the delegated subgraph. + * Encountering any other kind during deserialization is a fatal + * error. + */ +enum class InstructionKind : uint8_t { + Kernel = 0, +}; + +/** + * Lightweight reference to an instruction's location in the graph. + * Returned by the use-def analysis helpers (producer, users, find_ops). + */ +struct InstructionRef { + uint32_t chain_idx; + uint32_t instr_idx; +}; + +static constexpr InstructionRef kNoProducer = {UINT32_MAX, UINT32_MAX}; + +/** + * 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. + 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) { + 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); + } + } + + // Precompute use-def analysis indices (producer, users, op name index). + // Single pass over all instructions; CSR build for users. + producers_.assign(n_vals, kNoProducer); + std::vector user_counts(n_vals, 0); + + auto record_user = [&](uint32_t vid) { + if (vid < n_vals) + ++user_counts[vid]; + }; + auto record_producer = [&](uint32_t vid, InstructionRef ref) { + if (vid < n_vals) + producers_[vid] = ref; + }; + + // Pass 1: count users per value and record producers. + for (size_t ci = 0; ci < num_chains(); ++ci) { + auto instrs = plan_->chains()->Get(ci)->instructions(); + size_t n_instr = instrs ? instrs->size() : 0; + for (size_t ii = 0; ii < n_instr; ++ii) { + InstructionRef ref{ + static_cast(ci), static_cast(ii)}; + OperatorCall op = get_op(ci, ii); + for (size_t j = 0; j < op.num_inputs(); ++j) + record_user(op.input(j)); + for (size_t j = 0; j < op.num_outputs(); ++j) + record_producer(op.output(j), ref); + const char* op_name = op.name(); + if (op_name) + op_name_index_[op_name].push_back(ref); + } + } + + // Pass 2: build CSR from counts. + user_starts_.resize(n_vals + 1); + user_starts_[0] = 0; + for (uint32_t i = 0; i < n_vals; ++i) + user_starts_[i + 1] = user_starts_[i] + user_counts[i]; + user_entries_.resize(user_starts_[n_vals]); + + // Reuse user_counts as write cursors. + std::fill(user_counts.begin(), user_counts.end(), 0); + + auto emit_user = [&](uint32_t vid, InstructionRef ref) { + if (vid < n_vals) { + user_entries_[user_starts_[vid] + user_counts[vid]] = ref; + ++user_counts[vid]; + } + }; + + // Pass 3: fill user entries (same traversal as pass 1). + for (size_t ci = 0; ci < num_chains(); ++ci) { + auto instrs = plan_->chains()->Get(ci)->instructions(); + size_t n_instr = instrs ? instrs->size() : 0; + for (size_t ii = 0; ii < n_instr; ++ii) { + InstructionRef ref{ + static_cast(ci), static_cast(ii)}; + OperatorCall op = get_op(ci, ii); + for (size_t j = 0; j < op.num_inputs(); ++j) + emit_user(op.input(j), ref); + } + } + } + + //===------------------------------------------------------------------===// + // Version + //===------------------------------------------------------------------===// + + // cppcheck-suppress functionStatic + 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]; + } + + //===------------------------------------------------------------------===// + // Use-def analysis (precomputed at construction) + //===------------------------------------------------------------------===// + + // Which instruction produces this value? Returns nullptr if the value + // has no producer (graph input, constant, or mutable buffer placeholder). + const InstructionRef* producer(uint32_t value_id) const { + if (value_id >= producers_.size()) + return nullptr; + const auto& ref = producers_[value_id]; + if (ref.chain_idx == kNoProducer.chain_idx && + ref.instr_idx == kNoProducer.instr_idx) + return nullptr; + return &ref; + } + + // All instructions that consume this value (as an input, move source, + // jump condition, or free target). + ::executorch::runtime::Span users( + uint32_t value_id) const { + if (value_id >= producers_.size()) + return {}; + uint32_t start = user_starts_[value_id]; + uint32_t end = user_starts_[value_id + 1]; + return ::executorch::runtime::Span( + user_entries_.data() + start, end - start); + } + + size_t num_users(uint32_t value_id) const { + if (value_id >= producers_.size()) + return 0; + return user_starts_[value_id + 1] - user_starts_[value_id]; + } + + // All KernelCall instructions whose operator base name matches. + // Returns empty span if no matches. + ::executorch::runtime::Span find_ops( + const char* name) const { + auto it = op_name_index_.find(name); + if (it == op_name_index_.end()) + return {}; + return ::executorch::runtime::Span( + it->second.data(), it->second.size()); + } + + //===------------------------------------------------------------------===// + // 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; + } + + // cppcheck-suppress functionStatic + 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 + //===------------------------------------------------------------------===// + + 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; + ET_CHECK_MSG( + instr->instr_args_type() == IA::KernelCall, + "Graph: non-Kernel instruction at chain=%zu op_idx=%zu (type=%u). " + "Control flow is not supported in the native backend.", + chain_idx, + op_idx, + static_cast(instr->instr_args_type())); + return InstructionKind::Kernel; + } + + InstructionKind instruction_kind(size_t op_idx) const { + return instruction_kind(main_chain_idx(), op_idx); + } + + OperatorCall get_kernel_call(size_t chain_idx, size_t op_idx) const { + return get_op(chain_idx, op_idx); + } + + //===------------------------------------------------------------------===// + // 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_; + + // Use-def analysis indices (precomputed at construction). + std::vector producers_; // indexed by value_id + std::vector user_starts_; // CSR offsets, size = num_values + 1 + std::vector user_entries_; // CSR entries + std::unordered_map> op_name_index_; +}; + +// 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/test/CMakeLists.txt b/backends/native/ir/test/CMakeLists.txt new file mode 100644 index 00000000000..b99fe90b78f --- /dev/null +++ b/backends/native/ir/test/CMakeLists.txt @@ -0,0 +1,26 @@ +# 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. + +cmake_minimum_required(VERSION 3.19) + +include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + +et_cxx_test( + graph_types_test SOURCES graph_types_test.cpp EXTRA_LIBS program_schema +) + +target_include_directories( + graph_types_test + PRIVATE "${EXECUTORCH_ROOT}" "${CMAKE_INSTALL_PREFIX}/schema/include" + "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" +) + +set_property( + TEST graph_types_test + PROPERTY + ENVIRONMENT + "ET_NATIVE_TEST_DATA=${CMAKE_CURRENT_SOURCE_DIR}/testdata/linear_4x4.bin" +) diff --git a/backends/native/ir/test/generate_test_data.py b/backends/native/ir/test/generate_test_data.py new file mode 100644 index 00000000000..6f4a963c132 --- /dev/null +++ b/backends/native/ir/test/generate_test_data.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# 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. + +"""Generate test data (delegate flatbuffer blobs) for GraphTypes C++ tests.""" + +import pathlib + +import torch +import torch.nn as nn + +from executorch.backends.native.partitioner import NativePartitioner +from executorch.exir import to_edge_transform_and_lower + + +def _generate_linear_blob() -> bytes: + model = nn.Linear(4, 4) + ep = torch.export.export(model, (torch.randn(1, 4),)) + edge = to_edge_transform_and_lower(ep, partitioner=[NativePartitioner()]) + et = edge.to_executorch() + delegates = et.executorch_program.backend_delegate_data + assert len(delegates) == 1 + return bytes(delegates[0].data) + + +class _DiamondModel(nn.Module): + """x -> add(x,x) -> a; a -> mul(a,2) -> b; a -> add(a,1) -> c; add(b,c) -> out. + + 'a' has 2 users (mul and add), creating an interesting use-def graph. + """ + + def forward(self, x): + a = x + x + b = a * 2 + c = a + 1 + return b + c + + +def _generate_diamond_blob() -> bytes: + model = _DiamondModel() + ep = torch.export.export(model, (torch.randn(1, 4),)) + edge = to_edge_transform_and_lower(ep, partitioner=[NativePartitioner()]) + et = edge.to_executorch() + delegates = et.executorch_program.backend_delegate_data + assert len(delegates) == 1 + return bytes(delegates[0].data) + + +class _KVCacheModel(nn.Module): + """HF-style KV cache with index_copy_ (lowered to index_put). + + k_cache is a mutable buffer [1, max_seq=8, head_dim=4]. + forward updates the cache at position `pos` and returns a reduction. + """ + + def __init__(self): + super().__init__() + self.register_buffer("k_cache", torch.zeros(1, 8, 4)) + + def forward(self, k_new, pos): + self.k_cache.index_copy_(1, pos, k_new) + return self.k_cache.sum(dim=1) + + +def _generate_kv_cache_blob() -> bytes: + model = _KVCacheModel() + k_new = torch.randn(1, 1, 4) + pos = torch.tensor([0]) + ep = torch.export.export(model, (k_new, pos), strict=False) + edge = to_edge_transform_and_lower(ep, partitioner=[NativePartitioner()]) + et = edge.to_executorch() + delegates = et.executorch_program.backend_delegate_data + assert len(delegates) == 1 + return bytes(delegates[0].data) + + +def main(): + out_dir = pathlib.Path(__file__).parent / "testdata" + out_dir.mkdir(parents=True, exist_ok=True) + + blob = _generate_linear_blob() + out_path = out_dir / "linear_4x4.bin" + out_path.write_bytes(blob) + print(f"Wrote {len(blob)} bytes to {out_path}") + + blob = _generate_diamond_blob() + out_path = out_dir / "diamond.bin" + out_path.write_bytes(blob) + print(f"Wrote {len(blob)} bytes to {out_path}") + + blob = _generate_kv_cache_blob() + out_path = out_dir / "kv_cache.bin" + out_path.write_bytes(blob) + print(f"Wrote {len(blob)} bytes to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/backends/native/ir/test/graph_types_test.cpp b/backends/native/ir/test/graph_types_test.cpp new file mode 100644 index 00000000000..1df5579c9b0 --- /dev/null +++ b/backends/native/ir/test/graph_types_test.cpp @@ -0,0 +1,440 @@ +/* + * 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. + */ + +#include +#include + +#include + +#include +#include +#include +#include +#include + +using namespace executorch::backends::portable; + +namespace { + +std::vector load_file(const char* path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + EXPECT_TRUE(f.is_open()) << "Cannot open " << path; + auto size = f.tellg(); + f.seekg(0); + std::vector buf(size); + f.read(reinterpret_cast(buf.data()), size); + return buf; +} + +std::string testdata_path(const char* filename) { + const char* base = std::getenv("ET_NATIVE_TEST_DATA"); + EXPECT_NE(base, nullptr) << "Set ET_NATIVE_TEST_DATA env var"; + std::string p(base); + auto slash = p.rfind('/'); + EXPECT_NE(slash, std::string::npos); + return p.substr(0, slash + 1) + filename; +} + +std::unique_ptr load_graph( + const std::string& path, + std::vector& data_out, + const executorch_flatbuffer::Program** prog_out = nullptr) { + data_out = load_file(path.c_str()); + auto* program = + flatbuffers::GetRoot(data_out.data()); + if (prog_out) + *prog_out = program; + return std::make_unique(program->execution_plan()->Get(0), program); +} + +// Verifies that every op's outputs are recorded as produced by that op, +// and every op's inputs list that op as a user. +void verify_producer_user_consistency(const Graph& g) { + for (size_t ci = 0; ci < g.num_chains(); ++ci) { + for (size_t ii = 0; ii < g.num_ops_in_chain(ci); ++ii) { + OperatorCall op = g.get_op(ci, ii); + for (size_t k = 0; k < op.num_outputs(); ++k) { + auto* p = g.producer(op.output(k)); + ASSERT_NE(p, nullptr) << "output vid=" << op.output(k); + EXPECT_EQ(p->chain_idx, ci); + EXPECT_EQ(p->instr_idx, ii); + } + for (size_t k = 0; k < op.num_inputs(); ++k) { + auto u = g.users(op.input(k)); + bool found = false; + for (size_t j = 0; j < u.size(); ++j) { + if (u[j].chain_idx == ci && u[j].instr_idx == ii) { + found = true; + break; + } + } + EXPECT_TRUE(found) << "op(" << ci << "," << ii + << ") not in users of vid=" << op.input(k); + } + } + } +} + +} // namespace + +// ============================================================================= +// Linear model (nn.Linear(4,4)): single op, covers core API surface +// ============================================================================= + +class LinearGraphTest : public ::testing::Test { + protected: + void SetUp() override { + graph_ = load_graph(testdata_path("linear_4x4.bin"), data_, &program_); + } + std::vector data_; + const executorch_flatbuffer::Program* program_ = nullptr; + std::unique_ptr graph_; +}; + +TEST_F(LinearGraphTest, StructureAndMetadata) { + EXPECT_STREQ(graph_->version(), "1.0"); + EXPECT_GT(graph_->num_values(), 0u); + EXPECT_EQ(graph_->num_input_ids(), 1u); + EXPECT_EQ(graph_->num_output_ids(), 1u); + EXPECT_LT(graph_->input_id(0), graph_->num_values()); + EXPECT_LT(graph_->output_id(0), graph_->num_values()); + EXPECT_NE(graph_->input_id(0), graph_->output_id(0)); + EXPECT_GE(graph_->num_chains(), 1u); + EXPECT_GE(graph_->num_operators(), 1u); + EXPECT_EQ(graph_->main_chain_idx(), 0); +} + +TEST_F(LinearGraphTest, ValueKindsAndTypes) { + // Input + uint32_t in_vid = graph_->input_id(0); + EXPECT_EQ(graph_->value_kind(in_vid), ValueKind::INPUT); + EXPECT_EQ(graph_->value_type(in_vid), ValueType::Tensor); + EXPECT_FALSE(graph_->is_constant(in_vid)); + EXPECT_EQ(graph_->tensor_constant_data_key(in_vid), nullptr); + + // Output + uint32_t out_vid = graph_->output_id(0); + EXPECT_EQ(graph_->value_kind(out_vid), ValueKind::OUTPUT); + EXPECT_EQ(graph_->value_type(out_vid), ValueType::Tensor); + + // Constants exist (weight + bias) + bool has_constant = false; + for (uint32_t i = 0; i < graph_->num_values(); ++i) { + EXPECT_NE(graph_->value_meta(i), nullptr) << "vid=" << i; + auto kind = graph_->value_kind(i); + EXPECT_TRUE( + kind == ValueKind::INPUT || kind == ValueKind::OUTPUT || + kind == ValueKind::CONSTANT || kind == ValueKind::MUTABLE_BUFFER || + kind == ValueKind::INTERMEDIATE); + if (kind == ValueKind::CONSTANT) + has_constant = true; + } + EXPECT_TRUE(has_constant); +} + +TEST_F(LinearGraphTest, TensorAccessors) { + uint32_t vid = graph_->input_id(0); + EXPECT_EQ(graph_->tensor_dtype(vid), ::executorch::aten::ScalarType::Float); + + auto sizes = graph_->tensor_sizes(vid); + EXPECT_EQ(sizes.size(), 2u); + EXPECT_EQ(sizes[0], 1); + EXPECT_EQ(sizes[1], 4); + + auto dim_order = graph_->tensor_dim_order(vid); + EXPECT_EQ(dim_order.size(), 2u); + EXPECT_EQ(dim_order[0], 0); + EXPECT_EQ(dim_order[1], 1); + + EXPECT_EQ(graph_->tensor_nbytes_max(vid), 16u); // 1*4*sizeof(float) + EXPECT_EQ(graph_->tensor_nbytes_max(graph_->output_id(0)), 16u); + + graph_->tensor_shape_dynamism(vid); // shouldn't crash +} + +TEST_F(LinearGraphTest, MemObjId) { + for (uint32_t i = 0; i < graph_->num_values(); ++i) { + EXPECT_GE(graph_->mem_obj_id(i), -1); + } + for (size_t i = 0; i < graph_->num_mutable_buffer_ids(); ++i) { + EXPECT_LT(graph_->mutable_buffer_id(i), graph_->num_values()); + } +} + +TEST_F(LinearGraphTest, OperatorTable) { + bool found_linear = false; + for (size_t i = 0; i < graph_->num_operators(); ++i) { + const char* name = graph_->operator_name(i); + EXPECT_NE(name, nullptr); + graph_->operator_overload(i); // shouldn't crash + if (name && std::string(name).find("linear") != std::string::npos) + found_linear = true; + } + EXPECT_TRUE(found_linear); +} + +TEST_F(LinearGraphTest, InstructionAccess) { + EXPECT_EQ( + graph_->num_instructions(), + graph_->num_ops_in_chain(graph_->main_chain_idx())); + + for (size_t ci = 0; ci < graph_->num_chains(); ++ci) { + for (size_t ii = 0; ii < graph_->num_ops_in_chain(ci); ++ii) { + EXPECT_EQ(graph_->instruction_kind(ci, ii), InstructionKind::Kernel); + OperatorCall a = graph_->get_op(ci, ii); + OperatorCall b = graph_->get_kernel_call(ci, ii); + EXPECT_STREQ(a.name(), b.name()); + EXPECT_EQ(a.num_inputs(), b.num_inputs()); + } + } + + for (size_t ii = 0; ii < graph_->num_instructions(); ++ii) { + EXPECT_EQ(graph_->instruction_kind(ii), InstructionKind::Kernel); + OperatorCall a = graph_->get_instruction(ii); + OperatorCall b = graph_->get_op(graph_->main_chain_idx(), ii); + EXPECT_STREQ(a.name(), b.name()); + } +} + +TEST_F(LinearGraphTest, OperatorCallAPI) { + auto refs = graph_->find_ops("aten::linear"); + ASSERT_EQ(refs.size(), 1u); + OperatorCall op = graph_->get_op(refs[0].chain_idx, refs[0].instr_idx); + + EXPECT_STREQ(op.name(), "aten::linear"); + EXPECT_TRUE( + std::string(op.full_name()).find("aten::linear") != std::string::npos); + + EXPECT_GE(op.num_inputs(), 1u); + EXPECT_EQ(op.num_outputs(), 1u); + EXPECT_EQ(op.args().size(), op.inputs().size() + op.num_outputs()); + + for (size_t i = 0; i < op.num_inputs(); ++i) + EXPECT_LT(op.input(i), graph_->num_values()); + EXPECT_LT(op.output(0), graph_->num_values()); + + // linear's first input is graph input, output is graph output + EXPECT_EQ(op.input(0), graph_->input_id(0)); + EXPECT_EQ(op.output(0), graph_->output_id(0)); + + // node_id is mutable + EXPECT_EQ(op.node_id(), 0u); + OperatorCall op2 = graph_->get_instruction(0); + op2.set_node_id(42); + EXPECT_EQ(op2.node_id(), 42u); +} + +TEST_F(LinearGraphTest, FindOps) { + EXPECT_EQ(graph_->find_ops("aten::linear").size(), 1u); + EXPECT_EQ(graph_->find_ops("nonexistent::op").size(), 0u); + EXPECT_EQ(graph_->find_ops("").size(), 0u); +} + +TEST_F(LinearGraphTest, ProducerAndUsers) { + // Input: no producer, has users + uint32_t in_vid = graph_->input_id(0); + EXPECT_EQ(graph_->producer(in_vid), nullptr); + EXPECT_GT(graph_->num_users(in_vid), 0u); + + // Output: has producer (linear op), no users + uint32_t out_vid = graph_->output_id(0); + auto* p = graph_->producer(out_vid); + ASSERT_NE(p, nullptr); + auto refs = graph_->find_ops("aten::linear"); + EXPECT_EQ(p->chain_idx, refs[0].chain_idx); + EXPECT_EQ(p->instr_idx, refs[0].instr_idx); + EXPECT_EQ(graph_->num_users(out_vid), 0u); + + // Constants: no producer + for (uint32_t i = 0; i < graph_->num_values(); ++i) { + if (graph_->value_kind(i) == ValueKind::CONSTANT) + EXPECT_EQ(graph_->producer(i), nullptr) << "vid=" << i; + } + + // num_users == users().size() for all values + for (uint32_t i = 0; i < graph_->num_values(); ++i) + EXPECT_EQ(graph_->num_users(i), graph_->users(i).size()) << "vid=" << i; + + // User refs are in range + for (uint32_t i = 0; i < graph_->num_values(); ++i) { + auto u = graph_->users(i); + for (size_t j = 0; j < u.size(); ++j) { + EXPECT_LT(u[j].chain_idx, graph_->num_chains()); + EXPECT_LT(u[j].instr_idx, graph_->num_ops_in_chain(u[j].chain_idx)); + } + } +} + +TEST_F(LinearGraphTest, OutOfRangeSafety) { + EXPECT_EQ(graph_->value_meta(UINT32_MAX), nullptr); + EXPECT_EQ(graph_->value_type(UINT32_MAX), ValueType::None); + EXPECT_EQ(graph_->mem_obj_id(UINT32_MAX), -1); + EXPECT_EQ(graph_->tensor_nbytes_max(UINT32_MAX), 0u); + EXPECT_EQ(graph_->producer(UINT32_MAX), nullptr); + EXPECT_EQ(graph_->users(UINT32_MAX).size(), 0u); + EXPECT_EQ(graph_->num_users(UINT32_MAX), 0u); + EXPECT_EQ(graph_->operator_name(999), nullptr); + EXPECT_EQ(graph_->operator_overload(999), nullptr); +} + +TEST_F(LinearGraphTest, ConstructWithPlanOnly) { + Graph plan_only(program_->execution_plan()->Get(0)); + EXPECT_GT(plan_only.num_values(), 0u); + EXPECT_EQ(plan_only.num_input_ids(), 1u); + EXPECT_TRUE(plan_only.tensor_inline_data(plan_only.input_id(0)).empty()); +} + +TEST_F(LinearGraphTest, ProducerUserConsistency) { + verify_producer_user_consistency(*graph_); +} + +// ============================================================================= +// Diamond model: multi-user values and producer chains +// +// x -> add(x,x) -> a -> mul(a,2) -> b +// -> add(a,1) -> c -> add(b,c) -> out +// ============================================================================= + +class DiamondGraphTest : public ::testing::Test { + protected: + void SetUp() override { + graph_ = load_graph(testdata_path("diamond.bin"), data_); + } + std::vector data_; + std::unique_ptr graph_; +}; + +TEST_F(DiamondGraphTest, Structure) { + EXPECT_EQ(graph_->num_instructions(), 4u); + EXPECT_EQ(graph_->find_ops("aten::add").size(), 3u); + EXPECT_EQ(graph_->find_ops("aten::mul").size(), 1u); +} + +TEST_F(DiamondGraphTest, MultiUserAndProducerChain) { + // x (input 0) used twice by add[0] + uint32_t x_vid = graph_->input_id(0); + auto x_users = graph_->users(x_vid); + EXPECT_GE(x_users.size(), 2u); + EXPECT_EQ(x_users[0].instr_idx, 0u); + EXPECT_EQ(x_users[1].instr_idx, 0u); + + // 'a' (output of add[0]) has 2 distinct user instructions + OperatorCall first_add = graph_->get_instruction(0); + uint32_t a_vid = first_add.output(0); + EXPECT_GE(graph_->num_users(a_vid), 2u); + std::unordered_set a_user_instrs; + for (auto& ref : graph_->users(a_vid)) + a_user_instrs.insert(ref.instr_idx); + EXPECT_TRUE(a_user_instrs.count(1)); // mul + EXPECT_TRUE(a_user_instrs.count(2)); // add(a,1) + + // Trace from output back to input through the diamond + uint32_t out_vid = graph_->output_id(0); + auto* p_final = graph_->producer(out_vid); + ASSERT_NE(p_final, nullptr); + EXPECT_EQ(p_final->instr_idx, 3u); + + OperatorCall final_add = + graph_->get_op(p_final->chain_idx, p_final->instr_idx); + auto* p_b = graph_->producer(final_add.input(0)); + auto* p_c = graph_->producer(final_add.input(1)); + ASSERT_NE(p_b, nullptr); + ASSERT_NE(p_c, nullptr); + EXPECT_EQ(p_b->instr_idx, 1u); + EXPECT_EQ(p_c->instr_idx, 2u); + + // Both fan-in paths converge at 'a' + OperatorCall mul_op = graph_->get_op(p_b->chain_idx, p_b->instr_idx); + OperatorCall add_a1 = graph_->get_op(p_c->chain_idx, p_c->instr_idx); + EXPECT_EQ(mul_op.input(0), add_a1.input(0)); // same 'a' + EXPECT_EQ(graph_->producer(mul_op.input(0))->instr_idx, 0u); + + // Output has no users + EXPECT_EQ(graph_->num_users(out_vid), 0u); +} + +TEST_F(DiamondGraphTest, ProducerUserConsistency) { + verify_producer_user_consistency(*graph_); +} + +// ============================================================================= +// KV cache model: inplace ops and cache buffer pattern +// +// k_cache [1,8,4] updated via index_copy_ (lowered to index_put_), then +// reduced via sum. index_put_ aliases input[0] and output[0] (inplace). +// ============================================================================= + +class KVCacheGraphTest : public ::testing::Test { + protected: + void SetUp() override { + graph_ = load_graph(testdata_path("kv_cache.bin"), data_); + // Find the cache: intermediate tensor with allocation, consumed by 2+ ops + for (uint32_t i = 0; i < graph_->num_values(); ++i) { + if (graph_->value_kind(i) == ValueKind::INTERMEDIATE && + graph_->value_type(i) == ValueType::Tensor && + graph_->mem_obj_id(i) >= 0 && graph_->num_users(i) >= 2) { + cache_vid_ = i; + break; + } + } + } + std::vector data_; + std::unique_ptr graph_; + uint32_t cache_vid_ = UINT32_MAX; +}; + +TEST_F(KVCacheGraphTest, Structure) { + EXPECT_EQ(graph_->num_instructions(), 2u); + EXPECT_EQ(graph_->find_ops("aten::index_put_").size(), 1u); + EXPECT_EQ(graph_->find_ops("aten::sum").size(), 1u); +} + +TEST_F(KVCacheGraphTest, InplaceCachePattern) { + ASSERT_NE(cache_vid_, UINT32_MAX) << "cache tensor not found"; + + // Cache is not an input, output, or constant + EXPECT_EQ(graph_->value_kind(cache_vid_), ValueKind::INTERMEDIATE); + EXPECT_FALSE(graph_->is_constant(cache_vid_)); + EXPECT_GE(graph_->mem_obj_id(cache_vid_), 0); + + // Cache shape: [1, 8, 4], 128 bytes + auto sizes = graph_->tensor_sizes(cache_vid_); + EXPECT_EQ(sizes.size(), 3u); + EXPECT_EQ(sizes[0], 1); + EXPECT_EQ(sizes[1], 8); + EXPECT_EQ(sizes[2], 4); + EXPECT_EQ(graph_->tensor_nbytes_max(cache_vid_), 128u); + + // index_put_ is inplace: input[0] == output[0] == cache + auto refs = graph_->find_ops("aten::index_put_"); + ASSERT_EQ(refs.size(), 1u); + OperatorCall ip = graph_->get_op(refs[0].chain_idx, refs[0].instr_idx); + EXPECT_EQ(ip.input(0), ip.output(0)); + EXPECT_EQ(ip.output(0), cache_vid_); + + // Cache consumed by both index_put_ and sum + std::unordered_set user_instrs; + for (auto& ref : graph_->users(cache_vid_)) + user_instrs.insert(ref.instr_idx); + EXPECT_TRUE(user_instrs.count(0)); // index_put_ + EXPECT_TRUE(user_instrs.count(1)); // sum + + // Producer is index_put_ (inplace writes back to same value) + auto* p = graph_->producer(cache_vid_); + ASSERT_NE(p, nullptr); + EXPECT_STREQ( + graph_->get_op(p->chain_idx, p->instr_idx).name(), "aten::index_put_"); + + // Cache and output have different mem_obj_ids (128B vs 16B) + EXPECT_NE( + graph_->mem_obj_id(cache_vid_), graph_->mem_obj_id(graph_->output_id(0))); +} + +TEST_F(KVCacheGraphTest, ProducerUserConsistency) { + verify_producer_user_consistency(*graph_); +} diff --git a/backends/native/ir/test/testdata/diamond.bin b/backends/native/ir/test/testdata/diamond.bin new file mode 100644 index 00000000000..7ac0a05e362 Binary files /dev/null and b/backends/native/ir/test/testdata/diamond.bin differ diff --git a/backends/native/ir/test/testdata/kv_cache.bin b/backends/native/ir/test/testdata/kv_cache.bin new file mode 100644 index 00000000000..22b3626a1ea Binary files /dev/null and b/backends/native/ir/test/testdata/kv_cache.bin differ diff --git a/backends/native/ir/test/testdata/linear_4x4.bin b/backends/native/ir/test/testdata/linear_4x4.bin new file mode 100644 index 00000000000..21e740a48d1 Binary files /dev/null and b/backends/native/ir/test/testdata/linear_4x4.bin differ