XLA용 새 백엔드 개발

이 가이드는 XLA가 하드웨어를 효율적으로 타겟팅하는 프로그램을 출력하도록 하려는 시스템 엔지니어를 대상으로 합니다. 이 가이드는 단계별 설명이 아니며 LLVM, Bazel, XLA에 대한 지식이 있다고 가정합니다.

XLA는 새 아키텍처 또는 액셀러레이터가 XLA의 ML 프로그램 출력을 실행하기 위한 백엔드를 만들기 위해 구현할 수 있는 추상 인터페이스를 제공합니다. XLA 재타겟팅은 새로운 하드웨어에 대해 PyTorch 또는 TensorFlow와 같은 프런트엔드 프레임워크의 모든 기존 작업을 구현하는 것보다 훨씬 간단하고 확장성이 높습니다.

대부분의 구현은 다음 시나리오 중 하나에 속합니다.

1. 기존 LLVM 백엔드 유무와 관계없이 XLA에서 아직 공식적으로 지원되지 않는 기존 CPU 아키텍처
2. 기존 LLVM 백엔드가 있는 비 CPU 유사 하드웨어
3. 기존 LLVM 백엔드가 없는 비 CPU 유사 하드웨어

시나리오 1: XLA에서 아직 공식적으로 지원되지 않는 기존 CPU 아키텍처

이 시나리오에서는 먼저 기존 XLA CPU 백엔드를 살펴봅니다. XLA를 사용하면 LLVM을 사용하여 다양한 CPU를 쉽게 타겟팅할 수 있습니다. CPU의 XLA 백엔드 간의 주요 차이점은 LLVM에서 생성되는 코드이기 때문입니다.

하드웨어 공급업체가 자사 하드웨어를 위한 LLVM 백엔드를 보유하고 있다면 백엔드를 XLA로 빌드된 LLVM과 간단하게 연결할 수 있습니다. JIT 모드에서 XLA CPU 백엔드는 호스트 CPU의 코드를 내보냅니다. AOT 컴파일의 경우 xla::AotCompilationOptions는 대상 아키텍처를 구성하기 위한 LLVM 세 가지를 제공할 수 있습니다.

기존 LLVM 백엔드는 없지만 다른 종류의 코드 생성기가 있는 경우 대부분의 기존 CPU 백엔드를 재사용할 수 있습니다.

시나리오 2: 기존 LLVM 백엔드가 있는 CPU와 유사하지 않은 하드웨어

xla::CPUCompiler 및 xla::GPUCompiler 클래스는 이미 LLVM IR을 방출하기 때문에 기존 xla::CPUCompiler 및 xla::GPUCompiler 클래스에서 새로운 xla::Compiler 구현을 모델링할 수 있습니다. 하드웨어의 특성에 따라 LLVM IR 생성의 많은 측면을 변경해야 할 수 있지만 많은 코드를 기존 백엔드와 공유할 수 있습니다.

좋은 예는 XLA의 GPU 백엔드입니다. GPU 백엔드는 CPU와 유사한 ISA를 타겟팅하므로 코드 생성의 일부 측면은 GPU 도메인에 고유합니다. 업스트림 LLVM 백엔드가 있는 Hexagon과 같은 DSP와 같은 다른 종류의 하드웨어는 LLVM IR 방출 로직의 일부를 재사용할 수 있지만 다른 부분은 고유합니다.

시나리오 3: 기존 LLVM 백엔드가 없는 CPU와 유사한 하드웨어

LLVM을 활용할 수 없는 경우 최선의 옵션은 원하는 하드웨어에 XLA의 새 백엔드를 구현하는 것입니다. 이 옵션은 가장 큰 노력이 필요합니다. 구현해야 하는 클래스는 다음과 같습니다.

• StreamExecutor: 대부분의 기기에서 모든 StreamExecutor 메서드가 필요한 것은 아닙니다. 자세한 내용은 기존 StreamExecutor 구현을 참고하세요.
• xla::Compiler: 이 클래스는 HLO 계산의 컴파일을 xla::Executable로 캡슐화합니다.
• xla::Executable: 이 클래스는 플랫폼에서 컴파일된 계산을 실행하는 데 사용됩니다.
• xla::TransferManager: 이 클래스를 사용하면 백엔드가 지정된 기기 메모리 핸들에서 XLA 리터럴 데이터를 구성하기 위한 플랫폼별 메커니즘을 제공할 수 있습니다. 즉, 호스트에서 기기로, 그리고 그 반대로 데이터 전송을 캡슐화하는 데 도움이 됩니다.