Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance (22’)

Abs and Intro
오늘은 auto-tuners는 효율적인 tensor programs을 생성하지만 불투명한 하드웨어 세부 사항으로도 그렇게 한다. 이러한 성능은 수동으로 최적화하는 방식인 hardware-native libraries의 성능보다 뒤쳐진다. 뿐만 아니라 vendor library들은 지원되는 기능의 집합이 고정되어 있으며 auto-tuner에 의해 제공되는 customization기능 및 automation기능의 지원이 부족합니다. BOLT는 이 격차를 줄이고 hardware-native templated search를 사용함으로써 양쪽의 장점을 모두 구현합니다. 또한 Bolt는 graph, operator, model level에서 종단간 tensor optimization을 다시 생각하게 해준다.
Opaque device models는 검색 프로세스가 하드웨어 세부 정보에 의해 덜 알려지기 때문에 auto tuning 시간이 길어집니다. 이는 cache와 tuning log를 재사용하는 데이터베이스 개발을 이끌고 있지만 이러한 접근법은 갈 길이 멉니다. 모델의 dynamism이 커지고 dynamic data structure관점 뿐만 아니라 dynamic shape도 cache 만드는 것을 비효율적으로 만듭니다. Bolt는 새로운 방식의 computational graph-level optimization을 가능하게 하고 graph-level과 operator-level optimization을 동시에 할 수 있게 합니다. Hardware-native performance로 generation하는 방식은 상당히 turnaround 시간을 줄일 수 있습니다. 그 외로 operator fusion을 위한 지속적인 kernel fusion을 가능하게 할 뿐 아니라 light-weight performance profiler기능을 제공한다.

Implementation
Bolt는 TVM의 frontend를 사용해서 model을 relay graph로 만들고 이 graph에서 computational graph optimization을 수행한 뒤에 graph partition을 실행한다. 그 다음 각 연산에 대해 최고의 커널을 찾는 hardware-native profiling을 실행한다. 그리고 고성능의 CUDA code를 생성하고 컴파일한다. 생성된 code는 runtime file에 저장된다.

Enabling deeper operator fusion
여러 개의 GEMM과 Conv연산을 하나의 operator로 fusing하는 작업은 다음과 같은 이점이 있다. Inter-layer activation에서 storing과 loading하는데 Memory traffic 제거할 수 있고 작은 batch size로 생기는 launch latency를 제거한다. 그리고 명령어 스케줄링 옵션을 탐색하기 위해 컴파일러를 위한 optimization scope을 확장할 수 있다.
필수조건으로는 Epilogue fusion이 있고 지속적인 kernel fusion을 위해서 필요하다. Bolt는 우선 CUTLASS에서 제공되는 epilogue fusion을 통합하고 이는 GEMM/Conv kernel을 fusing함으로써 persistent kernel을 가능하게 한다.

Automating templated code generation
템플릿 라이브러리는 종단간 tensor program optimization를 위한 새로운 도전에 직면해 있다. 단순한 BYOC 방식으로는 전체적인 문제를 해결할 수 없다.

Designing system-friendly models
Graph-level optimization and operator-level optimization은 model-level optimization 기회를 살릴 수 있다. Model은 효율적인 inference가 가능하도록 설계될 수 있고 이를 system-model codesign이라고 부른다.