Skip to content

index.md

Latest commit

 

History

History
102 lines (74 loc) · 4.58 KB

index.md

File metadata and controls

102 lines (74 loc) · 4.58 KB
Raw

Core AI PyTorch Extensions (coreai-torch)

Bring PyTorch models to Core AI for on-device execution.

Overview

Core AI PyTorch Extensions (coreai-torch) is a Python package that bridges PyTorch and Core AI. It converts an existing PyTorch model — exported as a torch.export.ExportedProgram — into a Core AI AIProgram ready to run on Apple hardware, traversing the FX graph node-by-node and mapping ATen operators to Core AI operations. The package also supports authoring Core AI models directly from PyTorch by composing the library of composite ops in coreai_torch.composite_ops, authoring new ops via register_torch_lowering, and authoring inline Metal GPU kernels through TorchMetalKernel and register_custom_kernels — all expressed as PyTorch nn.Modules and lowered to Core AI IR that the compiler recognizes and optimizes natively.

The bring-up pipeline has three steps. First, export the PyTorch model with torch.export.export to capture the computation graph. Second, decompose the exported program with get_decomp_table(), which lowers composite ATen ops to the primitive set that TorchConverter can map while preserving the operations that TorchConverter lowers as composite ops. Third, call TorchConverter().add_exported_program(ep).to_coreai() to produce the AIProgram.

For authoring, coreai_torch.composite_ops exposes well-known building blocks — such as attention, RoPE embeddings, RMSNorm, and gather-matmul (the MoE primitive) — as PyTorch modules. Passing these modules to externalize_modules preserves each one's operation boundary as a named composite op that the compiler can recognize and optimize. When a PyTorch op has no built-in lowering rule, register a custom lowering function with register_torch_lowering. For compute-intensive custom operations, TorchMetalKernel lets authors write Metal kernel source; pass the resulting kernel objects to register_custom_kernels to wire them into the conversion pipeline.

Quick example

import torch
from coreai_torch import TorchConverter, get_decomp_table

model = MyModel().eval()
ep = torch.export.export(model, args=(torch.randn(1, 10),))
ep = ep.run_decompositions(get_decomp_table())
coreai_program = TorchConverter().add_exported_program(ep).to_coreai()
coreai_program.optimize()

Choosing a workflow

Use the following table to choose the conversion approach that matches the starting point.

Starting point Recommended approach
Already have a decomposed ExportedProgram TorchConverter().add_exported_program(ep).to_coreai()
Have an nn.Module, no externalization Either add_exported_program or add_pytorch_module
Have an nn.Module, need externalization add_pytorch_module(model, ..., externalize_modules=[...])

{doc}guides/externalization lets the Core AI compiler optimize submodules independently or hand them off to specialized backends. See {doc}guides/conversion-workflows for detailed code and a decision guide.

Next steps

  • New users: {doc}getting-started/installation and {doc}getting-started/quickstart cover setup and a first end-to-end bring-up.
  • Authoring Core AI models from PyTorch: {doc}guides/composite-ops covers the built-in composite op library, {doc}guides/custom-op-lowering shows how to author Core AI IR for new torch ops, and {doc}guides/custom-metal-kernels walks through authoring inline Metal GPU kernels.
  • Customizing bring-up: {doc}guides/conversion-workflows covers each bring-up workflow. {doc}guides/externalization covers preserving submodule boundaries as composite ops.
  • API reference: {doc}api/TorchConverter documents every method and parameter. {doc}api/composite-ops lists all built-in composite ops. {doc}api/TorchMetalKernel covers the Metal-kernel authoring API.

Links

Notices

PyTorch is a trademark of Meta Platforms, Inc.

:caption: Getting Started
:hidden:

getting-started/installation
getting-started/quickstart
whats-new

:caption: Guides
:hidden:

guides/conversion-workflows
guides/custom-op-lowering
guides/custom-metal-kernels
guides/composite-ops
guides/externalization

:caption: API Reference
:hidden:

api/TorchConverter
api/composite-ops
api/generate-composite-decl
api/ExternalizeSpec
api/TorchMetalKernel
api/supported-aten-ops
api/debugging

:caption: Core AI (coreai-core)
:hidden:

coreai-core/index

:caption: More
:hidden:

faq
contributing
resources
release-notes