FreeTensor入门指南：编写你的第一个张量计算程序

作为一款高性能张量计算框架，FreeTensor为开发者提供了简洁高效的编程接口。本文将带你了解FreeTensor的核心概念，并通过实例演示如何编写第一个FreeTensor程序。

向量加法示例

让我们从一个简单的向量加法程序开始：

import freetensor as ft
import numpy as np

n = 4

@ft.optimize
def test(a: ft.Var[(n,), "int32"], b: ft.Var[(n,), "int32"]):
    y = ft.empty((n,), "int32")
    for i in range(n):
        y[i] = a[i] + b[i]
    return y

result = test(np.array([1, 2, 3, 4], dtype="int32"),
             np.array([2, 3, 4, 5], dtype="int32")).numpy()
print(result)

这个示例展示了FreeTensor的基本工作流程：

1. 定义一个操作张量的Python函数
2. 使用@ft.optimize装饰器标记
3. 调用函数时，FreeTensor会自动生成并编译高效的原生代码

核心概念解析

张量的声明与定义

在FreeTensor中，所有张量（包括函数参数、中间变量和返回值）都需要明确定义：

• 函数参数使用ft.Var[shape, data_type]语法声明
• 中间变量可通过ft.empty、ft.var或ft.zeros创建
• 标量被视为0维张量

张量操作

操作张量的方式与常规Python代码类似，但需要注意：

• 使用for...in range()循环遍历张量元素
• 张量赋值必须使用切片语法（如y[i] = ...）
• 支持完整的算术运算和部分张量操作函数

动态与静态代码

FreeTensor代码分为两类：

动态代码：

• 实际参与编译的部分
• 包括张量操作和涉及张量的控制流语句

静态代码：

• 编译前执行的"元编程"部分
• 可以是任意Python代码

这种区分使得我们可以在保持高性能的同时，利用Python的灵活性构建复杂计算逻辑。

高级特性

JIT编译支持

当程序需要适应不同输入大小时，可以使用JIT编译：

@ft.optimize
def test(n: ft.JIT, a, b):
    a: ft.Var[(n,), "int32"]
    b: ft.Var[(n,), "int32"]
    y = ft.empty((n,), "int32")
    for i in range(n):
        y[i] = a[i] + b[i]
    return y

FreeTensor会自动为不同的n值生成优化代码，并缓存编译结果。

动态形状支持

对于运行时才能确定形状的情况，可以使用动态形状：

@ft.optimize
def test(n: ft.Var[(), "int32"], a, b):
    a: ft.Var[(n,), "int32"]
    b: ft.Var[(n,), "int32"]
    y = ft.empty((n,), "int32")
    for i in range(n):
        y[i] = a[i] + b[i]
    return y

这种方式只需编译一次，但某些优化可能无法应用。

PyTorch集成

FreeTensor可与PyTorch无缝协作：

@ft.optimize
def test(a: ft.Var[(n,), "int32"], b: ft.Var[(n,), "int32"]):
    y = ft.empty((n,), "int32")
    for i in range(n):
        y[i] = a[i] + b[i]
    return y

y = test(torch.tensor([1, 2, 3, 4], dtype=torch.int32),
         torch.tensor([2, 3, 4, 5], dtype=torch.int32)).torch()

还支持自动微分集成：

@ft.optimize_to_pytorch
def test(a: ft.Var[(n,), "float32"], b: ft.Var[(n,), "float32"]):
    y = ft.empty((n,), "float32")
    for i in range(n):
        y[i] = a[i] * b[i]
    return y

总结

FreeTensor通过简洁的Python接口提供了高性能的张量计算能力。本文介绍了基本编程模式、核心概念和高级特性，帮助你快速上手FreeTensor开发。通过合理运用静态/动态代码分离、JIT编译等特性，你可以构建既灵活又高效的数值计算程序。