zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-19 19:38:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

rebase

Browse Source
This commit is contained in:
CircleCI Docs
2025-01-09 21:56:20 +00:00
parent 04b749a588
commit d8d647015b
2642 changed files with 137687 additions and 70861 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
docs/build/html/backend_2metal_2kernels_2fft_8h_source.html vendored
View File

@@ -3,7 +3,7 @@
<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.12.0"/>
<meta name="generator" content="Doxygen 1.13.1"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>MLX: mlx/backend/metal/kernels/fft.h Source File</title>
<link href="tabs.css" rel="stylesheet" type="text/css"/>
@@ -11,11 +11,18 @@
<script type="text/javascript" src="dynsections.js"></script>
<script type="text/javascript" src="clipboard.js"></script>
<link href="navtree.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript" src="navtreedata.js"></script>
<script type="text/javascript" src="navtree.js"></script>
<script type="text/javascript" src="resize.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>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
$(function() { init_search(); });
/* @license-end */
</script>
<link href="doxygen.css" rel="stylesheet" type="text/css" />
</head>
<body>
@@ -28,12 +35,24 @@
<div id="projectname">MLX
</div>
</td>
<td> <div id="MSearchBox" class="MSearchBoxInactive">
<span class="left">
<span id="MSearchSelect" onmouseover="return searchBox.OnSearchSelectShow()" onmouseout="return searchBox.OnSearchSelectHide()">&#160;</span>
<input type="text" id="MSearchField" value="" placeholder="Search" accesskey="S"
onfocus="searchBox.OnSearchFieldFocus(true)"
onblur="searchBox.OnSearchFieldFocus(false)"
onkeyup="searchBox.OnSearchFieldChange(event)"/>
</span><span class="right">
<a id="MSearchClose" href="javascript:searchBox.CloseResultsWindow()"><img id="MSearchCloseImg" border="0" src="search/close.svg" alt=""/></a>
</span>
</div>
</td>
</tr>
</tbody>
</table>
</div>
<!-- end header part -->
<!-- Generated by Doxygen 1.12.0 -->
<!-- Generated by Doxygen 1.13.1 -->
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
var searchBox = new SearchBox("searchBox", "search/",'.html');
@@ -44,22 +63,23 @@ var searchBox = new SearchBox("searchBox", "search/",'.html');
$(function() { codefold.init(0); });
/* @license-end */
</script>
<script type="text/javascript" src="menudata.js"></script>
<script type="text/javascript" src="menu.js"></script>
</div><!-- top -->
<div id="side-nav" class="ui-resizable side-nav-resizable">
<div id="nav-tree">
<div id="nav-tree-contents">
<div id="nav-sync" class="sync"></div>
</div>
</div>
<div id="splitbar" style="-moz-user-select:none;"
class="ui-resizable-handle">
</div>
</div>
<script type="text/javascript">
/* @license magnet:?xt=urn:btih:d3d9a9a6595521f9666a5e94cc830dab83b65699&amp;dn=expat.txt MIT */
$(function() {
initMenu('',true,false,'search.php','Search',false);
$(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(){ initResizable(false); });
$(function(){initNavTree('backend_2metal_2kernels_2fft_8h_source.html',''); initResizable(true); });
/* @license-end */
</script>
<div id="doc-content">
<!-- window showing the filter options -->
<div id="MSearchSelectWindow"
onmouseover="return searchBox.OnSearchSelectShow()"
@@ -81,12 +101,6 @@ $(function(){ initResizable(false); });
</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> </ul>
</div>
</div><!-- top -->
<div id="doc-content">
<div class="header">
<div class="headertitle"><div class="title">fft.h</div></div>
</div><!--header-->
@@ -202,7 +216,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00106" name="l00106"></a><span class="lineno"> 106</span> thread float2* inputs,</div>
<div class="line"><a id="l00107" name="l00107"></a><span class="lineno"> 107</span> thread <span class="keywordtype">short</span>* indices,</div>
<div class="line"><a id="l00108" name="l00108"></a><span class="lineno"> 108</span> thread float2* values,</div>
<div class="line"><a id="l00109" name="l00109"></a><span class="lineno"> 109</span> threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>) {</div>
<div class="line"><a id="l00109" name="l00109"></a><span class="lineno"> 109</span> threadgroup float2* buf) {</div>
<div class="line"><a id="l00110" name="l00110"></a><span class="lineno"> 110</span> <span class="keywordtype">int</span> m_r = n / radix;</div>
<div class="line"><a id="l00111" name="l00111"></a><span class="lineno"> 111</span> <span class="comment">// When combining different sized radices, we have to do</span></div>
<div class="line"><a id="l00112" name="l00112"></a><span class="lineno"> 112</span> <span class="comment">// multiple butterflies in a single thread.</span></div>
@@ -220,7 +234,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00124" name="l00124"></a><span class="lineno"> 124</span> index = i + t * m;</div>
<div class="line"><a id="l00125" name="l00125"></a><span class="lineno"> 125</span> <span class="keywordflow">if</span> (index &lt; m_r) {</div>
<div class="line"><a id="l00126" name="l00126"></a><span class="lineno"> 126</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> r = 0; r &lt; radix; r++) {</div>
<div class="line"><a id="l00127" name="l00127"></a><span class="lineno"> 127</span> inputs[r] = <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[index + r * m_r];</div>
<div class="line"><a id="l00127" name="l00127"></a><span class="lineno"> 127</span> inputs[r] = buf[index + r * m_r];</div>
<div class="line"><a id="l00128" name="l00128"></a><span class="lineno"> 128</span> }</div>
<div class="line"><a id="l00129" name="l00129"></a><span class="lineno"> 129</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a278d980ed397e6841ce0af44b9aa4396">radix_butterfly&lt;radix, radix_func&gt;</a>(</div>
<div class="line"><a id="l00130" name="l00130"></a><span class="lineno"> 130</span> index, *p, inputs, indices + t * radix, values + t * radix);</div>
@@ -235,7 +249,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00139" name="l00139"></a><span class="lineno"> 139</span> <span class="keywordflow">if</span> (index &lt; m_r) {</div>
<div class="line"><a id="l00140" name="l00140"></a><span class="lineno"> 140</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> r = 0; r &lt; radix; r++) {</div>
<div class="line"><a id="l00141" name="l00141"></a><span class="lineno"> 141</span> r_index = t * radix + r;</div>
<div class="line"><a id="l00142" name="l00142"></a><span class="lineno"> 142</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[indices[r_index]] = values[r_index];</div>
<div class="line"><a id="l00142" name="l00142"></a><span class="lineno"> 142</span> buf[indices[r_index]] = values[r_index];</div>
<div class="line"><a id="l00143" name="l00143"></a><span class="lineno"> 143</span> }</div>
<div class="line"><a id="l00144" name="l00144"></a><span class="lineno"> 144</span> }</div>
<div class="line"><a id="l00145" name="l00145"></a><span class="lineno"> 145</span> }</div>
@@ -256,7 +270,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00157" name="l00157"></a><span class="lineno"> 157</span><span class="keyword">template</span> &lt;<span class="keywordtype">bool</span> rader = false&gt;</div>
<div class="line"><a id="l00158" name="l00158"></a><span class="lineno"> 158</span>METAL_FUNC <span class="keywordtype">void</span></div>
<div class="foldopen" id="foldopen00159" data-start="{" data-end="}">
<div class="line"><a id="l00159" name="l00159"></a><span class="lineno"><a class="line" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81"> 159</a></span><a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(<span class="keywordtype">int</span> fft_idx, thread <span class="keywordtype">int</span>* p, <span class="keywordtype">int</span> m, <span class="keywordtype">int</span> n, threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>) {</div>
<div class="line"><a id="l00159" name="l00159"></a><span class="lineno"><a class="line" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81"> 159</a></span><a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(<span class="keywordtype">int</span> fft_idx, thread <span class="keywordtype">int</span>* p, <span class="keywordtype">int</span> m, <span class="keywordtype">int</span> n, threadgroup float2* buf) {</div>
<div class="line"><a id="l00160" name="l00160"></a><span class="lineno"> 160</span> float2 inputs[<a class="code hl_define" href="backend_2metal_2kernels_2fft_8h.html#a7b6e56afa21f022c5e754b000955735a">MAX_RADIX</a>];</div>
<div class="line"><a id="l00161" name="l00161"></a><span class="lineno"> 161</span> <span class="keywordtype">short</span> indices[<a class="code hl_define" href="backend_2metal_2kernels_2fft_8h.html#a28d683cf067736d76f867f30c066317e">MAX_OUTPUT_SIZE</a>];</div>
<div class="line"><a id="l00162" name="l00162"></a><span class="lineno"> 162</span> float2 values[<a class="code hl_define" href="backend_2metal_2kernels_2fft_8h.html#a28d683cf067736d76f867f30c066317e">MAX_OUTPUT_SIZE</a>];</div>
@@ -310,9 +324,9 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00208" name="l00208"></a><span class="lineno"> 208</span> <span class="keywordtype">int</span> fft_idx = elem.z; <span class="comment">// Thread index in DFT</span></div>
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> <span class="keywordtype">int</span> m = grid.z; <span class="comment">// Threads per DFT</span></div>
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> <span class="keywordtype">int</span> tg_idx = elem.y * n; <span class="comment">// Index of this DFT in threadgroup</span></div>
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> = &amp;shared_in[tg_idx];</div>
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> threadgroup float2* buf = &amp;shared_in[tg_idx];</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> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>);</div>
<div class="line"><a id="l00213" name="l00213"></a><span class="lineno"> 213</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, buf);</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> read_writer.write();</div>
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span>}</div>
@@ -379,7 +393,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00275" name="l00275"></a><span class="lineno"> 275</span> </div>
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"> 276</span> <span class="keywordtype">int</span> fft_idx = elem.z;</div>
<div class="line"><a id="l00277" name="l00277"></a><span class="lineno"> 277</span> <span class="keywordtype">int</span> tg_idx = elem.y * n;</div>
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"> 278</span> threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> = &amp;shared_in[tg_idx];</div>
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"> 278</span> threadgroup float2* buf = &amp;shared_in[tg_idx];</div>
<div class="line"><a id="l00279" name="l00279"></a><span class="lineno"> 279</span> </div>
<div class="line"><a id="l00280" name="l00280"></a><span class="lineno"> 280</span> <span class="comment">// rader_m = n / rader_n;</span></div>
<div class="line"><a id="l00281" name="l00281"></a><span class="lineno"> 281</span> <span class="keywordtype">int</span> rader_m = <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ac309c77ef775a3ef13850c5287a86480">rader_m_</a>;</div>
@@ -391,7 +405,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00287" name="l00287"></a><span class="lineno"> 287</span> <span class="comment">// 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1</span></div>
<div class="line"><a id="l00288" name="l00288"></a><span class="lineno"> 288</span> <span class="keywordtype">short</span> x_0_index =</div>
<div class="line"><a id="l00289" name="l00289"></a><span class="lineno"> 289</span> <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx * <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a> / (rader_n - 1), rader_m - 1);</div>
<div class="line"><a id="l00290" name="l00290"></a><span class="lineno"> 290</span> float2 x_0[2] = {<a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[x_0_index], <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[x_0_index + 1]};</div>
<div class="line"><a id="l00290" name="l00290"></a><span class="lineno"> 290</span> float2 x_0[2] = {buf[x_0_index], buf[x_0_index + 1]};</div>
<div class="line"><a id="l00291" name="l00291"></a><span class="lineno"> 291</span> </div>
<div class="line"><a id="l00292" name="l00292"></a><span class="lineno"> 292</span> <span class="comment">// Do the Rader permutation in shared memory</span></div>
<div class="line"><a id="l00293" name="l00293"></a><span class="lineno"> 293</span> float2 temp[<a class="code hl_define" href="backend_2metal_2kernels_2fft_8h.html#a7b6e56afa21f022c5e754b000955735a">MAX_RADIX</a>];</div>
@@ -399,26 +413,26 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00295" name="l00295"></a><span class="lineno"> 295</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> e = 0; e &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; e++) {</div>
<div class="line"><a id="l00296" name="l00296"></a><span class="lineno"> 296</span> <span class="keywordtype">short</span> index = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx * <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a> + e, max_index);</div>
<div class="line"><a id="l00297" name="l00297"></a><span class="lineno"> 297</span> <span class="keywordtype">short</span> g_q = raders_g_q[index / rader_m];</div>
<div class="line"><a id="l00298" name="l00298"></a><span class="lineno"> 298</span> temp[e] = <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[rader_m + (g_q - 1) * rader_m + index % rader_m];</div>
<div class="line"><a id="l00298" name="l00298"></a><span class="lineno"> 298</span> temp[e] = buf[rader_m + (g_q - 1) * rader_m + index % rader_m];</div>
<div class="line"><a id="l00299" name="l00299"></a><span class="lineno"> 299</span> }</div>
<div class="line"><a id="l00300" name="l00300"></a><span class="lineno"> 300</span> </div>
<div class="line"><a id="l00301" name="l00301"></a><span class="lineno"> 301</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00302" name="l00302"></a><span class="lineno"> 302</span> </div>
<div class="line"><a id="l00303" name="l00303"></a><span class="lineno"> 303</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> e = 0; e &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; e++) {</div>
<div class="line"><a id="l00304" name="l00304"></a><span class="lineno"> 304</span> <span class="keywordtype">short</span> index = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx * <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a> + e, max_index);</div>
<div class="line"><a id="l00305" name="l00305"></a><span class="lineno"> 305</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[index + rader_m] = temp[e];</div>
<div class="line"><a id="l00305" name="l00305"></a><span class="lineno"> 305</span> buf[index + rader_m] = temp[e];</div>
<div class="line"><a id="l00306" name="l00306"></a><span class="lineno"> 306</span> }</div>
<div class="line"><a id="l00307" name="l00307"></a><span class="lineno"> 307</span> </div>
<div class="line"><a id="l00308" name="l00308"></a><span class="lineno"> 308</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00309" name="l00309"></a><span class="lineno"> 309</span> </div>
<div class="line"><a id="l00310" name="l00310"></a><span class="lineno"> 310</span> <span class="comment">// Rader FFT on x[rader_m:]</span></div>
<div class="line"><a id="l00311" name="l00311"></a><span class="lineno"> 311</span> <span class="keywordtype">int</span> p = 1;</div>
<div class="line"><a id="l00312" name="l00312"></a><span class="lineno"> 312</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>&lt;<span class="comment">/*rader=*/</span><span class="keyword">true</span>&gt;(fft_idx, &amp;p, m, n - rader_m, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> + rader_m);</div>
<div class="line"><a id="l00312" name="l00312"></a><span class="lineno"> 312</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>&lt;<span class="comment">/*rader=*/</span><span class="keyword">true</span>&gt;(fft_idx, &amp;p, m, n - rader_m, buf + rader_m);</div>
<div class="line"><a id="l00313" name="l00313"></a><span class="lineno"> 313</span> </div>
<div class="line"><a id="l00314" name="l00314"></a><span class="lineno"> 314</span> <span class="comment">// x_1 + ... + x_n is computed for us in the first FFT step so</span></div>
<div class="line"><a id="l00315" name="l00315"></a><span class="lineno"> 315</span> <span class="comment">// we save it in the first rader_m indices of the array for later.</span></div>
<div class="line"><a id="l00316" name="l00316"></a><span class="lineno"> 316</span> <span class="keywordtype">int</span> x_sum_index = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx, rader_m - 1);</div>
<div class="line"><a id="l00317" name="l00317"></a><span class="lineno"> 317</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[x_sum_index] = <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[rader_m + x_sum_index * (rader_n - 1)];</div>
<div class="line"><a id="l00317" name="l00317"></a><span class="lineno"> 317</span> buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)];</div>
<div class="line"><a id="l00318" name="l00318"></a><span class="lineno"> 318</span> </div>
<div class="line"><a id="l00319" name="l00319"></a><span class="lineno"> 319</span> float2 inv = {1.0f, -1.0f};</div>
<div class="line"><a id="l00320" name="l00320"></a><span class="lineno"> 320</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> e = 0; e &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; e++) {</div>
@@ -426,7 +440,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00322" name="l00322"></a><span class="lineno"> 322</span> <span class="keywordtype">short</span> interleaved_index =</div>
<div class="line"><a id="l00323" name="l00323"></a><span class="lineno"> 323</span> index / rader_m + (index % rader_m) * (rader_n - 1);</div>
<div class="line"><a id="l00324" name="l00324"></a><span class="lineno"> 324</span> temp[e] = <a class="code hl_function" href="radix_8h.html#a5bfc53b531214c9ce277bebc18aa67d6">complex_mul</a>(</div>
<div class="line"><a id="l00325" name="l00325"></a><span class="lineno"> 325</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[rader_m + interleaved_index],</div>
<div class="line"><a id="l00325" name="l00325"></a><span class="lineno"> 325</span> buf[rader_m + interleaved_index],</div>
<div class="line"><a id="l00326" name="l00326"></a><span class="lineno"> 326</span> raders_b_q[interleaved_index % (rader_n - 1)]);</div>
<div class="line"><a id="l00327" name="l00327"></a><span class="lineno"> 327</span> }</div>
<div class="line"><a id="l00328" name="l00328"></a><span class="lineno"> 328</span> </div>
@@ -434,25 +448,25 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00330" name="l00330"></a><span class="lineno"> 330</span> </div>
<div class="line"><a id="l00331" name="l00331"></a><span class="lineno"> 331</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> e = 0; e &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; e++) {</div>
<div class="line"><a id="l00332" name="l00332"></a><span class="lineno"> 332</span> <span class="keywordtype">short</span> index = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx * <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a> + e, max_index);</div>
<div class="line"><a id="l00333" name="l00333"></a><span class="lineno"> 333</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[rader_m + index] = temp[e] * inv;</div>
<div class="line"><a id="l00333" name="l00333"></a><span class="lineno"> 333</span> buf[rader_m + index] = temp[e] * inv;</div>
<div class="line"><a id="l00334" name="l00334"></a><span class="lineno"> 334</span> }</div>
<div class="line"><a id="l00335" name="l00335"></a><span class="lineno"> 335</span> </div>
<div class="line"><a id="l00336" name="l00336"></a><span class="lineno"> 336</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00337" name="l00337"></a><span class="lineno"> 337</span> </div>
<div class="line"><a id="l00338" name="l00338"></a><span class="lineno"> 338</span> <span class="comment">// Rader IFFT on x[rader_m:]</span></div>
<div class="line"><a id="l00339" name="l00339"></a><span class="lineno"> 339</span> p = 1;</div>
<div class="line"><a id="l00340" name="l00340"></a><span class="lineno"> 340</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>&lt;<span class="comment">/*rader=*/</span><span class="keyword">true</span>&gt;(fft_idx, &amp;p, m, n - rader_m, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> + rader_m);</div>
<div class="line"><a id="l00340" name="l00340"></a><span class="lineno"> 340</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>&lt;<span class="comment">/*rader=*/</span><span class="keyword">true</span>&gt;(fft_idx, &amp;p, m, n - rader_m, buf + rader_m);</div>
<div class="line"><a id="l00341" name="l00341"></a><span class="lineno"> 341</span> </div>
<div class="line"><a id="l00342" name="l00342"></a><span class="lineno"> 342</span> float2 rader_inv_factor = {1.0f / (rader_n - 1), -1.0f / (rader_n - 1)};</div>
<div class="line"><a id="l00343" name="l00343"></a><span class="lineno"> 343</span> </div>
<div class="line"><a id="l00344" name="l00344"></a><span class="lineno"> 344</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> e = 0; e &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; e++) {</div>
<div class="line"><a id="l00345" name="l00345"></a><span class="lineno"> 345</span> <span class="keywordtype">short</span> index = <a class="code hl_function" href="namespacemetal.html#a6653b28c9473087141eddce39878d4d3">metal::min</a>(fft_idx * <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a> + e, n - rader_m - 1);</div>
<div class="line"><a id="l00346" name="l00346"></a><span class="lineno"> 346</span> <span class="keywordtype">short</span> diff_index = index / (rader_n - 1) - x_0_index;</div>
<div class="line"><a id="l00347" name="l00347"></a><span class="lineno"> 347</span> temp[e] = <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[rader_m + index] * rader_inv_factor + x_0[diff_index];</div>
<div class="line"><a id="l00347" name="l00347"></a><span class="lineno"> 347</span> temp[e] = buf[rader_m + index] * rader_inv_factor + x_0[diff_index];</div>
<div class="line"><a id="l00348" name="l00348"></a><span class="lineno"> 348</span> }</div>
<div class="line"><a id="l00349" name="l00349"></a><span class="lineno"> 349</span> </div>
<div class="line"><a id="l00350" name="l00350"></a><span class="lineno"> 350</span> <span class="comment">// Use the sum of elements that was computed in the first FFT</span></div>
<div class="line"><a id="l00351" name="l00351"></a><span class="lineno"> 351</span> float2 x_sum = <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[x_0_index] + x_0[0];</div>
<div class="line"><a id="l00351" name="l00351"></a><span class="lineno"> 351</span> float2 x_sum = buf[x_0_index] + x_0[0];</div>
<div class="line"><a id="l00352" name="l00352"></a><span class="lineno"> 352</span> </div>
<div class="line"><a id="l00353" name="l00353"></a><span class="lineno"> 353</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00354" name="l00354"></a><span class="lineno"> 354</span> </div>
@@ -461,15 +475,15 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00357" name="l00357"></a><span class="lineno"> 357</span> <span class="keywordtype">short</span> g_q_index = index % (rader_n - 1);</div>
<div class="line"><a id="l00358" name="l00358"></a><span class="lineno"> 358</span> <span class="keywordtype">short</span> g_q = raders_g_minus_q[g_q_index];</div>
<div class="line"><a id="l00359" name="l00359"></a><span class="lineno"> 359</span> <span class="keywordtype">short</span> out_index = index - g_q_index + g_q + (index / (rader_n - 1));</div>
<div class="line"><a id="l00360" name="l00360"></a><span class="lineno"> 360</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[out_index] = temp[e];</div>
<div class="line"><a id="l00360" name="l00360"></a><span class="lineno"> 360</span> buf[out_index] = temp[e];</div>
<div class="line"><a id="l00361" name="l00361"></a><span class="lineno"> 361</span> }</div>
<div class="line"><a id="l00362" name="l00362"></a><span class="lineno"> 362</span> </div>
<div class="line"><a id="l00363" name="l00363"></a><span class="lineno"> 363</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[x_0_index * rader_n] = x_sum;</div>
<div class="line"><a id="l00363" name="l00363"></a><span class="lineno"> 363</span> buf[x_0_index * rader_n] = x_sum;</div>
<div class="line"><a id="l00364" name="l00364"></a><span class="lineno"> 364</span> </div>
<div class="line"><a id="l00365" name="l00365"></a><span class="lineno"> 365</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00366" name="l00366"></a><span class="lineno"> 366</span> </div>
<div class="line"><a id="l00367" name="l00367"></a><span class="lineno"> 367</span> p = rader_n;</div>
<div class="line"><a id="l00368" name="l00368"></a><span class="lineno"> 368</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>);</div>
<div class="line"><a id="l00368" name="l00368"></a><span class="lineno"> 368</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, buf);</div>
<div class="line"><a id="l00369" name="l00369"></a><span class="lineno"> 369</span> </div>
<div class="line"><a id="l00370" name="l00370"></a><span class="lineno"> 370</span> read_writer.write();</div>
<div class="line"><a id="l00371" name="l00371"></a><span class="lineno"> 371</span>}</div>
@@ -520,22 +534,22 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00414" name="l00414"></a><span class="lineno"> 414</span> <span class="keywordtype">int</span> fft_idx = elem.z; <span class="comment">// Thread index in DFT</span></div>
<div class="line"><a id="l00415" name="l00415"></a><span class="lineno"> 415</span> <span class="keywordtype">int</span> m = grid.z; <span class="comment">// Threads per DFT</span></div>
<div class="line"><a id="l00416" name="l00416"></a><span class="lineno"> 416</span> <span class="keywordtype">int</span> tg_idx = elem.y * n; <span class="comment">// Index of this DFT in threadgroup</span></div>
<div class="line"><a id="l00417" name="l00417"></a><span class="lineno"> 417</span> threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> = &amp;shared_in[tg_idx];</div>
<div class="line"><a id="l00417" name="l00417"></a><span class="lineno"> 417</span> threadgroup float2* buf = &amp;shared_in[tg_idx];</div>
<div class="line"><a id="l00418" name="l00418"></a><span class="lineno"> 418</span> </div>
<div class="line"><a id="l00419" name="l00419"></a><span class="lineno"> 419</span> <span class="comment">// fft</span></div>
<div class="line"><a id="l00420" name="l00420"></a><span class="lineno"> 420</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>);</div>
<div class="line"><a id="l00420" name="l00420"></a><span class="lineno"> 420</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, buf);</div>
<div class="line"><a id="l00421" name="l00421"></a><span class="lineno"> 421</span> </div>
<div class="line"><a id="l00422" name="l00422"></a><span class="lineno"> 422</span> float2 inv = float2(1.0f, -1.0f);</div>
<div class="line"><a id="l00423" name="l00423"></a><span class="lineno"> 423</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> t = 0; t &lt; <a class="code hl_variable" href="backend_2metal_2kernels_2fft_8h.html#ad395c11e6f2aee72cd1928fba93a35a3">elems_per_thread_</a>; t++) {</div>
<div class="line"><a id="l00424" name="l00424"></a><span class="lineno"> 424</span> <span class="keywordtype">int</span> index = fft_idx + t * m;</div>
<div class="line"><a id="l00425" name="l00425"></a><span class="lineno"> 425</span> <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[index] = <a class="code hl_function" href="radix_8h.html#a5bfc53b531214c9ce277bebc18aa67d6">complex_mul</a>(<a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>[index], w_q[index]) * inv;</div>
<div class="line"><a id="l00425" name="l00425"></a><span class="lineno"> 425</span> buf[index] = <a class="code hl_function" href="radix_8h.html#a5bfc53b531214c9ce277bebc18aa67d6">complex_mul</a>(buf[index], w_q[index]) * inv;</div>
<div class="line"><a id="l00426" name="l00426"></a><span class="lineno"> 426</span> }</div>
<div class="line"><a id="l00427" name="l00427"></a><span class="lineno"> 427</span> </div>
<div class="line"><a id="l00428" name="l00428"></a><span class="lineno"> 428</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00429" name="l00429"></a><span class="lineno"> 429</span> </div>
<div class="line"><a id="l00430" name="l00430"></a><span class="lineno"> 430</span> <span class="comment">// ifft</span></div>
<div class="line"><a id="l00431" name="l00431"></a><span class="lineno"> 431</span> p = 1;</div>
<div class="line"><a id="l00432" name="l00432"></a><span class="lineno"> 432</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>);</div>
<div class="line"><a id="l00432" name="l00432"></a><span class="lineno"> 432</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, buf);</div>
<div class="line"><a id="l00433" name="l00433"></a><span class="lineno"> 433</span> </div>
<div class="line"><a id="l00434" name="l00434"></a><span class="lineno"> 434</span> read_writer.write_padded(length, w_k);</div>
<div class="line"><a id="l00435" name="l00435"></a><span class="lineno"> 435</span>}</div>
@@ -566,7 +580,7 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00458" name="l00458"></a><span class="lineno"> 458</span> <span class="keywordtype">int</span> fft_idx = elem.z;</div>
<div class="line"><a id="l00459" name="l00459"></a><span class="lineno"> 459</span> </div>
<div class="line"><a id="l00460" name="l00460"></a><span class="lineno"> 460</span> threadgroup float2 shared_in[tg_mem_size];</div>
<div class="line"><a id="l00461" name="l00461"></a><span class="lineno"> 461</span> threadgroup float2* <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a> = &amp;shared_in[elem.y * n];</div>
<div class="line"><a id="l00461" name="l00461"></a><span class="lineno"> 461</span> threadgroup float2* buf = &amp;shared_in[elem.y * n];</div>
<div class="line"><a id="l00462" name="l00462"></a><span class="lineno"> 462</span> </div>
<div class="line"><a id="l00463" name="l00463"></a><span class="lineno"> 463</span> <span class="keyword">using </span>read_writer_t = <a class="code hl_struct" href="struct_read_writer.html">ReadWriter&lt;in_T, out_T, step, real&gt;</a>;</div>
<div class="line"><a id="l00464" name="l00464"></a><span class="lineno"> 464</span> read_writer_t read_writer = read_writer_t(</div>
@@ -588,12 +602,11 @@ $(function(){ initResizable(false); });
<div class="line"><a id="l00480" name="l00480"></a><span class="lineno"> 480</span> threadgroup_barrier(mem_flags::mem_threadgroup);</div>
<div class="line"><a id="l00481" name="l00481"></a><span class="lineno"> 481</span> </div>
<div class="line"><a id="l00482" name="l00482"></a><span class="lineno"> 482</span> <span class="keywordtype">int</span> p = 1;</div>
<div class="line"><a id="l00483" name="l00483"></a><span class="lineno"> 483</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, <a class="code hl_variable" href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a>);</div>
<div class="line"><a id="l00483" name="l00483"></a><span class="lineno"> 483</span> <a class="code hl_function" href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a>(fft_idx, &amp;p, m, n, buf);</div>
<div class="line"><a id="l00484" name="l00484"></a><span class="lineno"> 484</span> </div>
<div class="line"><a id="l00485" name="l00485"></a><span class="lineno"> 485</span> read_writer.write_strided(stride, overall_n);</div>
<div class="line"><a id="l00486" name="l00486"></a><span class="lineno"> 486</span>}</div>
</div>
<div class="ttc" id="abackend_2metal_2allocator_8h_html_a15aa5cc1baf29be08d55cca88509e697"><div class="ttname"><a href="backend_2metal_2allocator_8h.html#a15aa5cc1baf29be08d55cca88509e697">buf</a></div><div class="ttdeci">MTL::Buffer * buf</div><div class="ttdef"><b>Definition</b> allocator.h:39</div></div>
<div class="ttc" id="abackend_2metal_2kernels_2fft_8h_html_a020469d1aca557ff1b49414f121dedbb"><div class="ttname"><a href="backend_2metal_2kernels_2fft_8h.html#a020469d1aca557ff1b49414f121dedbb">rader_6_steps_</a></div><div class="ttdeci">static constant constexpr const int rader_6_steps_</div><div class="ttdef"><b>Definition</b> fft.h:43</div></div>
<div class="ttc" id="abackend_2metal_2kernels_2fft_8h_html_a050ead8fa5cacdaec13d68ca3c0dcb81"><div class="ttname"><a href="backend_2metal_2kernels_2fft_8h.html#a050ead8fa5cacdaec13d68ca3c0dcb81">perform_fft</a></div><div class="ttdeci">METAL_FUNC void perform_fft(int fft_idx, thread int *p, int m, int n, threadgroup float2 *buf)</div><div class="ttdef"><b>Definition</b> fft.h:159</div></div>
<div class="ttc" id="abackend_2metal_2kernels_2fft_8h_html_a0abc609e9756475800e996775a96a87e"><div class="ttname"><a href="backend_2metal_2kernels_2fft_8h.html#a0abc609e9756475800e996775a96a87e">bluestein_fft</a></div><div class="ttdeci">void bluestein_fft(const device in_T *in, device out_T *out, const device float2 *w_q, const device float2 *w_k, constant const int &amp;length, constant const int &amp;n, constant const int &amp;batch_size, uint3 elem, uint3 grid)</div><div class="ttdef"><b>Definition</b> fft.h:374</div></div>
@@ -649,10 +662,13 @@ $(function(){ initResizable(false); });
<div class="ttc" id="asteel_2defines_8h_html_a90b91c866313ffa46eff6d9cc944ad2b"><div class="ttname"><a href="steel_2defines_8h.html#a90b91c866313ffa46eff6d9cc944ad2b">STEEL_CONST</a></div><div class="ttdeci">#define STEEL_CONST</div><div class="ttdef"><b>Definition</b> defines.h:3</div></div>
<div class="ttc" id="astruct_read_writer_html"><div class="ttname"><a href="struct_read_writer.html">ReadWriter</a></div><div class="ttdef"><b>Definition</b> readwrite.h:35</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.12.0
</small></address>
</div><!-- doc-content -->
<!-- start footer part -->
<div id="nav-path" class="navpath"><!-- id is needed for treeview function! -->
<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="backend_2metal_2kernels_2fft_8h.html">fft.h</a></li>
<li class="footer">Generated by <a href="https://www.doxygen.org/index.html"><img class="footer" src="doxygen.svg" width="104" height="31" alt="doxygen"/></a> 1.13.1 </li>
</ul>
</div>
</body>
</html>