XLA의 앨리어싱

이 문서에서는 XLA 프로그램을 빌드할 때 입력 및 출력 버퍼 간의 앨리어싱.

컴파일 시간에 앨리어싱 정의

예를 들어 단순히 1를 입력에 추가하는 간단한 HLO 모듈을 가정해 보겠습니다.

HloModule increment

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

이 모듈은 4바이트 버퍼 2개를 할당합니다. 하나는 입력 %p용이고 다른 하나는 입력 %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 런타임에 전달됨) 버퍼가 실제로 제공되고 업데이트가 제자리에서 실행됩니다. compile-time aliasing API에 의해 요청됩니다.

그러나 컴파일 시 앨리어싱된 버퍼가 제공되지 않는 런타임 시 복사 방지가 시작됩니다. 추가 출력 버퍼 O가 할당됩니다. 별칭이 지정되도록 의도된 입력 버퍼 P의 콘텐츠가 복사됨 O로의 진입 (따라서 버퍼 O가 제공됩니다.