유형 추론

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

아래에 제안된 가이드라인은 StableHLO 작업에 고품질 인증 기능과 셰이프 함수를 구현하기 위한 것입니다.

제안 내용

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

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

spec은 StableHLO 연산의 모든 인증기 및 도형 함수에 관한 정보 소스입니다. 모든 작업의 기존 인증 기능과 도형 함수를 다시 재검토하여 사양과 완전히 일치시켜야 합니다. 사양 문서는 계속 발전합니다. 작업의 사양을 사용할 수 없는 경우에는 xla/service/shape_inference.cc 및 xla/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): 모든 인증 코드를 도형 함수에 넣고 인증자를 완전히 삭제합니다.

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

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

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

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

ODS가 적용되는 확인을 위해 테스트를 추가/유지해야 하나요?

Google에서는 그러지 않습니다. 테스트는 인증기와 도형 함수에 초점을 맞춰야 하지만, ODS를 변경할 때는 이 작업을 다시 살펴봐야 합니다.

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

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

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

infer_stablehlo.mlir는 hlo_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라는 별도의 테스트 파일을 추가하는 것이 좋습니다.