KaiSpace

Shelf

Articles, videos, and links worth keeping.

A curated shelf of external work, sometimes with personal notes attached and sometimes kept as direct links.

Uncategorized

torch.compile 技术剖析(三):Lowering

1. realize 是什么? 在 Inductor 里,realize 的意思是:把一个还没有真正落到内存里的中间表达式,强制变成一个真实 buffer。 例如: a + b 不一定会生成中间 tensor。Inductor 可以直接融合成: 这里 a+b 没有 realize,没有写 HBM。 但如果是: 如果 mm lowering 到 cuBLAS,cuBLAS 只能接受真实矩阵指针: 它不能接受表达式: 所以必须先: 因此: 不是因为“后面要读数据”,而是因为后面的 extern kernel 需要真实 storage pointer。 Reduction 不一定需要 realize,因为如果 reduction 是 Inductor 自己生成的 kernel,它可以把表达式 inline 进去: 可以生成: 不需要先生成 tmp =

Read note
0 comments0 likes
tech

torch.compile 2.x 学习笔记

torch.compile粒度 torch.compile 的编译入口通常是一个 Python function/callable,但它的实际 JIT 粒度并不等于“整个函数大小”。更准确地说,Dynamo 会在函数执行过程中捕获能够连续 trace 的 tensor computation,并把这段连续区域表示成 FX graph;因此 torch.compile 的实际编译粒度是 FX graph region。一个函数可能被完整捕获成一张 FX graph,也可能因为 graph break 被切成多个 FX graph region;同时,同一个 region 还可能因为不同的 shape、dtype、device、Python flag 或控制流路径生成多个 specialized graph。 例如: 这里是以整个f作为一整个compiled region如果第一次调用是 f(x, True),Dynamo 实际 trace 到的是 x add mul relu 这条路径,并生成一张带有 guard 的 specialized

Read note
0 comments0 likes