XLA의 에일리어싱

이 문서에서는 XLA 프로그램을 빌드할 때 입력 버퍼와 출력 버퍼 간에 앨리어싱을 지정할 수 있는 XLA 앨리어싱 API에 관해 설명합니다.

컴파일 시간에 별칭 정의

예를 들어 단순히 입력에 1를 추가하는 간단한 HLO 모듈을 생각해 보세요.

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

이 모듈은 두 개의 4바이트 버퍼를 할당합니다. 하나는 입력 %p용, 다른 하나는 출력 %out용입니다.

하지만 그 자리에서 업데이트를 수행하는 것이 바람직한 경우가 많습니다. 예를 들어 표현식을 생성하는 프런트엔드에서 계산 후에 입력 변수가 더 이상 존재하지 않는 경우 p++의 증분에서와 같이 더 이상 활성 상태가 아닙니다.

이러한 업데이트를 효율적으로 수행하려면 입력 앨리어싱을 지정하면 됩니다.

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

이 형식은 전체 출력 ({}로 표시됨)이 입력 매개변수 0으로 별칭이 지정되도록 지정합니다.

프로그래매틱 방식으로 별칭을 지정하려면 XlaBuilder::SetUpAlias API를 참고하세요.

런타임 시 앨리어싱 정의

이전 단계에서 정의한 앨리어싱은 컴파일 도중에 지정됩니다. 실행 중에 LocalClient::RunAsync API를 사용하여 버퍼 제공 여부를 선택할 수 있습니다.

프로그램의 입력 버퍼는 ExecutionInput에 래핑되며, 이 안에 MaybeOwningDeviceMemory 트리가 포함됩니다. 메모리가 소유로 지정된 경우 (버퍼의 소유권이 XLA 런타임에 전달됨) 버퍼가 실제로 제공되고 컴파일 시간 별칭 API에서 요청한 대로 업데이트가 실행됩니다.

하지만 컴파일 시간에 별칭이 지정된 버퍼가 런타임에 제공되지 않으면 복사 보호가 시작됩니다. 즉, 추가 출력 버퍼 O가 할당되고 별칭이 지정되어야 했던 입력 버퍼 P의 콘텐츠가 O에 복사됩니다. 따라서 마치 버퍼 O가 런타임에 제공된 것처럼 효과적으로 프로그램을 실행할 수 있습니다.