cutlass-learning

本文最后更新于：3 天前

Cutlass和CuTe作用

CuTe是用于内存排布和张量表示的模板工具,cutlass是从多个层面逐层分解逐层解耦的高性能模板库,CuTe有利于代码编写时不用过于关心一些数据偏移的计算,而cutlass则将高性能算子实现拆解,如果要单独实现一个特定的,就得按照它的代码规范实现一个,不然只能用它已有的模板,因此学习CuTe有利于我们很好的处理复杂的tensor,不用数格子数个半死,学习cutlass的example可以让我们实现编写一些已有模板的算子,并且接上后续的激活函数于epilogue(算子融合相当于),学习cutlass源码有利于我们遵循其规范编写自定义高性能算子

cutlass2.x用于[Volta,Hopper) (←这是个区间),cutlass3.x因为wgmma,tma等一堆新东西,重写且增加了很多新东西,但也完全兼容cutlass2.x,所以打好2.x的底座我个人感觉很有必要

Cutlass2.x

cutlass四大软件抽象

  • 模板推导: 类似于Gemm的DefaultGemm,通过偏特化层层推导
  • 数据移动: 采用非常抽象的**Iteartor**,这个Iterator是一个笼统的概念,它有gmem->reg,gmem->smem等好多个iterator,它需要用到如Layout(数据排布),TensorMap(每个thread的分工)一类的,还有个Fragment这个类,它是抽象寄存器来用,本质是个一维定长数组
  • 数据计算: 计算,比如MmaTensorOp和MmaSimt这些个类,用tensor core还是cuda core
  • 流水编排: 即是软流水,比如MmaPipelined和MmaMultiStage这些个类

cutlass源码目录分类

include/cutlass下支持的计算的目录:conv,epilogue,gemm,reduction;别的目录是计算公用的:

  • arch
  • conv
  • detail
    • help_macros.hpp
      • 一些源码中常见的宏定义,比如很常见的CUTLASS_HOST_DEVICE,就是强制内敛的__device__ __host__,循环展开的以及对c++17的宏等等
  • epilogue
  • gemm
  • layout
  • pipeline
  • platform
    • platform.h
      • 宏定义用于标注操作系统信息、标注一些空的关键字（以使得不支持的c++标准也能编译）、常量表达式需要的函数(max,min);
      • stl函数和类型的重新实现，以使得能编译在设备端(加个__device__关键字); 使用的namespace: platform
      • 为当前的c标准提供一些会用到的但是可能是后续比如c11/14/17…才提供的stl函数和类型 ; 使用的namespace: platform
  • reduction
  • thread
  • transform

在include/cutlass下的头文件的作用(慢慢看,有空looklook它们):

  • aligned_buffer.h

  • array.h

  • array_planar_complex.h

  • array_subbyte.h

  • barrier.h

  • bfloat16.h

  • blas3.h

  • blas3_types.h

  • block_striped.h

  • cluster_launch.hpp

  • complex.h

  • constants.h

  • coord.h

  • core_io.h

  • cuda_host_adapter.hpp

  • cutlass.h

    这个头文件做了cutlass::Stauts这个enum class的定义,还定义了warp、warp group的线程数目,也有些算lane_id,warp_id,warp_group_id的小函数(用shfl_sync作广播(warpId那些)又同步)

  • device_kernel.h

  • fast_math.h

  • float8.h

  • floating_point_nvrtc.h

  • functional.h

  • gemm_coord.h

  • gemm_coord.hpp

  • half.h

  • integer_subbyte.h

  • kernel_hardware_info.h

  • kernel_hardware_info.hpp

  • kernel_launch.h

  • matrix.h

  • trace.h debug跟踪代码用

点进去单个计算如gemm目录,有device->kernel->threadblock->warp->thread这几个目录(现在还有个collective).在这些目录下,文件名用default作前缀的是负责模板推导的文件

  • device层主要是调kernel,类似于host端的功能,problemsize那些在这里处理,然后launch kernel;
  • kernel层(进入到了device端)解耦gemm和epilogue,同时确保了splitK的数据依赖,注意splitk是多个threadblock负责同一个结果阵的部分,这个核函数做完之后对它们进行归约操作,sliceK则是多个warp负责同一个结果阵的部分,结束后做一次归约,而正常的是iterationK,即一个tile size(threadblock负责的大小)的k跟warp size的k是相同的,不同则为slickK,而splitK有不同的上层模板类;
  • threadblock层负责流水编排和各种数据移动,又大又杂;
  • warp/thread: 做数据计算的,此外里面会定义一些移动的iterator如smem->reg,然后由threadblock调

部分别的目录(include/cutlass下)的作用:

  • conv: 由于cutlass的conv是im2col + gemm实现的,所以warp/thread目录的代码大量复用gemm的;
  • epilogue: 只有threadblock以下的,毕竟你gemm完后的数据在reg即Fragment实例里,由kernel层将gemm的threadblock的Fragment对象传给尾处理,即此刻只有从threadblock开始的流程;
  • transform: 负责一些通用的iterator,如gmem->reg和reg->smem

以下二者流水编排(发生在threadblock层)策略的区别:描述区别

  • pipelined
  • multistage

参考文件

example代码测试加源码理解

源码解读之gemm系列

先把example看完再弄这个,一大坨

CuTe

Layout

在cute体系下,逻辑空间被称作domain,代表存储的物理空间被称作codomain

对于逻辑空间上元素对应的位置用coordinate表示,而它每个位置的元素之间的间隔用stride的表示,通过内积可以得到对应的一维物理空间的地址,即:

index=∑coordinate∗strideindex = \sum{coordinate * stride}

Layout是个函数，由shape和stride表达，是支持层次化(hierarchy)的Tensor，即一个shape可以嵌套表达,同时stride又与shape里的每个参数一一对应：

shape((内tensor行,外tensor行1,外tensor行2,...),(内tensor列，外tensor列1,外tensor列2,...),...)

这其中的一些属性如:

  • shape通过所有行列元素进行累乘得到的元素个数
  • rank表示行,列,…这些个维度的数目
  • depth表示嵌套的深度,即单个维度内比如行有几层嵌套
  • coshape
  • cosize物理空间的大小,如果数据不紧凑,则cosize会大于size

同时layout的编译器常量在CuTe中可以通过Int<N>{}使用,也可以通过_1一类的方式使用

支持切片,不只可以通过coordinate获取单个元素的物理位置,可以通过slice获取符合我们选择的一系列元素的物理位置,cute提供了对某个维度全选的UnderScore类型的变量_(python或matlab用的是:)

补集,因为Layout本质是个函数,从集合的角度来讲,就是定义了从domain到codomain的映射,数据紧凑则codomain没有空洞,否则的话会产生空洞,可以利用补集把空洞的位置填充了

乘法,本质上就是对tensor的重复,类似于python中的[2] * 3得到的是[2,2,2],即shape:(x,y)和shape:(z,w)得到的是shape:(x,y,z,w),说人话就是元素由x*y个乘上z*w个后变成了x*y*z*w个元素,那这个操作的毕竟是个高维数据,它要怎么变化,就分化了几种不同的乘法模式出来

除法, layout除法是对层次的划分,比如10 / 5 = 2;对于layout而言就是10 / 5 = (5,2) ; 其中小括号中的5表示的是被划分的大小,2是被划分的数目,具体操作过程就是把被除数的layout按列拉成大长条,把除数也按列主序展开,然后与除数对应的位置元素进行位置变换,变换后的layout按行排列,其中最终得到的layout的shape的rank0是除数的shape,rank1表示的是被划分的数目

复合函数和逆 暂不理解

Tensor

cute中的Tensor与深度学习框架中的Tensor不同,这个Tensor = Layout + Storage(指针表达的数据 / 栈上数据(GPU表现为寄存器)).我们**在cute中用Tensor大多数情况只是用它的逻辑语义和变换,不涉及数据的实际搬运**(比如由Tensor生成Tensor也只是Layout的表达变换,并不产生数据实体的搬运).Tensor是一个对数据表达进行了抽象的工具.Tensor的存储空间可以是栈上内存(寄存器,需要静态数据)也可以是堆上

MMA

MMA_tutorial

  • MMAOpearation与微架构相关,主要封装的是特定架构在特定MNK下特定精度的方法

  • MMA_Traits类似于C++特征萃取,用来提供一些信息,比如类型等,用作桥梁

  • MMA_Atom原子操作,硬件提供的能执行矩阵乘法的最小单位

  • TiledMMA在MMA_Atom上进行扩展,以提供更大的矩阵乘.这里的tiled体现在某些方向(MNK)上重复几次,可能重复的时候涉及到开线程(AtomLayoutMNK),可能涉及到开寄存器以重复计算(ValLayoutMNK以及后面替换了前者的PermutationMNK)

    当我们CopyOperation采用SM75_XXX_LDSM_X来进行数据s2r移动时,由MMA_Atom和ValLayoutMNK共同作用,看我们的ldmatrix是.x1 | .x2 | .x4

  • ThrMMA毕竟整个CUDA还是SIMT下的.这一整个任务(这里指的是TiledMMA,分块的矩阵)划分到某一个线程号下的任务是怎么样的,由这个Thread MMA领取它自己的计算任务

  • cute::gemm(tiledMMA,DThrMMA,AThrMMA,BThrMMA,CThrMMA)各个线程领到自己的计算任务后,一起调用这个函数,开始执行矩阵乘

Copy

COPY_tutorial

数据搬运表示为D=S,其中D,S都是Tensor

gpu-storage-hierarhy

这张图的主要一点,走1,2两条路到smem的,都是Ampere架构及其之后架构的事,之前架构的GPU都只能走3,再走4到smem.走3到reg的这种方式在CopyOperation里是被抽象为UniversalCopy类

  • CopyOperation封装了硬件相关的拷贝能力,类名以SMxx_dstBitxnums_op_trans大致组成,其中op就是跟架构相关的拷贝能力dstBit * nums体现的是目标存储的位数和数目

    • SM75_U32x1_LDSM_N代表了ldmatrix,以row-major,即无需开启.trans限定符,在.x1的情况下的数据移动(.x2则是U32x2,.x4则是U32x4);

    • SM75_U16x2_LDSM_T代表了ldmatrix,以col-major,即开启.trans,在.x1的情况下的数据移动(.x2则是U16x4,.x4则是U16x8);

      这时涉及到的行地址/列地址(.trans了),会在make_tiled_copy_A/B()中根据TiledMMA提供的layout和Copy_Atom自动计算

  • Copy_Traits封装了operation的线程数以及源数据和目标数据的Layout(里面的看起来是bit,不解)还有一个==RefLayout,暂不清楚怎么个线程级的retile==

  • Copy_Atom封装上面俩,继承Copy_Traits,并提供call()用于实现对底层指令的调用入口

  • TiledCopy封装了Copy_Atom的重复方式,以及提供了每个线程Layout描述的拷贝任务的获得方法get_slice()和get_thread_slice(),这俩函数returntype是ThrCopy

    其中这个重复方式包含了执行单元的重复ThrLayout,也包含了每个执行单元负责的tensor分块ValLayout,前者描述是描述线程排布的layout,后者是描述1个thread拷贝多少数据的layout

    二者组合得到的Tiler_MN的shape是TiledCopy执行一次copy时,从S/D copy的tensor shape,无关stride!

    此外还有个TiledLayout_TV的layout,实现根据thread的id,想知道某一个thread所负责的tensor分块的坐标在这个Tiler_MN中处在哪里,获取这个coordinate(coordinate -> offset的反方向).这需要根据TiledLayout_TV的layout算出它的offset,然后根据这个offset从Tiler_MN找到coordinate,其中这个layout的线程排布是列主序的(疑似是规则,这里的东西比较复杂,建议重看参考文件2的第2篇相关文章和实验)

  • ThrCopy提供partition_S()函数用以获取当前线程拷贝所需的src和dst,通过一个划分一个大的tensor到所需的tensor;retile_S()则是在这之上变换形状以匹配拷贝所要求的形状

  • cute::copy

感觉COPY比MMA难理解很多

Swizzle

在矩阵乘中,数据搬运的过程为gmem -> smem -> reg,这其中逻辑上来讲某一个C块,对应的A块和B块在gmem和在reg上是一样的(reg的矩阵块由warp中的reg整体表示),但是由于smem bank conflict的存在,数据存储在smem上就需要沿着bank的方向错开,即按照二维坐标来说是icol方向错开,那怎么错开,错开之后的offset咋取,还有怎么从一维变到二维以形成(irow,icol)?这些都可以交由Swizzle来解决

Swizzle与Layout一样,都是函数,Swizzle是嵌套Layout来使用,Layout提供的是offset/index,而Swizzle提供的是解决了smem bank conflict的offset,这里对bank conflict的解决是通过icol=irow⊕icolicol = irow \oplus icol实现的

Swizzle通过B,M,S三个参数构造出一个二维空间,一维坐标的元素可以按此空间的定义转换成二维坐标,B,S表示新空间的行列,M表示新空间的基本元素,由一维坐标下的连续元素组成,这三个参数都是以2为基底的,可以通过ldmatrix指令,smem bank来理解B,M,S三个参数(结合下图):

  • 首先设置M是为了构造新元素,这个新元素从ldmatrix的角度来理解,以m8n8.x1.b16来看,T0到T7各占一行的数据,一行的8个16位数据连续,T0到T7之间代表的该行首个元素地址不要求连续,一行8个元素,看作一个新元素,故M设置为3;
  • 紧接着新元素的大小是16Byte,smem的bank是4Byte一个,有32个,那此时的新bank大小是16Byte一个,有8个新bank,每个新bank可以放下一个新元素,支持并行访问8个新元素,即此刻新bank数目为8,构造的S从新bank数目来,故S设置为3;
  • 进行ldmatrix,讨论此刻发生的n-way bank conflict,此时可以设置B值用以表示irow,指代重映射的基本单位,注意 B≤SB \le S,这个B的取值实际上可以结合M,S将一行填满bank来看,观察此刻的bank conflict情况,从而取值

CuTe-Swizzle-BMS

此外除了解决smem bank conflict的Swizzle外,还有个用于线程块的Swizzle,可以用来增大L2Cache hit,个人感觉类似于执行配置参数中对grid的X,Y,Z进行位置变换,使之在调度block到SM上时对考虑L2 Cache的命中率

参考文件

Layout Tensor MMA Copy Swizzle cute cutlass2.x

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议 ，转载请注明出处！