StableHLO에서의 역동주의

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

동적 용어 및 지원 개요

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

동적 크기

동적 측정기준은 측정기준 크기를 알 수 없는 측정기준을 의미합니다. StableHLO에서는 ?(즉, tensor<16x?xf32>)를 사용하여 동적 측정기준을 나타냅니다.

제한된 동적

Bounded dynamism은 값이 알려진 상한값이 있는 동적 측정기준을 나타냅니다. 일반적으로 실행 중에 텐서를 패딩하는 데 유용합니다. StableHLO에서는 #stablehlo.bounds를 텐서 인코딩으로 사용하여 제한된 동적을 나타냅니다. 즉, 한 동적 측정기준이 16으로 제한되고 다른 측정기준은 제한되지 않은 2차 텐서는 tensor<?x?xf32, #stablehlo.bounds<16, ?>>로 나타낼 수 있습니다.

StableHLO는 제한된 동적성을 표현할 수 있지만, TensorFlow에서 시작된 제한적인 프레임워크 지원이 있으며 PyTorch/XLA에서 일부 지원됩니다.

무제한 동적

이름에서 알 수 있듯이 무제한 동적은 크기에 알려진 경계가 없는 동적 측정기준을 나타냅니다. 이러한 유형의 동적은 JAX, PyTorch/XLA, TF 지원이 포함된 StableHLO에서 매우 일반적이며 동적 일괄 크기 또는 시퀀스 길이가 있는 모델을 내보내는 데 자주 사용됩니다.

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>
}