StableHLO에서의 역동주의

동적 상태의 현재 상태는 동적 RFC에 더 공식적으로 설명되어 있습니다. 이 페이지에서는 RFC의 개요를 간략하게 제공하고 동적 프로그램과 상호작용하기 위한 중요한 API와 도구를 설명합니다.

동적 용어 및 지원 개요

먼저 이 문서에 표시되는 몇 가지 용어와 StableHLO에서의 지원에 관한 간략한 소개를 살펴보겠습니다.

동적 크기

동적 크기는 크기를 알 수 없는 모든 크기를 의미합니다. StableHLO에서는 ?, 즉 tensor<16x?xf32>를 사용하여 동적 차원을 나타냅니다.

제한된 동적

제한된 동적성은 값이 알려진 상한을 갖는 동적 차원을 의미합니다. 일반적으로 이는 실행 중에 텐서를 패딩하는 데 유용합니다. StableHLO에서는 #stablehlo.bounds를 텐서 인코딩으로 사용하여 바운드된 동적성을 나타냅니다. 즉, 16으로 바운드된 동적 차원이 하나 있고 바운드가 없는 다른 차원이 있는 순위 2 텐서는 tensor<?x?xf32, #stablehlo.bounds<16, ?>>로 나타낼 수 있습니다.

StableHLO는 제한된 동적 특성을 나타낼 수 있지만 TensorFlow에서 시작되고 PyTorch/XLA에서 일부 지원되는 제한된 프레임워크 지원이 있습니다.

무한한 역동성

이름에서 알 수 있듯이 무한한 동적성은 크기에 알려진 경계가 없는 동적 측정기준을 의미합니다. 이러한 유형의 동적 특성은 StableHLO에서 매우 일반적이며, JAX, PyTorch/XLA, TF 지원을 통해 동적 배치 크기 또는 시퀀스 길이로 모델을 내보내는 데 자주 사용됩니다.

StableHLO에서는 이러한 형태의 동적 특성에 대한 경계 인코딩을 간단히 생략합니다(즉, tensor<?x?xf32>).

모양 다형성

모양 다형성은 JAX에서 상속받은 용어입니다.

다형성을 형성하는 데는 두 가지 주요 의미가 있습니다.

  1. 프로그램의 모든 동적 요소는 입력 인수로 거슬러 올라갑니다.
  2. 모든 동적성은 텐서 모양에만 적용됩니다(데이터 종속 아님).

이 두 규칙을 사용하면 프로그램의 정적 모양이 알려진 후 동적 프로그램을 가져와 컴파일을 위해 정적 프로그램으로 완전히 세분화할 수 있습니다 ('동적 프로그램 세분화를 위한 컴파일러 패스' 참고).

일반적으로 모양 다형성은 무한한 동적 특성을 사용합니다. 알려진 인수 모양이 완전히 정적인 프로그램으로 이어질 수 있는 경우 값을 바인딩하는 방법을 추측할 필요가 없습니다.

데이터 종속 동적

데이터 종속 동적성은 텐서 내 데이터와 관련된 동적 측정기준 크기를 나타냅니다. 표준 예는 텐서 값에서 0인 모든 요소의 색인을 반환하는 nonzeros 함수입니다. 데이터를 평가하지 않고는 모양을 알 수 없지만, 바운드된 동적성을 사용하여 컴파일할 수 있는 경우가 많으며 잠재적 출력 텐서 크기에 추가 메모리를 사용합니다.

많은 데이터 종속 동적 작업은 텐서 크기의 상한이 지정되고 하드웨어가 일반적으로 텐서 패딩을 통해 이를 구현하는 제한된 동적 작업을 사용하여 모델링할 수 있습니다. 현재 PyTorch/XLA 및 TensorFlow에서는 데이터 종속 동적 동작을 일부 지원하지만 JAX는 데이터 종속 동적 동작으로 이어지는 작업을 추적하지 않습니다.

동적 측정기준이 있는 프로그램 내보내기

동적 배치 크기 또는 시퀀스 길이로 프로그램을 내보내는 방법은 StableHLO 튜토리얼을 참고하세요.

동적 프로그램을 개선하기 위한 컴파일러 패스

동적 패스 파이프라인 삭제

모양을 다듬는 데 유용한 몇 가지 패스가 있으며, 이 패스는 모두 createStablehloRemoveDynamismPipeline 패스 파이프라인에 번들로 제공됩니다.

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

동적 효과를 미세 조정하기 위한 개별 패스

개별적으로 모양 개선에 유용한 패스는 다음과 같습니다.

최신 정보와 예시는 연결된 문서를 참고하세요.

예: 동적 기능은 어떻게 유용하며 어떻게 사용할 수 있나요?

동적 기능은 다양한 용도로 사용됩니다. 여기서는 주로 모양 다형성의 일반적인 사용 사례인 유연한 내보낸 모델 표현을 만드는 데 중점을 둡니다. 이는 일반적으로 동적 배치 크기나 시퀀스 길이를 나타내는 데 사용됩니다.

정적 add_one 모델

이를 설명하기 위해 다음 간단한 add_one 모델을 사용합니다.

def add_one(x):
  return x + 1

tensor<4xf32>를 사용하여 추적하면 다음 StableHLO 프로그램이 표시됩니다.

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

이 모델은 tensor<4xf32> 모양이 있는 입력 인수에만 작동합니다. 배치 크기나 시퀀스 길이를 변경하는 경우 소스 코드를 다시 추적하고 StableHLO로 다시 낮춰야 하며, 소스 코드에 계속 액세스할 수 있다는 보장도 없습니다.

동적 add_one 모델

이때 모양 다형성 동적 특성이 사용됩니다. 대신 JAX와 PyTorch/XLA는 다음과 같이 동적 입력 모양과 일치하도록 상수를 브로드캐스트하는 동적으로 유효한 IR을 사용하여 add_one 모델을 내보낼 수 있습니다.

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

이 모델 표현은 훨씬 더 유연하며 배치 크기나 시퀀스 길이와 같은 값의 사양을 지연할 수 있습니다. 이 모델은 동적 모양을 지원하는 플랫폼 (예: AI Edge)에 배포하거나 이 문서에 언급된 동적 전달을 사용하여 미세 조정할 수 있습니다.

동적 모델 개선

예를 들어 다음 패스 순서를 사용하면 이 프로그램을 완전히 개선할 수 있습니다.

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

프로그램은 다음과 같이 점진적으로 변환됩니다.

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}