沒有 NVIDIA 顯卡也能練 AI？工程師怎麼用「翻譯層」讓 AMD 顯卡假裝自己是 CUDA

# -*- coding: utf-8 -*-
"""
shim_demo.py
玩具版「CUDA 翻譯層」示範：模擬 ZLUDA 如何把「寫給 NVIDIA 的呼叫」
轉譯成「AMD 顯卡看得懂的指令」，並用一個迷你神經網路前向計算
當作驗證，呼應素材裡用 PPO 網路做真實驗證的做法。
"""

import hashlib
import random
import time

# 素材裡驗證過的顯卡型號，只有這張是「保證能動」
VALIDATED_GPU = "AMD Radeon RX 9060 XT (gfx1200)"

# 目前你的「顯卡」（可以自己改這行，模擬換一張卡會發生什麼事）
CURRENT_GPU = "AMD Radeon RX 9060 XT (gfx1200)"


def rocblas_gemm(a, b):
    """AMD 官方數學運算庫做矩陣乘法，真正動手算的人。"""
    rows_a, cols_a = len(a), len(a[0])
    cols_b = len(b[0])
    result = [[0.0] * cols_b for _ in range(rows_a)]
    for i in range(rows_a):
        for j in range(cols_b):
            result[i][j] = sum(a[i][k] * b[k][j] for k in range(cols_a))
    return result


def rocblas_relu(matrix):
    return [[max(0.0, v) for v in row] for row in matrix]


# 翻譯表：左邊是程式「以為」自己在呼叫的 CUDA 函式，
# 右邊是真正負責執行的 AMD 函式（呼應素材裡的架構圖）
TRANSLATION_TABLE = {
    "cublasGemm": rocblas_gemm,
    "cublasRelu": rocblas_relu,
}


def cuda_call(cuda_func_name, *args):
    """
    這支函式假裝自己是 CUDA API。
    任何程式呼叫它，其實都會被偷偷轉發到 AMD 的實作。
    這就是 ZLUDA 在素材裡做的「翻譯層」角色。
    """
    if cuda_func_name not in TRANSLATION_TABLE:
        raise NotImplementedError(f"這個翻譯層還沒學會翻譯：{cuda_func_name}")

    backend_func = TRANSLATION_TABLE[cuda_func_name]
    start = time.perf_counter()
    result = backend_func(*args)
    elapsed_ms = (time.perf_counter() - start) * 1000
    print(f"[翻譯層] 收到 CUDA 呼叫 {cuda_func_name} -> 轉交給 AMD 後端執行，耗時 {elapsed_ms:.3f} ms")
    return result


def verify_file_integrity(fake_file_bytes, expected_hash):
    """模擬素材裡下載完 LibTorch 後核對 SHA-256 雜湊的步驟。"""
    actual_hash = hashlib.sha256(fake_file_bytes).hexdigest()
    return actual_hash == expected_hash


def toy_ppo_forward_pass():
    """
    迷你版「強化學習網路前向計算」，呼應素材裡
    2,216,347 個參數的 PPO 網路做 forward/inference 驗證的做法，
    這裡把參數量縮小到肉眼看得懂的程度。
    """
    random.seed(6)
    # 假裝這是「觀察到的環境狀態」，例如遊戲 NPC 看到的 3 個數值
    state = [[0.5, -0.2, 0.8]]
    weights_1 = [[random.uniform(-1, 1) for _ in range(4)] for _ in range(3)]
    weights_2 = [[random.uniform(-1, 1) for _ in range(2)] for _ in range(4)]

    hidden = cuda_call("cublasGemm", state, weights_1)
    hidden = cuda_call("cublasRelu", hidden)
    action_scores = cuda_call("cublasGemm", hidden, weights_2)
    return action_scores


def main():
    print(f"驗證用的顯卡型號：{VALIDATED_GPU}")
    print(f"目前偵測到的顯卡型號：{CURRENT_GPU}")

    if CURRENT_GPU != VALIDATED_GPU:
        print("警告：這張顯卡沒有在素材裡被驗證過，接下來的結果只能參考，不保證正確。")

    print("\n--- 第一步：模擬下載檔案，核對雜湊，防止檔案被動手腳 ---")
    fake_file = b"pretend-this-is-a-2.66GB-libtorch-package"
    expected_hash = hashlib.sha256(fake_file).hexdigest()
    ok = verify_file_integrity(fake_file, expected_hash)
    print(f"雜湊核對結果：{'通過' if ok else '失敗，中止安裝'}")

    print("\n--- 第二步：透過翻譯層執行一次迷你 PPO 網路前向計算 ---")
    scores = toy_ppo_forward_pass()
    print(f"網路輸出的行動分數：{scores}")

    print("\n--- 驗證結論 ---")
    print("所有翻譯呼叫都成功轉交給 AMD 後端執行，沒有丟出例外，")
    print("這就是素材裡「不只是能開機，是真的能訓練 AI」的縮小版證明。")


if __name__ == "__main__":
    main()