Qualcomm AI Engine Direct - [GenAI Pipeline Phase 2] PRB1 - Compilation foundation

backends/qualcomm/genai_pipeline/__init__.py

@@ -6,6 +6,21 @@
 
 __version__ = "1.0.0"
 
+from executorch.backends.qualcomm.genai_pipeline.artifact_keys import (
+    ALL_ARTIFACT_KEYS,
+    ARTIFACT_ATTENTION_SINK_EVICTOR,
+    ARTIFACT_AUDIO_ENCODER,
+    ARTIFACT_TEXT_DECODER,
+    ARTIFACT_TEXT_ENCODER,
+    ARTIFACT_TOK_EMBEDDING,
+    ARTIFACT_VISION_ENCODER,
+    DECODE_QDQ_FILENAME,
+)
+from executorch.backends.qualcomm.genai_pipeline.compilation import (
+    QnnCompileSpecBuilder,
+    resolve_backend_type,
+    resolve_soc_model,
+)
 from executorch.backends.qualcomm.genai_pipeline.configs import (
     CompilationInputConfig,
     CompilationOutputConfig,
@@ -16,6 +31,7 @@
     QuantizationInputConfig,
     QuantizationOutputConfig,
 )
+from executorch.backends.qualcomm.genai_pipeline.control_args import ControlArgs
 from executorch.backends.qualcomm.genai_pipeline.engine_proxy import EngineProxy
 from executorch.backends.qualcomm.genai_pipeline.exceptions import (
     ConfigValidationError,
@@ -24,6 +40,16 @@
     StageError,
 )
 from executorch.backends.qualcomm.genai_pipeline.genai_pipeline import GenAIPipeline
+from executorch.backends.qualcomm.genai_pipeline.graph_bundle import GraphBundle
+from executorch.backends.qualcomm.genai_pipeline.graph_names import (
+    DECODER_GRAPH_NAMES,
+    GRAPH_FORWARD,
+    GRAPH_KV_FORWARD,
+    GRAPH_PREFILL_FORWARD,
+    GRAPH_TOK_EMBEDDING_KV_FORWARD,
+    GRAPH_TOK_EMBEDDING_PREFILL_FORWARD,
+    TOK_EMBEDDING_GRAPH_NAMES,
+)
 from executorch.backends.qualcomm.genai_pipeline.pipeline_context import (
     PipelineContext,
     PipelineContextBuilder,
@@ -32,13 +58,29 @@
 from executorch.backends.qualcomm.genai_pipeline.pipeline_types import EngineType
 
 __all__ = [
+    "ALL_ARTIFACT_KEYS",
+    "ARTIFACT_ATTENTION_SINK_EVICTOR",
+    "ARTIFACT_AUDIO_ENCODER",
+    "ARTIFACT_TEXT_DECODER",
+    "ARTIFACT_TEXT_ENCODER",
+    "ARTIFACT_TOK_EMBEDDING",
+    "ARTIFACT_VISION_ENCODER",
     "CompilationInputConfig",
     "CompilationOutputConfig",
     "ConfigValidationError",
+    "ControlArgs",
+    "DECODER_GRAPH_NAMES",
+    "DECODE_QDQ_FILENAME",
     "EngineNotAvailableError",
     "EngineProxy",
     "EngineType",
     "GenAIPipeline",
+    "GRAPH_FORWARD",
+    "GRAPH_KV_FORWARD",
+    "GRAPH_PREFILL_FORWARD",
+    "GRAPH_TOK_EMBEDDING_KV_FORWARD",
+    "GRAPH_TOK_EMBEDDING_PREFILL_FORWARD",
+    "GraphBundle",
     "InferenceInputConfig",
     "InferenceOutputConfig",
     "ModelPreparationInputConfig",
@@ -47,7 +89,11 @@
     "PipelineContextBuilder",
     "PipelineError",
     "PipelineStage",
+    "QnnCompileSpecBuilder",
     "QuantizationInputConfig",
     "QuantizationOutputConfig",
+    "resolve_backend_type",
+    "resolve_soc_model",
     "StageError",
+    "TOK_EMBEDDING_GRAPH_NAMES",
 ]

backends/qualcomm/genai_pipeline/artifact_keys.py

@@ -0,0 +1,69 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# 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.
+
+"""Keys identifying the compiled artifacts a GenAI model produces.
+
+``CompilationOutputConfig.artifact_paths`` is keyed by these names, one entry
+per ``.pte`` file. The values match the on-device runner's ``pte_paths`` keys
+exactly, so a compiled bundle can be handed to the runner without translation.
+
+.. note::
+    These are **artifact** keys, not **graph** names. One artifact may hold
+    several methods: a hybrid text decoder lowers its AR-N prefill and AR-1
+    decode graphs into a single multi-method ``.pte`` so they can share
+    weights, and that one file is reached through
+    ``ARTIFACT_TEXT_DECODER``. Graph names (``kv_forward`` /
+    ``prefill_forward``) live in a separate key space used *inside* a lowering
+    call.
+
+    The same string constants exist in the legacy
+    ``examples/qualcomm/oss_scripts/llama/decoder_constants.py``. They are
+    duplicated here rather than imported so that this package does not depend
+    on the example scripts; ``tests/test_artifact_keys.py`` asserts the two
+    stay in agreement, and the legacy copy goes away with the legacy flow.
+"""
+
+from __future__ import annotations
+
+# Text decoder: the language model itself. Always present.
+ARTIFACT_TEXT_DECODER = "text_decoder"
+
+# Token embedding, lowered separately for multimodal models so that the
+# embedding lookup can run while modality features are being inserted.
+ARTIFACT_TOK_EMBEDDING = "tok_embedding"
+
+# Attention sink evictor, compiled only when the attention sink feature is on.
+ARTIFACT_ATTENTION_SINK_EVICTOR = "attention_sink_evictor"
+
+# Modality encoders. Present only for the corresponding multimodal model.
+ARTIFACT_AUDIO_ENCODER = "audio_encoder"
+ARTIFACT_TEXT_ENCODER = "text_encoder"
+ARTIFACT_VISION_ENCODER = "vision_encoder"
+
+# Every key this package may emit.
+#
+# Membership in ``artifact_paths`` is meaningful: the runner decides whether a
+# model is multimodal by testing for the encoder keys, so an artifact that was
+# not compiled must be **absent** rather than mapped to ``None``.
+ALL_ARTIFACT_KEYS = frozenset(
+    {
+        ARTIFACT_ATTENTION_SINK_EVICTOR,
+        ARTIFACT_AUDIO_ENCODER,
+        ARTIFACT_TEXT_DECODER,
+        ARTIFACT_TEXT_ENCODER,
+        ARTIFACT_TOK_EMBEDDING,
+        ARTIFACT_VISION_ENCODER,
+    }
+)
+
+# The decode graph's QDQ exported program, written by quantization before the
+# calibration graph is released and read back by the SQNR evaluation.
+#
+# A **filename**, not an artifact key: it is a ``.pt2`` exported program rather
+# than a lowered ``.pte``, the runner never sees it, and only the A-side stages
+# touch it. Hence deliberately absent from ``ALL_ARTIFACT_KEYS``. It lives here
+# because this module already owns the names the pipeline writes to disk.
+DECODE_QDQ_FILENAME = "decode_qdq.pt2"

backends/qualcomm/genai_pipeline/compilation/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# 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.
+
+from executorch.backends.qualcomm.genai_pipeline.compilation.compile_spec_builder import (
+    QnnCompileSpecBuilder,
+    resolve_backend_type,
+    resolve_soc_model,
+)
+
+__all__ = [
+    "QnnCompileSpecBuilder",
+    "resolve_backend_type",
+    "resolve_soc_model",
+]