mirror of
https://github.com/ml-explore/mlx.git
synced 2025-09-23 05:58:09 +08:00
docs update
This commit is contained in:
Awni Hannun
committed by
CircleCI Docs
CircleCI Docs
parent
3e724a7c98
commit
c620a28b16
229
docs/build/html/mma_8h_source.html
vendored
229
docs/build/html/mma_8h_source.html
vendored
@@ -296,87 +296,158 @@ $(function() { codefold.init(0); });
|
||||
<div class="line"><a id="l00199" name="l00199"></a><span class="lineno"> 199</span> }</div>
|
||||
</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> <span class="comment">/* Store results from simdgroup_matrix results into device memory */</span></div>
|
||||
<div class="foldopen" id="foldopen00202" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00202" name="l00202"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3"> 202</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3">store_result</a>(</div>
|
||||
<div class="line"><a id="l00203" name="l00203"></a><span class="lineno"> 203</span> device U* D,</div>
|
||||
<div class="line"><a id="l00204" name="l00204"></a><span class="lineno"> 204</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldd,</div>
|
||||
<div class="line"><a id="l00205" name="l00205"></a><span class="lineno"> 205</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00206" name="l00206"></a><span class="lineno"> 206</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00207" name="l00207"></a><span class="lineno"> 207</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00208" name="l00208"></a><span class="lineno"> 208</span> thread <span class="keyword">const</span> Epilogue& epilogue_op)<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> D += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldd + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>;</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> <span class="comment">// Loop over all simdgroup tiles</span></div>
|
||||
<div class="line"><a id="l00214" name="l00214"></a><span class="lineno"> 214</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00215" name="l00215"></a><span class="lineno"> 215</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00217" name="l00217"></a><span class="lineno"> 217</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00218" name="l00218"></a><span class="lineno"> 218</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00219" name="l00219"></a><span class="lineno"> 219</span> thread <span class="keyword">const</span> <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00220" name="l00220"></a><span class="lineno"> 220</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00221" name="l00221"></a><span class="lineno"> 221</span> <span class="keywordtype">int</span> offset_d = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldd + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>);</div>
|
||||
<div class="line"><a id="l00222" name="l00222"></a><span class="lineno"> 222</span> </div>
|
||||
<div class="line"><a id="l00223" name="l00223"></a><span class="lineno"> 223</span> <span class="comment">// Apply epilogue</span></div>
|
||||
<div class="line"><a id="l00224" name="l00224"></a><span class="lineno"> 224</span> U outs[2] = {</div>
|
||||
<div class="line"><a id="l00225" name="l00225"></a><span class="lineno"> 225</span> epilogue_op.apply(accum[0], C[offset_c]),</div>
|
||||
<div class="line"><a id="l00226" name="l00226"></a><span class="lineno"> 226</span> epilogue_op.apply(accum[1], C[offset_c + fdc])};</div>
|
||||
<div class="line"><a id="l00227" name="l00227"></a><span class="lineno"> 227</span> </div>
|
||||
<div class="line"><a id="l00228" name="l00228"></a><span class="lineno"> 228</span> <span class="comment">// Write out D</span></div>
|
||||
<div class="line"><a id="l00229" name="l00229"></a><span class="lineno"> 229</span> D[offset_d] = outs[0];</div>
|
||||
<div class="line"><a id="l00230" name="l00230"></a><span class="lineno"> 230</span> D[offset_d + 1] = outs[1];</div>
|
||||
<div class="line"><a id="l00231" name="l00231"></a><span class="lineno"> 231</span> }</div>
|
||||
<div class="line"><a id="l00232" name="l00232"></a><span class="lineno"> 232</span> }</div>
|
||||
<div class="line"><a id="l00233" name="l00233"></a><span class="lineno"> 233</span> }</div>
|
||||
<div class="line"><a id="l00201" name="l00201"></a><span class="lineno"> 201</span> <span class="comment">/* Apply epilogue */</span></div>
|
||||
<div class="line"><a id="l00202" name="l00202"></a><span class="lineno"> 202</span> <span class="keyword">template</span> <<span class="keyword">typename</span> BinaryEpilogue></div>
|
||||
<div class="foldopen" id="foldopen00203" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00203" name="l00203"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a823c56cbd2086f10272df7284a5247ae"> 203</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a823c56cbd2086f10272df7284a5247ae">apply_epilogue</a>(</div>
|
||||
<div class="line"><a id="l00204" name="l00204"></a><span class="lineno"> 204</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00205" name="l00205"></a><span class="lineno"> 205</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00206" name="l00206"></a><span class="lineno"> 206</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00207" name="l00207"></a><span class="lineno"> 207</span> thread <span class="keyword">const</span> BinaryEpilogue& epilogue_op) {</div>
|
||||
<div class="line"><a id="l00208" name="l00208"></a><span class="lineno"> 208</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> </div>
|
||||
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> <span class="comment">// Loop over all simdgroup tiles</span></div>
|
||||
<div class="line"><a id="l00212" name="l00212"></a><span class="lineno"> 212</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00213" name="l00213"></a><span class="lineno"> 213</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00214" name="l00214"></a><span class="lineno"> 214</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00215" name="l00215"></a><span class="lineno"> 215</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00217" name="l00217"></a><span class="lineno"> 217</span> thread <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00218" name="l00218"></a><span class="lineno"> 218</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00219" name="l00219"></a><span class="lineno"> 219</span> </div>
|
||||
<div class="line"><a id="l00220" name="l00220"></a><span class="lineno"> 220</span> <span class="comment">// Apply epilogue</span></div>
|
||||
<div class="line"><a id="l00221" name="l00221"></a><span class="lineno"> 221</span> accum[0] = epilogue_op.apply(accum[0], C[offset_c]);</div>
|
||||
<div class="line"><a id="l00222" name="l00222"></a><span class="lineno"> 222</span> accum[1] = epilogue_op.apply(accum[1], C[offset_c + fdc]);</div>
|
||||
<div class="line"><a id="l00223" name="l00223"></a><span class="lineno"> 223</span> }</div>
|
||||
<div class="line"><a id="l00224" name="l00224"></a><span class="lineno"> 224</span> }</div>
|
||||
<div class="line"><a id="l00225" name="l00225"></a><span class="lineno"> 225</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00234" name="l00234"></a><span class="lineno"> 234</span> </div>
|
||||
<div class="foldopen" id="foldopen00235" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00235" name="l00235"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391"> 235</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391">store_result_safe</a>(</div>
|
||||
<div class="line"><a id="l00236" name="l00236"></a><span class="lineno"> 236</span> device U* D,</div>
|
||||
<div class="line"><a id="l00237" name="l00237"></a><span class="lineno"> 237</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldd,</div>
|
||||
<div class="line"><a id="l00238" name="l00238"></a><span class="lineno"> 238</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00239" name="l00239"></a><span class="lineno"> 239</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00240" name="l00240"></a><span class="lineno"> 240</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00241" name="l00241"></a><span class="lineno"> 241</span> short2 dst_tile_dims,</div>
|
||||
<div class="line"><a id="l00242" name="l00242"></a><span class="lineno"> 242</span> thread <span class="keyword">const</span> Epilogue& epilogue_op)<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00243" name="l00243"></a><span class="lineno"> 243</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00244" name="l00244"></a><span class="lineno"> 244</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00245" name="l00245"></a><span class="lineno"> 245</span> D += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldd + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>;</div>
|
||||
<div class="line"><a id="l00246" name="l00246"></a><span class="lineno"> 246</span> dst_tile_dims -= short2(<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>, <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>);</div>
|
||||
<div class="line"><a id="l00247" name="l00247"></a><span class="lineno"> 247</span> </div>
|
||||
<div class="line"><a id="l00248" name="l00248"></a><span class="lineno"> 248</span> <span class="keywordflow">if</span> (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)</div>
|
||||
<div class="line"><a id="l00249" name="l00249"></a><span class="lineno"> 249</span> <span class="keywordflow">return</span>;</div>
|
||||
<div class="line"><a id="l00226" name="l00226"></a><span class="lineno"> 226</span> </div>
|
||||
<div class="line"><a id="l00227" name="l00227"></a><span class="lineno"> 227</span> <span class="comment">/* Apply epilogue */</span></div>
|
||||
<div class="line"><a id="l00228" name="l00228"></a><span class="lineno"> 228</span> <span class="keyword">template</span> <<span class="keyword">typename</span> BinaryEpilogue></div>
|
||||
<div class="foldopen" id="foldopen00229" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00229" name="l00229"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a9e48f2d51099ec00171506724faab54a"> 229</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a9e48f2d51099ec00171506724faab54a">apply_epilogue_safe</a>(</div>
|
||||
<div class="line"><a id="l00230" name="l00230"></a><span class="lineno"> 230</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00231" name="l00231"></a><span class="lineno"> 231</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00232" name="l00232"></a><span class="lineno"> 232</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00233" name="l00233"></a><span class="lineno"> 233</span> short2 dst_tile_dims,</div>
|
||||
<div class="line"><a id="l00234" name="l00234"></a><span class="lineno"> 234</span> thread <span class="keyword">const</span> BinaryEpilogue& epilogue_op) {</div>
|
||||
<div class="line"><a id="l00235" name="l00235"></a><span class="lineno"> 235</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00236" name="l00236"></a><span class="lineno"> 236</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00237" name="l00237"></a><span class="lineno"> 237</span> dst_tile_dims -= short2(<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>, <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>);</div>
|
||||
<div class="line"><a id="l00238" name="l00238"></a><span class="lineno"> 238</span> </div>
|
||||
<div class="line"><a id="l00239" name="l00239"></a><span class="lineno"> 239</span> <span class="keywordflow">if</span> (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)</div>
|
||||
<div class="line"><a id="l00240" name="l00240"></a><span class="lineno"> 240</span> <span class="keywordflow">return</span>;</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> <span class="comment">// Loop over all simdgroup tiles</span></div>
|
||||
<div class="line"><a id="l00243" name="l00243"></a><span class="lineno"> 243</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00244" name="l00244"></a><span class="lineno"> 244</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00245" name="l00245"></a><span class="lineno"> 245</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00246" name="l00246"></a><span class="lineno"> 246</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00247" name="l00247"></a><span class="lineno"> 247</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00248" name="l00248"></a><span class="lineno"> 248</span> thread <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00249" name="l00249"></a><span class="lineno"> 249</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00250" name="l00250"></a><span class="lineno"> 250</span> </div>
|
||||
<div class="line"><a id="l00251" name="l00251"></a><span class="lineno"> 251</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00252" name="l00252"></a><span class="lineno"> 252</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00253" name="l00253"></a><span class="lineno"> 253</span> <span class="keywordflow">if</span> (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a> < dst_tile_dims.y) {</div>
|
||||
<div class="line"><a id="l00254" name="l00254"></a><span class="lineno"> 254</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00255" name="l00255"></a><span class="lineno"> 255</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00256" name="l00256"></a><span class="lineno"> 256</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00257" name="l00257"></a><span class="lineno"> 257</span> thread <span class="keyword">const</span> <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00258" name="l00258"></a><span class="lineno"> 258</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00259" name="l00259"></a><span class="lineno"> 259</span> <span class="keywordtype">int</span> offset_d = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldd + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>);</div>
|
||||
<div class="line"><a id="l00251" name="l00251"></a><span class="lineno"> 251</span> <span class="comment">// Read C</span></div>
|
||||
<div class="line"><a id="l00252" name="l00252"></a><span class="lineno"> 252</span> U c_elems[2] = {0};</div>
|
||||
<div class="line"><a id="l00253" name="l00253"></a><span class="lineno"> 253</span> </div>
|
||||
<div class="line"><a id="l00254" name="l00254"></a><span class="lineno"> 254</span> <span class="keywordflow">if</span> ((j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a> + 1) < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00255" name="l00255"></a><span class="lineno"> 255</span> c_elems[0] = C[offset_c];</div>
|
||||
<div class="line"><a id="l00256" name="l00256"></a><span class="lineno"> 256</span> c_elems[1] = C[offset_c + fdc];</div>
|
||||
<div class="line"><a id="l00257" name="l00257"></a><span class="lineno"> 257</span> } <span class="keywordflow">else</span> <span class="keywordflow">if</span> ((j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00258" name="l00258"></a><span class="lineno"> 258</span> c_elems[0] = C[offset_c];</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> </div>
|
||||
<div class="line"><a id="l00261" name="l00261"></a><span class="lineno"> 261</span> <span class="comment">// Apply epilogue and output C</span></div>
|
||||
<div class="line"><a id="l00262" name="l00262"></a><span class="lineno"> 262</span> <span class="keywordflow">if</span> (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a> < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00263" name="l00263"></a><span class="lineno"> 263</span> D[offset_d] = epilogue_op.apply(accum[0], C[offset_c]);</div>
|
||||
<div class="line"><a id="l00264" name="l00264"></a><span class="lineno"> 264</span> }</div>
|
||||
<div class="line"><a id="l00265" name="l00265"></a><span class="lineno"> 265</span> </div>
|
||||
<div class="line"><a id="l00266" name="l00266"></a><span class="lineno"> 266</span> <span class="keywordflow">if</span> (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a> + 1 < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00267" name="l00267"></a><span class="lineno"> 267</span> D[offset_d + 1] = epilogue_op.apply(accum[1], C[offset_c + fdc]);</div>
|
||||
<div class="line"><a id="l00268" name="l00268"></a><span class="lineno"> 268</span> }</div>
|
||||
<div class="line"><a id="l00269" name="l00269"></a><span class="lineno"> 269</span> }</div>
|
||||
<div class="line"><a id="l00270" name="l00270"></a><span class="lineno"> 270</span> }</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> }</div>
|
||||
<div class="line"><a id="l00261" name="l00261"></a><span class="lineno"> 261</span> <span class="comment">// Apply epilogue</span></div>
|
||||
<div class="line"><a id="l00262" name="l00262"></a><span class="lineno"> 262</span> accum[0] = epilogue_op.apply(accum[0], c_elems[0]);</div>
|
||||
<div class="line"><a id="l00263" name="l00263"></a><span class="lineno"> 263</span> accum[1] = epilogue_op.apply(accum[1], c_elems[1]);</div>
|
||||
<div class="line"><a id="l00264" name="l00264"></a><span class="lineno"> 264</span> }</div>
|
||||
<div class="line"><a id="l00265" name="l00265"></a><span class="lineno"> 265</span> }</div>
|
||||
<div class="line"><a id="l00266" name="l00266"></a><span class="lineno"> 266</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00273" name="l00273"></a><span class="lineno"> 273</span>};</div>
|
||||
<div class="line"><a id="l00267" name="l00267"></a><span class="lineno"> 267</span> </div>
|
||||
<div class="line"><a id="l00268" name="l00268"></a><span class="lineno"> 268</span> <span class="comment">/* Store results from simdgroup_matrix results into device memory */</span></div>
|
||||
<div class="foldopen" id="foldopen00269" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00269" name="l00269"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3"> 269</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3">store_result</a>(</div>
|
||||
<div class="line"><a id="l00270" name="l00270"></a><span class="lineno"> 270</span> device U* D,</div>
|
||||
<div class="line"><a id="l00271" name="l00271"></a><span class="lineno"> 271</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldd,</div>
|
||||
<div class="line"><a id="l00272" name="l00272"></a><span class="lineno"> 272</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00273" name="l00273"></a><span class="lineno"> 273</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00274" name="l00274"></a><span class="lineno"> 274</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00275" name="l00275"></a><span class="lineno"> 275</span> thread <span class="keyword">const</span> Epilogue& epilogue_op)<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"> 276</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00277" name="l00277"></a><span class="lineno"> 277</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"> 278</span> D += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldd + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>;</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">// Loop over all simdgroup tiles</span></div>
|
||||
<div class="line"><a id="l00281" name="l00281"></a><span class="lineno"> 281</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00282" name="l00282"></a><span class="lineno"> 282</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00283" name="l00283"></a><span class="lineno"> 283</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00284" name="l00284"></a><span class="lineno"> 284</span> <span class="keywordflow">for</span> (<span class="keywordtype">short</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00285" name="l00285"></a><span class="lineno"> 285</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00286" name="l00286"></a><span class="lineno"> 286</span> thread <span class="keyword">const</span> <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00287" name="l00287"></a><span class="lineno"> 287</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00288" name="l00288"></a><span class="lineno"> 288</span> <span class="keywordtype">int</span> offset_d = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldd + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>);</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> <span class="comment">// Apply epilogue</span></div>
|
||||
<div class="line"><a id="l00291" name="l00291"></a><span class="lineno"> 291</span> U outs[2] = {</div>
|
||||
<div class="line"><a id="l00292" name="l00292"></a><span class="lineno"> 292</span> epilogue_op.apply(accum[0], C[offset_c]),</div>
|
||||
<div class="line"><a id="l00293" name="l00293"></a><span class="lineno"> 293</span> epilogue_op.apply(accum[1], C[offset_c + fdc])};</div>
|
||||
<div class="line"><a id="l00294" name="l00294"></a><span class="lineno"> 294</span> </div>
|
||||
<div class="line"><a id="l00295" name="l00295"></a><span class="lineno"> 295</span> <span class="comment">// Write out D</span></div>
|
||||
<div class="line"><a id="l00296" name="l00296"></a><span class="lineno"> 296</span> D[offset_d] = outs[0];</div>
|
||||
<div class="line"><a id="l00297" name="l00297"></a><span class="lineno"> 297</span> D[offset_d + 1] = outs[1];</div>
|
||||
<div class="line"><a id="l00298" name="l00298"></a><span class="lineno"> 298</span> }</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>
|
||||
<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="comment">// namespace steel</span></div>
|
||||
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"> 276</span>} <span class="comment">// namespace mlx</span></div>
|
||||
<div class="line"><a id="l00301" name="l00301"></a><span class="lineno"> 301</span> </div>
|
||||
<div class="foldopen" id="foldopen00302" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00302" name="l00302"></a><span class="lineno"><a class="line" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391"> 302</a></span> METAL_FUNC <span class="keywordtype">void</span> <a class="code hl_function" href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391">store_result_safe</a>(</div>
|
||||
<div class="line"><a id="l00303" name="l00303"></a><span class="lineno"> 303</span> device U* D,</div>
|
||||
<div class="line"><a id="l00304" name="l00304"></a><span class="lineno"> 304</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldd,</div>
|
||||
<div class="line"><a id="l00305" name="l00305"></a><span class="lineno"> 305</span> <span class="keyword">const</span> device U* C,</div>
|
||||
<div class="line"><a id="l00306" name="l00306"></a><span class="lineno"> 306</span> <span class="keyword">const</span> <span class="keywordtype">int</span> ldc,</div>
|
||||
<div class="line"><a id="l00307" name="l00307"></a><span class="lineno"> 307</span> <span class="keyword">const</span> <span class="keywordtype">int</span> fdc,</div>
|
||||
<div class="line"><a id="l00308" name="l00308"></a><span class="lineno"> 308</span> short2 dst_tile_dims,</div>
|
||||
<div class="line"><a id="l00309" name="l00309"></a><span class="lineno"> 309</span> thread <span class="keyword">const</span> Epilogue& epilogue_op)<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00310" name="l00310"></a><span class="lineno"> 310</span> <span class="comment">// Adjust for simdgroup and thread location</span></div>
|
||||
<div class="line"><a id="l00311" name="l00311"></a><span class="lineno"> 311</span> C += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldc + (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00312" name="l00312"></a><span class="lineno"> 312</span> D += (<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>) * ldd + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>;</div>
|
||||
<div class="line"><a id="l00313" name="l00313"></a><span class="lineno"> 313</span> dst_tile_dims -= short2(<a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ad7a2033983cfbb474c50c2878057d8f1">tn</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ade420e8b811d597345783c324c23a34a">sn</a>, <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">sm</a> + <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#ab84d70540d09ce115794f343849f016f">tm</a>);</div>
|
||||
<div class="line"><a id="l00314" name="l00314"></a><span class="lineno"> 314</span> </div>
|
||||
<div class="line"><a id="l00315" name="l00315"></a><span class="lineno"> 315</span> <span class="keywordflow">if</span> (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)</div>
|
||||
<div class="line"><a id="l00316" name="l00316"></a><span class="lineno"> 316</span> <span class="keywordflow">return</span>;</div>
|
||||
<div class="line"><a id="l00317" name="l00317"></a><span class="lineno"> 317</span> </div>
|
||||
<div class="line"><a id="l00318" name="l00318"></a><span class="lineno"> 318</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00319" name="l00319"></a><span class="lineno"> 319</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> i = 0; i < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#aba5f749fdf32d8bd9d9e29f2a9ae4591">TM</a>; i++) {</div>
|
||||
<div class="line"><a id="l00320" name="l00320"></a><span class="lineno"> 320</span> <span class="keywordflow">if</span> (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a> < dst_tile_dims.y) {</div>
|
||||
<div class="line"><a id="l00321" name="l00321"></a><span class="lineno"> 321</span> <a class="code hl_define" href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div>
|
||||
<div class="line"><a id="l00322" name="l00322"></a><span class="lineno"> 322</span> <span class="keywordflow">for</span> (<span class="keywordtype">int</span> j = 0; j < <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a>; j++) {</div>
|
||||
<div class="line"><a id="l00323" name="l00323"></a><span class="lineno"> 323</span> <span class="comment">// Get accumulated result and associated offset in C</span></div>
|
||||
<div class="line"><a id="l00324" name="l00324"></a><span class="lineno"> 324</span> thread <span class="keyword">const</span> <span class="keyword">auto</span>& accum = <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">results</a>[i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">TN</a> + j].thread_elements();</div>
|
||||
<div class="line"><a id="l00325" name="l00325"></a><span class="lineno"> 325</span> <span class="keywordtype">int</span> offset_c = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldc + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>) * fdc;</div>
|
||||
<div class="line"><a id="l00326" name="l00326"></a><span class="lineno"> 326</span> <span class="keywordtype">int</span> offset_d = (i * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a5b0029866f493363942133b55bff7307">TM_stride</a>) * ldd + (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a>);</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> <span class="comment">// Apply epilogue and output C</span></div>
|
||||
<div class="line"><a id="l00329" name="l00329"></a><span class="lineno"> 329</span> <span class="keywordflow">if</span> (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a> < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00330" name="l00330"></a><span class="lineno"> 330</span> D[offset_d] = epilogue_op.apply(accum[0], C[offset_c]);</div>
|
||||
<div class="line"><a id="l00331" name="l00331"></a><span class="lineno"> 331</span> }</div>
|
||||
<div class="line"><a id="l00332" name="l00332"></a><span class="lineno"> 332</span> </div>
|
||||
<div class="line"><a id="l00333" name="l00333"></a><span class="lineno"> 333</span> <span class="keywordflow">if</span> (j * <a class="code hl_variable" href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">TN_stride</a> + 1 < dst_tile_dims.x) {</div>
|
||||
<div class="line"><a id="l00334" name="l00334"></a><span class="lineno"> 334</span> D[offset_d + 1] = epilogue_op.apply(accum[1], C[offset_c + fdc]);</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> }</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> }</div>
|
||||
<div class="line"><a id="l00339" name="l00339"></a><span class="lineno"> 339</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00340" name="l00340"></a><span class="lineno"> 340</span>};</div>
|
||||
</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>} <span class="comment">// namespace steel</span></div>
|
||||
<div class="line"><a id="l00343" name="l00343"></a><span class="lineno"> 343</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_a5a5c3095b132a7589bc19cd5cb80e2c6"><div class="ttname"><a href="backend_2metal_2kernels_2steel_2utils_8h.html#a5a5c3095b132a7589bc19cd5cb80e2c6">STEEL_PRAGMA_UNROLL</a></div><div class="ttdeci">#define STEEL_PRAGMA_UNROLL</div><div class="ttdef"><b>Definition</b> utils.h:8</div></div>
|
||||
@@ -391,13 +462,15 @@ $(function() { codefold.init(0); });
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a6a2c2a6d5e767d52c41b42a9d36086b0"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a6a2c2a6d5e767d52c41b42a9d36086b0">mlx::steel::BlockMMA::mma</a></div><div class="ttdeci">METAL_FUNC void mma(const threadgroup T *As, const threadgroup T *Bs)</div><div class="ttdef"><b>Definition</b> mma.h:93</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a6dfb5ff1f174951ad7c0630ff190c95a"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a6dfb5ff1f174951ad7c0630ff190c95a">mlx::steel::BlockMMA::results</a></div><div class="ttdeci">simdgroup_matrix< AccumType, 8, 8 > results[TM *TN]</div><div class="ttdef"><b>Definition</b> mma.h:62</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a706ae779c1f8d2eb18f19c248567d424"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a706ae779c1f8d2eb18f19c248567d424">mlx::steel::BlockMMA::TN</a></div><div class="ttdeci">STEEL_CONST short TN</div><div class="ttdef"><b>Definition</b> mma.h:44</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a7b324c992750ed3aaa4c485f15b2f391"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391">mlx::steel::BlockMMA::store_result_safe</a></div><div class="ttdeci">METAL_FUNC void store_result_safe(device U *D, const int ldd, const device U *C, const int ldc, const int fdc, short2 dst_tile_dims, thread const Epilogue &epilogue_op) const</div><div class="ttdef"><b>Definition</b> mma.h:235</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a7cf757e9785e23997b1417e024559ed3"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3">mlx::steel::BlockMMA::store_result</a></div><div class="ttdeci">METAL_FUNC void store_result(device U *D, const int ldd, const device U *C, const int ldc, const int fdc, thread const Epilogue &epilogue_op) const</div><div class="ttdef"><b>Definition</b> mma.h:202</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a7b324c992750ed3aaa4c485f15b2f391"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a7b324c992750ed3aaa4c485f15b2f391">mlx::steel::BlockMMA::store_result_safe</a></div><div class="ttdeci">METAL_FUNC void store_result_safe(device U *D, const int ldd, const device U *C, const int ldc, const int fdc, short2 dst_tile_dims, thread const Epilogue &epilogue_op) const</div><div class="ttdef"><b>Definition</b> mma.h:302</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a7cf757e9785e23997b1417e024559ed3"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a7cf757e9785e23997b1417e024559ed3">mlx::steel::BlockMMA::store_result</a></div><div class="ttdeci">METAL_FUNC void store_result(device U *D, const int ldd, const device U *C, const int ldc, const int fdc, thread const Epilogue &epilogue_op) const</div><div class="ttdef"><b>Definition</b> mma.h:269</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a823c56cbd2086f10272df7284a5247ae"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a823c56cbd2086f10272df7284a5247ae">mlx::steel::BlockMMA::apply_epilogue</a></div><div class="ttdeci">METAL_FUNC void apply_epilogue(const device U *C, const int ldc, const int fdc, thread const BinaryEpilogue &epilogue_op)</div><div class="ttdef"><b>Definition</b> mma.h:203</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a85c57ad4092079f068683cb525fb08ee"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a85c57ad4092079f068683cb525fb08ee">mlx::steel::BlockMMA::store_result</a></div><div class="ttdeci">METAL_FUNC void store_result(device U *D, const int ldd) const</div><div class="ttdef"><b>Definition</b> mma.h:147</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a8b3690b383afd26563efb38f9c375e50"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a8b3690b383afd26563efb38f9c375e50">mlx::steel::BlockMMA::TN_stride</a></div><div class="ttdeci">STEEL_CONST short TN_stride</div><div class="ttdef"><b>Definition</b> mma.h:39</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a8fddaa78913cdc8eea5e1cf7d2776330"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a8fddaa78913cdc8eea5e1cf7d2776330">mlx::steel::BlockMMA::tile_stride_a</a></div><div class="ttdeci">STEEL_CONST short tile_stride_a</div><div class="ttdef"><b>Definition</b> mma.h:56</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a92ac4b483cb58c361d7a657e0738c070"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a92ac4b483cb58c361d7a657e0738c070">mlx::steel::BlockMMA::Asimd</a></div><div class="ttdeci">simdgroup_matrix< AccumType, 8, 8 > Asimd[TM]</div><div class="ttdef"><b>Definition</b> mma.h:60</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a92f6aeee432f53638447eac842f43eca"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a92f6aeee432f53638447eac842f43eca">mlx::steel::BlockMMA::Bs_offset</a></div><div class="ttdeci">short Bs_offset</div><div class="ttdef"><b>Definition</b> mma.h:73</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_a9e48f2d51099ec00171506724faab54a"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#a9e48f2d51099ec00171506724faab54a">mlx::steel::BlockMMA::apply_epilogue_safe</a></div><div class="ttdeci">METAL_FUNC void apply_epilogue_safe(const device U *C, const int ldc, const int fdc, short2 dst_tile_dims, thread const BinaryEpilogue &epilogue_op)</div><div class="ttdef"><b>Definition</b> mma.h:229</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_aa14406b7298456ac45d23dd3c4642dd8"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#aa14406b7298456ac45d23dd3c4642dd8">mlx::steel::BlockMMA::BlockMMA</a></div><div class="ttdeci">METAL_FUNC BlockMMA(ushort simd_group_id, ushort simd_lane_id)</div><div class="ttdef"><b>Definition</b> mma.h:76</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_aa85451edf6900fd6af164d4d50889ae3"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#aa85451edf6900fd6af164d4d50889ae3">mlx::steel::BlockMMA::sm</a></div><div class="ttdeci">short sm</div><div class="ttdef"><b>Definition</b> mma.h:69</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1steel_1_1_block_m_m_a_html_ab52d2644053c5bb28ab9b24fab8b194a"><div class="ttname"><a href="structmlx_1_1steel_1_1_block_m_m_a.html#ab52d2644053c5bb28ab9b24fab8b194a">mlx::steel::BlockMMA::simd_stride_a</a></div><div class="ttdeci">STEEL_CONST short simd_stride_a</div><div class="ttdef"><b>Definition</b> mma.h:47</div></div>
|
||||
|
||||
Reference in New Issue
Block a user