1. 編譯原理與什麼課程有聯系
與人工智慧、機器學習、操作系統、編譯器等都有關系，屬於計算機專業的核心課程，需要扎實掌握；
2. 開篇：XLA是什麼
XLA (Accelerated Linear Algebra)是專用於機器學習的編譯器，機器學習的運算中99%都是向量乘以矩陣、矩陣乘以矩陣的計算，XLA是專門用來優化這些計算的。
舉個例子，運行在GPU上的 model_fn 函數會順序調用 multiply 、 add 和 rece_sum 這三個op，而且 multiply ，也就是 y * z 的計算結果會先從GPU拷貝回host，再拷貝到device作為 add 的input，同樣的，add的計算結果也會以相同的方式傳遞給下一個op。
顯然，對於整個函數來說，將中間變數在host和device間來回倒騰是沒有意義的。因此，如果把函數看作一個op，那在計算中產生的中間結果就不必返回到host，少了數據傳輸的時間開銷，就可以大幅提升運算效率。
這種將多個op融合成一個op的方法就稱為 fuse ，當前fuse的技術路線有：
XLA的優化當然不只是fuse，還有對計算圖的優化，包括刪除無效指令、減少內存佔用、替換復雜指令等優化。下圖是官方提供的性能報告，經XLA優化過後，Tensorflow BERT MLPerf的訓練性能提升了~7倍。除了Tensorflow外，XLA還支持 JAX 、 Julia 、 PyTorch 和 Nx 等前端。
jit 是指在首次運行時將函數編譯成二進製程序，後續再調用該函數時直接運行先前編譯好的程序而非python code。 @tf.funciton 修飾的函數（包括它的子函數）會做 jit 。除非signature發生了變化，也就是input的shape或dtype和編譯時不同，否則 get_MSE 是不需要重復編譯的。
@tf.function 將函數內的ops替換成一組（ XlaCompile , XlaRun ) ops，在運行時前者負責編譯，並將編譯結果-- executable 保存到cache，後者負責運行executable。如果cache里已經有編譯好的程序就不需要編譯了，例如 get_MSE(tf.constant(3.0), tf.constant(4.0)) 。
XLA編譯器支持的語言（IR）是HLO（High Level Operations），顧名思義這些語言是由一個個op組成，因此，我們在編譯前需要先從python code中提取出所有ops，再將它們轉換成HLO。
JAX通過tracing的方式，從 @jax.jit 修飾的函數中提取ops，這些ops通過 jaxpr 來表示。然後再通過XLA client提供的API為ops生成相應的HLO。PyTorch/XLA也是採用類似的方法來生成HLO。
Tensorflow的 tf2xla 為每個 Op 創建了一個同名的 XlaOp 用於生成HLO， XlaOp 派生於 Op ，使用相同的注冊機制，因此，只要把要編譯的子圖根據拓撲排序運行一遍就能生成它的HLO。
HLO先經過一系列 pass 優化後再將HLO lowering成ISA，最後將編譯好的二進制封裝到 executable 。
除了二進製程序，它還包含運行該程序所需要的infos和options。調用 executable.run() 就可以執行計算圖。
3. AI編譯器技術剖析（一）-概述
AI編譯器是近年來在AI應用領域興起的關鍵技術，旨在解決AI模型從研發到多環境部署過程中所遇到的復雜性和繁瑣性問題。AI模型在智能家居、自動駕駛等領域的廣泛應用，促使它們在雲平台、專用計算設備以及物聯網感測器中的部署需求不斷增長。然而，這些模型往往包含非標准運算元，其部署在特定硬體上需要大量工程工作，這成為AI應用落地過程中的一個瓶頸。AI編譯器的誕生，旨在簡化這一過程，提供更高效、靈活的模型轉換與優化手段。
AI編譯器的產生背景與早期神經網路部署方式緊密相關。早期階段，神經網路的部署主要依賴框架和運算元庫，通過直觀的數據流圖進行表示。這種部署方式在一定程度上簡化了AI模型的運行，但隨著時間推移，其局限性逐漸顯現，特別是在復雜模型的優化和跨平台部署方面。AI編譯器的引入，旨在解決這些局限，通過將AI演算法從開發階段轉換為部署狀態，實現更高效、更靈活的模型優化與部署。
AI編譯器的目標在於優化程序性能，同時降低編程難度。與傳統編譯器相比，AI編譯器更側重於針對AI模型的特殊需求，如自動微分支持、運算元融合、內存優化等。它通過將機器學習演算法轉換為更優化的表示，並針對特定硬體進行編譯，以實現模型的高效執行。與傳統編譯器相似，兩者都涉及語言轉換，但AI編譯器更多關注於提升AI模型的性能和簡化開發過程。
AI編譯器的架構通常分為前端和後端兩個部分。前端負責解析用戶代碼，生成計算圖，並進行設備無關的優化；後端則負責將前端生成的計算圖進一步優化，使之更適應特定硬體，並選擇最適合的硬體運算元進行執行。這一過程包括計算圖優化、運算元選擇和內存分配等關鍵步驟，最終生成可在硬體上執行的任務序列。
AI編譯器的獨特之處在於其對自動微分的支持和前端優化策略。自動微分功能要求在傳統中間表示的基礎上設計新的中間表示結構，以支持更復雜的優化操作。此外，AI編譯器的前端還會執行一系列優化，如運算元融合、內存分配、常量折疊等，以提高程序的執行效率和減少能耗。
AI編譯器後端則專注於硬體相關優化，如運算元選擇和內存管理，以確保模型在特定硬體上的高效執行。它通過計算圖優化、運算元選擇和內存分配等步驟，生成一個完整的運算元序列，從而實現AI模型的高效執行。
綜上所述，AI編譯器通過提供從演算法到硬體執行的完整解決方案，簡化了AI模型的部署過程，提高了模型在不同環境下的執行效率。這一技術的不斷演進，將為AI應用的廣泛普及提供強大支持。
4. 機器學習編譯(MLC)基本思想
機器學習編譯（MLC）的基本思想是將機器學習演算法從開發階段優化和轉換，使其適應部署環境。機器學習演算法通常在通用框架如Paddle、PyTorch、TensorFlow或JAX中描述，並包含相應的權重。而部署形式涉及執行應用程序所需的代碼、資源管理（如內存）以及與開發環境的介面。MLC的目標是集成與最小化依賴，利用硬體加速，以及進行通用優化。
MLC通過將機器學習模型轉換為適配特定部署環境的形式，減少應用大小，並使應用能夠部署到更多環境。通過構建調用原生加速庫的代碼或生成利用原生指令（如TensorCore）的代碼，可以利用硬體加速。此外，MLC還包括對模型執行的等價變換，以優化內存使用和執行效率。
TVM是實現MLC的工具之一，它通過將外部模型轉換為Relay/Relax的高層IR表示，進一步轉換為底層tensor IR表示，並最終編譯為可執行代碼。IRMole是程序變換的核心組件，包含了一系列Function，如relay/relax::Function和tir::PrimFunc。運行時的Mole是IRMole的編譯結果，具備輸入、執行和輸出功能。
TensorIR是TVM中針對tensor計算過程的描述，如matmul + relu的表示方式與Numpy的對比。Schele是變換的基本組件，TVM提供多種方法進行變換。例如，split、reorder和reverse_compute_at等變換，都是對IRMole進行的優化，目的是使計算邏輯更加友好，以適應底層硬體。
訪問內存的效率是影響性能的關鍵因素。現代CPU採用緩存策略，獲取相鄰數據會更快。因此，優化內存訪問順序可以提高性能。MLC通過變換和優化計算圖，使得模型的執行更加高效。
TVM支持運算元自動融合，將多個運算元合並為一個，以優化模型執行。融合後的模型可以通過自動圖優化在計算圖之間進行變換，並將結果反饋到TensorIR中。此外，MLC可以集成現有的高度優化的矩陣計算庫，如MKL、CUBLAS、CUDNN和CUTLASS，以提高性能。
MLC的基本思想與手工優化類似，但通過計算圖融合和優化單個運算元性能來提升模型性能。編譯器採用計算調度分離的方式，通過自動程序優化尋找更優的運算元實現。在有先驗知識的情況下，可以結合手工優化來調整運算元實現，以進一步提升性能。
總結，機器學習編譯的基本思想旨在通過計算圖融合和運算元性能優化，以及集成現有優化庫，提高機器學習模型在不同部署環境中的執行效率。同時，MLC與手工優化相結合，可以進一步提升模型性能，實現更高效和靈活的機器學習應用部署。