Dive into Mojo🔥

Intermediate Representation (IR)

Intermediate Representation (IR)은 사람이 이해하기 쉬운 High-level 언어에서 기계어와 유사한 Low-level 언어로 변환하는 중간 단계의 표현 방식을 의미한다. 컴파일러나 인터프리터는 프로그램을 High-level 언어에서 IR로 변환하고, 이후 IR을 다시 기계어로 변환하여 실행 가능한 형태로 만든다.

• Low-level IR
  • Static Single Assignment (SSA) Form: 각 변수가 단 한 번만 정의되고 사용되는 형태로 변환하여 최적화를 용이하게 만듬
  • Three-Address Code: 각 명령어에 최대 세 개의 주소를 사용하는 형태로, 기계어에 가까운 표현
  • Control Flow Graph (CFG): 프로그램의 제어 흐름을 그래프로 표현하여 제어 흐름 분석을 수행
• Intermediate Representation (IR)
  • LLVM: 컴파일러 및 LLVM 프로젝트에서 사용되는 중간 표현 방식으로, 정적 타입을 가진 강력한 추상화를 제공
  • MLIR: 다양한 언어와 도구들 간의 상호 운용성과 최적화를 위한 중간 표현 방식
• High-level IR
  • Abstract Syntax Tree (AST): 프로그램의 구문 구조를 트리 형태로 표현하여 파싱한 결과를 나타냄
  • Typed High-level IR: 자료형에 대한 정보를 가지고 있는 고급 레벨의 중간 표현 방식
  • Source-level IR: 원시 코드와 유사한 고급 레벨의 중간 표현 방식으로, 디버깅이나 코드 이해에 유리

Low-level IR은 최적화와 코드 생성에 적합하며, High-level IR은 프로그램 구조를 이해하거나 변환할 때 유용하다

LLVM

• Architecture
  • Frontend: C, C++, Swift, Rust 등 다양한 언어를 지원하는 Frontend가 제공되어 해당 언어로 작성된 코드를 LLVM IR로 변환
  • Middle-End: 코드를 분석하고 중복 코드를 제거하거나 상수 폴딩(Constant Folding)과 같은 최적화를 수행
  • Backend: 목표 하드웨어 아키텍처에 맞는 기계어로 변환하는 역할

MLIR

• 다양한 Level의 중간 표현
  • MLIR은 다양한 수준의 중간 표현을 지원한다. 하위 레벨 IR부터 고급 레벨 IR까지 다양한 표현이 가능하다. 이는 하드웨어와 도구의 특성에 맞게 최적화를 수행하고, 각 언어와 도구에서 가장 적합한 표현을 선택할 수 있게 해줌
• 모듈화(Modularity)
  • MLIR은 모듈화된 디자인을 채택하여 각 언어와 도구에 특화된 기능을 제공
• 변환(Transformations)
  • MLIR은 중간 표현 간의 변환을 쉽게 수행할 수 있도록 지원
• 다양한 언어 및 도구와의 통
  • MLIR은 다양한 프로그래밍 언어와 도구를 통합하는데 사용

Modular AI Engine

Modular AI Engine은 Pytorch와 Tensorflow 같은 AI frameworks을 위한 고성능 inference engine이다.

https://www.modular.com/engine#performance

• Importer
  • serialized model을 MLIR graph representation으로 변환하는 역할
  • MLIR graph representation 변환 과정
    1. Parsing: AST로 파싱하여 추상적인 트리 구조로 변환
    2. AST를 MLIR 그래프로 변환: AST의 각 Node와 Edge를 MLIR 그래프와 패핑
• Optimizer
  • graph representation의 graph-level optimization(e.g, constant folding, shape inference)를 수행
• Kernel Generator
  • 다양한 기능을 수행하기 위한고성능 machine learning kernel(작은 코드 블록) 생성
  • 모듈화된 구조를 가지고 있어 kernel을 추가하거나 변경하는데 용이
• Runtime
  • optimized kernel graph를 실행

일반적으로 AI 알고리즘은 많은 연산이 필요하기 때문에 AI Engine의 kernel generator은 하드웨어별로 최적화된 machine learning operator 구현을 생성하는 역할을 한다. Modular에서는 처음에 이 코드를 MLIR에 바로 넣었었지만 생산성 문제가 있었다고 한다. 그래서 C/Rust와 같은 언어의 safety, predictability guarantee, Python의 usability를 갖추는 High-level language에서의 작업이 훨씬 더 생산적일거라 생각했다고 한다. 이처럼 AI Engine의 내부 개발에 의해서 탄생한 것이 Mojo🔥 이다.

Mojo🔥

Mojo🔥 는 앞서 말한 kernel을 위한 것만이 목적은 아니다.

• Python과 완벽하게 호환되는 것을 목표로 하고 있음
• 기존 Python 모듈을 가져와서 사용할 수 있음
• Python의 핵심 기능들(async/await, error handling 등)을 지원하지만 완벽하게 호환되지는 않음 (아직 개발 중)
• Python 보다 약 35,000배 빠른 속도로 동작
• 하드웨어에 대한 Low-level control을 직접 수행할 수 있도록 지원

• Mojo는 Python과 비슷한 문법을 가지지만 인터프리터를 가지는 Python과 달리 컴파일
• Mojo 컴파일러는 Mojo 코드를 MLIR로 낮추고 optimization을 수행하여 특정 하드웨어에 대한 코드를 생성

[Reference]

https://devocean.sk.com/blog/techBoardDetail.do?ID=165125 

https://docs.modular.com/mojo/why-mojo.html

https://mlir.llvm.org/
https://performance.modular.com/
https://www.modular.com/blog/whats-the-difference-between-the-ai-engine-and-mojo