[{"data":1,"prerenderedAt":38},["ShallowReactive",2],{"chapter:kernels\u002Fpatterns\u002Ffusion.json":3},{"project":4,"route":5,"title":6,"titleHtml":6,"navTitle":6,"part":7,"sourcePath":8,"editUrl":9,"html":10,"toc":11,"hasMermaid":31,"prev":32,"next":35},"kernels","\u002Fkernels\u002Fpatterns\u002Ffusion","Fusing an Epilogue","Real Patterns","patterns\u002Ffusion.md","https:\u002F\u002Fgithub.com\u002Fteenygrad\u002Fteenygrad\u002Fedit\u002Fmain\u002Fbooks\u002Fkernels\u002Fsrc\u002Fpatterns\u002Ffusion.md","\u003Cp>Chapter 1 claimed that fusion is the most common reason to write a kernel. This\nchapter is that claim, made concrete.\u003C\u002Fp>\n\u003Ch2 id=\"the-problem\">The problem\u003C\u002Fh2>\n\u003Cp>A convolution followed by batch normalisation followed by a SiLU activation is\none of the most common sequences in vision models. Called as three library\noperations, it does this:\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>conv       read x, write t1\nbatchnorm  read t1, write t2\nsilu       read t2, write y\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Six trips to memory for three operations. And the second and third operations\nare nothing: batch norm is a multiply and an add, SiLU is \u003Ccode>x * sigmoid(x)\u003C\u002Fcode>. All\nthe time goes into moving \u003Ccode>t1\u003C\u002Fcode> and \u003Ccode>t2\u003C\u002Fcode> in and out of memory that the kernel had\nalready read once.\u003C\u002Fp>\n\u003Cp>Fused, it is:\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>conv+bn+silu   read x, write y\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Two trips. The arithmetic is identical. You removed two thirds of the memory\ntraffic by doing the cheap work while the expensive data was still in registers.\u003C\u002Fp>\n\u003Ch2 id=\"what-an-epilogue-is\">What an epilogue is\u003C\u002Fh2>\n\u003Cp>The pattern has a name. The \u003Cstrong>epilogue\u003C\u002Fstrong> is the work you do to a result after\ncomputing it and before storing it, while it is still in registers.\u003C\u002Fp>\n\u003Cp>Anything element-wise is a candidate: activations, bias adds, scaling, dropout\nmasks, casts to a narrower dtype. The rule of thumb is that if an operation\ntouches each element once and needs no neighbours, it belongs in the epilogue of\nwhatever produced those elements.\u003C\u002Fp>\n\u003Ch2 id=\"the-shape-of-it\">The shape of it\u003C\u002Fh2>\n\u003Cp>Take the tiled matmul loop from Chapter 11. After the K loop finishes, \u003Ccode>acc\u003C\u002Fcode>\nholds the output tile in registers. The unfused version stores it and moves on.\nThe fused version does the extra work first:\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>acc = 0\nfor each k-chunk:\n    acc += A_tile @ B_tile\n\n# ── epilogue: the tile is in registers, use it ──\nacc = acc * bn_scale + bn_bias\nacc = acc * sigmoid(acc)\n\nstore acc\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>The library’s \u003Ccode>conv2d_bn_silu_gemm\u003C\u002Fcode> kernel is exactly this. Its K loop is the\none from Chapter 11, and immediately after it comes:\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#7F877D;font-style:italic\">\u002F\u002F ── BatchNorm epilog ──────────────────────────────────────────────────────\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> bn_off \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">arange\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_N\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> +\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> pid_n \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">*\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_N\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> bn_n_mask \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> bn_off\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">lt\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">C_OUT\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> bn_scale \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">load\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">bn_scale_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">add_offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">bn_off\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Some\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">bn_n_mask\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ...);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Note that the batch-norm parameters are loaded \u003Cem>inside\u003C\u002Fem> the same kernel. They\nare small — one scale and one bias per output channel — so loading them costs\nalmost nothing next to the tile they modify.\u003C\u002Fp>\n\u003Ch2 id=\"three-kernels-for-one-operation\">Three kernels for one operation\u003C\u002Fh2>\n\u003Cp>This tree ships three implementations of the same fused operation:\u003C\u002Fp>\n\u003Ctable>\n\u003Cthead>\n\u003Ctr>\n\u003Cth>Kernel\u003C\u002Fth>\n\u003Cth>Approach\u003C\u002Fth>\n\u003C\u002Ftr>\n\u003C\u002Fthead>\n\u003Ctbody>\n\u003Ctr>\n\u003Ctd>\u003Ccode>conv2d_bn_silu\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Scalar. One output element per lane.\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>conv2d_bn_silu_tiled\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Tiled over the output spatial extent.\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>conv2d_bn_silu_gemm\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Convolution as a GEMM, with \u003Ccode>T::dot\u003C\u002Fcode> and Tensor Cores.\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003C\u002Ftbody>\n\u003C\u002Ftable>\n\u003Cp>Three, because none of them wins everywhere. The GEMM version has the best\narithmetic intensity but its tiles are wasted on a convolution with few output\nchannels. The scalar version has no setup cost and wins on small shapes.\u003C\u002Fp>\n\u003Cp>So the lowering picks between them by shape, with thresholds hand-chosen and\nwritten down in \u003Ccode>kernels\u002Fteeny-kernels\u002Fsrc\u002Fgraph\u002Fmod.rs\u003C\u002Fcode>: the GEMM kernel for\n1×1 convolutions with at least 32 output channels, the tiled kernel above 16,\nthe scalar one otherwise.\u003C\u002Fp>\n\u003Cp>Two things worth taking from that.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>Hand-picked thresholds are normal.\u003C\u002Fstrong> Without an autotuner, somebody measures\nand writes the number down. The bench in \u003Ccode>benches\u002Fconv2d_bn_silu.rs\u003C\u002Fcode> exists\nspecifically to check those thresholds still hold, and its doc comment says so.\u003C\u002Fp>\n\u003Cp>The tile size \u003Cem>inside\u003C\u002Fem> the chosen kernel can now be derived from the layer shape\nand the target’s SM count instead of fixed — Chapter 16. Which kernel runs is\nstill one of these hand-picked thresholds.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>Shipping several kernels for one operation is normal too.\u003C\u002Fstrong> A single kernel\nthat is good at every shape is usually worse than three that are each good at\none.\u003C\u002Fp>\n\u003Ch2 id=\"when-not-to-fuse\">When not to fuse\u003C\u002Fh2>\n\u003Cp>Fusion is not free, and it is not always right.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>When the fused thing is also expensive.\u003C\u002Fstrong> Fusing two compute-bound operations\ngains you nothing on memory and may cost you registers.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>When it costs occupancy.\u003C\u002Fstrong> A longer kernel needs more registers. Past a\nthreshold the card runs fewer programs at once, and the loss can exceed the\nmemory saved. This is measurable, and only measurable.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>When the intermediate is needed anyway.\u003C\u002Fstrong> If \u003Ccode>t1\u003C\u002Fcode> is consumed by something\nelse too, fusing means computing it twice.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>When it makes the kernel unmaintainable.\u003C\u002Fstrong> A kernel fusing five operations has\nfive times the ways to be wrong, and Chapter 9’s MLIR is your only view into it.\u003C\u002Fp>\n\u003Cp>The honest test is the one Chapter 18 sets up: measure the fused version against\nthe sequence it replaces, on the shapes you actually run.\u003C\u002Fp>\n\u003Ch2 id=\"fusing-in-the-graph\">Fusing in the graph\u003C\u002Fh2>\n\u003Cp>Everything above is manual — you decide what to fuse and write one kernel that\ndoes it.\u003C\u002Fp>\n\u003Cp>There is a second kind, where the framework notices that two operations in a\ngraph could be merged and does it for you. teenygrad’s lowering does some of\nthis, and Part 5 covers where the seam is between what you fuse by hand and what\nthe graph fuses for you.\u003C\u002Fp>\n\u003Cp>Next: reductions beyond the sum.\u003C\u002Fp>\n",[12,16,19,22,25,28],{"id":13,"text":14,"level":15},"the-problem","The problem",2,{"id":17,"text":18,"level":15},"what-an-epilogue-is","What an epilogue is",{"id":20,"text":21,"level":15},"the-shape-of-it","The shape of it",{"id":23,"text":24,"level":15},"three-kernels-for-one-operation","Three kernels for one operation",{"id":26,"text":27,"level":15},"when-not-to-fuse","When not to fuse",{"id":29,"text":30,"level":15},"fusing-in-the-graph","Fusing in the graph",false,{"title":33,"titleHtml":33,"route":34},"Matrix Multiplication","\u002Fkernels\u002Fpatterns\u002Fmatmul",{"title":36,"titleHtml":36,"route":37},"Reductions and Scans","\u002Fkernels\u002Fpatterns\u002Freductions",1786271829885]