zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-04 02:28:13 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

docs

Browse Source
This commit is contained in:
Awni Hannun
2024-05-23 21:11:24 -07:00
committed by CircleCI Docs
parent a7d2c589b2
commit 4f6d88f179
1837 changed files with 6038 additions and 255894 deletions
Download Patch File Download Diff File Expand all files Collapse all files

413
docs/build/html/kernels_2steel_2gemm_2gemm_8h_source.html vendored
View File

@@ -1,413 +0,0 @@
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "https://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
<html xmlns="http://www.w3.org/1999/xhtml" lang="en-US">
<head>
<meta http-equiv="Content-Type" content="text/xhtml;charset=UTF-8"/>
<meta http-equiv="X-UA-Compatible" content="IE=11"/>
<meta name="generator" content="Doxygen 1.10.0"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>MLX: mlx/backend/metal/kernels/steel/gemm/gemm.h Source File</title>
<link href="tabs.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="jquery.js"></script>
<script type="text/javascript" src="dynsections.js"></script>
<script type="text/javascript" src="clipboard.js"></script>
<script type="text/javascript" src="cookie.js"></script>
<link href="search/search.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="search/searchdata.js"></script>
<script type="text/javascript" src="search/search.js"></script>
<link href="doxygen.css" rel="stylesheet" type="text/css" />
</head>
<body>
<div id="top"><!-- do not remove this div, it is closed by doxygen! -->
<div id="titlearea">
<table cellspacing="0" cellpadding="0">
<tbody>
<tr id="projectrow">
<td id="projectalign">
<div id="projectname">MLX
</div>
</td>
</tr>
</tbody>
</table>
</div>
<!-- end header part -->
<!-- Generated by Doxygen 1.10.0 -->
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
var searchBox = new SearchBox("searchBox", "search/",'.html');
/* @license-end */
</script>
<script type="text/javascript" src="menudata.js"></script>
<script type="text/javascript" src="menu.js"></script>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
$(function() {
initMenu('',true,false,'search.php','Search');
$(function() { init_search(); });
});
/* @license-end */
</script>
<div id="main-nav"></div>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
$(function() { codefold.init(0); });
/* @license-end */
</script>
<!-- window showing the filter options -->
<div id="MSearchSelectWindow"
onmouseover="return searchBox.OnSearchSelectShow()"
onmouseout="return searchBox.OnSearchSelectHide()"
onkeydown="return searchBox.OnSearchSelectKey(event)">
</div>
<!-- iframe showing the search results (closed by default) -->
<div id="MSearchResultsWindow">
<div id="MSearchResults">
<div class="SRPage">
<div id="SRIndex">
<div id="SRResults"></div>
<div class="SRStatus" id="Loading">Loading...</div>
<div class="SRStatus" id="Searching">Searching...</div>
<div class="SRStatus" id="NoMatches">No Matches</div>
</div>
</div>
</div>
</div>
<div id="nav-path" class="navpath">
<ul>
<li class="navelem"><a class="el" href="dir_938ab0ecf10b8b860ff766c820f665fd.html">mlx</a></li><li class="navelem"><a class="el" href="dir_1d446c9bd3c99228254c9484e0bc5c06.html">backend</a></li><li class="navelem"><a class="el" href="dir_d0c977ea65824390717cdb7efc36c157.html">metal</a></li><li class="navelem"><a class="el" href="dir_70a37effa88bcbd6b791977fa1e64356.html">kernels</a></li><li class="navelem"><a class="el" href="dir_76215a6c54e2b67053e723fc2395583c.html">steel</a></li><li class="navelem"><a class="el" href="dir_6768c99e6145fb9510ccdb40db8ede25.html">gemm</a></li> </ul>
</div>
</div><!-- top -->
<div class="header">
<div class="headertitle"><div class="title">gemm.h</div></div>
</div><!--header-->
<div class="contents">
<a href="kernels_2steel_2gemm_2gemm_8h.html">Go to the documentation of this file.</a><div class="fragment"><div class="line"><a id="l00001" name="l00001"></a><span class="lineno"> 1</span><span class="comment">// Copyright © 2024 Apple Inc.</span></div>
<div class="line"><a id="l00002" name="l00002"></a><span class="lineno"> 2</span> </div>
<div class="line"><a id="l00003" name="l00003"></a><span class="lineno"> 3</span><span class="preprocessor">#pragma once</span></div>
<div class="line"><a id="l00004" name="l00004"></a><span class="lineno"> 4</span> </div>
<div class="line"><a id="l00005" name="l00005"></a><span class="lineno"> 5</span><span class="preprocessor">#include &quot;<a class="code" href="gemm_2loader_8h.html">mlx/backend/metal/kernels/steel/gemm/loader.h</a>&quot;</span></div>
<div class="line"><a id="l00006" name="l00006"></a><span class="lineno"> 6</span><span class="preprocessor">#include &quot;<a class="code" href="mma_8h.html">mlx/backend/metal/kernels/steel/gemm/mma.h</a>&quot;</span></div>
<div class="line"><a id="l00007" name="l00007"></a><span class="lineno"> 7</span><span class="preprocessor">#include &quot;<a class="code" href="gemm_2params_8h.html">mlx/backend/metal/kernels/steel/gemm/params.h</a>&quot;</span></div>
<div class="line"><a id="l00008" name="l00008"></a><span class="lineno"> 8</span><span class="preprocessor">#include &quot;<a class="code" href="backend_2metal_2kernels_2steel_2gemm_2transforms_8h.html">mlx/backend/metal/kernels/steel/gemm/transforms.h</a>&quot;</span></div>
<div class="line"><a id="l00009" name="l00009"></a><span class="lineno"> 9</span><span class="preprocessor">#include &quot;<a class="code" href="backend_2metal_2kernels_2steel_2utils_8h.html">mlx/backend/metal/kernels/steel/utils.h</a>&quot;</span></div>
<div class="line"><a id="l00010" name="l00010"></a><span class="lineno"> 10</span> </div>
<div class="line"><a id="l00011" name="l00011"></a><span class="lineno"> 11</span><span class="keyword">using namespace </span><a class="code hl_namespace" href="namespacemetal.html">metal</a>;</div>
<div class="line"><a id="l00012" name="l00012"></a><span class="lineno"> 12</span> </div>
<div class="line"><a id="l00014" name="l00014"></a><span class="lineno"> 14</span><span class="comment">// GEMM kernel class</span></div>
<div class="line"><a id="l00016" name="l00016"></a><span class="lineno"> 16</span> </div>
<div class="line"><a id="l00017" name="l00017"></a><span class="lineno"> 17</span><span class="keyword">namespace </span><a class="code hl_namespace" href="namespacemlx.html">mlx</a> {</div>
<div class="line"><a id="l00018" name="l00018"></a><span class="lineno"> 18</span><span class="keyword">namespace </span>steel {</div>
<div class="line"><a id="l00019" name="l00019"></a><span class="lineno"> 19</span> </div>
<div class="line"><a id="l00020" name="l00020"></a><span class="lineno"> 20</span><span class="keyword">template</span> &lt;<span class="keywordtype">bool</span> M_aligned, <span class="keywordtype">bool</span> N_aligned, <span class="keywordtype">bool</span> K_aligned&gt;</div>
<div class="line"><a id="l00021" name="l00021"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_loop_alignment.html"> 21</a></span><span class="keyword">struct </span><a class="code hl_struct" href="structmlx_1_1steel_1_1_loop_alignment.html">LoopAlignment</a> {};</div>
<div class="line"><a id="l00022" name="l00022"></a><span class="lineno"> 22</span> </div>
<div class="line"><a id="l00023" name="l00023"></a><span class="lineno"> 23</span><span class="keyword">template</span> &lt;</div>
<div class="line"><a id="l00024" name="l00024"></a><span class="lineno"> 24</span> <span class="keyword">typename</span> T,</div>
<div class="line"><a id="l00025" name="l00025"></a><span class="lineno"> 25</span> <span class="keyword">typename</span> U,</div>
<div class="line"><a id="l00026" name="l00026"></a><span class="lineno"> 26</span> <span class="keywordtype">int</span> BM,</div>
<div class="line"><a id="l00027" name="l00027"></a><span class="lineno"> 27</span> <span class="keywordtype">int</span> BN,</div>
<div class="line"><a id="l00028" name="l00028"></a><span class="lineno"> 28</span> <span class="keywordtype">int</span> BK,</div>
<div class="line"><a id="l00029" name="l00029"></a><span class="lineno"> 29</span> <span class="keywordtype">int</span> WM,</div>
<div class="line"><a id="l00030" name="l00030"></a><span class="lineno"> 30</span> <span class="keywordtype">int</span> WN,</div>
<div class="line"><a id="l00031" name="l00031"></a><span class="lineno"> 31</span> <span class="keywordtype">bool</span> transpose_a,</div>
<div class="line"><a id="l00032" name="l00032"></a><span class="lineno"> 32</span> <span class="keywordtype">bool</span> transpose_b,</div>
<div class="line"><a id="l00033" name="l00033"></a><span class="lineno"> 33</span> <span class="keywordtype">bool</span> MN_aligned,</div>
<div class="line"><a id="l00034" name="l00034"></a><span class="lineno"> 34</span> <span class="keywordtype">bool</span> K_aligned,</div>
<div class="line"><a id="l00035" name="l00035"></a><span class="lineno"> 35</span> <span class="keyword">typename</span> AccumType = <span class="keyword">typename</span> <a class="code hl_typedef" href="structmlx_1_1steel_1_1_accum_helper.html#ae52abf69e7ba6af1a73d65d57182ed26">AccumHelper&lt;T&gt;::accum_type</a>,</div>
<div class="line"><a id="l00036" name="l00036"></a><span class="lineno"> 36</span> <span class="keyword">typename</span> Epilogue = <a class="code hl_struct" href="structmlx_1_1steel_1_1_transform_none.html">TransformNone&lt;U, AccumType&gt;</a>&gt;</div>
<div class="foldopen" id="foldopen00037" data-start="{" data-end="};">
<div class="line"><a id="l00037" name="l00037"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html"> 37</a></span><span class="keyword">struct </span><a class="code hl_struct" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html">GEMMKernel</a> {</div>
<div class="line"><a id="l00038" name="l00038"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0"> 38</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a> = 16 / <span class="keyword">sizeof</span>(T);</div>
<div class="line"><a id="l00039" name="l00039"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094"> 39</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a> = 16 / <span class="keyword">sizeof</span>(T);</div>
<div class="line"><a id="l00040" name="l00040"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ac00b149d76a903c2f91b0f477dc5037f"> 40</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ac00b149d76a903c2f91b0f477dc5037f">tgp_mem_size_a</a> =</div>
<div class="line"><a id="l00041" name="l00041"></a><span class="lineno"> 41</span> transpose_a ? BK * (BM + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a>) : BM * (BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a>);</div>
<div class="line"><a id="l00042" name="l00042"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a105af1069668028c6f1bc6d6dd162298"> 42</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a105af1069668028c6f1bc6d6dd162298">tgp_mem_size_b</a> =</div>
<div class="line"><a id="l00043" name="l00043"></a><span class="lineno"> 43</span> transpose_b ? BN * (BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a>) : BK * (BN + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a>);</div>
<div class="line"><a id="l00044" name="l00044"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a1ec583584e69dcbbb72106390a4fc5da"> 44</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a1ec583584e69dcbbb72106390a4fc5da">tgp_mem_size</a> = <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ac00b149d76a903c2f91b0f477dc5037f">tgp_mem_size_a</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a105af1069668028c6f1bc6d6dd162298">tgp_mem_size_b</a>;</div>
<div class="line"><a id="l00045" name="l00045"></a><span class="lineno"> 45</span> </div>
<div class="line"><a id="l00046" name="l00046"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a9058ddb73e30e83fb9c548ba22817d64"> 46</a></span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a> <span class="keywordtype">short</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a9058ddb73e30e83fb9c548ba22817d64">tgp_size</a> = WM * WN * 32;</div>
<div class="line"><a id="l00047" name="l00047"></a><span class="lineno"> 47</span> </div>
<div class="line"><a id="l00048" name="l00048"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#aa8a04ed74d2259f99b337d4662c64d83"> 48</a></span> <span class="keyword">using </span><a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_a_t</a> = <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">BlockLoader</a>&lt;</div>
<div class="line"><a id="l00049" name="l00049"></a><span class="lineno"> 49</span> T,</div>
<div class="line"><a id="l00050" name="l00050"></a><span class="lineno"> 50</span> transpose_a ? BK : BM,</div>
<div class="line"><a id="l00051" name="l00051"></a><span class="lineno"> 51</span> transpose_a ? BM : BK,</div>
<div class="line"><a id="l00052" name="l00052"></a><span class="lineno"> 52</span> transpose_a ? BM + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a> : BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a>,</div>
<div class="line"><a id="l00053" name="l00053"></a><span class="lineno"> 53</span> !transpose_a,</div>
<div class="line"><a id="l00054" name="l00054"></a><span class="lineno"> 54</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a9058ddb73e30e83fb9c548ba22817d64">tgp_size</a>&gt;;</div>
<div class="line"><a id="l00055" name="l00055"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#aa98f32278b5fd98c93ae5483c3596395"> 55</a></span> <span class="keyword">using </span><a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_b_t</a> = <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">BlockLoader</a>&lt;</div>
<div class="line"><a id="l00056" name="l00056"></a><span class="lineno"> 56</span> T,</div>
<div class="line"><a id="l00057" name="l00057"></a><span class="lineno"> 57</span> transpose_b ? BN : BK,</div>
<div class="line"><a id="l00058" name="l00058"></a><span class="lineno"> 58</span> transpose_b ? BK : BN,</div>
<div class="line"><a id="l00059" name="l00059"></a><span class="lineno"> 59</span> transpose_b ? BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a> : BN + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a>,</div>
<div class="line"><a id="l00060" name="l00060"></a><span class="lineno"> 60</span> transpose_b,</div>
<div class="line"><a id="l00061" name="l00061"></a><span class="lineno"> 61</span> <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a9058ddb73e30e83fb9c548ba22817d64">tgp_size</a>&gt;;</div>
<div class="line"><a id="l00062" name="l00062"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#add8c6a31011a4895667c2a94a5af3782"> 62</a></span> <span class="keyword">using </span><a class="code hl_struct" href="structmlx_1_1steel_1_1_block_m_m_a.html">mma_t</a> = <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_m_m_a.html">BlockMMA</a>&lt;</div>
<div class="line"><a id="l00063" name="l00063"></a><span class="lineno"> 63</span> T,</div>
<div class="line"><a id="l00064" name="l00064"></a><span class="lineno"> 64</span> U,</div>
<div class="line"><a id="l00065" name="l00065"></a><span class="lineno"> 65</span> BM,</div>
<div class="line"><a id="l00066" name="l00066"></a><span class="lineno"> 66</span> BN,</div>
<div class="line"><a id="l00067" name="l00067"></a><span class="lineno"> 67</span> BK,</div>
<div class="line"><a id="l00068" name="l00068"></a><span class="lineno"> 68</span> WM,</div>
<div class="line"><a id="l00069" name="l00069"></a><span class="lineno"> 69</span> WN,</div>
<div class="line"><a id="l00070" name="l00070"></a><span class="lineno"> 70</span> transpose_a,</div>
<div class="line"><a id="l00071" name="l00071"></a><span class="lineno"> 71</span> transpose_b,</div>
<div class="line"><a id="l00072" name="l00072"></a><span class="lineno"> 72</span> transpose_a ? BM + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a> : BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">tgp_padding_a</a>,</div>
<div class="line"><a id="l00073" name="l00073"></a><span class="lineno"> 73</span> transpose_b ? BK + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a> : BN + <a class="code hl_variable" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">tgp_padding_b</a>,</div>
<div class="line"><a id="l00074" name="l00074"></a><span class="lineno"> 74</span> AccumType,</div>
<div class="line"><a id="l00075" name="l00075"></a><span class="lineno"> 75</span> Epilogue&gt;;</div>
<div class="line"><a id="l00076" name="l00076"></a><span class="lineno"> 76</span> </div>
<div class="line"><a id="l00077" name="l00077"></a><span class="lineno"> 77</span> <span class="comment">/* Main kernel function */</span></div>
<div class="line"><a id="l00078" name="l00078"></a><span class="lineno"> 78</span> <span class="keyword">template</span> &lt;<span class="keywordtype">bool</span> M_aligned, <span class="keywordtype">bool</span> N_aligned, <span class="keywordtype">bool</span> K_aligned_&gt;</div>
<div class="foldopen" id="foldopen00079" data-start="{" data-end="}">
<div class="line"><a id="l00079" name="l00079"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a756d7bbcc96e2919cd65eec4bc135780"> 79</a></span> <span class="keyword">static</span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a756d7bbcc96e2919cd65eec4bc135780">gemm_loop</a>(</div>
<div class="line"><a id="l00080" name="l00080"></a><span class="lineno"> 80</span> threadgroup T* As [[threadgroup(0)]],</div>
<div class="line"><a id="l00081" name="l00081"></a><span class="lineno"> 81</span> threadgroup T* Bs [[threadgroup(1)]],</div>
<div class="line"><a id="l00082" name="l00082"></a><span class="lineno"> 82</span> <span class="keyword">const</span> <span class="keywordtype">int</span> gemm_k_iterations,</div>
<div class="line"><a id="l00083" name="l00083"></a><span class="lineno"> 83</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_a_t</a>&amp; loader_a,</div>
<div class="line"><a id="l00084" name="l00084"></a><span class="lineno"> 84</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_b_t</a>&amp; loader_b,</div>
<div class="line"><a id="l00085" name="l00085"></a><span class="lineno"> 85</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_m_m_a.html">mma_t</a>&amp; mma_op,</div>
<div class="line"><a id="l00086" name="l00086"></a><span class="lineno"> 86</span> thread <span class="keyword">const</span> <span class="keywordtype">short</span>&amp; tgp_bm,</div>
<div class="line"><a id="l00087" name="l00087"></a><span class="lineno"> 87</span> thread <span class="keyword">const</span> <span class="keywordtype">short</span>&amp; tgp_bn,</div>
<div class="line"><a id="l00088" name="l00088"></a><span class="lineno"> 88</span> thread <span class="keyword">const</span> <span class="keywordtype">short</span>&amp; lbk,</div>
<div class="line"><a id="l00089" name="l00089"></a><span class="lineno"> 89</span> <a class="code hl_struct" href="structmlx_1_1steel_1_1_loop_alignment.html">LoopAlignment&lt;M_aligned, N_aligned, K_aligned_&gt;</a> l = {}) {</div>
<div class="line"><a id="l00090" name="l00090"></a><span class="lineno"> 90</span> <span class="comment">// Appease the compiler</span></div>
<div class="line"><a id="l00091" name="l00091"></a><span class="lineno"> 91</span> (void)l;</div>
<div class="line"><a id="l00092" name="l00092"></a><span class="lineno"> 92</span> </div>
<div class="line"><a id="l00093" name="l00093"></a><span class="lineno"> 93</span> short2 tile_dims_A = transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);</div>
<div class="line"><a id="l00094" name="l00094"></a><span class="lineno"> 94</span> </div>
<div class="line"><a id="l00095" name="l00095"></a><span class="lineno"> 95</span> short2 tile_dims_B = transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);</div>
<div class="line"><a id="l00096" name="l00096"></a><span class="lineno"> 96</span> </div>
<div class="line"><a id="l00097" name="l00097"></a><span class="lineno"> 97</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> k = 0; k &lt; gemm_k_iterations; k++) {</div>
<div class="line"><a id="l00098" name="l00098"></a><span class="lineno"> 98</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00099" name="l00099"></a><span class="lineno"> 99</span> <span class="comment">// Load elements into threadgroup</span></div>
<div class="line"><a id="l00100" name="l00100"></a><span class="lineno"> 100</span> <span class="keywordflow">if</span> (M_aligned) {</div>
<div class="line"><a id="l00101" name="l00101"></a><span class="lineno"> 101</span> loader_a.load_unsafe();</div>
<div class="line"><a id="l00102" name="l00102"></a><span class="lineno"> 102</span> } <span class="keywordflow">else</span> {</div>
<div class="line"><a id="l00103" name="l00103"></a><span class="lineno"> 103</span> loader_a.load_safe(tile_dims_A);</div>
<div class="line"><a id="l00104" name="l00104"></a><span class="lineno"> 104</span> }</div>
<div class="line"><a id="l00105" name="l00105"></a><span class="lineno"> 105</span> </div>
<div class="line"><a id="l00106" name="l00106"></a><span class="lineno"> 106</span> <span class="keywordflow">if</span> (N_aligned) {</div>
<div class="line"><a id="l00107" name="l00107"></a><span class="lineno"> 107</span> loader_b.load_unsafe();</div>
<div class="line"><a id="l00108" name="l00108"></a><span class="lineno"> 108</span> } <span class="keywordflow">else</span> {</div>
<div class="line"><a id="l00109" name="l00109"></a><span class="lineno"> 109</span> loader_b.load_safe(tile_dims_B);</div>
<div class="line"><a id="l00110" name="l00110"></a><span class="lineno"> 110</span> }</div>
<div class="line"><a id="l00111" name="l00111"></a><span class="lineno"> 111</span> </div>
<div class="line"><a id="l00112" name="l00112"></a><span class="lineno"> 112</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00113" name="l00113"></a><span class="lineno"> 113</span> </div>
<div class="line"><a id="l00114" name="l00114"></a><span class="lineno"> 114</span> <span class="comment">// Multiply and accumulate threadgroup elements</span></div>
<div class="line"><a id="l00115" name="l00115"></a><span class="lineno"> 115</span> mma_op.mma(As, Bs);</div>
<div class="line"><a id="l00116" name="l00116"></a><span class="lineno"> 116</span> </div>
<div class="line"><a id="l00117" name="l00117"></a><span class="lineno"> 117</span> <span class="comment">// Prepare for next iteration</span></div>
<div class="line"><a id="l00118" name="l00118"></a><span class="lineno"> 118</span> loader_a.next();</div>
<div class="line"><a id="l00119" name="l00119"></a><span class="lineno"> 119</span> loader_b.next();</div>
<div class="line"><a id="l00120" name="l00120"></a><span class="lineno"> 120</span> }</div>
<div class="line"><a id="l00121" name="l00121"></a><span class="lineno"> 121</span> </div>
<div class="line"><a id="l00122" name="l00122"></a><span class="lineno"> 122</span> <span class="keywordflow">if</span> (!K_aligned_) {</div>
<div class="line"><a id="l00123" name="l00123"></a><span class="lineno"> 123</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00124" name="l00124"></a><span class="lineno"> 124</span> </div>
<div class="line"><a id="l00125" name="l00125"></a><span class="lineno"> 125</span> short2 tile_dims_A_last =</div>
<div class="line"><a id="l00126" name="l00126"></a><span class="lineno"> 126</span> transpose_a ? short2(tgp_bm, lbk) : short2(lbk, tgp_bm);</div>
<div class="line"><a id="l00127" name="l00127"></a><span class="lineno"> 127</span> short2 tile_dims_B_last =</div>
<div class="line"><a id="l00128" name="l00128"></a><span class="lineno"> 128</span> transpose_b ? short2(lbk, tgp_bn) : short2(tgp_bn, lbk);</div>
<div class="line"><a id="l00129" name="l00129"></a><span class="lineno"> 129</span> </div>
<div class="line"><a id="l00130" name="l00130"></a><span class="lineno"> 130</span> loader_a.load_safe(tile_dims_A_last);</div>
<div class="line"><a id="l00131" name="l00131"></a><span class="lineno"> 131</span> loader_b.load_safe(tile_dims_B_last);</div>
<div class="line"><a id="l00132" name="l00132"></a><span class="lineno"> 132</span> </div>
<div class="line"><a id="l00133" name="l00133"></a><span class="lineno"> 133</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00134" name="l00134"></a><span class="lineno"> 134</span> </div>
<div class="line"><a id="l00135" name="l00135"></a><span class="lineno"> 135</span> mma_op.mma(As, Bs);</div>
<div class="line"><a id="l00136" name="l00136"></a><span class="lineno"> 136</span> }</div>
<div class="line"><a id="l00137" name="l00137"></a><span class="lineno"> 137</span> }</div>
</div>
<div class="line"><a id="l00138" name="l00138"></a><span class="lineno"> 138</span> </div>
<div class="line"><a id="l00139" name="l00139"></a><span class="lineno"> 139</span> <span class="comment">/* Main kernel function */</span></div>
<div class="foldopen" id="foldopen00140" data-start="{" data-end="}">
<div class="line"><a id="l00140" name="l00140"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a00e55d4a161758350ed7310817d2d2a5"> 140</a></span> <span class="keyword">static</span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a00e55d4a161758350ed7310817d2d2a5">run</a>(</div>
<div class="line"><a id="l00141" name="l00141"></a><span class="lineno"> 141</span> <span class="keyword">const</span> device T* A [[buffer(0)]],</div>
<div class="line"><a id="l00142" name="l00142"></a><span class="lineno"> 142</span> <span class="keyword">const</span> device T* B [[buffer(1)]],</div>
<div class="line"><a id="l00143" name="l00143"></a><span class="lineno"> 143</span> device U* D [[buffer(2)]],</div>
<div class="line"><a id="l00144" name="l00144"></a><span class="lineno"> 144</span> <span class="keyword">const</span> constant <a class="code hl_struct" href="structmlx_1_1steel_1_1_g_e_m_m_params.html">GEMMParams</a>* params [[buffer(3)]],</div>
<div class="line"><a id="l00145" name="l00145"></a><span class="lineno"> 145</span> threadgroup T* As [[threadgroup(0)]],</div>
<div class="line"><a id="l00146" name="l00146"></a><span class="lineno"> 146</span> threadgroup T* Bs [[threadgroup(1)]],</div>
<div class="line"><a id="l00147" name="l00147"></a><span class="lineno"> 147</span> uint simd_lane_id [[thread_index_in_simdgroup]],</div>
<div class="line"><a id="l00148" name="l00148"></a><span class="lineno"> 148</span> uint simd_group_id [[simdgroup_index_in_threadgroup]],</div>
<div class="line"><a id="l00149" name="l00149"></a><span class="lineno"> 149</span> uint3 tid [[threadgroup_position_in_grid]],</div>
<div class="line"><a id="l00150" name="l00150"></a><span class="lineno"> 150</span> uint3 lid [[thread_position_in_threadgroup]]) {</div>
<div class="line"><a id="l00151" name="l00151"></a><span class="lineno"> 151</span> <span class="comment">// Pacifying compiler</span></div>
<div class="line"><a id="l00152" name="l00152"></a><span class="lineno"> 152</span> (void)lid;</div>
<div class="line"><a id="l00153" name="l00153"></a><span class="lineno"> 153</span> </div>
<div class="line"><a id="l00154" name="l00154"></a><span class="lineno"> 154</span> <span class="keyword">const</span> <span class="keywordtype">int</span> tid_y = ((tid.y) &lt;&lt; params-&gt;swizzle_log) +</div>
<div class="line"><a id="l00155" name="l00155"></a><span class="lineno"> 155</span> ((tid.x) &amp; ((1 &lt;&lt; params-&gt;swizzle_log) - 1));</div>
<div class="line"><a id="l00156" name="l00156"></a><span class="lineno"> 156</span> <span class="keyword">const</span> <span class="keywordtype">int</span> tid_x = (tid.x) &gt;&gt; params-&gt;swizzle_log;</div>
<div class="line"><a id="l00157" name="l00157"></a><span class="lineno"> 157</span> </div>
<div class="line"><a id="l00158" name="l00158"></a><span class="lineno"> 158</span> <span class="keywordflow">if</span> (params-&gt;tiles_n &lt;= tid_x || params-&gt;tiles_m &lt;= tid_y) {</div>
<div class="line"><a id="l00159" name="l00159"></a><span class="lineno"> 159</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00160" name="l00160"></a><span class="lineno"> 160</span> }</div>
<div class="line"><a id="l00161" name="l00161"></a><span class="lineno"> 161</span> </div>
<div class="line"><a id="l00162" name="l00162"></a><span class="lineno"> 162</span> threadgroup_barrier(mem_flags::mem_none);</div>
<div class="line"><a id="l00163" name="l00163"></a><span class="lineno"> 163</span> </div>
<div class="line"><a id="l00164" name="l00164"></a><span class="lineno"> 164</span> <span class="comment">// Find block in A, B, C</span></div>
<div class="line"><a id="l00165" name="l00165"></a><span class="lineno"> 165</span> <span class="keyword">const</span> <span class="keywordtype">int</span> c_row = tid_y * BM;</div>
<div class="line"><a id="l00166" name="l00166"></a><span class="lineno"> 166</span> <span class="keyword">const</span> <span class="keywordtype">int</span> c_col = tid_x * BN;</div>
<div class="line"><a id="l00167" name="l00167"></a><span class="lineno"> 167</span> <span class="keyword">const</span> <span class="keywordtype">size_t</span> c_row_long = size_t(c_row);</div>
<div class="line"><a id="l00168" name="l00168"></a><span class="lineno"> 168</span> <span class="keyword">const</span> <span class="keywordtype">size_t</span> c_col_long = size_t(c_col);</div>
<div class="line"><a id="l00169" name="l00169"></a><span class="lineno"> 169</span> </div>
<div class="line"><a id="l00170" name="l00170"></a><span class="lineno"> 170</span> A += transpose_a ? c_row_long : c_row_long * params-&gt;lda;</div>
<div class="line"><a id="l00171" name="l00171"></a><span class="lineno"> 171</span> B += transpose_b ? c_col_long * params-&gt;ldb : c_col_long;</div>
<div class="line"><a id="l00172" name="l00172"></a><span class="lineno"> 172</span> D += c_row_long * params-&gt;ldd + c_col_long;</div>
<div class="line"><a id="l00173" name="l00173"></a><span class="lineno"> 173</span> </div>
<div class="line"><a id="l00174" name="l00174"></a><span class="lineno"> 174</span> <span class="comment">// Prepare threadgroup loading operations</span></div>
<div class="line"><a id="l00175" name="l00175"></a><span class="lineno"> 175</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_a_t</a> loader_a(A, params-&gt;lda, As, simd_group_id, simd_lane_id);</div>
<div class="line"><a id="l00176" name="l00176"></a><span class="lineno"> 176</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_loader.html">loader_b_t</a> loader_b(B, params-&gt;ldb, Bs, simd_group_id, simd_lane_id);</div>
<div class="line"><a id="l00177" name="l00177"></a><span class="lineno"> 177</span> </div>
<div class="line"><a id="l00178" name="l00178"></a><span class="lineno"> 178</span> <span class="comment">// Prepare threadgroup mma operation</span></div>
<div class="line"><a id="l00179" name="l00179"></a><span class="lineno"> 179</span> thread <a class="code hl_struct" href="structmlx_1_1steel_1_1_block_m_m_a.html">mma_t</a> mma_op(simd_group_id, simd_lane_id);</div>
<div class="line"><a id="l00180" name="l00180"></a><span class="lineno"> 180</span> </div>
<div class="line"><a id="l00181" name="l00181"></a><span class="lineno"> 181</span> <span class="keywordtype">int</span> gemm_k_iterations = params-&gt;gemm_k_iterations_aligned;</div>
<div class="line"><a id="l00182" name="l00182"></a><span class="lineno"> 182</span> </div>
<div class="line"><a id="l00184" name="l00184"></a><span class="lineno"> 184</span> <span class="comment">// MNK aligned loop</span></div>
<div class="line"><a id="l00185" name="l00185"></a><span class="lineno"> 185</span> <span class="keywordflow">if</span> (MN_aligned) {</div>
<div class="line"><a id="l00186" name="l00186"></a><span class="lineno"> 186</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> k = 0; k &lt; gemm_k_iterations; k++) {</div>
<div class="line"><a id="l00187" name="l00187"></a><span class="lineno"> 187</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00188" name="l00188"></a><span class="lineno"> 188</span> <span class="comment">// Load elements into threadgroup</span></div>
<div class="line"><a id="l00189" name="l00189"></a><span class="lineno"> 189</span> loader_a.load_unsafe();</div>
<div class="line"><a id="l00190" name="l00190"></a><span class="lineno"> 190</span> loader_b.load_unsafe();</div>
<div class="line"><a id="l00191" name="l00191"></a><span class="lineno"> 191</span> </div>
<div class="line"><a id="l00192" name="l00192"></a><span class="lineno"> 192</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00193" name="l00193"></a><span class="lineno"> 193</span> </div>
<div class="line"><a id="l00194" name="l00194"></a><span class="lineno"> 194</span> <span class="comment">// Multiply and accumulate threadgroup elements</span></div>
<div class="line"><a id="l00195" name="l00195"></a><span class="lineno"> 195</span> mma_op.mma(As, Bs);</div>
<div class="line"><a id="l00196" name="l00196"></a><span class="lineno"> 196</span> </div>
<div class="line"><a id="l00197" name="l00197"></a><span class="lineno"> 197</span> <span class="comment">// Prepare for next iteration</span></div>
<div class="line"><a id="l00198" name="l00198"></a><span class="lineno"> 198</span> loader_a.next();</div>
<div class="line"><a id="l00199" name="l00199"></a><span class="lineno"> 199</span> loader_b.next();</div>
<div class="line"><a id="l00200" name="l00200"></a><span class="lineno"> 200</span> }</div>
<div class="line"><a id="l00201" name="l00201"></a><span class="lineno"> 201</span> </div>
<div class="line"><a id="l00202" name="l00202"></a><span class="lineno"> 202</span> threadgroup_barrier(mem_flags::mem_none);</div>
<div class="line"><a id="l00203" name="l00203"></a><span class="lineno"> 203</span> </div>
<div class="line"><a id="l00204" name="l00204"></a><span class="lineno"> 204</span> <span class="comment">// Loop tail</span></div>
<div class="line"><a id="l00205" name="l00205"></a><span class="lineno"> 205</span> <span class="keywordflow">if</span> (!K_aligned) {</div>
<div class="line"><a id="l00206" name="l00206"></a><span class="lineno"> 206</span> <span class="keywordtype">int</span> lbk = params-&gt;K - params-&gt;gemm_k_iterations_aligned * BK;</div>
<div class="line"><a id="l00207" name="l00207"></a><span class="lineno"> 207</span> short2 tile_dims_A = transpose_a ? short2(BM, lbk) : short2(lbk, BM);</div>
<div class="line"><a id="l00208" name="l00208"></a><span class="lineno"> 208</span> short2 tile_dims_B = transpose_b ? short2(lbk, BN) : short2(BN, lbk);</div>
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> </div>
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> loader_a.load_safe(tile_dims_A);</div>
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> loader_b.load_safe(tile_dims_B);</div>
<div class="line"><a id="l00212" name="l00212"></a><span class="lineno"> 212</span> </div>
<div class="line"><a id="l00213" name="l00213"></a><span class="lineno"> 213</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00214" name="l00214"></a><span class="lineno"> 214</span> </div>
<div class="line"><a id="l00215" name="l00215"></a><span class="lineno"> 215</span> mma_op.mma(As, Bs);</div>
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span> }</div>
<div class="line"><a id="l00217" name="l00217"></a><span class="lineno"> 217</span> </div>
<div class="line"><a id="l00218" name="l00218"></a><span class="lineno"> 218</span> <span class="comment">// Store results to device memory</span></div>
<div class="line"><a id="l00219" name="l00219"></a><span class="lineno"> 219</span> mma_op.store_result(D, params-&gt;ldd);</div>
<div class="line"><a id="l00220" name="l00220"></a><span class="lineno"> 220</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00221" name="l00221"></a><span class="lineno"> 221</span> </div>
<div class="line"><a id="l00222" name="l00222"></a><span class="lineno"> 222</span> }</div>
<div class="line"><a id="l00224" name="l00224"></a><span class="lineno"> 224</span> <span class="comment">// MN unaligned loop</span></div>
<div class="line"><a id="l00225" name="l00225"></a><span class="lineno"> 225</span> <span class="keywordflow">else</span> { <span class="comment">// Loop over K - unaligned case</span></div>
<div class="line"><a id="l00226" name="l00226"></a><span class="lineno"> 226</span> <span class="keywordtype">short</span> tgp_bm = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">min</a>(BM, params-&gt;M - c_row);</div>
<div class="line"><a id="l00227" name="l00227"></a><span class="lineno"> 227</span> <span class="keywordtype">short</span> tgp_bn = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">min</a>(BN, params-&gt;N - c_col);</div>
<div class="line"><a id="l00228" name="l00228"></a><span class="lineno"> 228</span> <span class="keywordtype">short</span> leftover_bk = params-&gt;K - params-&gt;gemm_k_iterations_aligned * BK;</div>
<div class="line"><a id="l00229" name="l00229"></a><span class="lineno"> 229</span> </div>
<div class="line"><a id="l00230" name="l00230"></a><span class="lineno"> 230</span> <span class="keywordflow">if</span> (tgp_bm == BM &amp;&amp; tgp_bn == BN) {</div>
<div class="line"><a id="l00231" name="l00231"></a><span class="lineno"> 231</span> gemm_loop&lt;true, true, K_aligned&gt;(</div>
<div class="line"><a id="l00232" name="l00232"></a><span class="lineno"> 232</span> As,</div>
<div class="line"><a id="l00233" name="l00233"></a><span class="lineno"> 233</span> Bs,</div>
<div class="line"><a id="l00234" name="l00234"></a><span class="lineno"> 234</span> gemm_k_iterations,</div>
<div class="line"><a id="l00235" name="l00235"></a><span class="lineno"> 235</span> loader_a,</div>
<div class="line"><a id="l00236" name="l00236"></a><span class="lineno"> 236</span> loader_b,</div>
<div class="line"><a id="l00237" name="l00237"></a><span class="lineno"> 237</span> mma_op,</div>
<div class="line"><a id="l00238" name="l00238"></a><span class="lineno"> 238</span> tgp_bm,</div>
<div class="line"><a id="l00239" name="l00239"></a><span class="lineno"> 239</span> tgp_bn,</div>
<div class="line"><a id="l00240" name="l00240"></a><span class="lineno"> 240</span> leftover_bk);</div>
<div class="line"><a id="l00241" name="l00241"></a><span class="lineno"> 241</span> </div>
<div class="line"><a id="l00242" name="l00242"></a><span class="lineno"> 242</span> mma_op.store_result(D, params-&gt;ldd);</div>
<div class="line"><a id="l00243" name="l00243"></a><span class="lineno"> 243</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00244" name="l00244"></a><span class="lineno"> 244</span> </div>
<div class="line"><a id="l00245" name="l00245"></a><span class="lineno"> 245</span> } <span class="keywordflow">else</span> <span class="keywordflow">if</span> (tgp_bn == BN) {</div>
<div class="line"><a id="l00246" name="l00246"></a><span class="lineno"> 246</span> gemm_loop&lt;false, true, K_aligned&gt;(</div>
<div class="line"><a id="l00247" name="l00247"></a><span class="lineno"> 247</span> As,</div>
<div class="line"><a id="l00248" name="l00248"></a><span class="lineno"> 248</span> Bs,</div>
<div class="line"><a id="l00249" name="l00249"></a><span class="lineno"> 249</span> gemm_k_iterations,</div>
<div class="line"><a id="l00250" name="l00250"></a><span class="lineno"> 250</span> loader_a,</div>
<div class="line"><a id="l00251" name="l00251"></a><span class="lineno"> 251</span> loader_b,</div>
<div class="line"><a id="l00252" name="l00252"></a><span class="lineno"> 252</span> mma_op,</div>
<div class="line"><a id="l00253" name="l00253"></a><span class="lineno"> 253</span> tgp_bm,</div>
<div class="line"><a id="l00254" name="l00254"></a><span class="lineno"> 254</span> tgp_bn,</div>
<div class="line"><a id="l00255" name="l00255"></a><span class="lineno"> 255</span> leftover_bk);</div>
<div class="line"><a id="l00256" name="l00256"></a><span class="lineno"> 256</span> </div>
<div class="line"><a id="l00257" name="l00257"></a><span class="lineno"> 257</span> mma_op.store_result_safe(D, params-&gt;ldd, short2(tgp_bn, tgp_bm));</div>
<div class="line"><a id="l00258" name="l00258"></a><span class="lineno"> 258</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00259" name="l00259"></a><span class="lineno"> 259</span> </div>
<div class="line"><a id="l00260" name="l00260"></a><span class="lineno"> 260</span> } <span class="keywordflow">else</span> <span class="keywordflow">if</span> (tgp_bm == BM) {</div>
<div class="line"><a id="l00261" name="l00261"></a><span class="lineno"> 261</span> gemm_loop&lt;true, false, K_aligned&gt;(</div>
<div class="line"><a id="l00262" name="l00262"></a><span class="lineno"> 262</span> As,</div>
<div class="line"><a id="l00263" name="l00263"></a><span class="lineno"> 263</span> Bs,</div>
<div class="line"><a id="l00264" name="l00264"></a><span class="lineno"> 264</span> gemm_k_iterations,</div>
<div class="line"><a id="l00265" name="l00265"></a><span class="lineno"> 265</span> loader_a,</div>
<div class="line"><a id="l00266" name="l00266"></a><span class="lineno"> 266</span> loader_b,</div>
<div class="line"><a id="l00267" name="l00267"></a><span class="lineno"> 267</span> mma_op,</div>
<div class="line"><a id="l00268" name="l00268"></a><span class="lineno"> 268</span> tgp_bm,</div>
<div class="line"><a id="l00269" name="l00269"></a><span class="lineno"> 269</span> tgp_bn,</div>
<div class="line"><a id="l00270" name="l00270"></a><span class="lineno"> 270</span> leftover_bk);</div>
<div class="line"><a id="l00271" name="l00271"></a><span class="lineno"> 271</span> </div>
<div class="line"><a id="l00272" name="l00272"></a><span class="lineno"> 272</span> mma_op.store_result_safe(D, params-&gt;ldd, short2(tgp_bn, tgp_bm));</div>
<div class="line"><a id="l00273" name="l00273"></a><span class="lineno"> 273</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00274" name="l00274"></a><span class="lineno"> 274</span> </div>
<div class="line"><a id="l00275" name="l00275"></a><span class="lineno"> 275</span> } <span class="keywordflow">else</span> {</div>
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"> 276</span> gemm_loop&lt;false, false, K_aligned&gt;(</div>
<div class="line"><a id="l00277" name="l00277"></a><span class="lineno"> 277</span> As,</div>
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"> 278</span> Bs,</div>
<div class="line"><a id="l00279" name="l00279"></a><span class="lineno"> 279</span> gemm_k_iterations,</div>
<div class="line"><a id="l00280" name="l00280"></a><span class="lineno"> 280</span> loader_a,</div>
<div class="line"><a id="l00281" name="l00281"></a><span class="lineno"> 281</span> loader_b,</div>
<div class="line"><a id="l00282" name="l00282"></a><span class="lineno"> 282</span> mma_op,</div>
<div class="line"><a id="l00283" name="l00283"></a><span class="lineno"> 283</span> tgp_bm,</div>
<div class="line"><a id="l00284" name="l00284"></a><span class="lineno"> 284</span> tgp_bn,</div>
<div class="line"><a id="l00285" name="l00285"></a><span class="lineno"> 285</span> leftover_bk);</div>
<div class="line"><a id="l00286" name="l00286"></a><span class="lineno"> 286</span> </div>
<div class="line"><a id="l00287" name="l00287"></a><span class="lineno"> 287</span> mma_op.store_result_safe(D, params-&gt;ldd, short2(tgp_bn, tgp_bm));</div>
<div class="line"><a id="l00288" name="l00288"></a><span class="lineno"> 288</span> <span class="keywordflow">return</span>;</div>
<div class="line"><a id="l00289" name="l00289"></a><span class="lineno"> 289</span> }</div>
<div class="line"><a id="l00290" name="l00290"></a><span class="lineno"> 290</span> }</div>
<div class="line"><a id="l00291" name="l00291"></a><span class="lineno"> 291</span> }</div>
</div>
<div class="line"><a id="l00292" name="l00292"></a><span class="lineno"> 292</span>};</div>
</div>
<div class="line"><a id="l00293" name="l00293"></a><span class="lineno"> 293</span> </div>
<div class="line"><a id="l00294" name="l00294"></a><span class="lineno"> 294</span>} <span class="comment">// namespace steel</span></div>
<div class="line"><a id="l00295" name="l00295"></a><span class="lineno"> 295</span>} <span class="comment">// namespace mlx</span></div>
<div class="ttc" id="abackend_2metal_2kernels_2steel_2gemm_2transforms_8h_html"><div class="ttname"><a href="backend_2metal_2kernels_2steel_2gemm_2transforms_8h.html">transforms.h</a></div></div>
<div class="ttc" id="abackend_2metal_2kernels_2steel_2utils_8h_html"><div class="ttname"><a href="backend_2metal_2kernels_2steel_2utils_8h.html">utils.h</a></div></div>
<div class="ttc" id="abackend_2metal_2kernels_2steel_2utils_8h_html_a90b91c866313ffa46eff6d9cc944ad2b"><div class="ttname"><a href="backend_2metal_2kernels_2steel_2utils_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a></div><div class="ttdeci">#define STEEL_CONST</div><div class="ttdef"><b>Definition</b> utils.h:7</div></div>
<div class="ttc" id="agemm_2loader_8h_html"><div class="ttname"><a href="gemm_2loader_8h.html">loader.h</a></div></div>
<div class="ttc" id="agemm_2params_8h_html"><div class="ttname"><a href="gemm_2params_8h.html">params.h</a></div></div>
<div class="ttc" id="amma_8h_html"><div class="ttname"><a href="mma_8h.html">mma.h</a></div></div>
<div class="ttc" id="anamespacemetal_html"><div class="ttname"><a href="namespacemetal.html">metal</a></div><div class="ttdef"><b>Definition</b> bf16.h:265</div></div>
<div class="ttc" id="anamespacemetal_html_a6653b28c9473087141eddce39878d4d3"><div class="ttname"><a href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a></div><div class="ttdeci">METAL_FUNC bfloat16_t min(bfloat16_t x, bfloat16_t y)</div><div class="ttdef"><b>Definition</b> bf16_math.h:234</div></div>
<div class="ttc" id="anamespacemlx_html"><div class="ttname"><a href="namespacemlx.html">mlx</a></div><div class="ttdef"><b>Definition</b> allocator.h:7</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_accum_helper_html_ae52abf69e7ba6af1a73d65d57182ed26"><div class="ttname"><a href="structmlx_1_1steel_1_1_accum_helper.html#ae52abf69e7ba6af1a73d65d57182ed26">mlx::steel::AccumHelper::accum_type</a></div><div class="ttdeci">float accum_type</div><div class="ttdef"><b>Definition</b> transforms.h:57</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_block_loader_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_loader.html">mlx::steel::BlockLoader</a></div><div class="ttdef"><b>Definition</b> loader.h:25</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html">mlx::steel::BlockMMA</a></div><div class="ttdef"><b>Definition</b> mma.h:35</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html">mlx::steel::GEMMKernel</a></div><div class="ttdef"><b>Definition</b> gemm.h:37</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_a00e55d4a161758350ed7310817d2d2a5"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a00e55d4a161758350ed7310817d2d2a5">mlx::steel::GEMMKernel::run</a></div><div class="ttdeci">static METAL_FUNC void run(const device T *A, const device T *B, device U *D, const constant GEMMParams *params, threadgroup T *As, threadgroup T *Bs, uint simd_lane_id, uint simd_group_id, uint3 tid, uint3 lid)</div><div class="ttdef"><b>Definition</b> gemm.h:140</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_a105af1069668028c6f1bc6d6dd162298"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a105af1069668028c6f1bc6d6dd162298">mlx::steel::GEMMKernel::tgp_mem_size_b</a></div><div class="ttdeci">STEEL_CONST short tgp_mem_size_b</div><div class="ttdef"><b>Definition</b> gemm.h:42</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_a1ec583584e69dcbbb72106390a4fc5da"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a1ec583584e69dcbbb72106390a4fc5da">mlx::steel::GEMMKernel::tgp_mem_size</a></div><div class="ttdeci">STEEL_CONST short tgp_mem_size</div><div class="ttdef"><b>Definition</b> gemm.h:44</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_a756d7bbcc96e2919cd65eec4bc135780"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a756d7bbcc96e2919cd65eec4bc135780">mlx::steel::GEMMKernel::gemm_loop</a></div><div class="ttdeci">static METAL_FUNC void gemm_loop(threadgroup T *As, threadgroup T *Bs, const int gemm_k_iterations, thread loader_a_t &amp;loader_a, thread loader_b_t &amp;loader_b, thread mma_t &amp;mma_op, thread const short &amp;tgp_bm, thread const short &amp;tgp_bn, thread const short &amp;lbk, LoopAlignment&lt; M_aligned, N_aligned, K_aligned_ &gt; l={})</div><div class="ttdef"><b>Definition</b> gemm.h:79</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_a9058ddb73e30e83fb9c548ba22817d64"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#a9058ddb73e30e83fb9c548ba22817d64">mlx::steel::GEMMKernel::tgp_size</a></div><div class="ttdeci">STEEL_CONST short tgp_size</div><div class="ttdef"><b>Definition</b> gemm.h:46</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_ac00b149d76a903c2f91b0f477dc5037f"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ac00b149d76a903c2f91b0f477dc5037f">mlx::steel::GEMMKernel::tgp_mem_size_a</a></div><div class="ttdeci">STEEL_CONST short tgp_mem_size_a</div><div class="ttdef"><b>Definition</b> gemm.h:40</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_ad1b03941e869017558423c08b08bc094"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad1b03941e869017558423c08b08bc094">mlx::steel::GEMMKernel::tgp_padding_b</a></div><div class="ttdeci">STEEL_CONST short tgp_padding_b</div><div class="ttdef"><b>Definition</b> gemm.h:39</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_kernel_html_ad547704ccbff6c2076abeffa6628c5a0"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_kernel.html#ad547704ccbff6c2076abeffa6628c5a0">mlx::steel::GEMMKernel::tgp_padding_a</a></div><div class="ttdeci">STEEL_CONST short tgp_padding_a</div><div class="ttdef"><b>Definition</b> gemm.h:38</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_g_e_m_m_params_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_g_e_m_m_params.html">mlx::steel::GEMMParams</a></div><div class="ttdef"><b>Definition</b> params.h:12</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_loop_alignment_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_loop_alignment.html">mlx::steel::LoopAlignment</a></div><div class="ttdef"><b>Definition</b> gemm.h:21</div></div>
<div class="ttc" id="astructmlx_1_1steel_1_1_transform_none_html"><div class="ttname"><a href="structmlx_1_1steel_1_1_transform_none.html">mlx::steel::TransformNone</a></div><div class="ttdef"><b>Definition</b> transforms.h:15</div></div>
</div><!-- fragment --></div><!-- contents -->
<!-- start footer part -->
<hr class="footer"/><address class="footer"><small>
Generated by&#160;<a href="https://www.doxygen.org/index.html"><img class="footer" src="doxygen.svg" width="104" height="31" alt="doxygen"/></a> 1.10.0
</small></address>
</body>
</html>