ARMONIC: Autonomous Optimization for Arm64 Cloud AI

• 

  Work Flow Blocks

• 

  REPO-structure

• 

  AWS-instance

• 

  ARM configuration(architecture and model)

• 

  working -ssh (cloud)

• 

  results-(complete run.....moderate load)

• 

  results-(complete run....heavy load)

Inspiration

We kept hitting the same wall: Python AI inference pipelines running on AWS Graviton3 were painfully slow, and nobody could tell us why. Generic profilers like cProfile or PyTorch Profiler would show a flat "87% time in function X" — but they were blind to the silicon underneath. They couldn't see Neoverse V1 memory stalls, L1D cache refills, or branch mispredictions. So developers did what developers do: throw @njit at everything, copy fixes from StackOverflow, and hope for the best. Hours wasted. Code broken. Performance barely moved.

We asked a simple question: what if the machine could profile itself, understand its own bottleneck, and fix itself? Not a suggestion. Not a report. An autonomous, closed-loop agent that goes from raw hardware counters to a validated, committed patch — without a human touching the code.

That question became ARMONIC.

---

What it does

ARMONIC is a fully autonomous performance optimization agent for Python AI workloads on Arm64 cloud infrastructure (AWS Graviton3, Ampere Altra, Neoverse V1/V2).

The pipeline is simple to invoke — make install && make run — but what happens underneath is anything but:

1. Profile — apx trace reads real Neoverse hardware counters: CPU_CYCLES, INST_RETIRED, L1D_CACHE_REFILL, BR_MIS_PRED, MEM_ACCESS.
2. Score — A weighted bottleneck score $B_s$ tuned for Neoverse V1 characteristics identifies whether the workload is CPU-bound, memory-bound, or cache-bound.
3. Reason — Telemetry is sent via MCP JSON-RPC 2.0 to Gemini, conditioned on Arm64-specific optimization patterns (NEON vectorization, fastmath, cache=True, orjson).
4. Patch — The LLM generates a targeted code change for the top hotspot function.
5. Validate — Syntax check (py_compile), AST check (ast.parse), import resolution, and score regression gating — if the optimized score isn't better than baseline, the patch is rejected.
6. Commit — Only passing patches are committed to an isolated git branch (armonic/auto-refactor-<timestamp>). Original code on main is never touched.

The result: a ~98.7% wall-time reduction on a real AI inference workload, validated against 24 APX samples, with zero manual code changes.

---

How we built it

Stack: Python 3.10+, Arm Performix (APX), Google Gemini API, GitPython, Docker.

• src/profiling/apx_wrapper.py — Interfaces with apx trace to collect hardware counters. Falls back to cProfile when APX is unavailable (cross-platform dev).
• src/scoring/bottleneck.py — Computes $B_s = 0.40C_s + 0.30M_s + 0.15L_s + 0.10I_s + 0.05P_s$, with weights tuned for Neoverse V1's memory-bound AI inference behavior.
• src/mcp_server/server.py — Exposes schema-validated telemetry over JSON-RPC 2.0, bridging raw silicon data to the reasoning layer.
• src/refactor_engine/llm_agent.py — Constructs a system prompt that conditions Gemini on Arm64 Python patterns. Dynamically targets the top hotspot function from telemetry — never hardcoded.
• armonic/run.py — Orchestrates the full loop: baseline profile → LLM advisory → patch application → re-profile → score comparison → accept/reject.
• Validation layer — py_compile + ast.parse + import smoke tests + score regression gating. Four gates, zero tolerance for bad patches.
• Git isolation — Every successful patch is committed to a timestamped branch. Rollback is one git checkout away.
• Docker + Makefile — One-command reproducible runs on any Arm64 Linux instance.

---

Challenges we ran into

1. Generic patches vs. Arm-specific claims Our first iterations produced @njit patches that worked on x86 just as well as Arm64. An Arm Staff Engineer would call that out immediately. We solved this by hardcoding Arm64-specific rules into the LLM prompt — fastmath=True for NEON, cache=True for Graviton recompilation overhead, orjson for Arm64-optimized JSON, and explicit bans on x86 libraries like Intel MKL. We also tuned the bottleneck scoring weights to reflect Neoverse V1's actual memory-bound behavior.

2. APX integration depth apx trace outputs raw hardware counter data that isn't trivial to parse. Building a robust wrapper that handles missing counters, malformed output, and graceful fallback to cProfile took multiple iterations. The MCP server abstraction was our answer — schema-validated telemetry that decouples the profiler from the reasoning engine.

3. Score validation credibility Our README claimed "rejected if opt_score >= base_score" but the orchestrator wasn't actually doing the re-profile comparison. We caught this during code review and added a full re-profiling loop in armonic/run.py that re-runs the workload, recalculates $B_s$, and reverts the file if the patch regresses.

4. LLM output reliability Gemini occasionally returned malformed JSON or omitted required keys. We added JSON schema validation with required key checks and a fallback retry mechanism.

5. Cloud instance variance Wall times on AWS Graviton3 fluctuate ±3% between runs due to scheduling noise. We moved from presenting single-run exact numbers to statistical ranges (~17.3s → ~0.22s, ~98.7%) to maintain scientific honesty.

---

Accomplishments that we're proud of

Accomplishments that we're proud of

• ~94% wall-time reduction on the ai_inference workload, from ~42.5s baseline to ~2.7s optimized, autonomously discovered and validated on AWS EC2 c7g.xlarge (Graviton3, Neoverse V1).
• Fully autonomous closed-loop pipeline — profile → score → LLM → patch → validate → commit — with zero human code changes, running end-to-end in a single make run.
• Production safety by design — deterministic correctness validation (SHA-256 hash matching), wall-time regression gating, and automatic git isolation on timestamped branches. A bad patch literally cannot reach main.
• Neoverse-tuned bottleneck scoring — weighted formula derived from actual Arm64 microarchitecture behavior (CPU cycles, memory stalls, cache refills), with a wall-time validation gate for environments without APX.
• One-command setup — make install && make run on any Arm64 Linux instance, with Docker for full reproducibility.
• Graceful dual-mode profiling — attempts Arm Performix APX first for hardware-accurate counters, then seamlessly falls back to cProfile with tottime-sorted hotspot detection. The agent works everywhere, shines on Arm.

---

What we learned

• Hardware-aware profiling is non-negotiable. Generic profilers tell you where time is spent. APX tells you why — cache misses vs. branch mispredictions vs. memory stalls. That "why" is what makes optimization possible.
• LLMs need constraints, not freedom. An unconstrained LLM will suggest numpy.vectorize or x86 intrinsics. A tightly conditioned prompt with Arm64-specific rules produces genuinely useful patches.
• Validation is harder than generation. Generating a patch is easy. Proving it doesn't break anything, doesn't regress performance, and doesn't introduce bad imports — that's where the engineering lives.
• Git is the ultimate safety net. Automated branching means you can be aggressive with optimization and conservative with risk at the same time.

---

What's next for ARMONIC

• Deeper Arm-specific optimizations — SVE intrinsics suggestions, cache-line-aware loop restructuring (64-byte Neoverse lines), prange tuning for Graviton core topology.
• CI/CD integration — GitHub Actions runner that auto-runs ARMONIC on pull requests, proposing performance patches as automated PR comments.
• Broader workload support — PyTorch inference graphs, ONNX runtime pipelines, agentic AI frameworks beyond pure Python.
• Local LLM fallback — Reduce API dependency by supporting quantized local models (Llama 3, CodeLlama) for offline optimization.
• Multi-objective optimization — Balance speed, memory, and power on Arm64 edge devices, not just cloud instances.

Built With

• amazon-web-services
• arm64
• cloud
• docker
• git
• json-rpc
• linux
• llm
• neoverse
• numba
• python
• ubuntu