🔢torch代码笔记

type
status
date
slug
summary
tags
category
icon
password
AI summary

🏹Lowering 总览

notion image
fx_codegen_and_compile: 编译fx图
  • GraphLowering.run(): 遍历每个fx node， lowering成inductor ir
  • GraphLowering.compile_to_fn(): 继续lowering成可行的function，包含系列优化 GraphLowering:
  • 内部由scheduler.generate完成codegen Scheduler: 负责inductor -> codegen， 会根据device选择不同codgen实现。
  • codegen():遍历node, 用对应device的backend进行codegen
  • BaseScheduling: 负责kernel code的codegen的backend , 比如cuda => CudaCombinedScheduling / TritonScheduling
    • 内部codegen时，具体的op会dispatch，调用到TritonKernel身上的函数去做
  • WrapperCodeGen: 负责wrapper code的codegen， 也就是调用kernel的那部分代码

🍇fx

graph.Graph

Node

  • Placeholder: graph input
  • get_attr: 从module身上拿一个属性，通常用来取weight
  • call_function: 调用函数
  • call_module：调用另一个module的forward函数
  • call_method: 在一个value身上调用其的一个method
  • Output: graph output

😺inductor

Ir. IRNode

IRNode基类
  • realize: 实例化到一个physical memorby.
    • 没有被realize的data, 包含了其计算的所有信息，具体在ir上表现为依然是一个Pointwise/Reduction 节点
    • 被realize的data，则是一个computed buffer，会对应一个physical memory
    • Poiontwise/Reduction的loader是 其inner_fn
    • computed_buffer的loader就是 ops.load
    • 所以realize实际上也就表示了终止可能的fuse
  • codegen_reference:?

Layout

基类，包含size, stride, offset，dtype, device信息

FixedLayout

make_indexer: 返回一个indexer函数 logical_coord => physical_index。 ([int] -> int)

FlexibleLayout

一个临时状态的layout， 可以根据需要改变stride

NonOwningLayout

Is a view into the storage of another tensor 是另一个tensor的alias

Loops

基类 tensor计算中，对循环的抽象，记录了range和inner_fn
  • range
  • innfer_fn: loop内的计算

Pointwise

继承自Loops， ElementWise的抽象
  • store_output: 返回一个op.store， 用于store output
  • make_loader: 返回inner_fn， 语义上loader就是一系列op，将自己load进reg/smem?

Scatter

继承自Loops，没看懂

Reduction

继承自Loops，reduction 类op的抽象

num_splits

根据reduction的规模，考虑到nv 硬件的特性，切分reduction，以获得更好的并行度。 返回类型： [ReductionHint, split_count]， 前者代表了切分的类型，后者代表具体切分的份数

Buffer

  • name
  • Layout
  • should_allocate: 标记当前buffer节点，是否需要新alloc memory， 比如view op的结果就不用，而matmul的结果就需要。
  • is_no_op: 是否没有实际操作。InputKernel/NopKernel没有实际操作，ComputedBuffer当 elements数量为0时，没有实际操作。
  • make_loader: 返回一个loader函数， logical_index -> element， 内部需要做logical_index -> phy_index的转换。 这里的index可以是symbol

InputBuffer

继承自Buffer，单纯的区别，没有额外信息

ConstantBuffer

ComputedBuffer

通过loop计算得到的buffer， 有一个类型为Loops的data成员
  • make_loader： 通常都是返回 ops.load
  • data: Loops

TritonTemplateBuffer

预先写好的template，实例化后的buffer。 Triton kernel的output tensor， 通过参数传进去。

MutableBox

对buffer的tensor 级别的抽象，抽象level更高 TensorBox / StorageBox allow in-place mutation of Tensors

StorageBox

  • mark_reuse: 判断是否应该把tensor放到memory中，否的话，每次使用都需要重计算, 判断的参考因素有：
    • User num
    • 计算此tensor是否复杂， 比如有 exp， 子op数量大于30
    • ...
  • realize: 将self.data用ComputedBuffer包起来

TensorBox

  1. Abstraction Level:
    1. TensorBox is a higher-level abstraction that represents a tensor and its operations. It's the primary interface for tensor manipulations in the inductor system.
    2. StorageBox is a mid-level abstraction that manages the actual storage and computation state of a tensor. It sits between TensorBox and lower-level representations like Buffer.
  1. Relationship:
    1. TensorBox typically wraps a StorageBox, as seen in the TensorBox.create method:
一个graph input的实例：

ExternKernel

Extern Python/cpp kerne调用，构造函数传入kernel_name。 没有返回值的extern kernel调用。 没有返回值，也就是除了kernel args外，不需要额外memory作为返回值的空间。

ExternKernelOut

需要预先alloc memory作为返回值buffer的 kernel
  • should_allocate: True

FallbackKernel

继承自ExternKernel 创建时，会把所有input都给realize了 output是一个MultiOutputLayout的东西，后续会用一个MultiOutput node从这个output中，把对应index的真正的buffer取出来， 这大概是因为整个inductor只能有一个返回值的背景下，一种hack的实现？

MultiOutput

从fallback kernel的output中，取出对应index的buffer
codegen后，不会有真正的操作， 是一条: buf2 = buf1

IndexingConstant

编译期已知的Index const

BaseView

  • make_indexer： 先从view_logical_index -> data_logica_index -> phy_logical->index
  • make_loader: 里面会做一次 index remap，从view_index -> data_index

ExpandView

Dependencies

Dep

基类，Dependency的抽象

MemoryDep

对memroy 依赖的抽象，代表对memory中部分区域/维度的访问， 记录了对memroy的access pattern， 对同一个memory，两种没有交集的access pattern的MemoryDep， 可能可以做一些事情？

StarDep

内部记录了依赖memory的名字。 与MemroyDep不同，StarDep代表对整个buffer的依赖。 Star代表了 *， 而 * 通常代表了全部， 比如正则表达式里的 *

WeakDep

两个node之间没有直接data dep， 但在执行顺序上依然可能有依赖，比如 write after read。
if A reads a buffer and B mutates it, B must be ordered after A

IndexExprDep

前面提到的Dep，依赖的对象都是buffer， IndexExprDep看起来是和index相关的依赖，依赖用于计算index的变量？

ReadWrites

读写操作的依赖集合, 可以merge, remove...

lowering.py

make_point_wise

inner_fn: 对所有input逐index， load, 然后调用fn

make_reduction_inner

里面会构造Reduction的inner_fn，Reduction的inner_fn，语义上是 (ouput_index, reduction_index) -> input_ele 内部做了个index的转换，包了下input的loader
这里保证了 new_index的维度，符合inner_loader的维度要求。

make_reduction

创建reduction， 如果reduction没有unroll， 则直接realize

GraphLowering

继承自 fx.Interpreter
  • run： 挨个run_node， 将每个node lowering成: ExternKernel, ComptuedBuffer.....
  • Codegen: inductor ir -> output_code.py, 内部核心：
    • scheduler.codegen()
    • scheduler.wrapper_code.generate()

Schedule

NodeUser
  • can_inplace

BaseSchedulerNode

SchedulerNode基类
  • users
  • last_usage: # buffers that won't be used after this kernel， 用于在合适位置del buffer
  • can_inplace: 判断函数'
  • codegen

decide_inplace_update

决定是否inplace,前提条件：
  • 当前node是最后一个user
  • node_user.can_inplace == True

SchedulerNode

与ir.node的比较： ir.Node:
  • 关注计算本身逻辑的抽象 SchedulerNode:
  • wrap了ir.Node
  • 关注计算如何被执行, scheduling
    • 执行顺序
    • 优化可能
    • ...
  • ir.Node:
    • Computation definition
    • Memory layout
    • Data dependencies
    • Operation implementation
  • SchedulerNode:
    • Execution ordering
    • Resource allocation
    • Performance optimization
    • Dependency management
    • Fusion decisions

can_inplace

不是template_buffer 只有一个write。并且是MemoryDep类型 read和write的index以及size一致

ExternKernelSchedulerNode

主要处理extern kernel，没法做fuse， 处理side effects

Scheduler

构造时传入Buffer list， 为每个buffer都会创建一个对应的Scheduler Node
  • can_fuse_vertical
  • compute_dependencies:考虑到mutaion和aliasing， 计算node间依赖. 会填写 SchedulerNode::users

init(self, ir.buffers)

  1. 为每个ir.Buffer，创建schedulerNode
  1. compute_dependeccies
  1. topological_sort_schedule
  1. dead_node_elimination ： Remove any nodes without users
  1. compute_ancestors： 计算每个node的所有祖先节点
  1. fuse_nodes

codegen

nodes -> lines, 不同type的IRNode有不同的codegen逻辑
  • Extern kernel: 调用node.codgen free buffer的timing点：
  • Codegen extern kernel后 Alloc buffer的timing点:
  • Extern kernel out, should alloc = true
  • SchedulerNode

can_fuse(node1, node2)

score_fusion_memory:

fuse后，减少的memory op的 size即：
  • 既是那些就是属于node1的input，又属于node2的input， fuse后只用load一次。

can_fuse_vertical

  • 访问模式相同
  • Index
  • size

fuse_nodes

尝试10次fuse_nodes_once

fuse_nodes_once

KernelFormatterHandler

对于Loops的inner_fn，进行format到str

🤖codegen

🛳Wrapper Code

  • allocate_groups： 对所有group排序后进行alloc

WrapperCodegen

  • allocate_groups： 对所有group排序后进行alloc

memory_plan_reuse

while 判断最后一条line是否还是 MemoryPlanningLine， 如果是，说明plan没有结束，继续下一轮plan。 每一轮plan，都是遍历所在MemoryPlanningLine, 执行其 plan 函数

memory_plan

  1. 初始化buffer_groups， AllocateLine会创建一个group, reuseLine加入其input在的group
  1. 将原本的line，替换为对应的PoolLine
  1. 计算buffer group 的live range

codegen_free

释放memory 对于input buffer直接codegen “del buf"，因为input buffer不可复用。 一些特殊buffer，不做free，如：constants... 对于其它buffer，则writeline: FreeIfNotReusedLine

codegen_allocation

分配memory, writeline: AllocatieLine 对于以下情况不做处理：
  • ir.ExternKernelAlloc: 会在kernel内部自己分配
  • ir.MultiOutput: 只是从之前的大buffer中，取出对应的子buffer
  • NonOwningLayout：view / ir.ReinterpretView

ReuseKey

= Tuple[device, dtype, storage_size]

Line

Wrapper code的一行

WrapperLine

Base class for all wrapper code line, empty mem funcs:
  • plan
  • codegen

MemoryPlanningLine

Base class for mem plan line

AllocateLine

Alloc a new mem
  • Plan: 从planning_pool里面reuse_key查找是否有符合需求的buffer. 有就返回一个ReuseLine

FreeIfNotReusedLine

Free a mem, 如果这个buffer后续没有被reuse，就free。 在创建这个line的时候，还不知道这个buffer在后面会不会被reuse
  • Plan: 将自己push到 planing_pool里面

ReuseLine

对于inplace的node， 需要先把einterpret_tensor， 把input转成output的layout
  • plan: 不做撒动作

PoolMemoryPlanningLine

memory pool操作的抽象
  • group: 所在group
  • timestep: live_range的index值

AllocFromPoolLine

Similar to AllocationLine, but takes memory from a pool"""

DeallocFromPoolLine

Similar to FreeIfNotReusedLine, but takes memory from a pool"""

memory_planning

AllocationTreeNode

Abstract base class for nodes in allocation pool.
  • get_live_ranges -> LiveRanges
  • get_size_hint -> int
  • get_symbolic_size -> sympy.Expr

Empty

Placeholder to represent empty space in the allocation pool. Only exists to get the size_hint correct in parent nodes. AllocationTree的叶子节点， 只是一个占位符号

Allocation

Represents memory allocated to a given node in the allocation pool. AllocationTree的叶子节点
  • offset
  • live_range

TemporalSplit

时间分片，包含一个allocation list， 每一项代表了同一块空间的mem, size相同，但是live range互不重叠。 Contains a list of allocations not overlapping in LiveRanges. 创建timing点：
  • 创建AllocationPool时，会带一个初始的temporalSplit作为root
  • expand root时，新增的部分也是一个temporalSplit, 新旧temporalSplit一起构成一个spatialSplit， 作为新root temporalSplit的唯一allocation
  • allocations: List[AllocationTreeNode]
  • get_size_hint: max of all allocations ， 实际每次新增时间分片时，保证了每个分片的size相同
  • get_live_rangs: join of ranges of each allocation
  • _allocate(self, block):
    • 判断有无overlap
    • 如果over lap的数量大于1，直接失败, 这里注释也说了，可以进一步判断，但是目前没实现
    • 如果over lap数量\==1， 尝试在与其overlapping的allocation上alloctate一次
    • 没有overlap
        1. 如果大小刚好合适 slot_size == block_size， 直接append(block)
        1. 如果slot_size > block_size: 新增一个 SpatialSplit
        1. expand？ 所以allocation list会在没有overlap，并且slot_size > block_size时，新增一个spatial split (block_size, slot_size - block_size) 所以allocation list的每一项都代表了对应pool的full memory， 但是每一项之间在live_range上不重叠。

SpatialSplit

空间切分，分片之间live_range有重叠，但不一定完全一致。size也不一定一致。 Contains two allocations, left and right, that do not overlap in space. Right will be allocated immediately after left in memory.
  • left: TemporalSplit
  • right: TemporalSplit
  • get_size_hint: left_size + right_size
  • get_live_range: left_ranges + right_ranges

AllocationPool

代表一次调用torch.empty / alloc 所创建的memory timing to create a new pool:
  • allocate一个block时，如果无法在现有的pool中的alloc一块mem
  • allocate
    • try alloc in root
    • expand root

AllocationPools

Collection of many AllocationPool objects grouped by device.
notion image

MemoryPlanner

利用live_range对memory 进行planning 具体参考：[[torch memory planning]]

allocate_groups

核心plan过程 分配顺序：按照(size, len(live_range))进行降序排序

🏭 Kernel Code

SIMDKernel

  • index_dtype: "tl.int32"...
  • numels: [4, 8192], 最后一维代表reduction维度？
  • inside_reduction: numels[-1] != 1

TritonKernel(SIMDKernel)

  • codegen_kerenl():
    • codegen meta信息，函数签名信息等
    • codegen_body(): 拼接 index_code, loads, compute, stores, suffix
  • 身上放了各种 inductor op -> triton op 的逻辑 {my-arrow}
    • load(name, index): codegen reductionA
    • to_dtype():
    • reduction():
    • ...

SIMDScheduling

对 numel和reduction numel感知
  • codegen_node_schedulte(): 可以分为下面两步
    • codegen_node_schedule_with_kernel(): 遍历node，进行codegen
      • kernel.split_and_set_ranges(node.get_ranges())： codegen 循环变量, r1 = rindex
      • node.codegen(index_vars): codegen 循环变量对应的循环体, 内部将循环体捕获成了一个fx graph， 然后通过 set op handler + run_node dispatch 到 TritonKernel::xxxOP 上面
    • kernel.codegen_kernel()

tuned_mm

预定义的多个config， config里定义了block_m/n/k, num_stage/num_wraps。lowering的时候，实际跑了看性能选择.

📮Other

Inductor ir case

fx.graph lowering to Inductor ir

fx_codegen_and_compile中完成： Lowering by run: 每个node type都register lowering方法，一边lowering一边做fuse

realize的timing点

  • 创建Reduction，并且没有被unroll，会将reduction realize
  • 创建Fallback kernel时，会将所有input realize
  • Fallback kernel的output
  • Lowering aten.mm时，会将input都realize了， 并且如果realize后，不是连续的，会copy一遍。 这里的连续指：某一维度的stride = 1
  • Template buffer， 本身就是一个buffer了。 tuned_mm创建triton template时，output就是个template buffer

tuned_mm

triton_codegen

inductor fuse的timming点

  • GraphLowering.run() 将fx_node lowering成inductor ir的时候
 

添加自定义c++算子

  1. 定义算子
  1. 为指定后端添加实现
  1. torch.compile里面，会用到fake_mode，进行shape推演，需要op注册fake kernel， fake kernel处理shape相关的逻辑，返回一个正确shape的fake tensor
 
Prev
torch memory planning
Next
Ansor论文笔记
Loading...