XLA (Accelerated Linear Algebra)는 선형 대수를 최적화하는 머신러닝 (ML) 컴파일러로 실행 속도와 메모리 사용량을 개선합니다. 이 페이지에서는 XLA 컴파일러의 목표와 아키텍처에 관해 간략히 설명합니다.
목표
현재 XLA는 여러 ML 프레임워크 프런트엔드 (PyTorch, TensorFlow, JAX 포함)를 지원하며 OpenXLA 프로젝트에 포함되어 있습니다. OpenXLA 프로젝트에 포함되어 있습니다. OpenXLA는 업계를 선도하는 ML 하드웨어 및 소프트웨어 조직에서 공동으로 개발한 ML용 오픈소스 컴파일러 기술 생태계입니다. OpenXLA 프로젝트가 만들어지기 전에 XLA는 TensorFlow 프로젝트 내에서 개발되었지만 기본 목표는 동일합니다.
실행 속도를 개선합니다. 하위 그래프를 컴파일하여 단기 운영 작업의 실행 시간을 줄이고 런타임의 오버헤드를 제거하고 파이프라인 작업을 융합하여 메모리 오버헤드를 줄이고 알려진 텐서 형태를 전문화하여 보다 적극적인 상수 전파를 허용합니다.
메모리 사용량을 개선합니다. 메모리 사용량을 분석하고 예약하여 여러 중간 저장소 버퍼를 제거합니다.
커스텀 작업에 대한 의존도를 줄입니다. 자동으로 융합된 하위 수준 작업의 성능을 개선하여 원래 수동으로 융합된 커스텀 작업의 성능을 개선함으로써 많은 커스텀 작업의 필요성을 없앱니다.
이동성 개선. 새로운 하드웨어를 위한 새 백엔드를 비교적 쉽게 작성할 수 있어 많은 ML 모델을 해당 하드웨어에서 수정되지 않은 상태로 실행할 수 있습니다. 이는 새 하드웨어에 개별 모놀리식 작업을 전문화하는 접근 방식과 대조됩니다. 이러한 작업을 사용하려면 모델을 재작성해야 합니다.
사용 방법
XLA 컴파일러는 StableHLO에 정의된 ML 프레임워크에서 모델 그래프를 가져와서 다양한 아키텍처에 대한 기계 명령어로 컴파일합니다. StableHLO는 ML 프레임워크와 컴파일러 간의 이동성 레이어를 제공하는 버전이 지정된 작업 세트 (HLO = 상위 수준 작업)를 정의합니다.
일반적으로 모델 그래프를 대상에 최적화된 실행 파일로 변환하는 컴파일 프로세스에는 다음 단계가 포함됩니다.
XLA는 타겟과 무관한 StableHLO 그래프에서 여러 가지 기본 제공 최적화 및 분석 패스를 실행합니다. 예를 들면 CSE, 타겟과 무관한 작업 융합, 계산을 위해 런타임 메모리를 할당하는 버퍼 분석 등을 예로 들 수 있습니다. 이 최적화 단계에서 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를 지원합니다.