Axon: 텐서 프로그램을 위한 합성 슈퍼옵티마이저 (Synthesizing Superoptimizer)

AI 가속기(AI accelerator)를 위한 고성능 커널(kernel)을 작성하는 것은 타일링 (tiling), 명령어 선택 (instruction selection), 데이터 레이아웃 (data layout), 그리고 연산자 융합 (operator fusion)에 대한 깊은 전문 지식을 요구하며, 이는 프로그래머에게 상당한 부담을 줍니다. 본 논문에서 우리는 타일 기반 AI 가속기 프로그램에 초점을 맞추어, 텐서 프로그램(tensor programs)을 위한 합성 슈퍼옵티마이저 (synthesizing superoptimizer)인 Axon을 제시합니다. Axon은 프로그램 합성 (program synthesis)을 사용하여 의미론적 명세 (semantics specifications)로부터 타겟 명령어 (target instructions)를 자동으로 생성하며, 경험적으로 가장 성능이 좋은 커널을 선택하기 위해 의미론적으로 동일한 프로그램 변형들을 탐색합니다. Axon은 계산 그래프 (computation graphs)를 통해 연산자 (operators)를 전파함으로써 대수적 변환 (algebraic transformations)을 발견하며, 수동으로 작성된 재작성 규칙 (rewrite rules) 없이도 모든 변환이 의미론을 보존함을 보장하기 위해 무제한 텐서 (unbounded tensors)에 대한 SMT를 사용합니다. 그 후 텐서 연산 (tensor operations)을 타겟 ISA 명령어로 낮추고 (lowers), 하드웨어 기술 (hardware descriptions)에 의해 제약되는 타일링 구성 (tiling configurations)을 탐색하며, 메모리 트래픽 (memory traffic)을 최소화하기 위해 연산자와 명령어를 융합합니다.