TensorIR Transformation
<p>In the previous <a href="https://www.bodunhu.com/blog/posts/dive-into-tensorir/">post</a>, we’ve explored how to write primitive functions in TensorIR. Here, we will see how to transform TensorIR into other (potentially more performant) variants. The content is drived from the <a href="https://mlc.ai/summer22/">mlc</a> course taught by <a href="https://tqchen.com/">Tianqi Chen</a>.</p>
<h2 id="batched-bmm-relu">Batched BMM ReLu</h2>
<p>A batched matrix multiplication followed by a ReLu operation can be expressed using numpy as:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">lnumpy_mm_relu_v2</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">B</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">C</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">((</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">"float32"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">if</span> <span class="n">k</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">+</span> <span class="n">A</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">C</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">],</span> <span class="mi">0</span><span class="p">)</span>
</span></span></code></pre></div><p>Translating the numpy code into TensorIR we get:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="nd">@tvm.script.ir_module</span>
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">MyBmmRule</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="nd">@T.prim_func</span>
</span></span><span class="line"><span class="cl"> <span class="k">def</span> <span class="nf">bmm_relu</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span>
</span></span><span class="line"><span class="cl"> <span class="n">W</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">]):</span>
</span></span><span class="line"><span class="cl"> <span class="n">T</span><span class="o">.</span><span class="n">func_attr</span><span class="p">({</span><span class="s2">"global_symbol"</span><span class="p">:</span> <span class="s2">"bmm_relu"</span><span class="p">,</span> <span class="s2">"tir.noalias"</span><span class="p">:</span> <span class="kc">True</span><span class="p">})</span>
</span></span><span class="line"><span class="cl"> <span class="c1"># we must to allocate the buffer here!</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y_</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">alloc_buffer</span><span class="p">([</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">"float32"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="n">n</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vi</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vk</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">init</span><span class="p">():</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y_</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y_</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">+=</span> <span class="n">A</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vk</span><span class="p">]</span> <span class="o">*</span> <span class="n">W</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vk</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"R"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="n">n</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vi</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">Y_</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">],</span> <span class="n">T</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
</span></span></code></pre></div><p>Our ultimate goal is to transform the TensorIR above to the following form:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="nd">@tvm.script.ir_module</span>
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">TargetModule</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="nd">@T.prim_func</span>
</span></span><span class="line"><span class="cl"> <span class="k">def</span> <span class="nf">bmm_relu</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">B</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">C</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">])</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="n">T</span><span class="o">.</span><span class="n">func_attr</span><span class="p">({</span><span class="s2">"global_symbol"</span><span class="p">:</span> <span class="s2">"bmm_relu"</span><span class="p">,</span> <span class="s2">"tir.noalias"</span><span class="p">:</span> <span class="kc">True</span><span class="p">})</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">alloc_buffer</span><span class="p">([</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">"float32"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">parallel</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i1</span><span class="p">,</span> <span class="n">i2_0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax0_init</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">vectorized</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_init"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">n</span><span class="p">,</span> <span class="n">i</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">i0</span><span class="p">,</span> <span class="n">i1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">j</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i2_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">ax0_init</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax1_0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">32</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax1_1</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">unroll</span><span class="p">(</span><span class="mi">4</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_update"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">n</span><span class="p">,</span> <span class="n">i</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">i0</span><span class="p">,</span> <span class="n">i1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">j</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i2_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">ax0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">k</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">ax1_0</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">+</span> <span class="n">ax1_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">+</span> <span class="n">A</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i2_1</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">vectorized</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"R"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">n</span><span class="p">,</span> <span class="n">i</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">i0</span><span class="p">,</span> <span class="n">i1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">j</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i2_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">i2_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">C</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">],</span> <span class="n">T</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
</span></span></code></pre></div><p>Before we perform the transformation, let’s understand what the transformed TensorIR is doing by looking at several loops here.</p>
<p>First, taking a look at</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="k">for</span> <span class="n">i1</span><span class="p">,</span> <span class="n">i2_0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax0_init</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">vectorized</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_init"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">n</span><span class="p">,</span> <span class="n">i</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">i0</span><span class="p">,</span> <span class="n">i1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">j</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i2_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">ax0_init</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span></span></code></pre></div><p>The code block is initializing the <code>Y</code> matrix to be 0. But it does so by initializing every 8 consecutive elements in each row of <code>Y</code> using a <em>vectorized</em> operation (which might be faster).</p>
<p>The next loop is bit tricky:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="k">for</span> <span class="n">ax1_0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">32</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax1_1</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">unroll</span><span class="p">(</span><span class="mi">4</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax0</span> <span class="ow">in</span> <span class="n">T</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_update"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">n</span><span class="p">,</span> <span class="n">i</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">i0</span><span class="p">,</span> <span class="n">i1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">j</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">i2_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">ax0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">k</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">ax1_0</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">+</span> <span class="n">ax1_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">=</span> <span class="n">Y</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span> <span class="o">+</span> <span class="n">A</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">j</span><span class="p">]</span>
</span></span></code></pre></div><p>This loop is actually performing the matrix multiplication of <code>A</code> and <code>B</code>. We mutiply a row in <code>A</code> with a column in <code>B</code> and sum up the result into a number.</p>
<p>Here, <code>i</code> is mapped to <code>i1</code>, which means we access <code>A</code> one row at a time.i <code>k = T.axis.reduce(128, ax1_0 * 4 + ax1_1)</code> means we access one row in matrix <code>A</code> and one column in matrix <code>B</code> sequentially duing mutiplying, while applying unrolling in hope for better access efficency (\(128 = 32\times 4))). <code>j = T.axis.spatial(128, i2_0 * 8 + ax0)</code> really just means accessing each column sequentially, nothing special.</p>
<h2 id="perform-transformation">Perform Transformation</h2>
<p>To perform tranformation on any TensorIP, it’s very <strong>important to follow the steps listed below</strong>:</p>
<ol>
<li>Get block</li>
<li>Get loops</li>
<li>Organize loops by split, reorder, compute_at/reverse_compute_at</li>
<li>Decompose reduction</li>
<li>vectorize/unroll/parallel</li>
</ol>
<p>Applying step 1, 2, and 3, we first get the block from the original TensorIR:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="n">sch</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">Schedule</span><span class="p">(</span><span class="n">MyBmmRule</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="c1"># Step 1. Get blocks</span>
</span></span><span class="line"><span class="cl"><span class="n">block_M</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_block</span><span class="p">(</span><span class="s2">"M"</span><span class="p">,</span> <span class="n">func_name</span><span class="o">=</span><span class="s2">"bmm_relu"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="c1"># Step 2. Get loops</span>
</span></span><span class="line"><span class="cl"><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">k</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_loops</span><span class="p">(</span><span class="n">block_M</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="c1"># Step 3. Organize loops</span>
</span></span><span class="line"><span class="cl"><span class="n">k0</span><span class="p">,</span> <span class="n">k1</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">factors</span><span class="o">=</span><span class="p">[</span><span class="mi">32</span><span class="p">,</span> <span class="mi">4</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"><span class="n">j0</span><span class="p">,</span> <span class="n">j1</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">factors</span><span class="o">=</span><span class="p">[</span><span class="mi">16</span><span class="p">,</span> <span class="mi">8</span><span class="p">])</span>
</span></span></code></pre></div><p>The reason we split <code>k</code> and <code>j</code> in such a way is: we already mentioned <code>k</code> dimension is accessed sequentially but with unrolling (4) applied; when matrix <code>Y</code> is initialized, a vectorized operation (applied on 8 elements) is applied to dimension <code>j</code>, or every 8 elements in one row(TVM is row-major, therefore might be faster).</p>
<p>But the next question is: how do we reorder the spitted loop? I spent a lot of time trying to figure that out. Turns out the simplest way is to write out the implementation in numpy and proceed from there. Remember, we’ve already splitted <code>k</code> and <code>j</code>, which are used during matrix multiplication, so our new matrix multipliation in numy would be:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="k">for</span> <span class="n">j0</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">k0</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">32</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">k1</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">4</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">j1</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="mi">8</span><span class="o">*</span><span class="n">j0</span><span class="o">+</span><span class="n">j1</span><span class="p">]</span> <span class="o">+=</span> <span class="n">A</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="mi">4</span><span class="o">*</span><span class="n">k0</span> <span class="o">+</span> <span class="n">k1</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="mi">4</span><span class="o">*</span><span class="n">k0</span><span class="o">+</span><span class="n">k1</span><span class="p">,</span> <span class="mi">8</span><span class="o">*</span><span class="n">j0</span><span class="o">+</span><span class="n">j1</span><span class="p">]</span>
</span></span></code></pre></div><p>Because we move the the next column in <code>B</code> after traversing the previous column, we will put <code>j1</code> at the innermost loop. Therefore, the transformation for TensorIR would be:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="n">sch</span><span class="o">.</span><span class="n">reorder</span><span class="p">(</span><span class="n">j0</span><span class="p">,</span> <span class="n">k0</span><span class="p">,</span> <span class="n">k1</span><span class="p">,</span> <span class="n">j1</span><span class="p">)</span>
</span></span></code></pre></div><p>We can print out the transformed TensorIR with <code>print(sch.mod.script())</code>:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="nd">@tvm.script.ir_module</span>
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">Module</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="nd">@tir.prim_func</span>
</span></span><span class="line"><span class="cl"> <span class="k">def</span> <span class="nf">bmm_relu</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">B</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">C</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">])</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">func_attr</span><span class="p">({</span><span class="s2">"global_symbol"</span><span class="p">:</span> <span class="s2">"bmm_relu"</span><span class="p">,</span> <span class="s2">"tir.noalias"</span><span class="p">:</span> <span class="kc">True</span><span class="p">})</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">alloc_buffer</span><span class="p">([</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">"float32"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">parallel</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j_0</span><span class="p">,</span> <span class="n">k_0</span><span class="p">,</span> <span class="n">k_1</span><span class="p">,</span> <span class="n">j_1</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span><span class="p">,</span> <span class="n">vi</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">j_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vk</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">k_0</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">+</span> <span class="n">k_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">reads</span><span class="p">(</span><span class="n">A</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vk</span><span class="p">],</span> <span class="n">B</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vk</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">writes</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">init</span><span class="p">():</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">+</span> <span class="n">A</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vk</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vk</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"R"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SSS"</span><span class="p">,</span> <span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">reads</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">writes</span><span class="p">(</span><span class="n">C</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">C</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">],</span> <span class="n">tir</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
</span></span></code></pre></div><p>Now, we just need to move the ReLu operation (<code>for n, i, j in tir.grid(16, 128, 128):</code>) into the loop above:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="n">block_M</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_block</span><span class="p">(</span><span class="s2">"M"</span><span class="p">,</span> <span class="n">func_name</span><span class="o">=</span><span class="s2">"bmm_relu"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="n">sch</span><span class="o">.</span><span class="n">reverse_compute_at</span><span class="p">(</span><span class="n">block_M</span><span class="p">,</span> <span class="n">j0</span><span class="p">)</span>
</span></span></code></pre></div><p>Step 4 involves seperating initialization and matrix multiplication, therefore we use <code>M_init = sch.decompose_reduction(block_M, k0)</code>, which results in:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="nd">@tvm.script.ir_module</span>
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">Module</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="nd">@tir.prim_func</span>
</span></span><span class="line"><span class="cl"> <span class="k">def</span> <span class="nf">bmm_relu</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">B</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">],</span> <span class="n">C</span><span class="p">:</span> <span class="n">tir</span><span class="o">.</span><span class="n">Buffer</span><span class="p">[(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span> <span class="s2">"float32"</span><span class="p">])</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span>
</span></span><span class="line"><span class="cl"> <span class="c1"># function attr dict</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">func_attr</span><span class="p">({</span><span class="s2">"global_symbol"</span><span class="p">:</span> <span class="s2">"bmm_relu"</span><span class="p">,</span> <span class="s2">"tir.noalias"</span><span class="p">:</span> <span class="kc">True</span><span class="p">})</span>
</span></span><span class="line"><span class="cl"> <span class="c1"># body</span>
</span></span><span class="line"><span class="cl"> <span class="c1"># with tir.block("root")</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">alloc_buffer</span><span class="p">([</span><span class="mi">16</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">"float32"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">parallel</span><span class="p">(</span><span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j_0</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">16</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">j_1_init</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_init"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span><span class="p">,</span> <span class="n">vi</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">j_1_init</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">reads</span><span class="p">()</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">writes</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">k_0</span><span class="p">,</span> <span class="n">k_1</span><span class="p">,</span> <span class="n">j_1</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="mi">32</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"M_update"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span><span class="p">,</span> <span class="n">vi</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">j_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">vk</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">k_0</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">+</span> <span class="n">k_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">reads</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">],</span> <span class="n">A</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vk</span><span class="p">],</span> <span class="n">B</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vk</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">writes</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">+</span> <span class="n">A</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vk</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vk</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span>
</span></span><span class="line"><span class="cl"> <span class="k">for</span> <span class="n">ax0</span> <span class="ow">in</span> <span class="n">tir</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="k">with</span> <span class="n">tir</span><span class="o">.</span><span class="n">block</span><span class="p">(</span><span class="s2">"R"</span><span class="p">):</span>
</span></span><span class="line"><span class="cl"> <span class="n">vn</span><span class="p">,</span> <span class="n">vi</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">remap</span><span class="p">(</span><span class="s2">"SS"</span><span class="p">,</span> <span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">vj</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">axis</span><span class="o">.</span><span class="n">spatial</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">j_0</span> <span class="o">*</span> <span class="mi">8</span> <span class="o">+</span> <span class="n">ax0</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">reads</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">tir</span><span class="o">.</span><span class="n">writes</span><span class="p">(</span><span class="n">C</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">])</span>
</span></span><span class="line"><span class="cl"> <span class="n">C</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">]</span> <span class="o">=</span> <span class="n">tir</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">Y</span><span class="p">[</span><span class="n">vn</span><span class="p">,</span> <span class="n">vi</span><span class="p">,</span> <span class="n">vj</span><span class="p">],</span> <span class="n">tir</span><span class="o">.</span><span class="n">float32</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
</span></span></code></pre></div><p>The final step is easy, just apply vectorize/parallel/unroll onto corresponding loop:</p>
<div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j_0</span><span class="p">,</span> <span class="n">j_1_init</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_loops</span><span class="p">(</span><span class="n">M_init</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="n">sch</span><span class="o">.</span><span class="n">vectorize</span><span class="p">(</span><span class="n">j_1_init</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j_0</span><span class="p">,</span> <span class="n">i2_1</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_loops</span><span class="p">(</span><span class="n">block_R</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="n">sch</span><span class="o">.</span><span class="n">vectorize</span><span class="p">(</span><span class="n">i2_1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="n">block_M_update</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_block</span><span class="p">(</span><span class="s2">"M_update"</span><span class="p">,</span> <span class="n">func_name</span><span class="o">=</span><span class="s2">"bmm_relu"</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="n">n</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j_0</span><span class="p">,</span> <span class="n">k_0</span><span class="p">,</span> <span class="n">k_1</span><span class="p">,</span> <span class="n">j_1</span> <span class="o">=</span> <span class="n">sch</span><span class="o">.</span><span class="n">get_loops</span><span class="p">(</span><span class="n">block_M_update</span><span class="p">)</span>
</span></span></code></pre></div><p>Print out the final TensorIR to find out its final form ( ͡❛ ͜ʖ ͡❛).</p>