기술 요약: AI 컴파일러 퍼징을 위한 데이터 기반 테스트 생성

문제 정의

AI 컴파일러(예: TVM, TensorRT, ONNXRuntime)는 다양한 하드웨어에 AI 모델을 배포하는 데 필수적이지만, 컴파일러의 신뢰성과 모델의 정확성을 저해하는 버그가 발생하기 쉽습니다. 컴파일 과정은 세 가지 뚜렷한 단계로 구성됩니다: (1) 모델 로딩(Model Loading): PyTorch나 Keras와 같은 라이브러리의 모델을 고수준 중간 표현(IR)으로 변환하는 단계; (2) 고수준 최적화(High-level Optimization): 하드웨어 독립적인 그래프 변환을 포함하는 단계; (3) 저수준 최적화(Low-level Optimization): 실행 가능한 코드를 생성하기 위해 하드웨어 특화된 변환을 적용하는 단계입니다.

기존의 테스트 기법들(예: NNSmith, HirGen, Tzer)은 주로 유효하고 다양한 입력을 생성하는 데 집중하며, 단계별로 포괄적인 커버리지를 제공하는 데는 실패합니다. 구체적으로 다음 사항들에 어려움을 겪습니다:

1. 다양한 프런트엔드 커버리지: 다양한 AI 라이브러리로부터의 모든 연산자(operator) 변환 로직이 모델 로딩 중에 올바르게 처리되는지 보장하는 문제.
2. 최적화 단계 인지(Optimization-Stage Awareness): 복잡하고 문맥 의존적인 고수준 최적화를 효과적으로 트리거할 수 있는 계산 그래프를 생성하는 문제.
3. 하드웨어 특화 검증: 저수준, 하드웨어 특화 최적화를 트리거하는 데 필요한 기능들을 테스트에 포함시키는 문제.

방법론

저자들은 AI 컴파일 파이프라인의 특정 단계를 각각 타겟팅하는 세 가지 상호 보완적인 구성 요소가 통합된 통합 데이터 기반 테스트 프레임워크인 OPERA를 제안합니다.

1. OPERA: 마이그레이션 기반 모델 로딩 테스트

OPERA는 AI 라이브러리의 기존 테스트를 마이그레이션하여 모델 로딩 단계를 다룹니다.

• 접근 방식: AI 라이브러리 소스 코드의 API를 인스트루먼테이션(instrumentation)하여 라이브러리 테스트를 실행하고 연산자 인스턴스를 추출합니다. 이 인스턴스들은 템플릿 기반 생성을 통해 단일 연산자 모델로 캡슐화됩니다.
• 최적화: 마이그레이션된 테스트 실행의 높은 비용을 관리하기 위해, OPERA는 **다양성 인식 우선순위 지정 알고리즘(diversity-aware prioritization algorithm)**을 채택하여 연산자를 의미론적으로 클러스터링함으로써 커버리지를 극대화하는 동시에 중복을 제거합니다.
• 목표: 라이브러리 테스트에 의해 이미 검증된 핵심 의미론을 활용하여 다양한 연산자 변환 로직의 커버리지를 극대화하는 것입니다.

2. OATest: 합성 기반 고수준 최적화 테스트

OATest는 최적화 인지형 계산 그래프를 합성함으로써 고수준 최적화를 타겟팅합니다.

• 접근 방식: 컴파일러 개발자들이 작성한 테스트로부터 최적화 패턴을 추출하기 위해 컴파일러 소스 코드를 인스트루먼테이션합니다. 이 과정은 최적화에 해당하는 계산 그래프를 식별하고, 최적화의 세밀도(granularity)에 따라 패턴을 도출하는 과정을 포함합니다.
• 합성: OATest는 추출된 패턴을 다양한 문맥(시드 그래프)과 결합합니다. 패턴을 무작위로 선택하고 시드를 선택한 뒤, 구조적 다양성을 높이기 위해 무작위 합성 지점에 해당 패턴을 삽입합니다.
• 유효성 보장: 정의되지 않은 입력으로 인한 유효하지 않은 테스트를 방지하기 위해, OATest는 두 가지 수정 전략을 사용합니다: (1) 기존 노드에서 호환 가능한 입력/출력을 재사용하는 방식, (2) 패턴과 문맥 사이의 올바른 연결을 설정하기 위해 새로운 노드를 생성하는 방식.

3. HARMONY: 변이 기반 저수준 최적화 테스트

HARMONY는 변이 기반 접근 방식을 사용하여 하드웨어 특화 최적화에 집중합니다.

• 시드 생성: 처음부터 테스트를 생성하는 대신, HARMONY는 여러 지식 소스에서 수집된 연산자 제약 조건에 의해 가이드되는 대규모 언어 모델(LLM)을 사용하여 유효하고 다양한 저수준 IR 시드를 생성합니다. 이러한 교차 검증은 오래된 문서와 관련된 리스크를 완화합니다.
• 변이: 분석용 LLM은 문서와 코드 예제에서 최적화 패턴을 추출합니다. 이후 코드 생성용 LLM은 이러한 패턴에 따라 기존 IR 시드에 대해 **경량 변이(lightweight mutations)**를 수행합니다.
• 목표: LLM의 환각(hallucination)으로 인한 유효하지 않은 출력의 위험을 최소화하면서, 타겟팅된 저수준 최적화(예: 메모리 레이턴시 은닉, 인트린직 매핑)를 트리거할 가능성이 높은 의미론적으로 유효한 테스트를 생성하는 것입니다.

주요 기여

본 논문은 하나의 통합된 프레임워크를 형성하는 세 가지 특정 기술을 제시합니다:

1. OPERA: 기존 AI 라이브러리 테스트를 활용하여 연산자 변환 로직의 버그를 찾아내는 마이그레이션 기반 접근 방식.
2. OATest: 개발자가 작성한 테스트에서 의미론적 패턴을 추출하여 문맥 의존적인 계산 그래프를 생성하는 최적화 인지형 합성 기술.
3. HARMONY: 문서와 LLM을 활용하여 하드웨어 특화 최적화 테스트를 위한 저수준 IR 시드를 생성하고 변이시키는 변이 기반 기술.

결과

본 프레임워크는 4개의 널리 사용되는 AI 컴파일러(TVM, TensorRT, OpenVINO, ONNXRuntime)에 대해 평가되었으며, 266개의 미발견 버그를 탐지했습니다. 구성 요소별 상세 내역은 다음과 같습니다:

• OPERA: 8개의 프런트엔드에 걸쳐 170개의 버그를 발견했습니다 (이 중 90개는 개발자에 의해 확인됨). 우선순위 지정 전략은 일반적인 전략에 비해 테스트 효율성을 11.9%~47.4% 향상시켰습니다. 버그의 근본 원인은 텐서 형태 문제, 타입 문제, 잘못된 코드 로직 등의 카테고리에 걸쳐 있었습니다.
• OATest: TVM 및 ONNXRuntime에서 56개의 버그를 발견했습니다 (이 중 42개 확인, 24개 수정). 이는 NNSmith 및 WhiteFox와 비교했을 때 60.2% 더 높은 브랜치 커버리지와 66.98% 더 높은 라인 커버리지를 달 기록했습니다.
• HARMONY: TVM에서 40개의 버그를 탐지했습니다 (이 중 26개 확인 및 수정). 이는 NNSmith, WhiteFox, Tzer와 같은 최신 퍼저(fuzzer)들과 비교하여 저수준 최적화 단계에서 실질적으로 더 높은 커버리지를 보여주었습니다.

의의 및 주장

본 논문은 이 연구가 컴파일러를 단일한 블랙박스로 취급하는 대신, 각 컴파일 단계의 특정 과제를 체계적으로 해결하는 포괄적이고 단계 인지적인 프레임워크를 제공한다고 주장합니다. 데이터(라이브러리 테스트, 개발자 패턴, 문서)에 의해 구동되는 마이그레이션, 합성, 변이 기술을 통합함으로써, 본 프레임워크는 테스트 커버리지와 효과성을 크게 향상시킵니다.

저자들은 이 작업이 전체 파이프라인에 걸쳐 AI 컴파일러의 신뢰성을 보장하기 위한 단계적 조치라고 규정합니다. 또한 현재의 결과가 매우 유의미하지만, 향후 연구를 통해 프레임워크를 신흥 컴파일러(예: Triton, MLC-LLM)로 확장하고, 개발자의 유지보수 노력을 더욱 줄이기 위해 자동 버그 위치 파악(localization) 및 수정(repair)을 조사할 것이라고 언급했습니다.