이 문서에서는 XLA High Level Optimizer (HLO) 모듈이 초기 상태에서 최종 실행 파일로 이동하는 과정을 간략하게 설명합니다. 때로는 '모듈'을 생략하고 'HLO'라고만 언급합니다.
사전 최적화 HLO
사전 최적화 HLO 모듈부터 시작합니다. 사전 최적화 HLO에는 fusion 또는 bitcast과 같이 XLA 내부로 간주되는 작업 (ops)이 포함되지 않습니다. 이 단계에서는 작업에 레이아웃이 없거나 레이아웃이 있어도 무시됩니다. 사전 최적화 HLO는 일반적으로 TensorFlow 및 JAX와 같은 상위 수준 프레임워크에서 생성됩니다. XLA 플래그 -xla_dump_to를 사용하면 사전 최적화 HLO가 파일 이름 접미사가 'before_optimizations.txt'인 파일에 덤프됩니다.
HLO 모듈 최적화
XLA:GPU 파이프라인은 일련의 패스를 실행하여 최적화 전 HLO를 최적화된 HLO로 변환합니다. 패스는 의미상으로 그룹화할 수 있으며 다음 순서로 실행할 수 있습니다.
샤딩 관련 패스
여기에는 Shardy Partitioner와 같은 패스나 SPMD 샤딩용 패스가 포함됩니다.
최적화 패스
여기에는 합법화 패스와 단순화 패스가 모두 포함될 수 있습니다.
집단 최적화 패스
최적화 패스와 비슷하지만 집단 작업에 중점을 둡니다.
레이아웃 할당 패스
각 HLO 작업에는 명령어 모양의 일부인 레이아웃이 할당됩니다. 레이아웃은 텐서가 메모리에 물리적으로 배치되는 방식을 제어합니다.
레이아웃이 있는 도형의 예:
f32[10,20,30]{2,0,1}
요소 유형 뒤에는 도형의 논리적 차원이 나오고 그 뒤에는 레이아웃 순열이 작은 순서에서 큰 순서로 나옵니다. 이 예에서 가장 작은 측정기준은 30이고 두 번째로 작은 측정기준은 10이며 가장 큰 측정기준은 20입니다.
레이아웃 할당의 목표는 그리디 전략을 사용하여 필요한 실제 전치의 수를 최소화하는 것입니다. 특정 레이아웃 제약 조건 (예: cuDNN/cuBLAS 라이브러리는 연속된 차원을 예상함)으로 시작하여 레이아웃을 HLO 그래프의 '아래'로 전파한 다음 '위'로 전파합니다. 레이아웃 전파가 끝나면 일부 명령어의 레이아웃이 충돌할 수 있습니다. 하나는 피연산자에서 전파되고 하나는 사용자로부터 전파됩니다. 이 충돌을 해결하기 위해 피연산자 레이아웃에서 명령어 레이아웃으로 레이아웃을 변경하는 copy HLO 명령어가 삽입됩니다.
레이아웃 정규화 패스
물리적 모양을 파악하기가 다소 어렵다는 점을 고려할 때 레이아웃 정규화는 기본 레이아웃 {rank-1, rank-2, …, 0}를 사용하도록 모양을 다시 작성하려고 시도합니다. 위의 예에서 정규화된 모양은 f32[20,10,30]{2,1,0}입니다. 레이아웃을 변경하는 복사 작업은 transpose와 bitcast의 조합으로 다시 작성됩니다. 현재 모든 작업을 정규화할 수 없으므로 기본 레이아웃이 아닌 레이아웃을 사용할 수 있는 작업이 여전히 있습니다. 특히 gather 및 dot이 그렇습니다. 정규화된 작업과 정규화되지 않은 작업 사이의 경계에는 전치를 나타내는 bitcast 작업이 있습니다. 즉, 물리적으로 아무 작업도 하지 않는 레이아웃이 할당된 전치입니다.
레이아웃 정규화는 일부 암시적 전치를 명시적으로 만드는데, 이는 코드 생성기가 전용 이미터로 명시적 전치를 처리할 수 있기 때문에 중요합니다. 예를 들어 reshape는 피연산자와 결과 간에 다른 물리적 레이아웃을 가질 수 있습니다 (예: 순위가 다르기 때문). 레이아웃 정규화 패스의 일부로 실행되는 ReshapeDecomposer 패스는 재형상을 transpose, 재형상 bitcast, transpose의 시퀀스로 변환합니다.
레이아웃 할당 최적화 패스 후
여기서 가장 중요한 패스는 Triton 융합 (GEMM 융합 + Softmax/Layernorm 융합) 또는 라이브러리 호출로 다시 작성하는 것입니다. 또한 이 단계에서는 자동 튜닝이 실행됩니다. 여기에서 XLA는 다양한 이미터 중에서 선택하고, 컨볼루션이나 점에 가장 적합한 알고리즘을 선택하고, Triton 이미터에서 처리하는 융합에 가장 적합한 타일링을 찾습니다.
Fusion 패스
두 가지 주요 패스는 PriorityFusion 및 Multi-Output 융합입니다.
PriorityFusion에서는 비용 모델에 따라 융합을 형성합니다. 융합 시 작업이 모든 사용자로 융합될 수 있는 경우 여러 사용자로 작업을 복제할 수 있습니다. 가능한 경우 기존 Triton Softmax 융합을 확장하는 것도 허용됩니다.
Multi-Output 융합은 피연산자를 공유하는 작업/융합을 융합할 수 있는 별도의 패스입니다. 또한 추가 출력을 추가하여 피연산자/피연산자 융합을 중복 없이 사용자에게 융합할 수 있으므로 융합할 작업의 다른 사용자를 이러한 출력으로 리디렉션할 수 있습니다. 이 패스는 HLO 그래프에 사이클이 도입되지 않도록 주의해야 합니다.
다중 출력 융합 후 공통 하위 표현식 제거 (HloCSE 패스)가 실행되어 이전에 중복된 작업이 동일한 융합으로 끝난 경우 다시 병합될 수 있습니다.
여러 융합 후 패스
집단과 관련된 여러 패스 (예: 비동기로 전환하거나 집단의 특정 상대적 순서를 강제 적용)
마지막으로 CopyInsertion를 실행하여 복사본을 추가합니다. 이렇게 하면 아직 다른 곳에서 필요한 데이터가 인플레이스 작업으로 덮어쓰이지 않습니다.
최적화가 끝나면 -xla_dump_to 플래그를 사용하는 경우 최적화된 HLO가 파일 이름 접미사가 'after_optimizations.txt'인 파일에 덤프됩니다. HloModule을 실제로 변경하는 중간 패스 후에 HLO를 덤프하려면 -xla_dump_hlo_pass_re=.* 플래그 (또는 특정 패스로 제한하는 정규 표현식)를 사용하면 됩니다.
예약
일정이 없는 HLO 모듈은 작업이 처리되는 순서에 있어 어느 정도의 자유도가 있습니다. 피연산자/결과 관계와 제어 종속성을 따르는 모든 위상 정렬이 유효합니다. 일정 예약은 사용할 특정 순서를 결정합니다. 이 단계에서 가장 우려되는 점은 텐서의 수명에 따라 달라지는 최대 메모리 소비입니다. 초기 단계에서는 다양한 스케줄러 알고리즘을 시도하고 최대 메모리 소비를 최소화하는 일정을 선택합니다. 이 시점에서는 아직 실제 버퍼를 사용하지 않으며('버퍼 할당'에서 사용) 메모리 사용량을 시뮬레이션합니다.
그런 다음 LatencyHidingScheduler 패스가 실행되고 컴퓨팅-커뮤니케이션 오버랩을 최대화하려고 시도합니다. 하지만 이렇게 하면 메모리 사용량이 다시 증가할 수 있습니다.
마지막으로, 최대 메모리 소비량이 사용 가능한 메모리 양보다 높은 경우 HloRematerialization를 실행합니다. 이 패스는 성능을 희생하여 메모리 사용량을 줄이려고 시도합니다. 예를 들어 일부 융합이 분할되고 일부 작업이 중복되어 버퍼 수명이 짧아질 수 있습니다. 재구체화가 발생하는 경우 모델 측의 메모리 요구사항을 줄이는 방법을 조사하는 것이 유용할 수 있습니다 (예: 더 작은 배치 크기 사용).
버퍼 할당
LLVM IR로 낮추기 직전에 HLO 그래프의 각 명령어에 버퍼 슬라이스를 할당하는 버퍼 할당 패스를 실행합니다. 버퍼 할당은 여러 단계로 실행됩니다.
HloDataflowAnalysis는HloValues(기본적으로 논리적 버퍼)를 명령어에 할당합니다. 인플레이스 작업의 경우 피연산자의HloValue를 재사용할 수 있습니다. 작업은 두 개 이상의HloValue를 정의할 수 있습니다 (예: 튜플 결과 모양 사용).HloAliasAnalysis는 앨리어싱 작업을 위해 버퍼를 결합하려고 시도하고HloValue에서HloBuffer로의 매핑을 계산합니다.BufferAssignment는 수명이 겹치는 서로 다른HloBuffers에 동일한 버퍼 슬라이스가 사용되지 않도록HloBuffers을 큰 버퍼 내의 버퍼 슬라이스에 매핑합니다. 별칭이 지정될 수 있는 작업의 경우 약간의 중복이 있어도 됩니다 (한HloBuffer의 종료 시간이 다른HloBuffer의 시작 시간과 일치할 수 있음).-xla_dump_to플래그를 사용하면 버퍼 할당에 관한 일부 정보가 이름 접미사가 'after_optimizations-buffer-assignment.txt'인 파일에 덤프됩니다.
thunk
HLO 그래프가 최적화되고 예약된 후 특정 백엔드 (CPU 또는 GPU)의 청크 선형 시퀀스로 낮아집니다.
XLA에서 Thunk는 런타임이 실행하는 독립 실행형 작업 단위의 추상화입니다. 컴파일된 커널 실행, 특정 작업, 라이브러리 호출, 제어 흐름 구조, 집단 통신 등이 될 수 있습니다. Thunk Sequence는 특정 백엔드의 전체 실행 파일을 나타냅니다.
Thunk 방출
예약된 HLO 계산을 청크 시퀀스로 변환하는 프로세스를 '청크 내보내기'라고 합니다. 이는 각 백엔드의 전용 이미터 클래스에 의해 처리됩니다.
GPU 백엔드의 경우 IrEmitterUnnested에서 처리합니다.
EmitHloComputation는 계산에서 예약된 HLO 명령어 목록을 반복하고 전문화된 Emit... 메서드 (예:
EmitFusion, EmitConvolutionThunk, EmitWhile)를 호출합니다. 이러한 각 메서드는 적절한 청크 객체를 구성하고 청크 시퀀스에 추가합니다.
CPU 백엔드의 경우 ThunkEmitter가 이 역할을 수행하며 유사한 방식으로 구성됩니다. 최종 ThunkSequence가 CpuExecutable에 삽입됩니다.
HLO 모듈의 항목 계산에 있는 각 명령어는 최종 청크 시퀀스에서 0개 (kTuple, kConstant 등), 1개 또는 여러 개의 청크 (예: 정렬 명령어)에 해당할 수 있습니다.
명령어 버퍼: GPU에서 실행 최적화
최신 GPU 하드웨어를 사용하면 GPU 작업 (커널 실행, 메모리 복사 등) 시퀀스를 한 번 기록한 후 CPU 오버헤드를 최소화하여 시퀀스를 여러 번 재생할 수 있습니다. 이는 특히 작고 빠르게 실행되는 커널이 많은 워크로드의 경우 중요한 성능 최적화입니다. XLA는 명령어 버퍼를 CUDA 그래프 또는 HIP 그래프의 추상화로 사용합니다. 핵심 인터페이스는 GpuCommandBuffer에 정의되어 있습니다.
명령어 버퍼는 CommandBufferThunk에 의해 썽크 시퀀스에 표시됩니다.
이미터는 HLO 명령에서 이 썽크를 직접 생성하지 않습니다. 대신 ThunkSequence 자체에서 실행되는 CommandBufferConversionPass에 의해 실행됩니다.
이 패스는 호환되는 청크의 연속된 하위 시퀀스 (예: KernelThunk 및 GemmThunk 시리즈)를 식별합니다. 그런 다음 발견된 하위 시퀀스를 단일 CommandBufferThunk로 바꿉니다. 새로운 썽크는 원래 썽크의 로직을 경량 CommandBufferCmd 객체 목록으로 캡슐화합니다.
CommandBufferThunk가 특정 GPU 스트림에서 처음 실행되면 명령어 시퀀스를 하드웨어 명령어 버퍼에 '기록'합니다. 이후 모든 실행에서 기록된 시퀀스를 '재생'하도록 GPU에 단일 명령어를 실행합니다. 이렇게 하면 각 개별 커널을 실행할 때 발생하는 CPU 오버헤드가 방지됩니다.
실행 가능
XLA 컴파일 파이프라인의 최종 제품은 자체 포함된 플랫폼별 실행 파일입니다. 이 객체는 CPU나 GPU와 같은 타겟 기기에서 컴파일된 프로그램을 실행하는 데 필요한 모든 정보를 캡슐화합니다. 컴파일러와 런타임 사이의 브리지입니다. PJRT와 같은 최신 런타임은 약간 더 높은 수준의 추상화를 사용하지만 (PjRtExecutable 참고) 궁극적으로 백엔드별 실행 파일을 래핑합니다.
Executable에는 컴파일 중에 생성된 여러 주요 정보가 포함되어 있습니다. 정확한 콘텐츠는 백엔드에 따라 다르지만 일반적으로 다음이 포함됩니다.
컴파일된 코드: 기기에서 실행되는 하위 수준 머신 코드입니다. CPU의 경우 일반적으로 하나 이상의 객체 파일입니다. GPU의 경우 이는 런타임에 GPU에 로드되는 PTX 또는 HSACO 형식의 컴파일된 기기 코드입니다.
실행 계획 (ThunkSequence): 런타임 로직의 핵심입니다. 이는 Thunk 객체의 선형 시퀀스입니다. 각 썽크는 커널 실행, 라이브러리 함수 (예: cuBLAS) 호출, 제어 흐름 처리와 같은 단일 작업 단위를 나타냅니다. 런타임은 이 시퀀스를 반복하여 프로그램을 실행합니다.
메모리 레이아웃 (BufferAssignment): BufferAssigner에 의해 생성되는 이 중요한 메타데이터는 계산의 전체 메모리 레이아웃을 설명합니다. 모든 버퍼의 크기와 매개변수, 출력, 임시 값에 메모리가 할당되고 재사용되는 방식을 지정합니다. 런타임은 이를 사용하여 기기 메모리를 할당하고 각 썽크에 올바른 포인터를 전달합니다.
(선택사항) HLO 모듈: 디버깅 및 프로파일링을 위해 실행 파일은 컴파일된 최종 최적화 HloModule에 대한 참조를 유지하는 경우가 많습니다.
최종 실행 파일의 생성은 각 특정 백엔드에 대해 컴파일러에 의해 조정됩니다. 컴파일러 구현의 RunBackend 메서드는 컴파일 프로세스의 마지막 단계로, 컴파일된 모든 아티팩트를 실행 파일 객체로 패키징합니다.
GpuCompiler와 CpuCompiler는 각각 GPU와 CPU를 타겟팅합니다.
사용자가 실행 파일에서 Execute...를 호출하면 런타임은 BufferAssignment를 사용하여 메모리를 할당한 다음 ThunkSequence를 호출하여 컴파일된 코드를 사용하여 기기에서 작업을 실행합니다.