XLA (Accelerated Linear Algebra)는 선형 대수를 최적화하는 머신러닝 (ML) 컴파일러로, 실행 속도와 메모리 사용량을 개선합니다. 이 페이지에서는 XLA 컴파일러의 목표와 아키텍처를 간략하게 설명합니다.
목표
현재 XLA는 여러 ML 프레임워크 프런트엔드 (PyTorch, TensorFlow, JAX 포함)를 지원하며 OpenXLA 프로젝트의 일부입니다. OpenXLA 프로젝트는 선도적인 ML 하드웨어 및 소프트웨어 조직이 공동으로 개발한 ML용 오픈소스 컴파일러 기술 생태계입니다. OpenXLA 프로젝트가 만들어지기 전에 XLA는 TensorFlow 프로젝트 내에서 개발되었지만 기본 목표는 동일합니다.
실행 속도 개선. 하위 그래프를 컴파일하여 단기 작업의 실행 시간을 줄이고 런타임의 오버헤드를 제거하고, 파이프라인 작업을 융합하여 메모리 오버헤드를 줄이고, 알려진 텐서 형태를 전문화하여 보다 적극적인 상수 전파를 허용합니다.
메모리 사용량을 개선합니다. 메모리 사용량을 분석하고 예약하여 여러 중간 저장소 버퍼를 제거합니다.
커스텀 작업에 대한 의존도를 줄입니다. 원래 수동으로 융합된 커스텀 작업의 성능과 일치하도록 자동으로 융합된 하위 수준 작업의 성능을 개선하여 많은 커스텀 작업의 필요성을 없애세요.
이동성 개선. 새로운 하드웨어를 위한 새 백엔드를 비교적 쉽게 작성하여 많은 ML 모델을 해당 하드웨어에서 수정되지 않은 상태로 실행할 수 있도록 합니다. 이는 새 하드웨어에 개별 모놀리식 작업을 특화하여 해당 작업을 사용하려면 모델을 재작성해야 하는 접근 방식과 대조됩니다.
작동 방식
XLA 컴파일러는 StableHLO에 정의된 ML 프레임워크의 모델 그래프를 가져와서 다양한 아키텍처의 머신 명령으로 컴파일합니다. StableHLO는 ML 프레임워크와 컴파일러 간의 이동성 레이어를 제공하는 버전이 지정된 작업 세트 (HLO = 상위 수준 작업)를 정의합니다.
일반적으로 모델 그래프를 대상에 최적화된 실행 파일로 변환하는 컴파일 프로세스에는 다음 단계가 포함됩니다.
XLA는 CSE, 대상 독립적 작업 융합, 계산에 런타임 메모리를 할당하기 위한 버퍼 분석과 같이 타겟에 독립적인 StableHLO 그래프에 대해 여러 가지 기본 제공 최적화 및 분석 패스를 실행합니다. 또한 이 최적화 단계에서 XLA는 StableHLO 언어를 내부 HLO 언어로 변환합니다.
XLA는 HLO 수준의 추가 최적화를 위해 HLO 계산을 백엔드로 보냅니다. 이번에는 대상별 정보와 요구사항을 염두에 둡니다. 예를 들어 GPU 백엔드는 특히 GPU 프로그래밍 모델에 유용한 작업 퓨전을 실행하고 계산을 스트림으로 분할하는 방법을 결정할 수 있습니다. 이 단계에서 백엔드는 특정 작업 또는 작업의 조합을 최적화된 라이브러리 호출에 패턴 일치시킬 수도 있습니다.
그러면 백엔드에서 타겟별 코드 생성을 수행합니다. XLA에 포함된 CPU 및 GPU 백엔드는 하위 수준 IR, 최적화, 코드 생성에 LLVM을 사용합니다. 이러한 백엔드는 HLO 계산을 효율적으로 나타내는 데 필요한 LLVM IR을 방출한 다음 LLVM을 호출하여 이 LLVM IR에서 네이티브 코드를 내보냅니다.
이 프로세스 내에서 XLA 컴파일러는 대체 백엔드에서 슬롯 몇 가지 새로운 HW 아키텍처를 타겟팅하기가 쉽다는 점에서 모듈식입니다. GPU 백엔드는 현재 LLVM NVPTX 백엔드를 통해 NVIDIA GPU를 지원합니다. CPU 백엔드는 여러 CPU ISA를 지원합니다.