821 lines
52 KiB
HTML
821 lines
52 KiB
HTML
<!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"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
|
||
<link rel="index" title="Index" href="../genindex.html" /><link rel="search" title="Search" href="../search.html" /><link rel="next" title="Contributing" href="../contribute.html" /><link rel="prev" title="Installation" href="../install.html" />
|
||
<link rel="canonical" href="docs/externals/00_basic_gemm.html" />
|
||
|
||
<!-- Generated with Sphinx 6.1.3 and Furo 2023.03.27 -->
|
||
<title>Basic example of using the CUTLASS Python interface - 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/nbsphinx-code-cells.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" 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 class="current">
|
||
<li class="toctree-l1"><a class="reference internal" href="../install.html">Installation</a></li>
|
||
<li class="toctree-l1 current"><a class="current reference internal" href="#">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 class="current">
|
||
<li class="toctree-l1 current has-children"><a class="reference internal" href="../examples.html">Examples</a><input checked="" 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 class="current">
|
||
<li class="toctree-l2 current current-page"><a class="current reference internal" href="#">Basic GEMM</a></li>
|
||
<li class="toctree-l2"><a class="reference internal" href="01_epilogue.html">Epilogue</a></li>
|
||
<li class="toctree-l2"><a class="reference internal" href="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" 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">
|
||
<section id="Basic-example-of-using-the-CUTLASS-Python-interface">
|
||
<h1>Basic example of using the CUTLASS Python interface<a class="headerlink" href="#Basic-example-of-using-the-CUTLASS-Python-interface" title="Permalink to this heading">#</a></h1>
|
||
<p>This notebook walks through a basic example of using the CUTLASS Python interface to declare, compile, and run GEMMs.</p>
|
||
<p><a class="reference external" href="https://colab.research.google.com/github/NVIDIA/cutlass/tree/master/examples/00_basic_gemm.ipynb"><img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg" /></a></p>
|
||
<p>We first import various packages needed for the example and construct the input and output tensors that will be used in our example.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[1]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
|
||
<span class="kn">import</span> <span class="nn">random</span>
|
||
|
||
<span class="kn">import</span> <span class="nn">cutlass</span>
|
||
|
||
<span class="c1"># This controls whether ther C++ GEMM declaration will be printed at each step. Set to `false` to</span>
|
||
<span class="c1"># omit this information.</span>
|
||
<span class="n">print_module</span> <span class="o">=</span> <span class="kc">True</span>
|
||
|
||
<span class="n">m</span> <span class="o">=</span> <span class="mi">128</span>
|
||
<span class="n">n</span> <span class="o">=</span> <span class="n">m</span>
|
||
<span class="n">k</span> <span class="o">=</span> <span class="n">m</span>
|
||
|
||
<span class="n">dtype</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span>
|
||
<span class="n">type_A</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span>
|
||
<span class="n">type_B</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span>
|
||
<span class="n">type_C</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span>
|
||
<span class="n">type_D</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float16</span>
|
||
|
||
<span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="mi">1234</span><span class="p">)</span>
|
||
<span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="mi">1234</span><span class="p">)</span>
|
||
<span class="n">scope_min</span> <span class="o">=</span> <span class="o">-</span><span class="mi">4</span>
|
||
<span class="n">scope_max</span> <span class="o">=</span> <span class="mi">4</span>
|
||
<span class="n">tensor_A</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">k</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_A</span><span class="p">))</span>
|
||
<span class="n">tensor_B</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">n</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_B</span><span class="p">))</span>
|
||
<span class="n">tensor_C</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">n</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_C</span><span class="p">))</span>
|
||
|
||
<span class="n">alpha</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="mf">1.</span><span class="p">)</span>
|
||
<span class="n">beta</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="mf">0.</span><span class="p">)</span>
|
||
|
||
<span class="n">tensor_D</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">tensor_C</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_D</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area stderr docutils container">
|
||
<div class="highlight"><pre>
|
||
/usr/local/lib/python3.8/dist-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
|
||
from .autonotebook import tqdm as notebook_tqdm
|
||
</pre></div></div>
|
||
</div>
|
||
<section id="Declaring-and-running-a-GEMM">
|
||
<h2>Declaring and running a GEMM<a class="headerlink" href="#Declaring-and-running-a-GEMM" title="Permalink to this heading">#</a></h2>
|
||
<p>To get started, one only needs to provide the tensors declared above to the <code class="docutils literal notranslate"><span class="pre">cutlass.op.Gemm</span></code> call. This sets up a default GEMM operation for the given device on which you are running.</p>
|
||
<p>Assuming that we are running on SM80, this default to using a GEMM that leverages FP16 Tensor Core operations.</p>
|
||
<p>Calling <code class="docutils literal notranslate"><span class="pre">plan.run()</span></code> will generate the CUTLASS C++ kernel in question, compile it, and run it on the tensors we previously passed in. By setting <code class="docutils literal notranslate"><span class="pre">print_module</span></code> to <code class="docutils literal notranslate"><span class="pre">true</span></code>, the C++ code that is emitted is printed.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[2]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># We specify `element_accumulator` here so as to match the kernel run by NumPy below. However,</span>
|
||
<span class="c1"># specifying `element_accumulator` is not required if it is the same as `element`</span>
|
||
<span class="n">plan</span> <span class="o">=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">Gemm</span><span class="p">(</span><span class="n">element</span><span class="o">=</span><span class="n">dtype</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="p">,</span> <span class="n">element_accumulator</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="n">plan</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">tensor_A</span><span class="p">,</span> <span class="n">tensor_B</span><span class="p">,</span> <span class="n">tensor_C</span><span class="p">,</span> <span class="n">tensor_D</span><span class="p">,</span> <span class="n">print_module</span><span class="o">=</span><span class="n">print_module</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
|
||
// Gemm operator cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8
|
||
using cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_base =
|
||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor,
|
||
float,
|
||
cutlass::arch::OpClassTensorOp,
|
||
cutlass::arch::Sm80,
|
||
cutlass::gemm::GemmShape<256, 128, 64>,
|
||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||
cutlass::gemm::GemmShape<16, 8, 16>,
|
||
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 8, float, float>,
|
||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>,
|
||
3,
|
||
cutlass::arch::OpMultiplyAdd
|
||
>::GemmKernel;
|
||
|
||
// Define named type
|
||
struct cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_type :
|
||
public cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_base { };
|
||
|
||
</pre></div></div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[2]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
<cutlass.backend.gemm_operation.GemmArguments2x at 0x7f79cc556070>
|
||
</pre></div></div>
|
||
</div>
|
||
<p>There are many other ways to construct a plan from <code class="docutils literal notranslate"><span class="pre">cutlass.op.Gemm</span></code> (e.g., by specifiying they types and layouts of each operand, by providing representative tensors as inputs). For more details on these, see the documentation in the <code class="docutils literal notranslate"><span class="pre">cutlass.op.Gemm</span></code> constructor.</p>
|
||
<p>We then compare the output to running the GEMM using NumPy.</p>
|
||
<div class="nbinput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[3]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">tensor_D_numpy</span> <span class="o">=</span> <span class="p">(</span><span class="n">alpha</span> <span class="o">*</span> <span class="p">(</span><span class="n">tensor_A</span> <span class="o">@</span> <span class="n">tensor_B</span><span class="p">))</span> <span class="o">+</span> <span class="p">(</span><span class="n">beta</span> <span class="o">*</span> <span class="n">tensor_C</span><span class="p">)</span>
|
||
<span class="n">np</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_array_equal</span><span class="p">(</span><span class="n">tensor_D</span><span class="p">,</span> <span class="n">tensor_D_numpy</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<p>Note that one could use the same kernel just declared for tensors provided by other frameworks beyond NumPy, such as PyTorch or CuPy.</p>
|
||
</section>
|
||
<section id="Changing-operation-modes">
|
||
<h2>Changing operation modes<a class="headerlink" href="#Changing-operation-modes" title="Permalink to this heading">#</a></h2>
|
||
<p>By default, the CUTLASS Python interface will try to use Tensor Core operations whenever possible. If the configuration provided to <code class="docutils literal notranslate"><span class="pre">cutlass.op.Gemm</span></code> is not supported on Tensor Cores, the interface will fall back to using a SIMT kernel.</p>
|
||
<p>The operation mode currently in use can be returned via the <code class="docutils literal notranslate"><span class="pre">plan.opclass</span></code> property. In this case Tensor Core operations.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[4]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="nb">print</span><span class="p">(</span><span class="n">plan</span><span class="o">.</span><span class="n">opclass</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
OpcodeClass.TensorOp
|
||
</pre></div></div>
|
||
</div>
|
||
<p>Suppose that we don’t want to use Tensor Cores for this GEMM. One can change to using CUTLASS’s SIMT GEMMs by setting the plan’s <code class="docutils literal notranslate"><span class="pre">opclass</span></code> field.</p>
|
||
<p>As is shown in the printed output, the emitted kernel uses template parameters that fit CUTLASS’s SIMT GEMMs.</p>
|
||
<p>Also notice that, this time around, we provided tensor parameters to <code class="docutils literal notranslate"><span class="pre">plan.run()</span></code>. One is free to provide different parameters to <code class="docutils literal notranslate"><span class="pre">plan.run()</span></code> than were passed in at the initial call to <code class="docutils literal notranslate"><span class="pre">cutlass.op.Gemm</span></code>, provided that the passed-in tensors have the same data type and layout as those passed in on intialization.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[5]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">tensor_D_simt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">tensor_C</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_D</span><span class="p">)</span>
|
||
<span class="n">plan</span><span class="o">.</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="n">plan</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">tensor_A</span><span class="p">,</span> <span class="n">tensor_B</span><span class="p">,</span> <span class="n">tensor_C</span><span class="p">,</span> <span class="n">tensor_D_simt</span><span class="p">,</span> <span class="n">alpha</span><span class="p">,</span> <span class="n">beta</span><span class="p">,</span> <span class="n">print_module</span><span class="o">=</span><span class="n">print_module</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
|
||
// Gemm operator cutlass_sm80_simt_f16_sgemm_f16_1x1x1_128x128_8x2_tt_align1
|
||
using cutlass_sm80_simt_f16_sgemm_f16_1x1x1_128x128_8x2_tt_align1_base =
|
||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1,
|
||
cutlass::half_t, cutlass::layout::RowMajor,
|
||
float,
|
||
cutlass::arch::OpClassSimt,
|
||
cutlass::arch::Sm80,
|
||
cutlass::gemm::GemmShape<128, 128, 8>,
|
||
cutlass::gemm::GemmShape<32, 64, 8>,
|
||
cutlass::gemm::GemmShape<1, 1, 1>,
|
||
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 1, float, float>,
|
||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>,
|
||
2,
|
||
cutlass::arch::OpMultiplyAdd
|
||
>::GemmKernel;
|
||
|
||
// Define named type
|
||
struct cutlass_sm80_simt_f16_sgemm_f16_1x1x1_128x128_8x2_tt_align1_type :
|
||
public cutlass_sm80_simt_f16_sgemm_f16_1x1x1_128x128_8x2_tt_align1_base { };
|
||
|
||
</pre></div></div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[5]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
<cutlass.backend.gemm_operation.GemmArguments2x at 0x7f7b3075abe0>
|
||
</pre></div></div>
|
||
</div>
|
||
<p>If we compare the output of the Tensor Core and SIMT GEMMs we just ran we see that they are equal.</p>
|
||
<div class="nbinput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[6]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">np</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_array_equal</span><span class="p">(</span><span class="n">tensor_D</span><span class="p">,</span> <span class="n">tensor_D_simt</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
</section>
|
||
<section id="Running-cached-kernels">
|
||
<h2>Running cached kernels<a class="headerlink" href="#Running-cached-kernels" title="Permalink to this heading">#</a></h2>
|
||
<p>You may have noticed that the <code class="docutils literal notranslate"><span class="pre">plan.run()</span></code> calls for the previous two kernels took some time to execute. This is because the kernel being emitted had not yet been compiled.</p>
|
||
<p>CUTLASS caches compiled binaries so that recompilation isn’t necessary every time a kernel is run. For example, if we change modes back to using Tensor Cores and call <code class="docutils literal notranslate"><span class="pre">plan.run()</span></code> again (with a different set of tensor parameters), you’ll find the call to return much faster.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[7]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">m</span> <span class="o">=</span> <span class="mi">2400</span>
|
||
<span class="n">n</span> <span class="o">=</span> <span class="mi">3232</span>
|
||
<span class="n">k</span> <span class="o">=</span> <span class="mi">4096</span>
|
||
|
||
<span class="n">tensor_A</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">k</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_A</span><span class="p">))</span>
|
||
<span class="n">tensor_B</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">n</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_B</span><span class="p">))</span>
|
||
<span class="n">tensor_C</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">low</span><span class="o">=</span><span class="n">scope_min</span><span class="p">,</span> <span class="n">high</span><span class="o">=</span><span class="n">scope_max</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">n</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_C</span><span class="p">))</span>
|
||
<span class="n">tensor_D</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">tensor_C</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">type_D</span><span class="p">)</span>
|
||
|
||
<span class="n">alpha</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="mf">1.</span><span class="p">)</span>
|
||
<span class="n">beta</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="mf">2.</span><span class="p">)</span>
|
||
|
||
<span class="n">plan</span><span class="o">.</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="n">plan</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">tensor_A</span><span class="p">,</span> <span class="n">tensor_B</span><span class="p">,</span> <span class="n">tensor_C</span><span class="p">,</span> <span class="n">tensor_D</span><span class="p">,</span> <span class="n">alpha</span><span class="p">,</span> <span class="n">beta</span><span class="p">,</span> <span class="n">print_module</span><span class="o">=</span><span class="n">print_module</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
|
||
// Gemm operator cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8
|
||
using cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_base =
|
||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor,
|
||
float,
|
||
cutlass::arch::OpClassTensorOp,
|
||
cutlass::arch::Sm80,
|
||
cutlass::gemm::GemmShape<256, 128, 64>,
|
||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||
cutlass::gemm::GemmShape<16, 8, 16>,
|
||
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 8, float, float>,
|
||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>,
|
||
3,
|
||
cutlass::arch::OpMultiplyAdd
|
||
>::GemmKernel;
|
||
|
||
// Define named type
|
||
struct cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_type :
|
||
public cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_256x128_64x3_tt_align8_base { };
|
||
|
||
</pre></div></div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[7]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
<cutlass.backend.gemm_operation.GemmArguments2x at 0x7f7b30fb9880>
|
||
</pre></div></div>
|
||
</div>
|
||
</section>
|
||
<section id="Running-non-default-GEMMs">
|
||
<h2>Running non-default GEMMs<a class="headerlink" href="#Running-non-default-GEMMs" title="Permalink to this heading">#</a></h2>
|
||
<p>The previous examples showed how it is simple to get started running a default GEMM kernel in CUTLASS. But, what do you do if you want a bit more control over the parameters to the GEMM?</p>
|
||
<p>Under the hood, CUTLASS enumerates the different GEMM configuration parameters possible for this kernel from the CUTLASS profiler. The code below shows how one can access the tile descriptions for the kernels (e.g., cluster, threadblock, and warp shape).</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[8]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">tiles</span> <span class="o">=</span> <span class="n">plan</span><span class="o">.</span><span class="n">tile_descriptions</span><span class="p">()</span>
|
||
<span class="nb">print</span><span class="p">(</span><span class="s1">'</span><span class="si">{}</span><span class="s1"> tile descriptions returned'</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">tiles</span><span class="p">)))</span>
|
||
<span class="n">num_print</span> <span class="o">=</span> <span class="mi">10</span>
|
||
<span class="nb">print</span><span class="p">(</span><span class="s1">'First </span><span class="si">{}</span><span class="s1"> tile descriptions are:'</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">num_print</span><span class="p">))</span>
|
||
<span class="k">for</span> <span class="n">td</span> <span class="ow">in</span> <span class="n">tiles</span><span class="p">[:</span><span class="n">num_print</span><span class="p">]:</span>
|
||
<span class="nb">print</span><span class="p">(</span><span class="n">td</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
132 tile descriptions returned
|
||
First 10 tile descriptions are:
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [256, 128, 64]
|
||
WarpCount: [4, 2, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [128, 256, 64]
|
||
WarpCount: [2, 4, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [256, 128, 64]
|
||
WarpCount: [4, 2, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [128, 256, 64]
|
||
WarpCount: [2, 4, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [256, 128, 32]
|
||
WarpCount: [4, 2, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [128, 256, 32]
|
||
WarpCount: [2, 4, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [256, 64, 64]
|
||
WarpCount: [4, 1, 1]
|
||
Stages: 4
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [64, 256, 64]
|
||
WarpCount: [1, 4, 1]
|
||
Stages: 4
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [128, 128, 64]
|
||
WarpCount: [2, 2, 1]
|
||
Stages: 4
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [256, 64, 64]
|
||
WarpCount: [4, 1, 1]
|
||
Stages: 3
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
</pre></div></div>
|
||
</div>
|
||
<p>Next, we’ll pick one of these configurations at random and compile and run it.</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[9]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">idx</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">tiles</span><span class="p">)</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
|
||
<span class="n">td</span> <span class="o">=</span> <span class="n">tiles</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span>
|
||
<span class="nb">print</span><span class="p">(</span><span class="s1">'Tile description </span><span class="si">{}</span><span class="s1"> is: </span><span class="si">{}</span><span class="s1">'</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">td</span><span class="p">))</span>
|
||
<span class="n">plan</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">td</span><span class="p">)</span>
|
||
<span class="n">plan</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">tensor_A</span><span class="p">,</span> <span class="n">tensor_B</span><span class="p">,</span> <span class="n">tensor_C</span><span class="p">,</span> <span class="n">tensor_D</span><span class="p">,</span> <span class="n">alpha</span><span class="p">,</span> <span class="n">beta</span><span class="p">,</span> <span class="n">print_module</span><span class="o">=</span><span class="n">print_module</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
Tile description 112 is:
|
||
{
|
||
ClusterShape: [1, 1, 1]
|
||
ThreadblockShape: [128, 128, 32]
|
||
WarpCount: [2, 2, 1]
|
||
Stages: 4
|
||
Kernel schedule: ScheduleAuto
|
||
}
|
||
|
||
// Gemm operator cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8
|
||
using cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_base =
|
||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor,
|
||
float,
|
||
cutlass::arch::OpClassTensorOp,
|
||
cutlass::arch::Sm80,
|
||
cutlass::gemm::GemmShape<128, 128, 32>,
|
||
cutlass::gemm::GemmShape<64, 64, 32>,
|
||
cutlass::gemm::GemmShape<16, 8, 16>,
|
||
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 8, float, float>,
|
||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>,
|
||
4,
|
||
cutlass::arch::OpMultiplyAdd
|
||
>::GemmKernel;
|
||
|
||
// Define named type
|
||
struct cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_type :
|
||
public cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_base { };
|
||
|
||
</pre></div></div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[9]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
<cutlass.backend.gemm_operation.GemmArguments2x at 0x7f79cc58de20>
|
||
</pre></div></div>
|
||
</div>
|
||
<p>One can also change the swizzling function used by the kernel. For example, one can modify the kernel to use the stream K feature of CUTLASS via:</p>
|
||
<div class="nbinput docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[10]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># Stream K is only supported pre-SM90 (at least when this example was written)</span>
|
||
<span class="k">if</span> <span class="n">plan</span><span class="o">.</span><span class="n">cc</span> <span class="o">!=</span> <span class="mi">90</span><span class="p">:</span>
|
||
<span class="n">plan</span><span class="o">.</span><span class="n">swizzling_functor</span> <span class="o">=</span> <span class="n">cutlass</span><span class="o">.</span><span class="n">swizzle</span><span class="o">.</span><span class="n">ThreadblockSwizzleStreamK</span>
|
||
<span class="n">plan</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="n">tensor_A</span><span class="p">,</span> <span class="n">tensor_B</span><span class="p">,</span> <span class="n">tensor_C</span><span class="p">,</span> <span class="n">tensor_D</span><span class="p">,</span> <span class="n">alpha</span><span class="p">,</span> <span class="n">beta</span><span class="p">,</span> <span class="n">print_module</span><span class="o">=</span><span class="n">print_module</span><span class="p">)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
<div class="nboutput nblast docutils container">
|
||
<div class="prompt empty docutils container">
|
||
</div>
|
||
<div class="output_area docutils container">
|
||
<div class="highlight"><pre>
|
||
|
||
// Gemm operator cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8
|
||
using cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_base =
|
||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8,
|
||
cutlass::half_t, cutlass::layout::RowMajor,
|
||
float,
|
||
cutlass::arch::OpClassTensorOp,
|
||
cutlass::arch::Sm80,
|
||
cutlass::gemm::GemmShape<128, 128, 32>,
|
||
cutlass::gemm::GemmShape<64, 64, 32>,
|
||
cutlass::gemm::GemmShape<16, 8, 16>,
|
||
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 8, float, float>,
|
||
cutlass::gemm::threadblock::ThreadblockSwizzleStreamK,
|
||
4,
|
||
cutlass::arch::OpMultiplyAdd
|
||
>::GemmKernel;
|
||
|
||
// Define named type
|
||
struct cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_type :
|
||
public cutlass_sm80_tensorop_f16_s16x8x16gemm_f16_1x1x1_128x128_32x4_tt_align8_base { };
|
||
|
||
</pre></div></div>
|
||
</div>
|
||
</section>
|
||
<section id="Handling-errors">
|
||
<h2>Handling errors<a class="headerlink" href="#Handling-errors" title="Permalink to this heading">#</a></h2>
|
||
<p>The CUTLASS Python interface attempts to catch runtime and compilation errors in Python so as to provide more understandable error messages.</p>
|
||
<p>Here’s an example in which we try to use too many stages for a given GEMM kernel. Normally, this would result in a runtime error due to the GPU having insufficient shared memory to launch the kernel with 8 stages. The CUTLASS Python interface is able to detect this issue before compiling the kernel, and reports it back to the user.</p>
|
||
<div class="nbinput nblast docutils container">
|
||
<div class="prompt highlight-none notranslate"><div class="highlight"><pre><span></span>[11]:
|
||
</pre></div>
|
||
</div>
|
||
<div class="input_area highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># td = tiles[0]</span>
|
||
<span class="c1"># td.stages = 8</span>
|
||
<span class="c1"># plan.compile(td)</span>
|
||
</pre></div>
|
||
</div>
|
||
</div>
|
||
</section>
|
||
</section>
|
||
|
||
</article>
|
||
</div>
|
||
<footer>
|
||
|
||
<div class="related-pages">
|
||
<a class="next-page" href="../contribute.html">
|
||
<div class="page-info">
|
||
<div class="context">
|
||
<span>Next</span>
|
||
</div>
|
||
<div class="title">Contributing</div>
|
||
</div>
|
||
<svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
|
||
</a>
|
||
<a class="prev-page" href="../install.html">
|
||
<svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
|
||
<div class="page-info">
|
||
<div class="context">
|
||
<span>Previous</span>
|
||
</div>
|
||
|
||
<div class="title">Installation</div>
|
||
|
||
</div>
|
||
</a>
|
||
</div>
|
||
<div class="bottom-of-page">
|
||
<div class="left-details">
|
||
<div class="copyright">
|
||
Copyright © 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">
|
||
|
||
|
||
<div class="toc-sticky toc-scroll">
|
||
<div class="toc-title-container">
|
||
<span class="toc-title">
|
||
On this page
|
||
</span>
|
||
</div>
|
||
<div class="toc-tree-container">
|
||
<div class="toc-tree">
|
||
<ul>
|
||
<li><a class="reference internal" href="#">Basic example of using the CUTLASS Python interface</a><ul>
|
||
<li><a class="reference internal" href="#Declaring-and-running-a-GEMM">Declaring and running a GEMM</a></li>
|
||
<li><a class="reference internal" href="#Changing-operation-modes">Changing operation modes</a></li>
|
||
<li><a class="reference internal" href="#Running-cached-kernels">Running cached kernels</a></li>
|
||
<li><a class="reference internal" href="#Running-non-default-GEMMs">Running non-default GEMMs</a></li>
|
||
<li><a class="reference internal" href="#Handling-errors">Handling errors</a></li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
|
||
</div>
|
||
</div>
|
||
</div>
|
||
|
||
|
||
</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>
|
||
<script>window.MathJax = {"tex": {"inlineMath": [["$", "$"], ["\\(", "\\)"]], "processEscapes": true}, "options": {"ignoreHtmlClass": "tex2jax_ignore|mathjax_ignore|document", "processHtmlClass": "tex2jax_process|mathjax_process|math|output_area"}}</script>
|
||
<script defer="defer" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||
</body>
|
||
</html> |