【技术尝鲜】Mojo 1.0 Beta 来了！写起来像 Python，跑起来像 C++

┌─────────────────────────────────────────────────────────┐
│                    Mojo Source Code                      │
│              (Python-like syntax + types)                │
└──────────────────────┬──────────────────────────────────┘
                       │
                       v
┌─────────────────────────────────────────────────────────┐
│                  Mojo Frontend                          │
│         (Parsing, Type Checking, Semantic Analysis)     │
└──────────────────────┬──────────────────────────────────┘
                       │
                       v
┌─────────────────────────────────────────────────────────┐
│              MLIR Dialects (Multi-Level IR)              │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐              │
│  │  Mojo    │  │  Linalg  │  │  GPU     │   ...        │
│  │  Dialect │  │  Dialect │  │  Dialect │              │
│  └──────────┘  └──────────┘  └──────────┘              │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐              │
│  │  SIMD    │  │  Tensor  │  │  MemRef  │   ...        │
│  │  Dialect │  │  Dialect │  │  Dialect │              │
│  └──────────┘  └──────────┘  └──────────┘              │
└──────────────────────┬──────────────────────────────────┘
                       │
                       v
┌─────────────────────────────────────────────────────────┐
│              Target-Specific Code Generation             │
│    ┌─────────┐    ┌─────────┐    ┌─────────┐           │
│    │  x86    │    │  ARM64  │    │  NVIDIA │           │
│    │  LLVM   │    │  LLVM   │    │  PTX    │  ...      │
│    └─────────┘    └─────────┘    └─────────┘           │
└─────────────────────────────────────────────────────────┘

传统 CUDA/CUTLASS 内核开发:
┌──────────────────────────────────────┐
│         单一巨大函数 (3000+ 行)        │
│  数据加载 + 计算 + 同步 + 存储         │
│  所有逻辑耦合在一起，改一处动全身       │
└──────────────────────────────────────┘

Mojo Structured Kernel:
┌──────────────────────────────────────┐
│  TileIO     - 内存层级间数据搬运       │
├──────────────────────────────────────┤
│  TilePipeline - 流水线编排/共享内存管理 │
├──────────────────────────────────────┤
│  TileOp     - 实际计算操作             │
└──────────────────────────────────────┘
关注点分离：添加新功能只需修改对应模块

# 1.0 Beta 之前
fn strict_func(x: Int) -> Int:
    return x * 2

def flexible_func(x):
    return x * 2

# 1.0 Beta 之后：统一用 def
def func(x: Int) -> Int:
    return x * 2

# 无状态闭包自动提升为顶层函数（兼容 FFI 回调）
def process_data(callback: fn(Int) -> Int):
    callback(42)

# 使用 ref 捕获约定
def make_counter() -> fn() -> Int:
    var count = 0
    def increment() -> Int:
        count += 1
        return count
    return increment

# thin 函数效果：纯函数指针，无闭包上下文
def apply[func: fn(Int) -> Int](x: Int) -> Int:
    return func(x)

# 非空指针（默认）
var ptr: UnsafePointer[Int] = ...  # 保证非空

# 可空指针
var maybe_ptr: Optional[UnsafePointer[Int]] = None  # 显式可空

var arr = [1, 2, 3, 4, 5]

# ❌ 编译错误：不支持负索引
# print(arr[-1])

# ✅ 正确做法
print(arr[len(arr) - 1])

# 边界检查默认开启（CPU 上）
# arr[10]  # 运行时错误：Index out of bounds

传统张量:                    TileTensor:
┌─────────────────┐         ┌─────────────────┐
│ Tensor          │         │ Tensor[float32,  │
│  - data         │         │   layout=ROW_MAJOR│
│  - shape (运行时)│         │   swizzle=NONE   │
│  - stride (运行时)│        │ ]                 │
│  - layout (运行时)│        │ 编译期已知所有布局  │
└─────────────────┘         └─────────────────┘

# Python + NumPy
import numpy as np

def matmul(a, b):
    return np.dot(a, b)

# 性能：依赖底层 BLAS 库，约 10-100 GFLOPS
# 代码量：1 行（但黑箱，无法自定义优化）

// Rust + ndarray
use ndarray::Array2;

fn matmul(a: &Array2<f32>, b: &Array2<f32>) -> Array2<f32> {
    a.dot(b)
}

// 性能：接近 BLAS 水平
// 代码量：1 行 + 大量类型声明
// 学习曲线：陡峭（所有权、生命周期）

# Mojo
def matmul[
    M: Int, N: Int, K: Int
](a: Matrix[M, K], b: Matrix[K, N]) -> Matrix[M, N]:
    var result = Matrix[M, N]()
    for i in range(M):
        for j in range(N):
            var sum = SIMD[DType.float32, 0](0)
            for k in range(K):
                sum += a[i, k] * b[k, j]
            result[i, j] = sum
    return result

# 性能：通过 SIMD 向量化 + MLIR 优化，接近手写汇编
# 代码量：约 10 行
# 学习曲线：平滑（Python 语法 + 可选类型标注）

# 你可能会发现很多 Python 库在 Mojo 中不可用
import requests  # ❌ 不支持，这是 CPython 库

# 先用 Python 写主体逻辑
# 对性能关键的函数用 Mojo 重写
from python import Python

def optimized_kernel(data: PythonObject) raises:
    # 用 Mojo 的 SIMD 能力加速
    ...

# ❌ 旧语法（1.0 Beta 之前）
fn compute(x: Int) -> Int:
    return x * 2

# ✅ 新语法（1.0 Beta）
def compute(x: Int) -> Int:
    return x * 2

# 生产环境可考虑禁用边界检查（通过编译选项）
# mojoc --disable-bounds-checking my_module.mojo

# GPU 内核中不能随意 print（虽然 Apple Metal 已支持）
# 建议：先在 CPU 上验证逻辑，再移植到 GPU
def kernel_cpu(a: ..., b: ..., result: ...):
    # CPU 版本，可以用 print 调试
    ...

def kernel_gpu(a: ..., b: ..., result: ...):
    # GPU 版本，移植时注意边界条件
    ...

# macOS / Linux 安装（推荐使用 uv）
curl -LsSf https://astral.sh/uv/install.sh | sh
uv init hello-mojo
cd hello-mojo
uv add mojo

# 验证安装
mojo --version

# hello.mojo
def main():
    print("Hello, Mojo!")

    # SIMD 向量化计算
    var a = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0)
    var b = SIMD[DType.float32, 4](5.0, 6.0, 7.0, 8.0)
    print(a + b)  # [6.0, 8.0, 10.0, 12.0]

    # 编译期元编程
    comptime width = simd_width_of[DType.float32]()
    print("SIMD width:", width)

mojo hello.mojo

from python import Python

def main():
    # 调用 Python 库
    var np = Python.import_module("numpy")
    var arr = np.array([1, 2, 3, 4, 5])
    print("NumPy array:", arr)
    print("Sum:", np.sum(arr))

    # 在 Mojo 中操作 NumPy 数组的底层数据
    var ptr = arr.ctypes.data.unsafe_get_as_pointer[DType.float64]()
    for i in range(5):
        print(ptr[i])