유형 추론

StableHLO는 원래 MHLO 언어에서 부트스트랩되었으며 유형 추론의 MHLO 구현을 상속했습니다. 구현 진행 상황은 status.md에서 추적됩니다.

아래 제안된 가이드라인은 StableHLO 연산을 위한 고품질 인증기 및 셰이프 함수를 구현할 수 있도록 하기 위한 것입니다.

제안 내용

이러한 제안은 기존 구현을 재방문하고 포괄적인 적용 범위가 될 때까지 새로운 작업을 달성하는 데 모두 적용됩니다.

(P1) StableHLO 사양을 정보 소스로 사용

spec은 StableHLO 연산의 모든 인증자와 셰이프 함수에 관한 정보 소스입니다. 모든 작업의 기존 인증자와 셰이프 함수를 다시 재검토하여 사양에 맞게 완벽하게 조정해야 합니다. 사양 문서는 계속 바뀝니다. 작업의 사양을 사용할 수 없는 경우 xla/service/shape_inference.ccxla/service/hlo_verifier.cc를 포함하여 XLA 구현을 정보 소스로 대신 사용해야 합니다. XLA 구현은 제한되지 않은 동적 특성을 다루지 않습니다. 따라서 제한되지 않은 동적 기능의 경우 동적 RFC를 사용할 수 있을 때까지 상식을 적용합니다.

(P2) ODS를 최대한 활용하세요

ODS 파일 (예: StablehloOps.td)은 각 피연산자/속성/결과의 특성과 유형으로 작업을 정의하고 확인을 실행합니다. 따라서 ODS에서 이미 보장한 속성의 인증자 또는 도형 함수에는 인증 코드가 필요하지 않습니다. ODS와 중복된 인증 코드는 트리거되지 않으므로 삭제합니다.

ODS의 제약 조건 테스트를 추가해야 하나요? 테스트 가이드라인 수립을 참고하세요.

(P3) 인증기 및 셰이프 함수에서 인증 코드 유지

둘 다:

  • 인증자: Op::verify()에 의해 구현됩니다.
  • 도형 함수: Op::inferReturnTypes() 또는 Op::inferReturnTypeComponents와 같은 InferTypeOpInterface로 구현됩니다.

피연산자/속성/결과를 확인하기 위한 인증 코드가 있을 수 있습니다. 초기 분할은 다음과 같을 수 있습니다. 검증자가 피연산자/속성을 확인하도록 한 다음 셰이프 함수가 추론된 결과 유형만 계산하고 실제 결과 유형과의 호환성을 확인하도록 합니다. 하지만 실제로 이러한 분할에는 몇 가지 문제가 있습니다.

  • 도형 함수는 확인자를 먼저 호출하지 않고 자동 생성된 build() 함수에 의해 호출될 수 있습니다. 따라서 관련 입력도 셰이프 함수에서 확인해야 합니다.
  • 코드 중복: 예를 들어 인증기에서는 피연산자에 몇 가지 처리를 수행한 후 중간 결과를 확인합니다. 그런 다음 셰이프 함수에서 이러한 중간 결과는 최종 결과를 추론하는 데 유용합니다. 이러한 중간 결과는 두 번 계산해야 합니다.
  • 유지보수 부담: 작업 확인은 두 가지 방법으로 포함됩니다.

솔루션은 다음과 같습니다.

  1. 리전이 없는 대부분의 작업 (예: PadOp): 모든 인증 코드를 도형 함수에 입력하고 인증자를 완전히 삭제해 보세요. 반환 유형을 추론할 수 없어(예: ReshapeOp 또는 BroadcastInDimOp 사용) 확인할 수 없는 경우 필요한 확인 로직을 포함하는 확인자를 만드세요. 일반적으로 AddOp와 같이 추론 가능한 오퍼레이션은 유형/도형 추론 메서드에서 액세스할 수 없는 제공된 반환 유형의 제약 조건을 확인하므로 추가 검증을 실행하기 위해 인증기가 여전히 필요할 수 있습니다.

  2. 리전이 있는 오퍼레이션의 경우 (예: ReduceOp/IfOp, 전체 목록은 여기 참고): 자동 생성된 빌더는 리전을 매개변수로 사용하지 않으므로 이러한 빌더에 유형 추론이 포함된 경우 셰이프 함수가 빈 영역으로 호출됩니다 (이 예 참고).

    1. 리전이 유형 추론에 필요하지 않은 경우 (예: ReduceOp) 도형 함수 대신 영역 관련 확인 로직을 인증자에 배치합니다. 불가피한 경우 일부 코드를 복제합니다.

    2. 유형 추론 (IfOp/CaseOp/MapOp)에 리전이 필요한 경우 추가로 셰이프 함수는 ODS가 작업 정의에 이미 존재한다고 보장하더라도 리전이 비어 있지 않은지를 명시적으로 확인해야 합니다.

(P4) 테스트 가이드라인 수립

ODS가 적용되는 인증에 대한 테스트를 추가/유지해야 하나요?

Google은 그렇지 않습니다. 테스트는 인증자와 셰이프 함수에 중점을 두어야 하며 ODS 변경 시 이 작업을 다시 살펴봐야 합니다.

하지만 누락된 부분에 주의하세요. 예를 들어 작업에 도형만 확인하고 요소 유형은 검사하지 않는 특성 SameOperandsAndResultShape가 포함되어 있는 경우 피연산자/결과의 요소 유형 확인에는 여전히 테스트가 필요합니다.

확인자 및 유형 추론 테스트는 어디에 배치하나요?

ops_stablehlo.mlir에는 작업 양성 사례와 모든 검증 오류에 대해 최소 1개의 음성 테스트가 포함됩니다. 추론된 반환 유형이 실제 결과 유형과 호환되는지 (같지 않음)도 확인할 수 있습니다.

infer_stablehlo.mlirhlo_test_infer.get_return_type_components"(%x):...를 사용하여 작업의 셰이프 함수가 있는지 확인하고 추론된 유형이 예상과 정확히 일치하는지 확인합니다. 일반적으로 작업당 양성 테스트 1건이 필요합니다.

필요한 조치

작업의 인증 함수 및/또는 셰이프 함수를 구현하거나 재검토하는 경우:

  1. 양성 사례와 음성 사례를 모두 ops_stablehlo.mlir에 넣습니다.

  2. infer_stablehlo.mlir에 단일 양성 테스트를 추가하여 인터페이스를 테스트합니다.

  3. (선택사항) 작업이 복잡하고 많은 테스트가 포함될 수 있는 경우 동일한 폴더 내에 verify_<op_name>.mlir 또는 verify_<your_topic>.mlir라는 별도의 테스트 파일을 추가하는 것이 좋습니다.