leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

CUTLASS 3.1 Python interface documentation (#917)

Browse Source 
* Add 12.1 Dockerfile

* Add 3.1 docs
This commit is contained in:
Jack Kosaian
2023-04-18 15:11:35 -04:00
committed by GitHub
parent 54bebe417d
commit 9a83bd3381
81 changed files with 18997 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

476
python/docs/_modules/cutlass/utils/check.html Normal file
View File

@@ -0,0 +1,476 @@
<!doctype html>
<html class="no-js" lang="en">
<head><meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1"/>
<meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../genindex.html" /><link rel="search" title="Search" href="../../../search.html" />
<link rel="canonical" href="docs/_modules/cutlass/utils/check.html" />
<!-- Generated with Sphinx 6.1.3 and Furo 2023.03.27 -->
<title>cutlass.utils.check - CUTLASS Python</title>
<link rel="stylesheet" type="text/css" href="../../../_static/pygments.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
<link rel="stylesheet" type="text/css" href="../../../_static/copybutton.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/tabs.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
<style>
body {
--color-code-background: #eeffcc;
--color-code-foreground: black;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
@media not print {
body[data-theme="dark"] {
--color-code-background: #272822;
--color-code-foreground: #f8f8f2;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
@media (prefers-color-scheme: dark) {
body:not([data-theme="light"]) {
--color-code-background: #272822;
--color-code-foreground: #f8f8f2;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
}
}
</style></head>
<body>
<script>
document.body.dataset.theme = localStorage.getItem("theme") || "auto";
</script>
<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
<symbol id="svg-toc" viewBox="0 0 24 24">
<title>Contents</title>
<svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
<path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
</svg>
</symbol>
<symbol id="svg-menu" viewBox="0 0 24 24">
<title>Menu</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
<line x1="3" y1="12" x2="21" y2="12"></line>
<line x1="3" y1="6" x2="21" y2="6"></line>
<line x1="3" y1="18" x2="21" y2="18"></line>
</svg>
</symbol>
<symbol id="svg-arrow-right" viewBox="0 0 24 24">
<title>Expand</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
<polyline points="9 18 15 12 9 6"></polyline>
</svg>
</symbol>
<symbol id="svg-sun" viewBox="0 0 24 24">
<title>Light mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
<circle cx="12" cy="12" r="5"></circle>
<line x1="12" y1="1" x2="12" y2="3"></line>
<line x1="12" y1="21" x2="12" y2="23"></line>
<line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
<line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
<line x1="1" y1="12" x2="3" y2="12"></line>
<line x1="21" y1="12" x2="23" y2="12"></line>
<line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
<line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
</svg>
</symbol>
<symbol id="svg-moon" viewBox="0 0 24 24">
<title>Dark mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
<path stroke="none" d="M0 0h24v24H0z" fill="none" />
<path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
</svg>
</symbol>
<symbol id="svg-sun-half" viewBox="0 0 24 24">
<title>Auto light/dark mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
<path stroke="none" d="M0 0h24v24H0z" fill="none"/>
<circle cx="12" cy="12" r="9" />
<path d="M13 12h5" />
<path d="M13 15h4" />
<path d="M13 18h1" />
<path d="M13 9h4" />
<path d="M13 6h1" />
</svg>
</symbol>
</svg>
<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
<label class="overlay sidebar-overlay" for="__navigation">
<div class="visually-hidden">Hide navigation sidebar</div>
</label>
<label class="overlay toc-overlay" for="__toc">
<div class="visually-hidden">Hide table of contents sidebar</div>
</label>
<div class="page">
<header class="mobile-header">
<div class="header-left">
<label class="nav-overlay-icon" for="__navigation">
<div class="visually-hidden">Toggle site navigation sidebar</div>
<i class="icon"><svg><use href="#svg-menu"></use></svg></i>
</label>
</div>
<div class="header-center">
<a href="../../../index.html"><div class="brand">CUTLASS Python</div></a>
</div>
<div class="header-right">
<div class="theme-toggle-container theme-toggle-header">
<button class="theme-toggle">
<div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
<svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
<svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
<svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
</button>
</div>
<label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
<div class="visually-hidden">Toggle table of contents sidebar</div>
<i class="icon"><svg><use href="#svg-toc"></use></svg></i>
</label>
</div>
</header>
<aside class="sidebar-drawer">
<div class="sidebar-container">
<div class="sidebar-sticky"><a class="sidebar-brand" href="../../../index.html">
<div class="sidebar-logo-container">
<img class="sidebar-logo only-light" src="../../../_static/cutlass-logo-small.png" alt="Light Logo"/>
<img class="sidebar-logo only-dark" src="../../../_static/cutlass-logo-small.png" alt="Dark Logo"/>
</div>
<span class="sidebar-brand-text">CUTLASS Python</span>
</a><form class="sidebar-search-container" method="get" action="../../../search.html" role="search">
<input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
<input type="hidden" name="check_keywords" value="yes">
<input type="hidden" name="area" value="default">
</form>
<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../index.html">Home</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Getting Started:</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../install.html">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../externals/00_basic_gemm.html">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../contribute.html">Contributing</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Python Documentation:</span></p>
<ul>
<li class="toctree-l1 has-children"><a class="reference internal" href="../../../modules.html">CUTLASS Python API</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l2 has-children"><a class="reference internal" href="../../../cutlass.html">CUTLASS</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.emit.html">Emitters</a></li>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.op.html">Operations</a></li>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.utils.html">Utilities</a></li>
</ul>
</li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Examples and Tutorials:</span></p>
<ul>
<li class="toctree-l1 has-children"><a class="reference internal" href="../../../examples.html">Examples</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/00_basic_gemm.html">Basic GEMM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/01_epilogue.html">Epilogue</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/02_pytorch_extension_grouped_gemm.html">PyTorch Extension</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Reference:</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://github.com/NVIDIA/cutlass">Github</a></li>
</ul>
</div>
</div>
</div>
</div>
</aside>
<div class="main">
<div class="content">
<div class="article-container">
<a href="#" class="back-to-top muted-link">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
</svg>
<span>Back to top</span>
</a>
<div class="content-icon-container">
<div class="theme-toggle-container theme-toggle-content">
<button class="theme-toggle">
<div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
<svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
<svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
<svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
</button>
</div>
<label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
<div class="visually-hidden">Toggle table of contents sidebar</div>
<i class="icon"><svg><use href="#svg-toc"></use></svg></i>
</label>
</div>
<article role="main">
<h1>Source code for cutlass.utils.check</h1><div class="highlight"><pre>
<span></span><span class="c1">#################################################################################################</span>
<span class="c1">#</span>
<span class="c1"># Copyright (c) 2023 - 2023 NVIDIA CORPORATION &amp; AFFILIATES. All rights reserved.</span>
<span class="c1"># SPDX-License-Identifier: BSD-3-Clause</span>
<span class="c1">#</span>
<span class="c1"># Redistribution and use in source and binary forms, with or without</span>
<span class="c1"># modification, are permitted provided that the following conditions are met:</span>
<span class="c1">#</span>
<span class="c1"># 1. Redistributions of source code must retain the above copyright notice, this</span>
<span class="c1"># list of conditions and the following disclaimer.</span>
<span class="c1">#</span>
<span class="c1"># 2. Redistributions in binary form must reproduce the above copyright notice,</span>
<span class="c1"># this list of conditions and the following disclaimer in the documentation</span>
<span class="c1"># and/or other materials provided with the distribution.</span>
<span class="c1">#</span>
<span class="c1"># 3. Neither the name of the copyright holder nor the names of its</span>
<span class="c1"># contributors may be used to endorse or promote products derived from</span>
<span class="c1"># this software without specific prior written permission.</span>
<span class="c1">#</span>
<span class="c1"># THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS &quot;AS IS&quot;</span>
<span class="c1"># AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE</span>
<span class="c1"># IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE</span>
<span class="c1"># DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE</span>
<span class="c1"># FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL</span>
<span class="c1"># DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR</span>
<span class="c1"># SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER</span>
<span class="c1"># CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,</span>
<span class="c1"># OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE</span>
<span class="c1"># OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.</span>
<span class="c1">#</span>
<span class="c1">#################################################################################################</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Utility functions for checking constraints on kernels and calculating kernel attributes</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">ctypes</span>
<span class="kn">import</span> <span class="nn">cutlass_bindings</span>
<span class="kn">import</span> <span class="nn">cutlass</span>
<span class="kn">from</span> <span class="nn">cutlass.backend.library</span> <span class="kn">import</span> <span class="n">DataTypeSize</span><span class="p">,</span> <span class="n">TileDescription</span>
<div class="viewcode-block" id="calculate_smem_usage_per_stage"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.calculate_smem_usage_per_stage">[docs]</a><span class="k">def</span> <span class="nf">calculate_smem_usage_per_stage</span><span class="p">(</span><span class="n">tile_description</span><span class="p">,</span> <span class="n">operation_kind</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Returns the amount of shared memory in bytes consumed in a single stage of a kernel.</span>
<span class="sd"> :return: number of bytes of shared memory consumed by a single stage</span>
<span class="sd"> :rtype: int</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">m</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">k</span> <span class="o">=</span> <span class="n">tile_description</span><span class="o">.</span><span class="n">threadblock_shape</span>
<span class="k">if</span> <span class="n">operation_kind</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">OperationKind</span><span class="o">.</span><span class="n">Gemm</span><span class="p">:</span>
<span class="n">stage_barrier_bytes</span> <span class="o">=</span> <span class="mi">32</span>
<span class="k">return</span> <span class="p">(</span>
<span class="p">(</span><span class="n">DataTypeSize</span><span class="p">[</span><span class="n">tile_description</span><span class="o">.</span><span class="n">math_instruction</span><span class="o">.</span><span class="n">element_a</span><span class="p">]</span> <span class="o">*</span> <span class="n">m</span> <span class="o">*</span> <span class="n">k</span> <span class="o">//</span> <span class="mi">8</span><span class="p">)</span>
<span class="o">+</span> <span class="p">(</span><span class="n">DataTypeSize</span><span class="p">[</span><span class="n">tile_description</span><span class="o">.</span><span class="n">math_instruction</span><span class="o">.</span><span class="n">element_b</span><span class="p">]</span> <span class="o">*</span> <span class="n">k</span> <span class="o">*</span> <span class="n">n</span> <span class="o">//</span> <span class="mi">8</span><span class="p">)</span>
<span class="o">+</span> <span class="n">stage_barrier_bytes</span>
<span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No available shared memory calculation for operation kind </span><span class="si">{</span><span class="n">operation</span><span class="o">.</span><span class="n">operation_kind</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="calculate_smem_usage"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.calculate_smem_usage">[docs]</a><span class="k">def</span> <span class="nf">calculate_smem_usage</span><span class="p">(</span><span class="n">operation</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Returns the amount of shared memory in bytes consumed by a kernel.</span>
<span class="sd"> :return: number of bytes of shared memory consumed by the operation</span>
<span class="sd"> :return: int</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">_per_stage</span> <span class="o">=</span> <span class="n">calculate_smem_usage_per_stage</span><span class="p">(</span><span class="n">operation</span><span class="o">.</span><span class="n">tile_description</span><span class="p">,</span> <span class="n">operation</span><span class="o">.</span><span class="n">operation_kind</span><span class="p">)</span>
<span class="k">return</span> <span class="n">_per_stage</span> <span class="o">*</span> <span class="n">operation</span><span class="o">.</span><span class="n">tile_description</span><span class="o">.</span><span class="n">stages</span></div>
<div class="viewcode-block" id="valid_stage_count"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.valid_stage_count">[docs]</a><span class="k">def</span> <span class="nf">valid_stage_count</span><span class="p">(</span><span class="n">cc</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">td</span><span class="p">:</span> <span class="n">TileDescription</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Checks whether a device with `cc` supports the number of stages within `tile_description`, both</span>
<span class="sd"> based on raw limits on the number of stages and based on shared memory capacity</span>
<span class="sd"> :param cc: compute capability of device in question</span>
<span class="sd"> :type cc: int</span>
<span class="sd"> :param td: tile description to check</span>
<span class="sd"> :type td: TileDescription</span>
<span class="sd"> :return: tuple with the first element indicating whether the provided tile description is</span>
<span class="sd"> valid for the provided device and the second element being an error message</span>
<span class="sd"> :rtype: tuple</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">cc</span> <span class="o">==</span> <span class="mi">90</span> <span class="ow">and</span> <span class="p">(</span><span class="n">td</span><span class="o">.</span><span class="n">stages</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">td</span><span class="o">.</span><span class="n">stages</span> <span class="o">==</span> <span class="mi">0</span><span class="p">):</span>
<span class="c1"># Stage count of None or 0 for SM90 indicates that the CollectiveBuilder automatically</span>
<span class="c1"># determines the stage count to use. Thus, all settings are valid in these scenarios.</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">True</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="n">td</span><span class="o">.</span><span class="n">stages</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;Stage counts must be positive integers. Tile description has stage count of </span><span class="si">{</span><span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="n">cc</span> <span class="o">&lt;</span> <span class="mi">80</span> <span class="ow">and</span> <span class="n">td</span><span class="o">.</span><span class="n">stages</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;Tile description has stage count of </span><span class="si">{</span><span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="si">}</span><span class="s2">, &quot;</span>
<span class="sa">f</span><span class="s2">&quot;but only 2 stages are supported on SM</span><span class="si">{</span><span class="n">cc</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="n">smem_per_stage</span> <span class="o">=</span> <span class="n">calculate_smem_usage_per_stage</span><span class="p">(</span><span class="n">td</span><span class="p">,</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">OperationKind</span><span class="o">.</span><span class="n">Gemm</span><span class="p">)</span>
<span class="n">smem_arch</span> <span class="o">=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">SharedMemPerCC</span><span class="p">[</span><span class="n">cc</span><span class="p">]</span> <span class="o">&lt;&lt;</span> <span class="mi">10</span>
<span class="k">if</span> <span class="p">(</span><span class="n">smem_per_stage</span> <span class="o">*</span> <span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">smem_arch</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span> <span class="kc">False</span><span class="p">,</span>
<span class="s2">&quot;Configuration uses too much shared memory. Consider reducing stage count or tile shape.</span><span class="se">\n</span><span class="s2">&quot;</span>
<span class="sa">f</span><span class="s2">&quot;Details: configuration uses </span><span class="si">{</span><span class="n">smem_per_stage</span><span class="si">}</span><span class="s2"> bytes of shared memory per stage, and &quot;</span>
<span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="si">}</span><span class="s2"> stages for a total of </span><span class="si">{</span><span class="n">smem_per_stage</span><span class="w"> </span><span class="o">*</span><span class="w"> </span><span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="si">}</span><span class="s2"> bytes.</span><span class="se">\n</span><span class="s2">&quot;</span>
<span class="sa">f</span><span class="s2">&quot;The maxmium amoung of shared memory that can be used per block on CC </span><span class="si">{</span><span class="n">cc</span><span class="si">}</span><span class="s2"> is </span><span class="si">{</span><span class="n">smem_arch</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">True</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="valid_cluster_shape"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.valid_cluster_shape">[docs]</a><span class="k">def</span> <span class="nf">valid_cluster_shape</span><span class="p">(</span><span class="n">cc</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">cluster_shape</span><span class="p">:</span> <span class="nb">list</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Checks whether a device with `cc` supports a thread block cluster of shape `cluster_shape`.</span>
<span class="sd"> :param cc: compute capability of device in question</span>
<span class="sd"> :type cc: int</span>
<span class="sd"> :param cluster_shape: dimensions of thread block cluster shape to check</span>
<span class="sd"> :type cluster_shape: list</span>
<span class="sd"> :return: tuple with the first element indicating whether the provided cluster shape is</span>
<span class="sd"> valid for the provided device and the second element being an error message</span>
<span class="sd"> :rtype: tuple</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">cc</span> <span class="o">&lt;</span> <span class="mi">90</span><span class="p">:</span>
<span class="k">if</span> <span class="n">cluster_shape</span> <span class="o">!=</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span>
<span class="sa">f</span><span class="s2">&quot;Cluster shape for pre-SM90 architectures must be [1, 1, 1]. Received cluster shape of &quot;</span>
<span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">cluster_shape</span><span class="si">}</span><span class="s2"> for SM</span><span class="si">{</span><span class="n">cc</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">True</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">cluster_shape</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">3</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span>
<span class="sa">f</span><span class="s2">&quot;Cluster shapes must be rank-3. Received </span><span class="si">{</span><span class="n">cluster_shape</span><span class="si">}</span><span class="s2"> (rank </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">cluster_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="n">cluster_shape</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span>
<span class="s2">&quot;CUTLASS kernels currently require the third dimension of cluster shape to be 1. &quot;</span>
<span class="sa">f</span><span class="s2">&quot;Received cluster shape of </span><span class="si">{</span><span class="n">cluster_shape</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="c1"># The CUDA programming guide currently defines a maximum of 8 thread blocks per cluster</span>
<span class="c1"># as being portably supported (https://docs.nvidia.com/cuda/cuda-c-programming-guide/#thread-block-clusters).</span>
<span class="c1"># Current CUTLASS kernels only have non-unit cluster dimensions within the first two dimensions,</span>
<span class="c1"># so we check that the first two dimensions of the cluster shape do not exceed 8 thread blocks in total.</span>
<span class="n">blocks_in_2d</span> <span class="o">=</span> <span class="n">cluster_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">cluster_shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
<span class="k">if</span> <span class="n">blocks_in_2d</span> <span class="o">&gt;</span> <span class="mi">8</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span>
<span class="sa">f</span><span class="s2">&quot;Thread block clusters with more than 8 thread blocks are currently unsupported on SM</span><span class="si">{</span><span class="n">cc</span><span class="si">}</span><span class="s2">. &quot;</span>
<span class="sa">f</span><span class="s2">&quot;Received cluster shape </span><span class="si">{</span><span class="n">cluster_shape</span><span class="si">}</span><span class="s2">, which has </span><span class="si">{</span><span class="n">blocks_in_2d</span><span class="si">}</span><span class="s2"> thread blocks.&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">True</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="valid_kernel_schedule"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.valid_kernel_schedule">[docs]</a><span class="k">def</span> <span class="nf">valid_kernel_schedule</span><span class="p">(</span><span class="n">cc</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">kernel_schedule</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">KernelScheduleType</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Checks whether a device with ``cc`` supports ``kernel_schedule``.</span>
<span class="sd"> :param cc: compute capability of device in question</span>
<span class="sd"> :type cc: int</span>
<span class="sd"> :param kernel_schedule: kernel schedule type</span>
<span class="sd"> :type KernelScheduleType: cutlass.KernelScheduleType</span>
<span class="sd"> :return: tuple with the first element indicating whether the provided kernel schedule is</span>
<span class="sd"> valid for the provided device and the second element being an error message</span>
<span class="sd"> :rtype: tuple</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">kernel_schedule</span> <span class="o">!=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">KernelScheduleType</span><span class="o">.</span><span class="n">ScheduleAuto</span> <span class="ow">and</span> <span class="n">cc</span> <span class="o">&lt;</span> <span class="mi">90</span><span class="p">:</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="s2">&quot;Non-default kernel schedules are only supported on SM90 and beyond&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="kc">True</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="alignment_or_default"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.check.alignment_or_default">[docs]</a><span class="k">def</span> <span class="nf">alignment_or_default</span><span class="p">(</span><span class="n">alignment_provided</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">default_alignment</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Returns `alignment_provided` if it is set, otherwise `default_alignment` and checks</span>
<span class="sd"> that `alignment_provided` does not exceed `default_alignment`.</span>
<span class="sd"> :param alignment_provided: alignment preference specified. Can be None.</span>
<span class="sd"> :type alignment_provided: int</span>
<span class="sd"> :param default_alignment: alignment to use if `alignment_provided` is None</span>
<span class="sd"> :type default_alignment: int</span>
<span class="sd"> :return: alignment to use</span>
<span class="sd"> :rtype: int</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">alignment_provided</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="k">if</span> <span class="n">alignment_provided</span> <span class="o">&gt;</span> <span class="n">default_alignment</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Alignment </span><span class="si">{</span><span class="n">alignment_provided</span><span class="si">}</span><span class="s2"> exceeds the maximum supported of </span><span class="si">{</span><span class="n">default_alignment</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="n">alignment_provided</span>
<span class="k">return</span> <span class="n">default_alignment</span></div>
</pre></div>
</article>
</div>
<footer>
<div class="related-pages">
</div>
<div class="bottom-of-page">
<div class="left-details">
<div class="copyright">
Copyright &#169; 2023, NVIDIA
</div>
Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
<a href="https://github.com/pradyunsg/furo">Furo</a>
</div>
<div class="right-details">
<div class="icons">
<a class="muted-link " href="https://github.com/NVIDIA/cutlass" aria-label="GitHub">
<svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
<path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
</svg>
</a>
</div>
</div>
</div>
</footer>
</div>
<aside class="toc-drawer no-toc">
</aside>
</div>
</div><script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
<script src="../../../_static/doctools.js"></script>
<script src="../../../_static/sphinx_highlight.js"></script>
<script src="../../../_static/scripts/furo.js"></script>
<script src="../../../_static/clipboard.min.js"></script>
<script src="../../../_static/copybutton.js"></script>
<script src="../../../_static/tabs.js"></script>
<script crossorigin="anonymous" integrity="sha256-Ae2Vz/4ePdIu6ZyI/5ZGsYnb+m0JlOmKPjt6XZ9JJkA=" src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js"></script>
</body>
</html>

623
python/docs/_modules/cutlass/utils/datatypes.html Normal file
View File

@@ -0,0 +1,623 @@
<!doctype html>
<html class="no-js" lang="en">
<head><meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1"/>
<meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../genindex.html" /><link rel="search" title="Search" href="../../../search.html" />
<link rel="canonical" href="docs/_modules/cutlass/utils/datatypes.html" />
<!-- Generated with Sphinx 6.1.3 and Furo 2023.03.27 -->
<title>cutlass.utils.datatypes - CUTLASS Python</title>
<link rel="stylesheet" type="text/css" href="../../../_static/pygments.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
<link rel="stylesheet" type="text/css" href="../../../_static/copybutton.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/tabs.css" />
<link rel="stylesheet" type="text/css" href="../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
<style>
body {
--color-code-background: #eeffcc;
--color-code-foreground: black;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
@media not print {
body[data-theme="dark"] {
--color-code-background: #272822;
--color-code-foreground: #f8f8f2;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
@media (prefers-color-scheme: dark) {
body:not([data-theme="light"]) {
--color-code-background: #272822;
--color-code-foreground: #f8f8f2;
--color-brand-primary: #76B900;
--color-brand-content: #76B900;
}
}
}
</style></head>
<body>
<script>
document.body.dataset.theme = localStorage.getItem("theme") || "auto";
</script>
<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
<symbol id="svg-toc" viewBox="0 0 24 24">
<title>Contents</title>
<svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
<path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
</svg>
</symbol>
<symbol id="svg-menu" viewBox="0 0 24 24">
<title>Menu</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
<line x1="3" y1="12" x2="21" y2="12"></line>
<line x1="3" y1="6" x2="21" y2="6"></line>
<line x1="3" y1="18" x2="21" y2="18"></line>
</svg>
</symbol>
<symbol id="svg-arrow-right" viewBox="0 0 24 24">
<title>Expand</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
<polyline points="9 18 15 12 9 6"></polyline>
</svg>
</symbol>
<symbol id="svg-sun" viewBox="0 0 24 24">
<title>Light mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
<circle cx="12" cy="12" r="5"></circle>
<line x1="12" y1="1" x2="12" y2="3"></line>
<line x1="12" y1="21" x2="12" y2="23"></line>
<line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
<line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
<line x1="1" y1="12" x2="3" y2="12"></line>
<line x1="21" y1="12" x2="23" y2="12"></line>
<line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
<line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
</svg>
</symbol>
<symbol id="svg-moon" viewBox="0 0 24 24">
<title>Dark mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
<path stroke="none" d="M0 0h24v24H0z" fill="none" />
<path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
</svg>
</symbol>
<symbol id="svg-sun-half" viewBox="0 0 24 24">
<title>Auto light/dark mode</title>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
<path stroke="none" d="M0 0h24v24H0z" fill="none"/>
<circle cx="12" cy="12" r="9" />
<path d="M13 12h5" />
<path d="M13 15h4" />
<path d="M13 18h1" />
<path d="M13 9h4" />
<path d="M13 6h1" />
</svg>
</symbol>
</svg>
<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
<label class="overlay sidebar-overlay" for="__navigation">
<div class="visually-hidden">Hide navigation sidebar</div>
</label>
<label class="overlay toc-overlay" for="__toc">
<div class="visually-hidden">Hide table of contents sidebar</div>
</label>
<div class="page">
<header class="mobile-header">
<div class="header-left">
<label class="nav-overlay-icon" for="__navigation">
<div class="visually-hidden">Toggle site navigation sidebar</div>
<i class="icon"><svg><use href="#svg-menu"></use></svg></i>
</label>
</div>
<div class="header-center">
<a href="../../../index.html"><div class="brand">CUTLASS Python</div></a>
</div>
<div class="header-right">
<div class="theme-toggle-container theme-toggle-header">
<button class="theme-toggle">
<div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
<svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
<svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
<svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
</button>
</div>
<label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
<div class="visually-hidden">Toggle table of contents sidebar</div>
<i class="icon"><svg><use href="#svg-toc"></use></svg></i>
</label>
</div>
</header>
<aside class="sidebar-drawer">
<div class="sidebar-container">
<div class="sidebar-sticky"><a class="sidebar-brand" href="../../../index.html">
<div class="sidebar-logo-container">
<img class="sidebar-logo only-light" src="../../../_static/cutlass-logo-small.png" alt="Light Logo"/>
<img class="sidebar-logo only-dark" src="../../../_static/cutlass-logo-small.png" alt="Dark Logo"/>
</div>
<span class="sidebar-brand-text">CUTLASS Python</span>
</a><form class="sidebar-search-container" method="get" action="../../../search.html" role="search">
<input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
<input type="hidden" name="check_keywords" value="yes">
<input type="hidden" name="area" value="default">
</form>
<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../index.html">Home</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Getting Started:</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../install.html">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../externals/00_basic_gemm.html">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../contribute.html">Contributing</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Python Documentation:</span></p>
<ul>
<li class="toctree-l1 has-children"><a class="reference internal" href="../../../modules.html">CUTLASS Python API</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l2 has-children"><a class="reference internal" href="../../../cutlass.html">CUTLASS</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.emit.html">Emitters</a></li>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.op.html">Operations</a></li>
<li class="toctree-l3"><a class="reference internal" href="../../../cutlass.utils.html">Utilities</a></li>
</ul>
</li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Examples and Tutorials:</span></p>
<ul>
<li class="toctree-l1 has-children"><a class="reference internal" href="../../../examples.html">Examples</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/00_basic_gemm.html">Basic GEMM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/01_epilogue.html">Epilogue</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../../externals/02_pytorch_extension_grouped_gemm.html">PyTorch Extension</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Reference:</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://github.com/NVIDIA/cutlass">Github</a></li>
</ul>
</div>
</div>
</div>
</div>
</aside>
<div class="main">
<div class="content">
<div class="article-container">
<a href="#" class="back-to-top muted-link">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
</svg>
<span>Back to top</span>
</a>
<div class="content-icon-container">
<div class="theme-toggle-container theme-toggle-content">
<button class="theme-toggle">
<div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
<svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
<svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
<svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
</button>
</div>
<label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
<div class="visually-hidden">Toggle table of contents sidebar</div>
<i class="icon"><svg><use href="#svg-toc"></use></svg></i>
</label>
</div>
<article role="main">
<h1>Source code for cutlass.utils.datatypes</h1><div class="highlight"><pre>
<span></span><span class="c1">#################################################################################################</span>
<span class="c1">#</span>
<span class="c1"># Copyright (c) 2023 - 2023 NVIDIA CORPORATION &amp; AFFILIATES. All rights reserved.</span>
<span class="c1"># SPDX-License-Identifier: BSD-3-Clause</span>
<span class="c1">#</span>
<span class="c1"># Redistribution and use in source and binary forms, with or without</span>
<span class="c1"># modification, are permitted provided that the following conditions are met:</span>
<span class="c1">#</span>
<span class="c1"># 1. Redistributions of source code must retain the above copyright notice, this</span>
<span class="c1"># list of conditions and the following disclaimer.</span>
<span class="c1">#</span>
<span class="c1"># 2. Redistributions in binary form must reproduce the above copyright notice,</span>
<span class="c1"># this list of conditions and the following disclaimer in the documentation</span>
<span class="c1"># and/or other materials provided with the distribution.</span>
<span class="c1">#</span>
<span class="c1"># 3. Neither the name of the copyright holder nor the names of its</span>
<span class="c1"># contributors may be used to endorse or promote products derived from</span>
<span class="c1"># this software without specific prior written permission.</span>
<span class="c1">#</span>
<span class="c1"># THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS &quot;AS IS&quot;</span>
<span class="c1"># AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE</span>
<span class="c1"># IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE</span>
<span class="c1"># DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE</span>
<span class="c1"># FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL</span>
<span class="c1"># DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR</span>
<span class="c1"># SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER</span>
<span class="c1"># CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,</span>
<span class="c1"># OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE</span>
<span class="c1"># OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.</span>
<span class="c1">#</span>
<span class="c1">#################################################################################################</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Utility functions for converting between frontend datatypes and CUTLASS datatypes</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">cutlass_bindings</span>
<span class="kn">import</span> <span class="nn">cutlass</span>
<span class="kn">from</span> <span class="nn">cutlass.backend.library</span> <span class="kn">import</span> <span class="p">(</span>
<span class="n">DataTypeSize</span><span class="p">,</span>
<span class="n">MathInstruction</span><span class="p">,</span>
<span class="n">MathOperation</span><span class="p">,</span>
<span class="n">ShortLayoutTypeNames</span><span class="p">,</span>
<span class="n">TileDescription</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">try</span><span class="p">:</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="n">numpy_available</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">_library_to_numpy_dict</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s8</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s32</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span>
<span class="p">}</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
<span class="n">numpy_available</span> <span class="o">=</span> <span class="kc">False</span>
<span class="n">_library_to_numpy_dict</span> <span class="o">=</span> <span class="p">{}</span>
<div class="viewcode-block" id="numpy_library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.numpy_library_type">[docs]</a><span class="k">def</span> <span class="nf">numpy_library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">:</span>
<span class="k">if</span> <span class="n">numpy_available</span><span class="p">:</span>
<span class="k">if</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">float64</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">int8</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s8</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s32</span>
<span class="k">return</span> <span class="kc">None</span></div>
<div class="viewcode-block" id="numpy_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.numpy_type">[docs]</a><span class="k">def</span> <span class="nf">numpy_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">return</span> <span class="n">_library_to_numpy_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>
<span class="k">try</span><span class="p">:</span>
<span class="kn">import</span> <span class="nn">cupy</span> <span class="k">as</span> <span class="nn">cp</span>
<span class="n">cupy_available</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">_library_to_cupy_dict</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">:</span> <span class="n">cp</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">:</span> <span class="n">cp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">:</span> <span class="n">cp</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s8</span><span class="p">:</span> <span class="n">cp</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s32</span><span class="p">:</span> <span class="n">cp</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span>
<span class="p">}</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
<span class="n">cupy_available</span> <span class="o">=</span> <span class="kc">False</span>
<span class="n">_library_to_cupy_dict</span> <span class="o">=</span> <span class="p">{}</span>
<div class="viewcode-block" id="cupy_library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.cupy_library_type">[docs]</a><span class="k">def</span> <span class="nf">cupy_library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">:</span>
<span class="k">if</span> <span class="n">cupy_available</span><span class="p">:</span>
<span class="k">if</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">cp</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">cp</span><span class="o">.</span><span class="n">float32</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span>
<span class="k">elif</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">cp</span><span class="o">.</span><span class="n">float64</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span>
<span class="k">return</span> <span class="kc">None</span></div>
<div class="viewcode-block" id="cupy_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.cupy_type">[docs]</a><span class="k">def</span> <span class="nf">cupy_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">return</span> <span class="n">_library_to_cupy_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>
<span class="k">try</span><span class="p">:</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="n">torch_available</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">_torch_to_library_dict</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">torch</span><span class="o">.</span><span class="n">half</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">,</span>
<span class="n">torch</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">,</span>
<span class="n">torch</span><span class="o">.</span><span class="n">float</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">,</span>
<span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">,</span>
<span class="n">torch</span><span class="o">.</span><span class="n">double</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">,</span>
<span class="n">torch</span><span class="o">.</span><span class="n">float64</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">,</span>
<span class="p">}</span>
<span class="n">_library_to_torch_dict</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">half</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">float</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">double</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="p">}</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
<span class="n">torch_available</span> <span class="o">=</span> <span class="kc">False</span>
<span class="n">_torch_to_library_dict</span> <span class="o">=</span> <span class="p">{}</span>
<span class="n">_library_to_torch_dict</span> <span class="o">=</span> <span class="p">{}</span>
<div class="viewcode-block" id="torch_library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.torch_library_type">[docs]</a><span class="k">def</span> <span class="nf">torch_library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">:</span>
<span class="k">return</span> <span class="n">_torch_to_library_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>
<div class="viewcode-block" id="torch_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.torch_type">[docs]</a><span class="k">def</span> <span class="nf">torch_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">return</span> <span class="n">_library_to_torch_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>
<span class="k">try</span><span class="p">:</span>
<span class="kn">import</span> <span class="nn">bfloat16</span>
<span class="n">bfloat16_available</span> <span class="o">=</span> <span class="kc">True</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
<span class="n">bfloat16_available</span> <span class="o">=</span> <span class="kc">False</span>
<div class="viewcode-block" id="bfloat16_library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.bfloat16_library_type">[docs]</a><span class="k">def</span> <span class="nf">bfloat16_library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">:</span>
<span class="k">if</span> <span class="n">bfloat16_available</span><span class="p">:</span>
<span class="k">if</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">bfloat16</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">bf16</span></div>
<div class="viewcode-block" id="bfloat16_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.bfloat16_type">[docs]</a><span class="k">def</span> <span class="nf">bfloat16_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">bfloat16</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">:</span>
<span class="k">if</span> <span class="n">bfloat16_available</span><span class="p">:</span>
<span class="k">if</span> <span class="n">inp</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">bf16</span><span class="p">:</span>
<span class="k">return</span> <span class="n">bfloat16</span><span class="o">.</span><span class="n">bfloat16</span></div>
<span class="c1"># Mapping from library data type to Python-bound CUTLASS data type</span>
<span class="n">library_to_binding_dict</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s8</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s32</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">bf16</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">tf32</span><span class="p">:</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">tfloat32</span><span class="p">,</span>
<span class="p">}</span>
<span class="c1"># Mapping from Python-bound CUTLASS data type to library data type</span>
<span class="n">binding_to_library</span> <span class="o">=</span> <span class="p">{</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">int8</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s8</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">int32</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">s32</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f16</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">bf16</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float32</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f32</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">float64</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">f64</span><span class="p">,</span>
<span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">tfloat32</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="o">.</span><span class="n">tf32</span><span class="p">,</span>
<span class="p">}</span>
<div class="viewcode-block" id="binding_library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.binding_library_type">[docs]</a><span class="k">def</span> <span class="nf">binding_library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">if</span> <span class="n">inp</span> <span class="ow">in</span> <span class="n">binding_to_library</span><span class="p">:</span>
<span class="k">return</span> <span class="n">binding_to_library</span><span class="p">[</span><span class="n">inp</span><span class="p">]</span>
<span class="k">return</span> <span class="kc">None</span></div>
<div class="viewcode-block" id="has_binding_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.has_binding_type">[docs]</a><span class="k">def</span> <span class="nf">has_binding_type</span><span class="p">(</span><span class="n">inp</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">):</span>
<span class="k">return</span> <span class="n">inp</span> <span class="ow">in</span> <span class="n">library_to_binding_dict</span></div>
<div class="viewcode-block" id="library_to_binding"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.library_to_binding">[docs]</a><span class="k">def</span> <span class="nf">library_to_binding</span><span class="p">(</span><span class="n">inp</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataType</span><span class="p">):</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">has_binding_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No available conversion from library type </span><span class="si">{</span><span class="n">inp</span><span class="si">}</span><span class="s2"> to Python-bound CUTLASS type&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="n">library_to_binding_dict</span><span class="p">[</span><span class="n">inp</span><span class="p">]</span></div>
<div class="viewcode-block" id="library_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.library_type">[docs]</a><span class="k">def</span> <span class="nf">library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">if</span> <span class="n">inp</span> <span class="ow">in</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">DataTypeSize</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
<span class="k">return</span> <span class="n">inp</span>
<span class="k">for</span> <span class="n">cvt_fn</span> <span class="ow">in</span> <span class="p">[</span>
<span class="n">bfloat16_library_type</span><span class="p">,</span>
<span class="n">cupy_library_type</span><span class="p">,</span>
<span class="n">numpy_library_type</span><span class="p">,</span>
<span class="n">torch_library_type</span><span class="p">,</span>
<span class="n">binding_library_type</span><span class="p">,</span>
<span class="p">]:</span>
<span class="n">out</span> <span class="o">=</span> <span class="n">cvt_fn</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span>
<span class="k">if</span> <span class="n">out</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="k">return</span> <span class="n">out</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No available conversion from type </span><span class="si">{</span><span class="n">inp</span><span class="si">}</span><span class="s2"> to a library type.&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="library_layout"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.library_layout">[docs]</a><span class="k">def</span> <span class="nf">library_layout</span><span class="p">(</span><span class="n">layout</span><span class="p">):</span>
<span class="k">if</span> <span class="n">layout</span> <span class="ow">in</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutTag</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
<span class="k">return</span> <span class="n">layout</span>
<span class="c1"># Convert Python-bound CUTLASS layout to profiler library layout</span>
<span class="k">if</span> <span class="n">layout</span> <span class="o">==</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">RowMajor</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">RowMajor</span>
<span class="k">elif</span> <span class="n">layout</span> <span class="o">==</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">ColumnMajor</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">ColumnMajor</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No conversion available for layout </span><span class="si">{</span><span class="n">layout</span><span class="si">}</span><span class="s2"> to library layout.&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="binding_type"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.binding_type">[docs]</a><span class="k">def</span> <span class="nf">binding_type</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span>
<span class="k">if</span> <span class="n">inp</span> <span class="ow">in</span> <span class="n">DataTypeSize</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
<span class="k">return</span> <span class="n">inp</span>
<span class="n">libtype</span> <span class="o">=</span> <span class="n">library_type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span>
<span class="k">return</span> <span class="n">library_to_binding</span><span class="p">(</span><span class="n">libtype</span><span class="p">)</span></div>
<div class="viewcode-block" id="binding_layout"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.binding_layout">[docs]</a><span class="k">def</span> <span class="nf">binding_layout</span><span class="p">(</span><span class="n">layout</span><span class="p">):</span>
<span class="k">if</span> <span class="n">layout</span> <span class="ow">in</span> <span class="n">ShortLayoutTypeNames</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
<span class="k">return</span> <span class="n">layout</span>
<span class="k">elif</span> <span class="n">layout</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">RowMajor</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">RowMajor</span>
<span class="k">elif</span> <span class="n">layout</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">ColumnMajor</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">ColumnMajor</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No conversion available for layout </span><span class="si">{</span><span class="n">layout</span><span class="si">}</span><span class="s2"> to Python-bound CUTLASS layout.&quot;</span><span class="p">)</span></div>
<span class="k">def</span> <span class="nf">_tensor_from_numpy</span><span class="p">(</span><span class="n">np_tensor</span><span class="p">):</span>
<span class="n">dtype</span> <span class="o">=</span> <span class="n">library_type</span><span class="p">(</span><span class="n">np_tensor</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
<span class="k">if</span> <span class="n">np_tensor</span><span class="o">.</span><span class="n">flags</span><span class="o">.</span><span class="n">c_contiguous</span><span class="p">:</span>
<span class="n">layout</span> <span class="o">=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">RowMajor</span>
<span class="k">elif</span> <span class="n">np_tensor</span><span class="o">.</span><span class="n">flags</span><span class="o">.</span><span class="n">f_contiguous</span><span class="p">:</span>
<span class="n">layout</span> <span class="o">=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">ColumnMajor</span>
<span class="k">return</span> <span class="p">(</span><span class="n">dtype</span><span class="p">,</span> <span class="n">layout</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_tensor_from_torch</span><span class="p">(</span><span class="n">pt_tensor</span><span class="p">):</span>
<span class="n">dtype</span> <span class="o">=</span> <span class="n">library_type</span><span class="p">(</span><span class="n">pt_tensor</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="n">dtype</span><span class="p">,</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">LayoutType</span><span class="o">.</span><span class="n">RowMajor</span><span class="p">)</span>
<div class="viewcode-block" id="get_datatype_and_layout"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.get_datatype_and_layout">[docs]</a><span class="k">def</span> <span class="nf">get_datatype_and_layout</span><span class="p">(</span><span class="n">tensor</span><span class="p">):</span>
<span class="k">if</span> <span class="p">(</span><span class="n">numpy_available</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">))</span> <span class="ow">or</span> <span class="p">(</span>
<span class="n">cupy_available</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">cp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">)</span>
<span class="p">):</span>
<span class="k">return</span> <span class="n">_tensor_from_numpy</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
<span class="k">elif</span> <span class="n">torch_available</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
<span class="k">return</span> <span class="n">_tensor_from_torch</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Unable to convert tensor of type </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span><span class="si">}</span><span class="s2"> to Python-bound CUTLASS datatype and layout.&quot;</span><span class="p">)</span></div>
<div class="viewcode-block" id="binding_opclass"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.binding_opclass">[docs]</a><span class="k">def</span> <span class="nf">binding_opclass</span><span class="p">(</span><span class="n">opclass</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">OpcodeClass</span><span class="p">):</span>
<span class="k">if</span> <span class="n">opclass</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">OpcodeClass</span><span class="o">.</span><span class="n">TensorOp</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">OpClass</span><span class="o">.</span><span class="n">TensorOp</span>
<span class="k">elif</span> <span class="n">opclass</span> <span class="o">==</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">OpcodeClass</span><span class="o">.</span><span class="n">Simt</span><span class="p">:</span>
<span class="k">return</span> <span class="n">cutlass_bindings</span><span class="o">.</span><span class="n">OpClass</span><span class="o">.</span><span class="n">Simt</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Unable to convert opcode class of type </span><span class="si">{</span><span class="n">opclass</span><span class="si">}</span><span class="s2"> to Python-bound CUTLASS opcode class.&quot;</span><span class="p">)</span></div>
<span class="n">_math_operation_value_map</span> <span class="o">=</span> <span class="p">{</span><span class="n">x</span><span class="o">.</span><span class="n">value</span><span class="p">:</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">MathOperation</span><span class="p">}</span>
<div class="viewcode-block" id="backend_math_operation"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.backend_math_operation">[docs]</a><span class="k">def</span> <span class="nf">backend_math_operation</span><span class="p">(</span><span class="n">math_op</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">MathOperation</span><span class="p">):</span>
<span class="k">if</span> <span class="n">math_op</span><span class="o">.</span><span class="n">value</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_math_operation_value_map</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Unable to convert math operation of type </span><span class="si">{</span><span class="n">math_op</span><span class="si">}</span><span class="s2"> to backend math operation.&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="n">_math_operation_value_map</span><span class="p">[</span><span class="n">math_op</span><span class="o">.</span><span class="n">value</span><span class="p">]</span></div>
<div class="viewcode-block" id="construct_backend_td"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.construct_backend_td">[docs]</a><span class="k">def</span> <span class="nf">construct_backend_td</span><span class="p">(</span><span class="n">td</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">TileDescription</span><span class="p">,</span>
<span class="n">kernel_schedule</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">KernelScheduleType</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">TileDescription</span><span class="p">:</span>
<span class="n">mi</span> <span class="o">=</span> <span class="n">td</span><span class="o">.</span><span class="n">math_instruction</span>
<span class="n">backend_mi</span> <span class="o">=</span> <span class="n">MathInstruction</span><span class="p">(</span>
<span class="n">mi</span><span class="o">.</span><span class="n">instruction_shape</span><span class="p">,</span>
<span class="n">binding_type</span><span class="p">(</span><span class="n">mi</span><span class="o">.</span><span class="n">element_a</span><span class="p">),</span>
<span class="n">binding_type</span><span class="p">(</span><span class="n">mi</span><span class="o">.</span><span class="n">element_b</span><span class="p">),</span>
<span class="n">binding_type</span><span class="p">(</span><span class="n">mi</span><span class="o">.</span><span class="n">element_accumulator</span><span class="p">),</span>
<span class="n">binding_opclass</span><span class="p">(</span><span class="n">mi</span><span class="o">.</span><span class="n">opcode_class</span><span class="p">),</span>
<span class="n">backend_math_operation</span><span class="p">(</span><span class="n">mi</span><span class="o">.</span><span class="n">math_operation</span><span class="p">)</span>
<span class="p">)</span>
<span class="k">return</span> <span class="n">TileDescription</span><span class="p">(</span><span class="n">td</span><span class="o">.</span><span class="n">threadblock_shape</span><span class="p">,</span> <span class="n">td</span><span class="o">.</span><span class="n">stages</span><span class="p">,</span> <span class="n">td</span><span class="o">.</span><span class="n">warp_count</span><span class="p">,</span>
<span class="n">backend_mi</span><span class="p">,</span> <span class="n">td</span><span class="o">.</span><span class="n">cluster_shape</span><span class="p">,</span> <span class="n">kernel_schedule</span><span class="p">)</span></div>
<div class="viewcode-block" id="td_from_profiler_op"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.td_from_profiler_op">[docs]</a><span class="k">def</span> <span class="nf">td_from_profiler_op</span><span class="p">(</span><span class="n">op</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">TileDescription</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Converts the profiler&#39;s TileDescription in ``op`` into the backend TileDescription</span>
<span class="sd"> :param op: profiler Operation</span>
<span class="sd"> :returns: backend TileDescription</span>
<span class="sd"> :rtype: cutlass.backend.TileDescription</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">schedule</span> <span class="o">=</span> <span class="n">op</span><span class="o">.</span><span class="n">kernel_schedule</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="s1">&#39;kernel_schedule&#39;</span><span class="p">)</span> <span class="k">else</span> <span class="kc">None</span>
<span class="k">return</span> <span class="n">construct_backend_td</span><span class="p">(</span><span class="n">op</span><span class="o">.</span><span class="n">tile_description</span><span class="p">,</span> <span class="n">schedule</span><span class="p">)</span></div>
<div class="viewcode-block" id="td_from_profiler_td"><a class="viewcode-back" href="../../../cutlass.utils.html#cutlass.utils.datatypes.td_from_profiler_td">[docs]</a><span class="k">def</span> <span class="nf">td_from_profiler_td</span><span class="p">(</span><span class="n">td</span><span class="p">:</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">backend</span><span class="o">.</span><span class="n">TileDescription</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">TileDescription</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Converts the profiler&#39;s TileDescription into the backend TileDescription</span>
<span class="sd"> :param td: profiler TileDescription</span>
<span class="sd"> :type td: cutlass.TileDescription</span>
<span class="sd"> :returns: backend TileDescription</span>
<span class="sd"> :rtype: cutlass.backend.TileDescription</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">return</span> <span class="n">construct_backend_td</span><span class="p">(</span><span class="n">td</span><span class="p">,</span> <span class="n">kernel_schedule</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span></div>
</pre></div>
</article>
</div>
<footer>
<div class="related-pages">
</div>
<div class="bottom-of-page">
<div class="left-details">
<div class="copyright">
Copyright &#169; 2023, NVIDIA
</div>
Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
<a href="https://github.com/pradyunsg/furo">Furo</a>
</div>
<div class="right-details">
<div class="icons">
<a class="muted-link " href="https://github.com/NVIDIA/cutlass" aria-label="GitHub">
<svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
<path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
</svg>
</a>
</div>
</div>
</div>
</footer>
</div>
<aside class="toc-drawer no-toc">
</aside>
</div>
</div><script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
<script src="../../../_static/doctools.js"></script>
<script src="../../../_static/sphinx_highlight.js"></script>
<script src="../../../_static/scripts/furo.js"></script>
<script src="../../../_static/clipboard.min.js"></script>
<script src="../../../_static/copybutton.js"></script>
<script src="../../../_static/tabs.js"></script>
<script crossorigin="anonymous" integrity="sha256-Ae2Vz/4ePdIu6ZyI/5ZGsYnb+m0JlOmKPjt6XZ9JJkA=" src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js"></script>
</body>
</html>