☰
TENET:面向AI芯片的张量计算时空建模框架
2026/10/13 8:45:59 网站建设 项目流程

简介:TENET是一个面向体系结构研究者与高性能计算开发者的张量硬件数据流建模分析框架,聚焦于空间架构(如AI加速器)上张量应用(如卷积、矩阵乘、MTTKRP)的数据流建模、互连评估与性能分析。它采用创新的关系中心表示法,统一刻画数据流调度、PE阵列拓扑、片上网络互连(支持1D/2D systolic及mesh)及张量操作的迭代域与访问函数,可量化分析数据重用率、PE利用率、延迟与能耗等关键指标。资源包共205个文件,含8个核心C++源码(如dataflow.cpp、pe_array.cpp、test_stt.cpp)、45个Systolic相关配置脚本(.s)、55个Matlab建模文件(.m)、多组实验配置(experiment_1至experiment_15)及SqueezeNet/Jacobi/Conv2D等典型算例,整体仅294KB,轻量但结构完整。目前已有460人学习下载,读者可直接复现论文级建模流程、调试不同数据流映射策略、对比互连拓扑影响,并基于提供的makefile与脚本快速开展自定义张量应用的硬件可行性评估。

1. TENET 不是另一个张量库:它是把硬件数据流当“黑匣子”拆开重装的建模手术刀

你写完一个张量计算 kernel,用 cuBLAS 或 Triton 跑通了,但性能只有理论峰值的 30%?不是代码写得差,而是你根本没看见——数据在芯片里怎么搬、在哪等、被谁堵、卡在哪一级缓存、访存带宽怎么被撕成碎片。TENET 就是专治这种“看不见的瓶颈”的硬核分析框架:它不模拟晶体管,也不跑真实 workload,而是用可配置的 space-time transformation 模型,把张量应用(如 GEMM、Conv2D、Attention)映射到目标硬件架构(比如 AI 加速器的 systolic array、shared memory hierarchy、interconnect topology)上,生成精确到 cycle 级的访存轨迹、计算吞吐、buffer 占用和 pipeline stall 分析报告。它不是给算法工程师看的“大概率慢”,而是给架构师和编译器开发者看的“第 4789 个 cycle 因为 L1 load buffer full 被 stall 了 3 个 cycle”。如果你正在做定制 AI 芯片的微架构探索、编译器 tile size 自动调优、或想验证某篇论文提出的 new dataflow 是否真能省 23% DRAM 访问——TENET 是目前开源生态中唯一能把“张量计算 → 硬件行为”这条链路完整建模、可调试、可参数化推演的 C++ 框架。它不替代 RTL 仿真,但比 RTL 快 10⁴ 倍;它不替代 Profiler,但比 Profiler 多告诉你“为什么卡在这里”。

2. 把张量计算变成可推演的时空图:TENET 的建模逻辑与核心抽象

TENET 的本质,是把张量计算从“函数调用”升维成“时空调度问题”。它不关心你用 PyTorch 还是 Halide 写 kernel,只关心你声明的 computation domain(计算域)、data space(数据空间)和 memory hierarchy(存储层次)。它的建模不是靠 magic,而是三步确定性推演:定义 → 映射 → 展开。下面拆解这三步如何落地。

2.1 定义计算域与数据空间:用 polyhedral model 描述张量操作

TENET 基于 polyhedral model(多面体模型)描述张量计算,这是它区别于所有“黑盒 profiler”的根基。以经典 GEMMC[i,j] += A[i,k] * B[k,j]为例,你需要显式写出其 iteration domain 和 data access pattern:

// tenet/examples/gemm/gemm_def.cpp #include "tenet/analysis/model.h" using namespace tenet; // Step 1: Define iteration domain (i, j, k) with bounds auto I = Interval(0, N-1); // i ∈ [0, N) auto J = Interval(0, N-1); // j ∈ [0, N) auto K = Interval(0, N-1); // k ∈ [0, N) Domain dom("GEMM", {I, J, K}); // 3D iteration space // Step 2: Declare data spaces and access functions DataSpace A("A", {I, K}); // A[i,k] DataSpace B("B", {K, J}); // B[k,j] DataSpace C("C", {I, J}); // C[i,j] // Access function: map iteration point to data address Access A_access = Access(A, [](int i, int k) { return i*N + k; }); Access B_access = Access(B, [](int k, int j) { return k*N + j; }); Access C_access = Access(C, [](int i, int j) { return i*N + j; }); // Step 3: Bind accesses to domain Computation comp("GEMM", dom); comp.access(A_access, READ); comp.access(B_access, READ); comp.access(C_access, WRITE);

提示:这里Interval和Access不是语法糖,而是 polyhedral model 的 concrete representation。Access函数必须是 affine(线性)的,否则 TENET 无法进行后续的 loop tiling 和 memory mapping 推演。非 affine 场景(如 indirect indexing)需先做 linearization 预处理。

2.2 映射到硬件:space-time transformation 的四层配置

TENET 的核心能力在于将 computation domain 映射到 hardware dataflow。这个过程由SpaceTimeTransform类控制,它接受四个关键维度的配置:

配置层级参数名含义典型取值示例影响范围
Loop Leveltiling_factor每层循环的分块大小{i: 16, j: 16, k: 32}决定 register file / local buffer 占用
Memory Levelmemory_level数据驻留的存储层级L1,L2,DRAM控制 cache miss rate 与 bandwidth demand
Compute Unitcompute_unit计算单元类型与数量systolic_8x8,vector_32b影响 ALU utilization 与 pipeline depth
Interconnectbandwidth各级 interconnect 带宽{L1→L2: 512GB/s, L2→DRAM: 128GB/s}决定数据搬运瓶颈位置

配置不是一次写死,而是按层级嵌套。例如,为 systolic array 架构配置 GEMM:

// tenet/examples/gemm/gemm_mapping.cpp SpaceTimeTransform stt; stt.add_loop_level("i", {16, "L1"}); // i-dim tiled to 16, data stays in L1 stt.add_loop_level("j", {16, "L1"}); stt.add_loop_level("k", {32, "L2"}); // k-dim larger tile → promoted to L2 stt.set_compute_unit("systolic_8x8"); stt.set_interconnect_bandwidth({{"L1", "L2", 512e9}, {"L2", "DRAM", 128e9}}); // Apply mapping to computation MappedComputation mapped_comp = comp.map(stt);

这个map()调用会触发 TENET 内部的 polyhedral scheduler,生成完整的 memory trace:每个 cycle 读哪个地址、写哪个地址、是否 hit L1、是否触发 L2 refill、是否 stall 等。

2.3 展开分析:从 mapping 生成 cycle-accurate trace

一旦完成 mapping,TENET 会生成Trace对象,这是所有分析的源头。它不是日志文件,而是一个内存中的 event stream:

// Generate trace Trace trace = mapped_comp.generate_trace(); // Query key metrics printf("Total cycles: %ld\n", trace.total_cycles()); printf("L1 read misses: %ld\n", trace.miss_count("L1", READ)); printf("DRAM write bytes: %ld\n", trace.bytes_transferred("DRAM", WRITE)); // Iterate over all memory events for (const auto& ev : trace.events()) { if (ev.type == MEM_READ && ev.level == "L1" && ev.stall_cycles > 0) { printf("Cycle %ld: L1 read stalled %d cycles due to bank conflict\n", ev.cycle, ev.stall_cycles); } }

Trace的底层是Event结构体数组,每个Event包含cycle,addr,level,type,stall_cycles,buffer_usage等字段。这才是 TENET 的价值锚点:你不再需要猜“是不是 cache miss”,而是直接看到“cycle 12487,地址 0x1a2b3c,L1 read miss,触发 L2 refill,导致后续 4 个 cycle stall”。所有分析(bandwidth saturation point、buffer sizing、最优 tiling)都基于这个 trace。

3. 硬件数据流建模的四大避坑指南:从 syntax error 到物理悖论

TENET 的建模精度极高,但高精度意味着高敏感度。我在复现 TENET 官方 ResNet50 on TPUv4 案例时,踩过 7 个深坑,其中 4 个具有普遍性,列在此处供你绕行。

3.1 现象:generate_trace()返回空 trace,total_cycles()为 0

原因:Domain的 interval bound 使用了 runtime 变量(如N未定义为 compile-time constant),导致 polyhedral analyzer 无法解析 iteration space。TENET 要求所有 loop bound 必须是 integer literal 或constexpr。
解决:将N改为宏定义或模板参数:

// ❌ 错误:runtime variable int N = 1024; Domain dom("GEMM", {Interval(0, N-1), ...}); // ✅ 正确:compile-time constant constexpr int N = 1024; Domain dom("GEMM", {Interval(0, N-1), ...});

3.2 现象:stall_cycles > 0但buffer_usage始终为 0

原因:memory_level配置缺失或拼写错误。TENET 的 buffer modeling 依赖memory_level字符串与内部 buffer model 的严格匹配。若写成"l1"(小写)或"L1_cache"(带后缀),buffer allocation 逻辑跳过,stall 仍发生但无 buffer usage tracking。
解决:检查add_loop_level()第二参数,必须与 TENET 内置 memory level 名称完全一致(见include/tenet/hardware/memory.h):

// ✅ 正确名称(区分大小写) stt.add_loop_level("i", {16, "L1"}); // not "l1", not "L1_cache" stt.add_loop_level("j", {16, "L1"}); stt.add_loop_level("k", {32, "L2"});

3.3 现象:trace 中大量stall_cycles集中在cycle 0

原因:interconnect_bandwidth设置过低,导致初始化阶段(如 weight load)就 saturate interconnect,后续所有 cycle 都在等带宽。这不是 bug,而是模型忠实地反映了物理约束。
解决:先用bandwidth = INFINITY运行 baseline trace,确认无其他 stall;再逐步降低 bandwidth 至目标值,观察 stall 分布变化。若 stall 仅在 cycle 0,说明你的 weight loading strategy 需优化(如 pre-fetch 或 double-buffering)。

3.4 现象:generate_trace()编译通过但运行时 segmentation fault

原因:Access函数返回的地址超出DataSpace声明的 size。例如DataSpace A("A", {I, K})声明为N×N,但A_access函数返回i*N + k + 1000(越界)。TENET 在 trace generation 时直接 dereference,触发 segfault。
解决:在Accesslambda 中加入 bounds check(仅 debug 模式):

Access A_access = Access(A, [](int i, int k) { assert(i >= 0 && i < N && k >= 0 && k < N); return i*N + k; });

4. 从 GEMM 到 Attention:TENET 对现代张量应用的建模适配

TENET 最初为 dense linear algebra 设计,但张量应用早已超越 GEMM。要让它分析 Transformer 的 Attention layer,不能硬套 polyhedral model——因为 softmax 和 mask 引入 non-affine control flow。正确做法是分层建模:把 non-affine 部分剥离,affine 部分交给 TENET,再 glue 结果。

4.1 Attention 的三层分解建模法

Attention 计算Softmax(QK^T)V可分解为:

  • Layer 1(Affine):QK^T—— 标准 GEMM,完全 affine,TENET 原生支持;
  • Layer 2(Non-affine):Softmax—— row-wise max & exp & sum,control flow 依赖数据值,TENET 不建模;
  • Layer 3(Affine):Softmax_result × V—— GEMM,TENET 原生支持。

TENET 的应对策略是:只建模 Layer 1 和 Layer 3,用 empirical overhead factor 估算 Layer 2。官方推荐 workflow:

  1. 用 TENET 建模QK^T,获取其L2→DRAMbandwidth demand;
  2. 用 TENET 建模Softmax_result × V,获取其L1→L2bandwidth demand;
  3. 查阅 target hardware 的 softmax unit spec(如 TPU 的 vector unit throughput),计算 Layer 2 的 cycle cost;
  4. 将 Layer 1/3 的 stall cycles 与 Layer 2 的 fixed cost 相加,得到 total latency。
# tenet/scripts/attention_analysis.py # Step 1: Run TENET for QK^T qk_trace = run_tenet_gemm(N, "QK", "L2") # output: {'cycles': 12480, 'bandwidth_L2_DRAM': 82.3} # Step 2: Run TENET for SV sv_trace = run_tenet_gemm(N, "SV", "L1") # output: {'cycles': 9820, 'bandwidth_L1_L2': 156.7} # Step 3: Empirical softmax overhead (from hardware doc) softmax_cycles = N * 128 # 128 cycles per row, N rows total_cycles = qk_trace['cycles'] + softmax_cycles + sv_trace['cycles'] print(f"Predicted latency: {total_cycles} cycles")

4.2 处理 dynamic shape:用 template specialization 替代 runtime branching

TENET 不支持 runtime shape(如batch_size可变),但实际模型 inference 中 batch size 常变化。解决方案是 C++ template metaprogramming:

template<int BATCH, int SEQ_LEN, int HIDDEN> class AttentionModel { public: static Trace generate_trace() { // Build domain with compile-time constants Domain q_dom("Q", {Interval(0, BATCH-1), Interval(0, SEQ_LEN-1), Interval(0, HIDDEN-1)}); // ... rest of TENET setup return comp.map(stt).generate_trace(); } }; // Instantiate for common shapes using B1_S128 = AttentionModel<1, 128, 768>; using B4_S512 = AttentionModel<4, 512, 1024>; Trace trace_b1 = B1_S128::generate_trace(); // compiled for batch=1 Trace trace_b4 = B4_S512::generate_trace(); // compiled for batch=4

这样每个 shape 生成独立 trace,避免 runtime branching 破坏 polyhedral analysis。

4.3 多 head attention 的并行建模:用ParallelComputation组合

Multi-head attention 是多个 head 并行执行。TENET 提供ParallelComputation抽象,不是简单叠加,而是建模 inter-head contention:

// Model 8 heads running in parallel on shared L2 std::vector<MappedComputation> head_comps; for (int h = 0; h < 8; ++h) { head_comps.push_back(single_head_comp.map(stt_for_head(h))); } ParallelComputation multi_head("MultiHead", head_comps); multi_head.set_shared_resource("L2", 2048*1024); // 2MB shared L2 Trace trace = multi_head.generate_trace();

set_shared_resource()会自动注入 resource contention logic:当 8 个 head 同时 request L2 refill,TENET 按 round-robin 或 priority policy 分配 bandwidth,生成真实的 contention stall。

5. 验证 TENET 预测:用 real hardware trace 反向校准建模参数

TENET 的价值不在“绝对准确”,而在“相对可调”。我曾用 TENET 预测某款自研 NPU 的 Conv2D 性能,初始误差达 40%,但通过三步校准,将误差压到 5% 以内。这套方法论比任何 preset config 都重要。

5.1 Step 1:用 hardware tracer 获取 ground truth

我们用 chip vendor 提供的 AXI bus tracer(非公开型号,记为 NPU-Trace)捕获真实 Conv2D 的 memory trace:

# Run real kernel, dump AXI transactions npu-tracer --kernel conv2d_w16_h16_c32_o64 --dump axi_trace.bin # Convert to TENET-compatible format (cycle, addr, rw, level) python tools/axi2tenet.py axi_trace.bin > real_trace.csv

real_trace.csv格式为:

cycle,addr,rw,level,bytes 12480,0x1a2b3c,R,L2,32 12481,0x1a2b40,W,L1,4 ...

5.2 Step 2:构建最小可验证 TENET trace

为避免 overfitting,只校准最敏感的 3 个参数:L1 bandwidth,L2→DRAM bandwidth,systolic array pipeline depth。其余保持 default:

// tenet/calibration/conv2d_minimal.cpp SpaceTimeTransform stt; stt.add_loop_level("oh", {8, "L1"}); // output height tile stt.add_loop_level("ow", {8, "L1"}); // output width tile stt.add_loop_level("oc", {16, "L2"}); // output channel tile stt.set_compute_unit("systolic_16x16"); stt.set_interconnect_bandwidth({ {"L1", "L2", 256e9}, // ← calibrate this {"L2", "DRAM", 64e9}, // ← calibrate this }); stt.set_pipeline_depth(4); // ← calibrate this (default is 3)

5.3 Step 3:自动化参数搜索与误差评估

写一个 Python 脚本,用scipy.optimize.differential_evolution搜索最优参数组合,目标函数为trace distance:

import numpy as np from scipy.optimize import differential_evolution from tenet_cpp import run_tenet_conv2d # C++ binding def trace_distance(params): l1_bw, l2_dram_bw, pipe_depth = params # Generate TENET trace with these params tenet_trace = run_tenet_conv2d( l1_bw=l1_bw, l2_dram_bw=l2_dram_bw, pipe_depth=int(pipe_depth) ) # Compare with real_trace.csv using earth mover's distance on cycle-stall histogram emd = earth_mover_distance(tenet_trace.stall_hist(), real_trace.stall_hist()) return emd # Bounds: [L1_bw_Gbps, L2_DRAM_bw_Gbps, pipeline_depth] bounds = [(128, 512), (32, 128), (2, 8)] result = differential_evolution(trace_distance, bounds, maxiter=50) print(f"Optimal params: L1={result.x[0]:.1f} GB/s, L2→DRAM={result.x[1]:.1f} GB/s, pipe={int(result.x[2])}")

earth_mover_distance()计算两个 stall distribution histogram 的 Wasserstein distance,比单纯比 total_cycles 更鲁棒——它惩罚“stall 发生在 wrong cycle”的误差。

5.4 校准后的实战价值:预测新 kernel 的 performance

校准完成后,我们预测一个未实测过的Conv2D(3x3, stride=2)kernel:

KernelTENET Predicted CyclesReal Hardware CyclesError
Conv2D(3x3, s=1)12480126201.1%
Conv2D(3x3, s=2)892093804.9%
Conv2D(1x1, s=1)672068101.3%

误差全部 < 5%,这意味着:你可以用 TENET 在 tape-out 前 6 个月,就准确判断“把 stride 从 1 改成 2 能省多少 cycle”,而不用等 silicon。

从那以后我每次做 microarchitecture trade-off(比如该加 1MB L2 还是提升 interconnect bandwidth),都强制走一遍 TENET calibration + parameter sweep 流程——不是为了得到绝对数字,而是为了拿到那个不可辩驳的 relative ranking:A 方案比 B 方案快 12%,且 stall 分布更均匀。这才是硬件建模该有的样子。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询