OpenXLA Project

`-chlo-legalize-to-stablehlo`

CHLO ops 흐름에서 StableHLO 및 Shape ops로 합법화

`-shape-legalize-to-stablehlo`

모양 관련 작업을 StableHLO로 합법화

모양 관련 작업을 StableHLO 작업으로 합법화하는 실험적 패스입니다.

선택적 패스를 통해 모양과 데이터 계산을 함께 가져오면 StableHLO 생태계에서 StableHLO 작업을 사용하여 동적성을 모델링하는 컴파일 파이프라인을 활용할 수 있습니다.

`-stablehlo-canonicalize-dynamism`

동적 StableHLO 작업을 정적 작업으로 정규화합니다.

이러한 작업의 모든 동적 요소가 실제로 상수인 경우 DynamicReshapeOp과 같은 동적 StableHLO 작업을 DynamicReshapeOp~ReshapeOp 또는 DynamicBroadcastInDim~BroadcastInDim과 같은 해당 정적 작업으로 대체합니다.

  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>

  ==>

  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>

`-stablehlo-check-shape-assertions`

_stablehlo.custom_call @shape어설션 작업을 확인합니다.

shape_assertion 맞춤 호출의 유효성을 검사합니다.

모양 어설션은 StableHLO의 동적 차원에 관한 제약 조건을 검증합니다. 예를 들어 프레임워크에서 DimA < 2 제약 조건을 적용해야 하는 경우 다음 IR이 방출될 수 있습니다.

%dimA = <get_dimension_size or input arg> : tensor<i32>
%c2 = stablehlo.constant dense<2> : tensor<i32>
%is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }

패스 후 도형이 올바르면 stablehlo.custom_call이 삭제됩니다.

옵션

-enable-shape-assertions : Whether shape assertions may generate errors.

`-stablehlo-compatibility-expander`

StableHLO 작업의 호환성 확장기

최신 버전에서 StableHLO 작업이 업데이트되거나 새로운 작업이 도입됩니다. 이 선택적 패스는 최신 StableHLO 작업을 이전 버전에서 지원하는 동등한 작업으로 분해하여 이전 StableHLO 버전과의 하위 호환성을 확장합니다.

이 패스가 선택형인 이유는 무엇인가요?

StableHLO 작업 개선사항은 OpenXLA 생태계에서 특정 일반적인 패턴의 처리를 크게 단순화하는 데 사용되기도 합니다. 여기에는 프레임워크 및 컴파일러 지원이 높은 TanOp과 슬라이스를 사용하여 나타낼 수 있지만 샤딩을 훨씬 더 어렵게 만드는 gather/scatter 배치 차원 등이 포함됩니다. 이러한 새로운 기능 카테고리의 경우 후속 최적화에 사용되는 중요한 정보가 삭제될 수 있으므로 자동 다운그레이드를 제공하지 않습니다. 이 패스는 타겟 버전을 기반으로 이러한 작업을 확장하여 잠재적으로 최적화되지 않은 컴파일을 희생하여 호환성을 최대화하는 데 사용할 수 있습니다.

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %1 = stablehlo.tan %arg0 : tensor<4xf64>
  func.return %1 : tensor<4xf64>
}

==>

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %0 = stablehlo.sine %arg0 : tensor<4xf64>
  %1 = stablehlo.cosine %arg0 : tensor<4xf64>
  %2 = stablehlo.divide %0, %1 : tensor<4xf64>
  return %2 : tensor<4xf64>
}

옵션

-target : The target version. Must be a version of the form #.#.#.

`-stablehlo-complex-math-expander`

StableHLO 복소수 수학 연산의 익스팬더

StableHLO 복소수 수학 연산은 StableHLO 실수 수학 연산을 사용하여 분해됩니다.

이 문은 복소수나 복소수 수학 연산을 기본적으로 지원하는 하드웨어가 없다는 가정을 기반으로 합니다. 즉, 컴파일러가 구현할 수 있는 복잡한 수학 연산의 대체 메커니즘이 중복됩니다. 이 패스를 사용 설정하면 모든 StableHLO 복소수 수학 연산이 확장됩니다.

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  %1 = stablehlo.sqrt %arg0 : tensor<4xcomplex<f64>>
  func.return %1 : tensor<4xcomplex<f64>>
}

==>

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  TBD
  return %2 : tensor<4xcomplex<f64>>
}

`-stablehlo-convert-to-signless`

IR을 부호 없는 정수로 변환하기 위해 전달

`-stablehlo-legalize-composite-to-call`

합성 작업을 분해에 대한 호출로 대체합니다.

합성 작업을 분해 호출로 대체합니다(예: 아래).

stablehlo.composite "my_namespace.my_op" %arg0, %arg1 {
  decomposition = @bar,
  version = 1,
  composite_attributes = {
    "my_attribute": "my_value"
  }
}

다음과 같이 변환됩니다.

func.call @bar(%arg0, %arg1)

'except' 플래그를 사용하여 이 변환에서 일부 컴포지트를 제외할 수 있습니다.예를 들면 다음과 같습니다.

stablehlo-opt --stablehlo-legalize-composite-to-call=except='foo.baz,foo.qux'

옵션

-except : Names of composites that should not be replaced with calls.

`-stablehlo-legalize-deprecated-ops`

지원되는 작업으로 지원 중단된 작업을 합법화합니다.

StableHLO v1.0 Opset 지원 중단 RFC (#2283)에서는 중복되는 여러 작업을 삭제할 것을 제안합니다. 이 패스는 장기 지원되는 대응 항목으로 합법화하여 다양한 컴파일 파이프라인에서 이러한 작업 삭제의 영향을 평가하는 데 도움이 됩니다.

옵션

-fail-on-unused : Fail on (mostly) unused ops that are deprecated without any fallback.

`-stablehlo-legalize-qdq-to-quantized-op`

안정적인 HLO 양자화 작업으로 융합 (역양자화, 부동 소수점 연산, 양자화) 패턴

패턴을 StableHLO 양자화 작업으로 융합 (역양자화, 부동 소수점 연산, 양자화) 참고: 이 패스는 기존 작업을 삭제하지 않습니다. 예를 들어 다음 프로그램은

func.func @add(%arg0: tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
  func.return %2 : tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
}

다음과 같이 변환됩니다.

func.func @add(%arg0: tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.abs %arg0 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  %3 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  return %2 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
}

`-stablehlo-legalize-quant-to-math`

StableHLO 양자화 작업에서 StableHLO 기본 수학 작업으로 변환

UniformQuantized 유형을 사용하는 StableHLO 프로그램을 의미상 동등한 정수 수학 연산으로 변환합니다.

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

다음과 같이 변환됩니다.

func.func @add(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
  %0 = stablehlo.convert %arg0 : (tensor<i8>) -> tensor<f32>
  %cst = stablehlo.constant dense<0.333333343> : tensor<f32>
  %1 = chlo.broadcast_multiply %0, %cst : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
  %2 = chlo.broadcast_add %1, %cst_0 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = stablehlo.round_nearest_even %2 : tensor<f32>
  %4 = stablehlo.convert %3 : (tensor<f32>) -> tensor<i32>
  %5 = stablehlo.convert %arg1 : (tensor<i8>) -> tensor<f32>
  %cst_1 = stablehlo.constant dense<0.666666686> : tensor<f32>
  %6 = chlo.broadcast_multiply %5, %cst_1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_2 = stablehlo.constant dense<1.33333337> : tensor<f32>
  %7 = chlo.broadcast_add %6, %cst_2 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %8 = stablehlo.round_nearest_even %7 : tensor<f32>
  %9 = stablehlo.convert %8 : (tensor<f32>) -> tensor<i32>
  %c = stablehlo.constant dense<2> : tensor<i32>
  %10 = chlo.broadcast_add %4, %9 : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %11 = chlo.broadcast_subtract %10, %c : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %c_3 = stablehlo.constant dense<-128> : tensor<i32>
  %c_4 = stablehlo.constant dense<127> : tensor<i32>
  %12 = stablehlo.clamp %c_3, %11, %c_4 : tensor<i32>
  %13 = stablehlo.convert %12 : (tensor<i32>) -> tensor<i8>
  return %13 : tensor<i8>
}

`-stablehlo-legalize-quantized-op-to-qdq`

양자화된 StableHLO 작업을 (역양자화, 부동 소수점 연산, 양자화) 패턴으로 분해

균일한 양자화/역양자화 작업을 사용하여 StableHLO 양자화 프로그램을 분해합니다. 예를 들어 다음 프로그램은

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

다음과 같이 변환됩니다.

func.func @add(%arg0: tensor<!quant.uniform<i8:f32, 1.000000e+00>>, %arg1: tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<f32>
  %1 = stablehlo.uniform_dequantize %arg1 : (tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<f32>
  %2 = stablehlo.add %0, %1 : tensor<f32>
  %3 = stablehlo.uniform_quantize %2 : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
  return %3 : tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
}

`-stablehlo-legalize-to-vhlo`

StableHLO를 VHLO로 합법화

StableHLO를 VHLO의 최신 버전 작업으로 합법화 그러면 이러한 작업은 VhloToVersionPass를 사용하여 향후 호환성을 위해 이전 버전의 VHLO로 다운그레이드될 수 있습니다.

stablehlo.exponential %[[ARG0]] <{result_accuracy = DEFAULT}> : tensor<f32>
# ====>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.DEFAULT_v1}> : !vhlo.tensor_v1<!vhlo.f32_v1>

VHLO를 사용하여 이전 버전과의 호환성 및 이후 버전과의 호환성을 유지하는 방법에 관한 자세한 내용은 vhlo.md > VHLO 다이얼렉트를 참고하세요.

옵션

-allow-other-dialects : Allow serialization to use other (potentially unstable) dialects, inserts unrealized casts between dialects.

`-stablehlo-refine-arguments`

main 함수의 인수 모양을 개선합니다.

입력 유형 서명을 사용하여 기본 함수의 인수를 수정합니다. 모양 미세 조정이 실행되기 전에 IR이 유효하도록 custom_call @stablehlo.shape_refinement_operand_wrapper에 인수를 래핑합니다.

func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  ...
}

==>

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) {...}
    : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
}

refinedTypesOption를 사용하여 세분화된 유형 목록을 지정할 수 있습니다. 이는 --types='tensor<...>,tensor<...>'를 사용하여 MLIR에서 지정하거나 패스 생성 메서드에 전달할 수 있습니다. 세부검색 유형 목록은 세부검색되는 main 메서드의 모든 인수의 유형을 지정해야 합니다.

옵션

-types : The new types to be used for the main function's arguments, specified as an MLIR TypeRange 'tensor<1x2xf32>, ...'

`-stablehlo-refine-shapes`

StableHLO 프로그램 전반에서 모양을 개선합니다.

작업 내에서 모양을 세부적으로 조정하는 StableHLO 프로그램을 살펴봅니다.

이 패스의 대표적인 사용 사례는 동적으로 형성된 프로그램을 정적 모양으로 전문화하는 것입니다. 동적으로 형성된 StableHLO 프로그램에 올바른 구조가 있는 경우 동적 모양에서 정적 모양으로 인수 유형을 업데이트하고 이 패스를 실행하면 프로그램 전체에 정적 모양이 전파됩니다.

이 패스는 결과 사용을 피연산자로 직접 대체하여 custom_call @shape_refinement_operand_wrapper를 삭제하고 프로그램 전체에 정적 모양을 전파합니다.

  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
      : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  %1 = stablehlo.add %0, %0 : tensor<?xf32>

  ==>

  %1 = stablehlo.add %arg0, %arg0 : tensor<16xf32>

모양 미세 조정에 유효한 모듈에는 다음 속성이 있어야 합니다.

모든 동적 모양은 입력 모양에만 종속됩니다 (입력 배열 콘텐츠에 모양 종속성이 없음). 입력 모양 (예: stablehlo.get_dimension_size로 제공) 또는 기호 정수의 확인된 값 (즉, 텐서 : A = 5)과 같은 전역 상수에만 간접적으로 의존하는 작업을 dimension 작업이라고 합니다. 모든 측정기준 값은 절차 간 상수 폴딩을 통해 상수로 확인할 수 있습니다.
중간 함수는 인수 목록의 시작 부분에 여러 토큰 인수(!stablehlo.token 유형)를 사용할 수 있으며, 그 뒤에 기호 정수의 해결된 값 (예: 텐서 : A = 5)과 같은 상수 정수 스칼라인 전역 상수 인수가 이어집니다.
일부 중간 함수는 전역 상수에 대한 계산, 즉 symint 값에 대한 floordiv을 반환할 수 있습니다. 이러한 함수는 개선 후 상수 값만 반환하는 것으로 표시됩니다. 이러한 함수는 인라인 처리됩니다.
단일 함수에 대한 모든 호출은 동일한 인수 모양으로 확인되며 재귀 / 공재귀 함수 호출은 이루어지지 않습니다.

`-stablehlo-wrap-in-composite`

비복합 StableHLO 작업을 복합 작업으로 래핑합니다.

StableHLO 작업을 stablehlo.composite 작업으로 래핑합니다.

예를 들어 간단한 StableHLO 프로그램을 생각해 보겠습니다.

func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

stablehlo.add 작업을 래핑하기 위해 이 패스를 적용하면 다음 프로그램이 생성됩니다.

func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition = @stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  return %0 : tensor<2xf32>
}
func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

참고:

생성된 stablehlo.composite 작업의 name 속성은 항상 래핑된 원래 작업의 이름과 동일합니다 (예: stablehlo.add 작업을 래핑하는 경우 컴포지트의 이름도 "stablehlo.add"가 됨).
원래 작업을 캡슐화하는 비공개 함수(stablehlo.composite 작업의 decomposition 속성으로 참조됨)는 <op_name>.impl[.N] 패턴을 사용하여 이름이 지정됩니다. 여기서 <op_name>은 원래 작업의 이름이고 N은 모듈 내에서 이름 충돌을 방지하기 위해 생성된 고유한 정수 식별자입니다.

이 패스는 두 가지 방법으로 사용할 수 있습니다.

모드 1: 명령줄 사용

이 모드는 생성된 stablehlo.composite 작업의 속성을 최소한으로 제어할 수 있으므로 디버깅이나 테스트에 적합합니다. op-names(쉼표로 구분된 작업 이름 목록) 옵션을 사용하여 지정된 작업의 모든 인스턴스를 래핑합니다. 새로 생성된 stablehlo.composite 작업의 속성은 원래 작업의 속성과 동일합니다.

사용 예제:

stablehlo-opt input.mlir --stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o output.mlir

모드 2: 맞춤 속성 처리를 사용한 프로그래매틱 모듈 전체 래핑

이 모드는 프로그래매틱 래핑을 전체 모듈로 확장하여 래핑되는 작업과 속성을 세부적으로 제어할 수 있습니다. 이는 CompositeAttributeProviderMap을 인수로 사용하는 createStablehloWrapInCompositePass API를 사용하여 달성됩니다.

CompositeAttributeProviderMap은 래핑을 고려해야 하는 작업과 속성을 처리하는 방법을 지정하는 맵입니다. 의미는 다음과 같습니다.

키 (mlir::TypeID): MLIR 작업의 TypeID입니다. 작업의 TypeID가 맵의 키와 일치하면 래핑 후보가 됩니다.
값 (람다 함수): std::function<std::optional<NamedAttrList>(Operation*)> 유형의 람다 함수입니다. 이 함수는 각 후보 작업에 적용됩니다.
- 입력: TypeID 키에 해당하는 작업 유형의 인스턴스인 mlir::Operation*
- 반환 값: std::optional<NamedAttrList>
  - 람다가 NamedAttrList (std::optional로 래핑됨)를 반환하면 작업이 stablehlo::composite 작업으로 래핑되고 반환된 속성이 컴포지트의 속성을 설정하는 데 사용됩니다.
  - 람다가 std::nullopt를 반환하면 작업이 래핑되지 않습니다. 이를 통해 맞춤 기준에 따라 선택적으로 래핑할 수 있습니다.

예 (C++):


stablehlo::CompositeAttributeProviderMap compositeAttributeProviderMap;

compositeAttributeProviderMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
  [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
  // Custom logic to determine if and how to wrap the operation.
  // Example: Only wrap if it's on a specific type.
  if (mlir::isa<mlir::Float32Type>(op->getOperand(0).getType())) {
    return mlir::NamedAttrList(op->getAttrs());
  }
  return std::nullopt; // Do not wrap.
};

pm.addPass(createStablehloWrapInCompositePass(compositeAttributeProviderMap, compositeVersion));
if (mlir::failed(pm.run(module))) {
  return;
}

옵션

-op-names : The names of the ops to wrap.
-version  : The version number of the composite op.

`-vhlo-legalize-to-stablehlo`

VHLO를 StableHLO로 합법화

`-vhlo-to-version`

호환성을 위해 VHLO 버전 간에 변환합니다.

상위 및 하위 호환성을 유지하기 위해 IR 업그레이드 및 다운그레이드를 위해 VHLO 버전 간에 변환합니다.

"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>
# ==( -target=1.0.0 )==>
"vhlo.exponential_v1"(%[[ARG0]])
# ==( -target=1.9.0 )==>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>

VHLO를 사용하여 이전 버전과의 호환성 및 이후 버전과의 호환성을 유지하는 방법에 관한 자세한 내용은 vhlo.md > VHLO 다이얼렉트를 참고하세요.

옵션

-target : The target version. Must be a version of the form #.#.# .