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

v4.5 tag update (#3202)

Browse Source 
* Python DSL examples reorganization.

* v4.5 tag update.
This commit is contained in:
Junkai-Wu
2026-05-05 20:55:27 -04:00
committed by GitHub
parent f74fea9ce3
commit cb37157db5
351 changed files with 36688 additions and 8117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/cute/notebooks/README.md
+31
View File
@@ -0,0 +1,31 @@
# Copyright
Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
SPDX-License-Identifier: BSD-3-Clause
```
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
```
examples/python/CuTeDSL/cute/notebooks/async_pipeline.ipynb
+599
View File
@@ -0,0 +1,599 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"# Tutorial: Warp Specialization with Async Pipeline in CuTe DSL\n",
"\n",
"This tutorial explores advanced CUDA programming techniques for implementing efficient producer-consumer \n",
"patterns using asynchronous communication primitives in the CuTe Domain Specific Language (DSL).\n",
"\n",
"## Foundation: Inter-Warp Communication Basics\n",
"\n",
"### Understanding CUDA Warps and Shared Memory\n",
"\n",
"A **warp** is the fundamental execution unit in CUDA, consisting of 32 threads that execute instructions in Single Instruction, \n",
"Multiple Thread (SIMT) fashion on a Streaming Multiprocessor (SM). Understanding warp-level programming is crucial for \n",
"achieving optimal GPU performance.\n",
"\n",
"**Key Concepts:**\n",
"- Warps execute in lockstep, making them ideal for SIMD operations\n",
"- Multiple warps within a thread block (CTA) can cooperate through shared memory\n",
"- Shared memory provides low-latency, high-bandwidth communication between threads\n",
"\n",
"### Shared Memory Architecture\n",
"\n",
"**Shared memory** serves as a programmer-managed cache with several important characteristics:\n",
"\n",
"- **Speed**: ~100x faster than global memory access\n",
"- **Scope**: Accessible by all threads within the same thread block\n",
"- **Organization**: Divided into banks (typically 32) to enable parallel access\n",
"- **Conflicts**: Bank conflicts occur when multiple threads access the same bank simultaneously\n",
"\n",
"### Traditional Synchronous Communication\n",
"\n",
"The conventional approach for inter-warp communication relies on explicit synchronization barriers. The following sequence diagram \n",
"illustrates the typical producer-consumer pattern:\n",
"\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer Warp\n",
" participant SMEM as Shared Memory\n",
" participant W1 as Consumer Warp\n",
" \n",
" W0->>SMEM: Write data\n",
" critical Synchronization Barrier\n",
" W0-->W1: __syncthreads()\n",
" SMEM->>W1: Read data\n",
" W0-->W1: __syncthreads()\n",
" end\n",
"```\n",
"\n",
"**Limitations of Synchronous Communication:**\n",
"- All warps must wait at synchronization points\n",
"- No opportunity for overlapped computation\n",
"- Reduced overall throughput due to forced serialization"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def synced_producer_consumer(SharedStorage: cutlass.Constexpr, res: cute.Tensor):\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(cute.make_layout(1))\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" for i in cutlass.range(cute.size(res)):\n",
" if warp_idx == 0:\n",
" staging_smem[0] = i * 1.0\n",
" # mark enter of critical region\n",
" cute.arch.sync_threads()\n",
" if warp_idx == 1:\n",
" res[i] = staging_smem[0]\n",
" # mark exit of critical region\n",
" cute.arch.sync_threads()\n",
"\n",
"\n",
"@cute.jit\n",
"def run_synced_producer_consumer(res: cute.Tensor):\n",
" @cute.struct\n",
" class SharedStorage:\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, 1], 1024\n",
" ]\n",
"\n",
" synced_producer_consumer(SharedStorage, res).launch(\n",
" grid=(1, 1, 1), block=(64, 1, 1)\n",
" )\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"run_synced_producer_consumer(from_dlpack(res))"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0')"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res"
]
},
{
"cell_type": "markdown",
"metadata": {
"editable": true,
"slideshow": {
"slide_type": ""
},
"tags": []
},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"## Asynchronous Communication: Breaking the Synchronization Bottleneck\n",
"\n",
"### The Problem with Synchronous Patterns\n",
"\n",
"The previous example demonstrates traditional synchronized communication between warps. While functional, this approach \n",
"has significant performance limitations:\n",
"\n",
"**Critical Section Analysis:**\n",
"- **First `__syncthreads()`**: Ensures data is written and ready for consumption\n",
"- **Second `__syncthreads()`**: Guarantees data has been consumed and memory can be safely overwritten\n",
"\n",
"**Performance Impact:**\n",
"- All warps are forced into lockstep execution\n",
"- No computational overlap between producer and consumer operations\n",
"- Wasted cycles as warps wait at synchronization barriers\n",
"\n",
"### Hopper Architecture: Enabling Asynchronous Primitives\n",
"\n",
"Starting with the Hopper architecture, CUDA introduced sophisticated asynchronous communication primitives that enable \n",
"**warp specialization**—allowing different warps to perform distinct, specialized roles while maintaining loose coupling.\n",
"\n",
"**Key Benefits:**\n",
"- **Overlapped Execution**: Producer and consumer warps can perform computations concurrently\n",
"- **Reduced Latency**: Eliminates unnecessary synchronization stalls\n",
"- **Better Resource Utilization**: Maximizes SM occupancy and throughput\n",
"\n",
"### Async Pipeline Communication Pattern\n",
"\n",
"The async pipeline abstraction provides a elegant solution for producer-consumer communication without rigid synchronization constraints:\n",
"\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer Warp\n",
" participant Pipeline as Async Pipeline\n",
" participant SMEM as Shared Memory \n",
" participant W1 as Consumer Warp\n",
" \n",
" W0->>Pipeline: Acquire (request write slot)\n",
" activate W1\n",
" Pipeline-->>W0: Grant access\n",
" deactivate W1\n",
" \n",
" W1->>Pipeline: Wait (for data availability)\n",
" activate Pipeline\n",
" \n",
" W0->>SMEM: Write data\n",
" W0->>Pipeline: Commit (signal data ready)\n",
" \n",
" Pipeline-->>W1: Data available\n",
" deactivate Pipeline\n",
" \n",
" activate W0\n",
" SMEM->>W1: Read data\n",
" deactivate W0\n",
" W1->>Pipeline: Release (mark slot available)\n",
"```\n",
"\n",
"**Async Pipeline Advantages:**\n",
"- **Non-blocking Operations**: Warps can perform other work while waiting\n",
"- **Fine-grained Control**: Explicit control over data readiness and consumption\n",
"- **Scalable**: Supports multiple producer-consumer pairs efficiently"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Async Pipeline API Reference\n",
"\n",
"The `PipelineAsync` abstraction in CuTe DSL provides a comprehensive set of primitives for implementing efficient producer-consumer patterns:\n",
"\n",
"#### Producer Operations\n",
"- **`PipelineProducer.acquire()`**: Blocks until a write slot becomes available (released by consumer)\n",
" - Returns with a handle pointing to a available slot immediately if there is\n",
" - Enables backpressure control to prevent buffer overflow\n",
" - **`PipelineProducer.acquire_and_advance()`** additionally moves the producer's write index to the next buffer slot\n",
"\n",
"- **`PipelineProducer.commit(PipelineProducer.ImmutableProducerHandle)`** / **`PipelineProducer.ImmutableProducerHandle.commit()`**: Signals that data has been written to the handle-pointed slot and is ready for consumption\n",
" - Triggers waiting consumers\n",
" - Maintains data consistency guarantees\n",
" - If no assigned handle, **`PipelineConsumerHandle.release()`** tracks its internal maintained handle (pointed to the last one it acquires)\n",
"\n",
"#### Consumer Operations \n",
"- **`PipelineConsumer.wait()`**: Blocks until data becomes available for reading\n",
" - Returns with a handle pointing to a committed slot when producer commits new data\n",
" - Supports timeout and polling variants\n",
" - **`PipelineConsumer.wait_and_advance()`** additionally moves the consumer's read index to the next buffer slot\n",
"\n",
"- **`PipelineConsumerHandle.release(PipelineConsumer.ImmutableConsumerHandle)`** / **`PipelineConsumer.ImmutableConsumerHandle.release()`**: Marks data as consumed and the handle-pointed slot as consumed and available for reuse\n",
" - Enables producers to acquire released slots\n",
" - Critical for preventing deadlock in circular buffers\n",
" - If no assigned handle, **`PipelineConsumerHandle.release()`** tracks its internal maintained handle (pointed to the last one it waits for)\n",
"\n",
"#### Disclaimer\n",
"\n",
"The `pipeline` APIs provided abstractions for developers to manage synchornization between warps, thread-blocks, etc. It doesn't provide deadlock-free guarantee. It's still developer's responsibility to write correct code to avoid deadlock.\n",
"\n",
"#### Performance Characteristics\n",
"\n",
"**Computational Overlap**: This asynchronous communication pattern enables limited but significant computational overlap:\n",
"- **Producer**: Can perform preprocessing, data transformation, or prefetching while consumer processes previous data\n",
"- **Consumer**: Can execute post-processing, result computation, or output operations while producer prepares next data\n",
"\n",
"**Memory Efficiency**: Explicit slot management ensures optimal memory utilization without unnecessary copying or buffering."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def async_pipeline_kernel(res: cute.Tensor):\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" @cute.struct\n",
" class SharedStorage:\n",
" tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, 1], 1024\n",
" ]\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" # Warp 0\n",
" producer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
" # Warp 1\n",
" consumer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
"\n",
" pipeline = cutlass.pipeline.PipelineAsync.create(\n",
" num_stages=1,\n",
" producer_group=producer_group,\n",
" consumer_group=consumer_group,\n",
" barrier_storage=storage.tma_mbar_ptr.data_ptr(),\n",
" )\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(cute.make_layout(1))\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" producer, consumer = pipeline.make_participants()\n",
"\n",
" # Producer warp\n",
" if warp_idx == 0:\n",
" for i in cutlass.range(cute.size(res)):\n",
" # Producer: Wait for data buffer is available\n",
" handle = producer.acquire_and_advance()\n",
" # Producer: Write data to shared memory\n",
" staging_smem[handle.index] = 1.0 * i\n",
" # Producer: Signal data is ready for consumption\n",
" handle.commit()\n",
" producer.tail()\n",
"\n",
" # Consumer warp\n",
" if warp_idx == 1:\n",
" for i in cutlass.range(cute.size(res)):\n",
" # Consumer: Wait for producer to signal when data is available for use\n",
" handle = consumer.wait_and_advance()\n",
" # Conumer: consumes data\n",
" res[i] = staging_smem[handle.index]\n",
" # Conumer: Signal data buffer is ready for write\n",
" handle.release()\n",
"\n",
"\n",
"@cute.jit\n",
"def async_pipeline(res: cute.Tensor):\n",
" # Launch kernel with two warps: producer and consumer\n",
" async_pipeline_kernel(res).launch(grid=(1, 1, 1), block=(64, 1, 1))\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"async_pipeline(from_dlpack(res))"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0')"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"## Advanced Pattern: Staged Async Pipeline with Circular Buffering\n",
"\n",
"### Limitations of Single-Stage Pipelines\n",
"\n",
"While async communication provides significant improvements over synchronous patterns, single-stage pipelines \n",
"still exhibit serialization bottlenecks:\n",
"\n",
"**Dependency Chain Analysis:**\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer\n",
" participant Pipeline as Pipeline\n",
" participant W1 as Consumer\n",
" \n",
" W0->>Pipeline: Acquire\n",
" Note over W0,W1: Producer waits here\n",
" W1->>Pipeline: Release\n",
" Pipeline-->>W0: Granted\n",
"```\n",
"\n",
"**Performance Bottleneck**: The producer must wait for the consumer to complete processing and release the buffer \n",
"before acquiring the next write slot. This creates a serialization point that limits overall throughput.\n",
"\n",
"### Multi-Stage Pipeline Architecture\n",
"\n",
"The **staged async pipeline** implements a circular buffer managed by an array of synchronization barriers, \n",
"enabling much higher degrees of parallelism:\n",
"\n",
"#### Core Concepts\n",
"\n",
"**Circular Buffer Management:**\n",
"- **Multiple Stages**: Support for N concurrent buffer slots (typically 2-8 stages)\n",
"- **Independent Indexing**: Producer and consumer maintain separate advancement indices\n",
"- **Barrier Array**: Each stage has an associated memory barrier for fine-grained synchronization\n",
"\n",
"#### Enhanced API Operations\n",
"\n",
"- **`PipelineProducer.advance()`**: Moves the producer's write index to the next buffer slot\n",
" - Enables round-robin buffer allocation\n",
" - Allows producer to continue without waiting for all previous data to be consumed\n",
" - Can be conducted implicitly when calling **`PipelineProducer.require_and_advance()`**\n",
"\n",
"- **`PipelineConsumer.advance()`**: Moves the consumer's read index to the next buffer slot\n",
" - Maintains proper ordering of data consumption\n",
" - Signals availability of processed slots\n",
" - Can be conducted implicitly when calling **`PipelineConsumer.wait_and_advance()`**\n",
"\n",
"- **`PipelineProducer.ImmutableResourceHandle.index`** / **`PipelineConsumer.ImmutableResourceHandle.index`**: Returns pointed buffer slot index\n",
" - Used for addressing specific staging buffer locations\n",
" - Enables direct slot-based data access\n",
"\n",
"### Circular Buffer State Visualization\n",
"\n",
"```\n",
"Legend:\n",
" W: Currently being written (producer active)\n",
" D: Data ready for consumption \n",
" R: Currently being read (consumer active)\n",
" X: Empty slot available for writing\n",
" \n",
" Advance Direction\n",
" <-------------------\n",
"\n",
" Producer Consumer\n",
" | ^\n",
" V |\n",
" +-----------------+\n",
" --|X|X|W|D|D|D|D|R|X|<-.\n",
" / +-----------------+ \\\n",
" | |\n",
" `------------------------' \n",
"```\n",
"\n",
"**Key Advantages:**\n",
"- **Increased Throughput**: Producer can stay ahead of consumer by multiple stages\n",
"- **Latency Hiding**: Consumer processing latency is hidden by buffered data\n",
"- **Better Resource Utilization**: Both warps can maintain high activity levels\n",
"- **Scalable Design**: Buffer depth can be tuned based on workload characteristics\n",
"\n",
"The following implementation demonstrates efficient multi-stage pipeline communication with proper circular buffer management:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def async_pipeline_staged_kernel(\n",
" SharedStorage: cutlass.Constexpr, res: cute.Tensor, staging: cute.Tensor\n",
"):\n",
" stages = cute.size(staging)\n",
"\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" # Warp 0\n",
" producer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
" # Warp 1\n",
" consumer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
"\n",
" pipeline = cutlass.pipeline.PipelineAsync.create(\n",
" num_stages=stages,\n",
" producer_group=producer_group,\n",
" consumer_group=consumer_group,\n",
" barrier_storage=storage.tma_mbar_ptr.data_ptr(),\n",
" )\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(staging.layout)\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" producer, consumer = pipeline.make_participants()\n",
"\n",
" # Producer warp\n",
" if warp_idx == 0:\n",
" for i in cutlass.range(cute.size(res)):\n",
" handle = producer.acquire_and_advance()\n",
" staging_smem[handle.index] = 1.0 * i\n",
" handle.commit() # or producer.commit(handle)\n",
"\n",
" # prevents CTA0 from retiring until it receives all expected arrives.\n",
" producer.tail()\n",
"\n",
" # Consumer warp\n",
" if warp_idx == 1:\n",
" for i in cutlass.range(cute.size(res)):\n",
" handle = consumer.wait_and_advance()\n",
" res[i] = staging_smem[handle.index]\n",
" handle.release() # or consumer.release(handle)\n",
"\n",
" tidx, _, _ = cute.arch.thread_idx()\n",
" if tidx == 0:\n",
" staging.store(staging_smem.load())\n",
"\n",
"\n",
"@cute.jit\n",
"def async_pipeline_staged(res: cute.Tensor, staging: cute.Tensor):\n",
" stages = cute.size(staging)\n",
"\n",
" @cute.struct\n",
" class SharedStorage:\n",
" tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, stages * 2]\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, stages], 1024\n",
" ]\n",
"\n",
" async_pipeline_staged_kernel(SharedStorage, res, staging).launch(\n",
" grid=(1, 1, 1), block=(64, 1, 1)\n",
" )\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"staging = torch.zeros((5,), device=\"cuda\")\n",
"async_pipeline_staged(from_dlpack(res), from_dlpack(staging))\n",
"torch.cuda.synchronize()"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0'),\n",
" tensor([5., 6., 7., 3., 4.], device='cuda:0'))"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res, staging"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Try Acquire/Wait\n",
"\n",
"In some circumstances, developers may want to just check status of pipeline state without blocking. This could benefit some cases that we have independent instructions to hide latency of checking pipeline state. We provided `try_aquire` or `try_wait` which are non-blocking APIs. "
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.10"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/benchmark_autotune.ipynb
+460
View File
@@ -0,0 +1,460 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"import torch\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute\n",
"import cutlass.cute.testing as testing\n",
"import cutlass.torch as cutlass_torch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## The Usage of Benchmark and Autotune Utilities in CuTe DSL\n",
"\n",
"CuTe DSL provides autotune and benchmark utilities to help users evaluate and optimize kernel performance. This notebook demonstrates how to use these tools.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"### Autotune\n",
"\n",
"We provides two kinds of autotune utilities for users: `autotune.jit` decorator and the `tune` function. The former is used as a decorator used on top of `@cute.jit` while the latter is used as an individual function.\n",
"\n",
"#### @autotune.jit\n",
"\n",
"We take the `elementwise_add_kernel` as an example. After writing the jit host function and kernel, we could add the `@autotune_jit` decorator on top of the jit host function to enable autotune. \n",
"```python\n",
"@testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"],\n",
" warmup_iterations=100,\n",
" iterations=100,\n",
")\n",
"```\n",
"\n",
"The `autotune_jit` decorator provides several parameters to control the autotuning process:\n",
"\n",
"- params_dict: A dictionary containing the parameters to be tuned and their possible values\n",
"- update_on_change: A list of argument names that trigger re-tuning when their values change\n",
"- warmup_iterations: Number of warmup iterations before timing\n",
"- iterations: Number of iterations for timing each parameter combination\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"@cute.kernel\n",
"def elementwise_add_kernel(\n",
" gA: cute.Tensor,\n",
" gB: cute.Tensor,\n",
" gC: cute.Tensor,\n",
" cC: cute.Tensor, # coordinate tensor\n",
" shape: cute.Shape,\n",
" thr_layout: cute.Layout,\n",
" val_layout: cute.Layout,\n",
"):\n",
" tidx, _, _ = cute.arch.thread_idx()\n",
" bidx, _, _ = cute.arch.block_idx()\n",
"\n",
" # slice for CTAs\n",
" # logical id -> address\n",
" blk_coord = ((None, None), bidx)\n",
" blkA = gA[blk_coord] # (TileM,TileN)\n",
" blkB = gB[blk_coord] # (TileM,TileN)\n",
" blkC = gC[blk_coord] # (TileM,TileN)\n",
" blkCrd = cC[blk_coord] # (TileM, TileN)\n",
"\n",
" # # declare the atoms which will be used later for memory copy\n",
" copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)\n",
" copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type)\n",
"\n",
" tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)\n",
" tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)\n",
" tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout)\n",
"\n",
" thr_copy_A = tiled_copy_A.get_slice(tidx)\n",
" thr_copy_B = tiled_copy_B.get_slice(tidx)\n",
" thr_copy_C = tiled_copy_C.get_slice(tidx)\n",
"\n",
" thrA = thr_copy_A.partition_S(blkA)\n",
" thrB = thr_copy_B.partition_S(blkB)\n",
" thrC = thr_copy_C.partition_S(blkC)\n",
"\n",
" # allocate fragments for gmem->rmem\n",
" frgA = cute.make_fragment_like(thrA)\n",
" frgB = cute.make_fragment_like(thrB)\n",
" frgC = cute.make_fragment_like(thrC)\n",
"\n",
" thrCrd = thr_copy_C.partition_S(blkCrd)\n",
" frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)\n",
"\n",
" for i in range(0, cute.size(frgPred), 1):\n",
" val = cute.elem_less(thrCrd[i], shape)\n",
" frgPred[i] = val\n",
"\n",
" ##########################################################\n",
" # Move data to reg address space\n",
" ##########################################################\n",
"\n",
" cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)\n",
" cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)\n",
"\n",
" # Load data before use. The compiler will optimize the copy and load\n",
" # operations to convert some memory ld/st into register uses.\n",
" result = frgA.load() + frgB.load()\n",
"\n",
" # Save the results back to registers. Here we reuse b's registers.\n",
" frgC.store(result)\n",
"\n",
" # Copy the results back to c\n",
" cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)\n",
"\n",
"\n",
"@testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"],\n",
" warmup_iterations=100,\n",
" iterations=100,\n",
")\n",
"@cute.jit\n",
"def elementwise_add_autotune(mA, mB, mC, M, N, copy_bits: cutlass.Constexpr = 128):\n",
" dtype = mA.element_type\n",
" vector_size = copy_bits // dtype.width\n",
"\n",
" thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))\n",
" val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))\n",
" tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)\n",
"\n",
" gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" idC = cute.make_identity_tensor(mC.shape)\n",
" cC = cute.zipped_divide(idC, tiler=tiler_mn)\n",
"\n",
" elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(\n",
" grid=[cute.size(gC, mode=[1]), 1, 1],\n",
" block=[cute.size(tv_layout, mode=[0]), 1, 1],\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When we run the jit funciton `elementwise_add_autotune`, the CuTe DSL will help us tune the kernels by looping the specified configs and run the kernel with the best config.\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"\n",
"M, N = 1024, 1024\n",
"dtype = cutlass.Float32\n",
"skip_ref_check = False\n",
"\n",
"print(f\"\\nRunning Elementwise Add test with:\")\n",
"print(f\"Tensor dimensions: [{M}, {N}]\")\n",
"print(f\"Input and Output Data type: {dtype}\")\n",
"\n",
"torch_dtype = cutlass_torch.dtype(dtype)\n",
"\n",
"a = torch.randn(M, N, device=torch.device(\"cuda\"), dtype=torch_dtype)\n",
"b = torch.randn(M, N, device=torch.device(\"cuda\"), dtype=torch_dtype)\n",
"\n",
"c = torch.zeros_like(a)\n",
"\n",
"print(f\"Input tensor shapes:\")\n",
"print(f\"a: {a.shape}, dtype: {a.dtype}\")\n",
"print(f\"b: {b.shape}, dtype: {b.dtype}\")\n",
"print(f\"c: {c.shape}, dtype: {c.dtype}\\n\")\n",
"\n",
"elementwise_add_autotune(a, b, c, M, N)\n",
"\n",
"if not skip_ref_check:\n",
" print(\"Verifying results for autotuned function ...\")\n",
" torch.testing.assert_close(a + b, c)\n",
" print(\"Results verified successfully!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The output is as follows:\n",
"\n",
"```\n",
"Running Elementwise Add test with:\n",
"Tensor dimensions: [1024, 1024]\n",
"Input and Output Data type: Float32\n",
"Input tensor shapes:\n",
"a: torch.Size([1024, 1024]), dtype: torch.float32\n",
"b: torch.Size([1024, 1024]), dtype: torch.float32\n",
"c: torch.Size([1024, 1024]), dtype: torch.float32\n",
"Verifying results for autotuned function ...\n",
"Results verified successfully!\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"To monitor the autotuning process in detail, you can enable logging by setting the environment variable `CUTE_DSL_LOG_AUTOTUNE`. \n",
"```shell\n",
"export CUTE_DSL_LOG_AUTOTUNE=1\n",
"```\n",
"This will display comprehensive information including:\n",
"- Each configuration being evaluated and its corresponding execution time\n",
"- The optimal configuration that was selected\n",
"- Total time spent on tuning\n",
"- Cache hit/miss statistics\n",
"\n",
"\n",
"Below is a sample output showing the autotuning process with different configurations:\n",
"```python\n",
"2025-07-23 06:17:03,978 - cutlass.cute.testing_Autotune - INFO - Tuning configuration: {'copy_bits': 64}\n",
"2025-07-23 06:17:04,519 - cutlass.cute.testing_Autotune - INFO - Execution time: 0.010857919985428453 us\n",
"2025-07-23 06:17:04,519 - cutlass.cute.testing_Autotune - INFO - Tuning configuration: {'copy_bits': 128}\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Execution time: 0.011117440033704042 us\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Best configuration: {'copy_bits': 64}, execution time: 0.010857919985428453 us\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Total tuning time: 0.7053244113922119 s\n",
"...\n",
"2025-07-23 06:17:04,700 - cutlass.cute.testing_Autotune - INFO - Using cached best configuration: {'copy_bits': 64}\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### tune\n",
"\n",
"We also provide a `tune` funtion. The interface of the `tune` function is as follows:\n",
"\n",
"```python\n",
"def tune(\n",
" func: Callable[[Any], Callable[[], Any]],\n",
" params_dict: Dict[str, List[Any]] = None,\n",
" kernel_arguments: JitArguments = JitArguments(),\n",
" warmup_iterations=10,\n",
" iterations=100,\n",
" stream: Optional[cuda_driver.CUstream] = None,\n",
") -> Dict[str, Any]:\n",
"```\n",
"\n",
"The `tune` function takes the following parameters:\n",
"\n",
"- func: A callable that takes configuration parameters and returns a kernel function\n",
"- params_dict: Dictionary mapping parameter names to lists of possible values to tune\n",
"- kernel_arguments: Arguments to pass to the kernel for tuning\n",
"- warmup_iterations: Number of warmup iterations before timing (default: 10)\n",
"- iterations: Number of timing iterations per configuration (default: 100)\n",
"- stream: Optional CUDA stream to use for execution. defaults to default CUDA stream. The stream parameter must match the stream passed to the kernel, mismatched streams will result in an error.\n",
"\n",
"It returns a dictionary containing the best kernel configuration found.\n",
"\n",
"\n",
"Here is an example to use the `tune` function:\n",
"\n",
"1. First remove the `@testing.autotune_jit` decorator from the `elementwise_add_autotune` function:\n",
" ```python\n",
" @testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"], \n",
" warmup_iterations=100,\n",
" iterations=100,\n",
" )\n",
" ```\n",
"\n",
" 2. Define a `tune_func` that:\n",
" - Takes input tensors (a, b, c), dimensions (M, N) and tuning parameter copy_bits\n",
" - Compiles the `elementwise_add_autotune` function using `cute.compile()`\n",
" - Returns a lambda function that executes the compiled kernel\n",
"\n",
" 3. Pass `tune_func` to `testing.tune` function along with:\n",
" - Parameter space to explore (copy_bits values)\n",
" - Kernel arguments wrapped in JitArguments\n",
" - The `tune` function will find optimal parameters automatically\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"def tune_func(a, b, c, M, N, copy_bits=128):\n",
" compiled_func = cute.compile(elementwise_add_autotune, a, b, c, M, N, copy_bits=128)\n",
" return lambda: compiled_func(a, b, c, M, N)\n",
"\n",
"params = testing.tune(\n",
" tune_func,\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" kernel_arguments=testing.JitArguments(a, b, c, M, N),\n",
")\n",
"print(f\"The best kernel configs found: {params}\")\n",
"\n",
"# run the kernel with the best config\n",
"compiled_func = cute.compile(elementwise_add_autotune, a, b, c, M, N, **params)\n",
"compiled_func(a, b, c, M, N)\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### benchmark\n",
"\n",
"In CuTe DSL, the benchmark utility can be used to measure kernel execution time. The interface of benchmark routine is as follows:\n",
"\n",
"```python\n",
"def benchmark(\n",
" callable: Callable,\n",
" *,\n",
" warmup_iterations: int = 10,\n",
" iterations: int = 100,\n",
" stream: Optional[cuda_driver.CUstream] = None,\n",
" kernel_arguments: Optional[JitArguments] = None,\n",
" workspace_generator: Optional[Callable[[], JitArguments]] = None,\n",
" workspace_count: int = 1,\n",
" use_cuda_graphs: bool = False,\n",
") -> float:\n",
"```\n",
"\n",
"The benchmark utility exposes several key configuration parameters to control profiling behavior:\n",
"\n",
"- callable: The function to be benchmarked\n",
"- warmup_iterations: Controls the number of initial warmup iterations before measurement begins (default: 10)\n",
"- iterations: Specifies how many iterations to profile for performance measurement (default: 100)\n",
"- stream: Designates which CUDA stream to execute the kernel on (default: default stream) \n",
"- use_cuda_graphs: Whether enables CUDA graph for the callable function to minimize kernel launch overhead (default: False)\n",
"- workspace_generator: Provides a function that generates fresh kernel arguments each iteration to avoid caching effects\n",
"- workspace_count: Determines how many different workspaces to cycle through during profiling (default: 1)\n",
"\n",
"When benchmarking, there are several key parameters that can be configured:\n",
"\n",
"1. Core parameters:\n",
" - The function to profile (callable)\n",
" - Number of warmup iterations before measurement\n",
" - Number of profiling iterations for measurement\n",
"\n",
"2. Stream configuration:\n",
" - For kernels running in non-default streams, the stream must be specified\n",
" - The stream parameter must match the stream passed to the kernel, mismatched streams will result in an error\n",
"\n",
"3. Cache effects mitigation:\n",
" - To prevent L2 cache effects from skewing results, multiple workspaces can be cycled through\n",
" - This is controlled via workspace_count and workspace_generator parameters\n",
" - Each workspace provides fresh kernel arguments\n",
"\n",
"4. CUDA Graph support:\n",
" - Enables measuring kernel execution time without host overhead\n",
" - Requires the callable to be decorated with @cute.jit\n",
" - Must use a non-default CUDA stream when using graphs\n",
"\n",
"This function will return the execution time of the callable in microseconds. As GPU frequency can vary dynamically, we could fix the SM and memory frequencies to get more stable and reproducible benchmark results. This can be done by setting the GPU clocks using nvidia-smi before running the benchmark. In the next, let's use the benchmark function to get the execution time of the above elementwise_add kernel."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"def generate_kernel_arguments():\n",
" a = torch.randn(\n",
" M, N, device=torch.device(\"cuda\"), dtype=torch_dtype\n",
" )\n",
" b = torch.randn(\n",
" M, N, device=torch.device(\"cuda\"), dtype=torch_dtype\n",
" )\n",
"\n",
" c = torch.zeros_like(a)\n",
"\n",
" return testing.JitArguments(a, b, c, M, N)\n",
"\n",
"avg_time_us = testing.benchmark(\n",
" elementwise_add_autotune,\n",
" workspace_generator=generate_kernel_arguments,\n",
" workspace_count=10,\n",
" warmup_iterations=10,\n",
" iterations=100,\n",
")\n",
"\n",
"# Print execution results\n",
"print(\n",
" f\"Kernel execution time for cute.jit kernel with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms\"\n",
")\n",
"print(\n",
" f\"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After running the code, we will get output similar to the following:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```\n",
"Kernel execution time for cute.jit kernel with M=1024, N=1024: 0.0403 ms\n",
"Achieved memory throughput for M=1024, N=1024: 312.37 GB/s\n",
"```"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
examples/python/CuTeDSL/cute/notebooks/composed_layout.ipynb
+225
View File
@@ -0,0 +1,225 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0c7cf795",
"metadata": {},
"source": [
"# Composed Layout in CuTe\n",
"\n",
"A **Composed Layout** is a powerful abstraction in CuTe that enables complex data transformations through \n",
"the composition of layouts and transformations. It provides a flexible way to manipulate memory layouts \n",
"and coordinate systems.\n",
"\n",
"## Components\n",
"\n",
"A Composed Layout consists of three key components:\n",
"\n",
"1. **Inner Layout/Transformation** (`inner`):\n",
" - Can be a layout, swizzle, or custom transformation function\n",
" - Applies the final transformation to the coordinates\n",
" - Supports arbitrary coordinate manipulations\n",
"\n",
"2. **Offset** (`offset`):\n",
" - Typically represented as an integer tuple\n",
" - Adds a constant displacement to coordinates\n",
" - Enables fine-grained control over data positioning\n",
"\n",
"3. **Outer Layout** (`outer`):\n",
" - The layout visible to the user\n",
" - Defines the initial coordinate transformation\n",
" - Determines the shape and organization of the data structure\n",
"\n",
"## Mathematical Representation\n",
"\n",
"The mathematical composition of these components is defined as:\n",
"\n",
"$\n",
"R(c) := (inner \\circ offset \\circ outer)(c) := inner(offset + outer(c))\n",
"$\n",
"\n",
"Where:\n",
"- $c$ represents the input coordinates\n",
"- $\\circ$ denotes function composition\n",
"- The transformation is applied from right to left\n",
"\n",
"## Usage in Python\n",
"\n",
"To create a Composed Layout in Python, use the `make_composed_layout` function:\n",
"\n",
"```python\n",
"layout = cute.make_composed_layout(inner, offset, outer)\n",
"```\n",
"\n",
"## Key Benefits\n",
"\n",
"1. **Flexibility**: Supports complex transformations that direct composition cannot handle\n",
"2. **Modularity**: Separates different aspects of the transformation\n",
"3. **Performance**: Enables optimized memory access patterns for GPU computations\n",
"4. **Compatibility**: Works with various types of transformations and layouts"
]
},
{
"cell_type": "markdown",
"id": "24448f7d",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": [
"## Custom Transformation Example\n",
"\n",
"This example demonstrates how to create a Composed Layout with a custom transformation function. We'll create a simple transformation that:\n",
"\n",
"1. Takes a 2D coordinate input `(x, y)`\n",
"2. Increments the y-coordinate by 1\n",
"3. Combines this with an offset and identity layout\n",
"\n",
"The example shows how to:\n",
"- Define a custom transformation function\n",
"- Create a composed layout with the transformation\n",
"- Apply the layout to coordinates\n",
"- Print the results for verification"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "184f30e6",
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack, make_ptr\n",
"\n",
"\n",
"@cute.jit\n",
"def customized_layout():\n",
" def inner(c):\n",
" x, y = c\n",
" return x, y + 1\n",
"\n",
" layout = cute.make_composed_layout(\n",
" inner, (1, 0), cute.make_identity_layout(shape=(8, 4))\n",
" )\n",
" print(layout)\n",
" cute.printf(layout(0))\n",
"\n",
"\n",
"customized_layout()"
]
},
{
"cell_type": "markdown",
"id": "c897187f",
"metadata": {},
"source": [
"## Gather/Scatter Operations with Composed Layout\n",
"\n",
"Gather and Scatter operations are fundamental data access patterns in parallel computing and GPU programming. In CuTe, we can implement these operations elegantly using Composed Layout.\n",
"\n",
"### Gather Operation\n",
"A gather operation collects elements from a source array using an index array (also called an indirection array). It's defined as:\n",
"```python\n",
"output[i] = source[index[i]]\n",
"```\n",
"\n",
"#### Components in CuTe Implementation:\n",
"1. **Offset Tensor**: Contains the indices for gathering (`offset_tensor`)\n",
"2. **Data Pointer**: Points to the source data array (`data_ptr`)\n",
"3. **Shape**: Defines the shape of logic tensor viewed by user (`shape`)\n",
"\n",
"### How it Works\n",
"1. The inner transformation function reads from the offset tensor:\n",
" ```python\n",
" def inner(c):\n",
" return offset_tensor[c] # Returns the gather index\n",
" ```\n",
"2. The composed layout maps input coordinates through the offset tensor:\n",
" ```python\n",
" gather_layout = cute.make_composed_layout(inner, 0, cute.make_layout(shape))\n",
" ```\n",
"3. This creates an indirect access pattern where:\n",
" - Input coordinate `i` → `offset_tensor[i]` → `data_ptr[offset_tensor[i]]`\n",
"\n",
"4. notably, layout operations like slice, partition can still be applied on `outer` layout\n",
"\n",
"### Use Cases\n",
"- **Sparse Operations**: Accessing non-contiguous memory efficiently\n",
"- **Graph Processing**: Following edge connections in graph algorithms\n",
"- **Feature Embedding**: Looking up embeddings for discrete tokens\n",
"- **Irregular Data Access**: Any pattern requiring indirect memory access\n",
"\n",
"### Example Output Interpretation\n",
"The example code prints pairs of numbers `i -> j` where:\n",
"- `i` is the output index\n",
"- `j` is the gathered source index from `offset_tensor`\n",
"\n",
"This demonstrates how the composed layout transforms coordinates for indirect memory access.\n",
"\n",
"Note: Scatter operations (writing to indirect locations) can be implemented similarly by reversing the data flow direction.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d68f9476",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"\n",
"@cute.jit\n",
"def gather_tensor(\n",
" offset_tensor: cute.Tensor, data_ptr: cute.Pointer, shape: cute.Shape\n",
"):\n",
" def inner(c):\n",
" return offset_tensor[c]\n",
"\n",
" gather_layout = cute.make_composed_layout(inner, 0, cute.make_layout(shape))\n",
" for i in cutlass.range_constexpr(cute.size(shape)):\n",
" cute.printf(\"%d -> %d\", i, gather_layout(i))\n",
"\n",
" # TODO: support in future\n",
" # gather_tensor = cute.make_tensor(data_ptr, gather_layout)\n",
" # cute.printf(gather_tensor[0])\n",
"\n",
"\n",
"shape = (16,)\n",
"offset_tensor = torch.randint(0, 256, shape, dtype=torch.int32)\n",
"data_tensor = torch.arange(0, 256, dtype=torch.int32)\n",
"\n",
"\n",
"gather_tensor(\n",
" from_dlpack(offset_tensor),\n",
" make_ptr(cutlass.Int32, data_tensor.data_ptr(), cute.AddressSpace.generic),\n",
" shape,\n",
")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv3_12",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
examples/python/CuTeDSL/cute/notebooks/cuda_graphs.ipynb
+651
View File
@@ -0,0 +1,651 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0e95f0df-4d1a-4e2e-92ff-90539bb4c517",
"metadata": {},
"source": [
"# Example 06: CUDA Graphs\n",
"\n",
"In this example we demonstrate how to use CUDA graphs through PyTorch with CuTe DSL.\n",
"The process of interacting with PyTorch's CUDA graph implementation requires exposing PyTorch's CUDA streams to CUTLASS.\n",
"\n",
"To use CUDA graphs with Blackwell requires a version of PyTorch that supports Blackwell.\n",
"This can be obtained through:\n",
"- The [PyTorch NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)\n",
"- [PyTorch 2.7 with CUDA 12.8 or later](https://pytorch.org/) (e.g., `pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128`)\n",
"- Building PyTorch directly with your version of CUDA."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "46b8fb6f-9ac5-4a3d-b765-b6476f182bf7",
"metadata": {},
"outputs": [],
"source": [
"# import torch for CUDA graphs\n",
"import torch\n",
"import cutlass.cute as cute\n",
"\n",
"# import CUstream type from the cuda driver bindings\n",
"from cuda.bindings.driver import CUstream\n",
"\n",
"# import the current_stream function from torch\n",
"from torch.cuda import current_stream"
]
},
{
"cell_type": "markdown",
"id": "bcf5e06e-1f5b-4d72-ad73-9b36efb78ca0",
"metadata": {},
"source": [
"## Kernel Creation\n",
"\n",
"We create a kernel which prints \"Hello world\" as well as a host function to launch the kernel.\n",
"We then compile the kernel for use in our graph, by passing in a default stream.\n",
"\n",
"Kernel compilation before graph capture is required since CUDA graphs cannot JIT compile kernels during graph execution."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "0c2a6ca8-98d7-4837-b91f-af769ca8fcd8",
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def hello_world_kernel():\n",
" \"\"\"\n",
" A kernel that prints hello world\n",
" \"\"\"\n",
" cute.printf(\"Hello world\")\n",
"\n",
"\n",
"@cute.jit\n",
"def hello_world(stream: CUstream):\n",
" \"\"\"\n",
" Host function that launches our (1,1,1), (1,1,1) grid in stream\n",
" \"\"\"\n",
" hello_world_kernel().launch(grid=[1, 1, 1], block=[1, 1, 1], stream=stream)\n",
"\n",
"\n",
"# Grab a stream from PyTorch, this will also initialize our context\n",
"# so we can omit cutlass.cuda.initialize_cuda_context()\n",
"stream = current_stream()\n",
"hello_world_compiled = cute.compile(hello_world, CUstream(stream.cuda_stream))"
]
},
{
"cell_type": "markdown",
"id": "ecc850af-09f8-4a29-9c93-ff31fbb9326f",
"metadata": {},
"source": [
"## Creating and replaying a CUDA Graph\n",
"\n",
"We create a stream through torch as well as a graph.\n",
"When we create the graph we can pass the stream we want to capture to torch. We similarly run the compiled kernel with the stream passed as a CUstream.\n",
"\n",
"Finally we can replay our graph and synchronize."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "f673e5ae-42bb-44d0-b652-3280606181c4",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Hello world\n",
"Hello world\n"
]
}
],
"source": [
"# Create a CUDA Graph\n",
"g = torch.cuda.CUDAGraph()\n",
"# Capture our graph\n",
"with torch.cuda.graph(g):\n",
" # Turn our torch Stream into a cuStream stream.\n",
" # This is done by getting the underlying CUstream with .cuda_stream\n",
" graph_stream = CUstream(current_stream().cuda_stream)\n",
" # Run 2 iterations of our compiled kernel\n",
" for _ in range(2):\n",
" # Run our kernel in the stream\n",
" hello_world_compiled(graph_stream)\n",
"\n",
"# Replay our graph\n",
"g.replay()\n",
"# Synchronize all streams (equivalent to cudaDeviceSynchronize() in C++)\n",
"torch.cuda.synchronize()"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "db76d9c3-7617-4bf2-b326-11982e6803bf",
"metadata": {},
"source": [
"Our run results in the following execution when viewed in NSight Systems:\n",
"\n",
"![Image of two hello world kernels run back to back in a CUDA graph](images/cuda_graphs_image.png)\n",
"\n",
"We can observe the launch of the two kernels followed by a `cudaDeviceSynchronize()`.\n",
"\n",
"Now we can confirm that this minimizes some launch overhead:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "3ebe15bf-dc97-42e9-913c-224ecfb472e8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n",
"Hello world\n"
]
}
],
"source": [
"# Get our CUDA stream from PyTorch\n",
"stream = CUstream(current_stream().cuda_stream)\n",
"\n",
"# Create a larger CUDA Graph of 100 iterations\n",
"g = torch.cuda.CUDAGraph()\n",
"# Capture our graph\n",
"with torch.cuda.graph(g):\n",
" # Turn our torch Stream into a cuStream stream.\n",
" # This is done by getting the underlying CUstream with .cuda_stream\n",
" graph_stream = CUstream(current_stream().cuda_stream)\n",
" # Run 2 iterations of our compiled kernel\n",
" for _ in range(100):\n",
" # Run our kernel in the stream\n",
" hello_world_compiled(graph_stream)\n",
"\n",
"# Create CUDA events for measuring performance\n",
"start = torch.cuda.Event(enable_timing=True)\n",
"end = torch.cuda.Event(enable_timing=True)\n",
"\n",
"# Run our kernel to warm up the GPU\n",
"for _ in range(100):\n",
" hello_world_compiled(stream)\n",
"\n",
"# Record our start time\n",
"start.record()\n",
"# Run 100 kernels\n",
"for _ in range(100):\n",
" hello_world_compiled(stream)\n",
"# Record our end time\n",
"end.record()\n",
"# Synchronize (cudaDeviceSynchronize())\n",
"torch.cuda.synchronize()\n",
"\n",
"# Calculate the time spent when launching kernels in a stream\n",
"# Results are in ms\n",
"stream_time = start.elapsed_time(end)\n",
"\n",
"# Warmup our GPU again\n",
"g.replay()\n",
"# Record our start time\n",
"start.record()\n",
"# Run our graph\n",
"g.replay()\n",
"# Record our end time\n",
"end.record()\n",
"# Synchronize (cudaDeviceSynchronize())\n",
"torch.cuda.synchronize()\n",
"\n",
"# Calculate the time spent when launching kernels in a graph\n",
"# units are ms\n",
"graph_time = start.elapsed_time(end)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "12b8151a-46b3-4c99-9945-301f6b628131",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"8.94% speedup when using CUDA graphs for this kernel!\n"
]
}
],
"source": [
"# Print out speedup when using CUDA graphs\n",
"percent_speedup = (stream_time - graph_time) / graph_time\n",
"print(f\"{percent_speedup * 100.0:.2f}% speedup when using CUDA graphs for this kernel!\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
examples/python/CuTeDSL/cute/notebooks/cute_layout_algebra.ipynb
+1030
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/cute/notebooks/data_types.ipynb
+270
View File
@@ -0,0 +1,270 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Understanding data structure in CuTe DSL\n",
"\n",
"In most cases, data structures in CuTe DSL work the same as Python data structures with the notable difference that Python data structures in most cases are considered as static data which are interpreted by the DSL compiler embedded inside Python interpreter.\n",
"\n",
"To differentiate between compile-time and runtime values, CuTe DSL introduces primitive types that \n",
"represent dynamic values in JIT-compiled code.\n",
"\n",
"CuTe DSL provides a comprehensive set of primitive numeric types for representing dynamic values at \n",
"runtime. These types are formally defined within the CuTe DSL typing system:\n",
"\n",
"### Integer Types\n",
"- `Int8` - 8-bit signed integer\n",
"- `Int16` - 16-bit signed integer \n",
"- `Int32` - 32-bit signed integer\n",
"- `Int64` - 64-bit signed integer\n",
"- `Int128` - 128-bit signed integer\n",
"- `Uint8` - 8-bit unsigned integer\n",
"- `Uint16` - 16-bit unsigned integer\n",
"- `Uint32` - 32-bit unsigned integer\n",
"- `Uint64` - 64-bit unsigned integer\n",
"- `Uint128` - 128-bit unsigned integer\n",
"\n",
"### Floating Point Types\n",
"- `Float16` - 16-bit floating point\n",
"- `Float32` - 32-bit floating point \n",
"- `Float64` - 64-bit floating point\n",
"- `BFloat16` - Brain Floating Point format (16-bit)\n",
"- `TFloat32` - Tensor Float32 format (reduced precision format used in tensor operations)\n",
"- `Float8E4M3` - 8-bit floating point with 4-bit exponent and 3-bit mantissa\n",
"- `Float8E5M2` - 8-bit floating point with 5-bit exponent and 2-bit mantissa\n",
"\n",
"These specialized types are designed to represent dynamic values in CuTe DSL code that will be \n",
"evaluated at runtime, in contrast to Python's built-in numeric types which are evaluated during \n",
"compilation.\n",
"\n",
"### Example usage:\n",
"\n",
"```python\n",
"x = cutlass.Int32(5) # Creates a 32-bit integer\n",
"y = cutlass.Float32(3.14) # Creates a 32-bit float\n",
"\n",
"@cute.jit\n",
"def foo(a: cutlass.Int32): # annotate `a` as 32-bit integer passed to jit function via ABI\n",
" ...\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"a(static) = ?\n",
"b(static) = ?\n",
"a(dynamic) = 3.140000\n",
"b(dynamic) = 5\n"
]
}
],
"source": [
"@cute.jit\n",
"def bar():\n",
" a = cutlass.Float32(3.14)\n",
" print(\"a(static) =\", a) # prints `a(static) = ?`\n",
" cute.printf(\"a(dynamic) = {}\", a) # prints `a(dynamic) = 3.140000`\n",
"\n",
" b = cutlass.Int32(5)\n",
" print(\"b(static) =\", b) # prints `b(static) = 5`\n",
" cute.printf(\"b(dynamic) = {}\", b) # prints `b(dynamic) = 5`\n",
"\n",
"\n",
"bar()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Type Conversion API\n",
"\n",
"CUTLASS numeric types provide type conversion through the `to()` method available on all Numeric types. This allows you to convert between different numeric data types at runtime.\n",
"\n",
"Syntax:\n",
"\n",
"```python\n",
"new_value = value.to(target_type)\n",
"```\n",
"\n",
"The `to()` method supports conversion between:\n",
"- Integer types (Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64)\n",
"- Floating point types (Float16, Float32, Float64, BFloat16)\n",
"- Mixed integer/floating point conversions\n",
"\n",
"Note that when converting from floating point to integer types, the decimal portion is truncated. When converting between types with different ranges, values may be clamped or lose precision if they exceed the target type's representable range."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Int32(42) => Float32(42.000000)\n",
"Float32(3.140000) => Int32(3)\n",
"Int32(127) => Int8(127)\n",
"Int32(300) => Int8(44) (truncated due to range limitation)\n"
]
}
],
"source": [
"@cute.jit\n",
"def type_conversion():\n",
" # Convert from Int32 to Float32\n",
" x = cutlass.Int32(42)\n",
" y = x.to(cutlass.Float32)\n",
" cute.printf(\"Int32({}) => Float32({})\", x, y)\n",
"\n",
" # Convert from Float32 to Int32\n",
" a = cutlass.Float32(3.14)\n",
" b = a.to(cutlass.Int32)\n",
" cute.printf(\"Float32({}) => Int32({})\", a, b)\n",
"\n",
" # Convert from Int32 to Int8\n",
" c = cutlass.Int32(127)\n",
" d = c.to(cutlass.Int8)\n",
" cute.printf(\"Int32({}) => Int8({})\", c, d)\n",
"\n",
" # Convert from Int32 to Int8 with value exceeding Int8 range\n",
" e = cutlass.Int32(300)\n",
" f = e.to(cutlass.Int8)\n",
" cute.printf(\"Int32({}) => Int8({}) (truncated due to range limitation)\", e, f)\n",
"\n",
"\n",
"type_conversion()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Operator Overloading\n",
"\n",
"CUTLASS numeric types support Python's built-in operators, allowing you to write natural mathematical expressions. The operators work with both CUTLASS numeric types and Python native numeric types.\n",
"\n",
"Supported operators include:\n",
"- Arithmetic: `+`, `-`, `*`, `/`, `//`, `%`, `**`\n",
"- Comparison: `<`, `<=`, `==`, `!=`, `>=`, `>`\n",
"- Bitwise: `&`, `|`, `^`, `<<`, `>>`\n",
"- Unary: `-` (negation), `~` (bitwise NOT)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"a: Int32(10), b: Int32(3)\n",
"x: Float32(5.500000)\n",
"\n",
"a + b = 13\n",
"x * 2 = 11.000000\n",
"a + x = 15.500000 (Int32 + Float32 promotes to Float32)\n",
"a / b = 3.333333\n",
"x / 2.0 = 2.750000\n",
"a > b = 1\n",
"a & b = 2\n",
"-a = -10\n",
"~a = -11\n"
]
}
],
"source": [
"@cute.jit\n",
"def operator_demo():\n",
" # Arithmetic operators\n",
" a = cutlass.Int32(10)\n",
" b = cutlass.Int32(3)\n",
" cute.printf(\"a: Int32({}), b: Int32({})\", a, b)\n",
"\n",
" x = cutlass.Float32(5.5)\n",
" cute.printf(\"x: Float32({})\", x)\n",
"\n",
" cute.printf(\"\")\n",
"\n",
" sum_result = a + b\n",
" cute.printf(\"a + b = {}\", sum_result)\n",
"\n",
" y = x * 2 # Multiplying with Python native type\n",
" cute.printf(\"x * 2 = {}\", y)\n",
"\n",
" # Mixed type arithmetic (Int32 + Float32) that integer is converted into float32\n",
" mixed_result = a + x\n",
" cute.printf(\"a + x = {} (Int32 + Float32 promotes to Float32)\", mixed_result)\n",
"\n",
" # Division with Int32 (note: integer division)\n",
" div_result = a / b\n",
" cute.printf(\"a / b = {}\", div_result)\n",
"\n",
" # Float division\n",
" float_div = x / cutlass.Float32(2.0)\n",
" cute.printf(\"x / 2.0 = {}\", float_div)\n",
"\n",
" # Comparison operators\n",
" is_greater = a > b\n",
" cute.printf(\"a > b = {}\", is_greater)\n",
"\n",
" # Bitwise operators\n",
" bit_and = a & b\n",
" cute.printf(\"a & b = {}\", bit_and)\n",
"\n",
" neg_a = -a\n",
" cute.printf(\"-a = {}\", neg_a)\n",
"\n",
" not_a = ~a\n",
" cute.printf(\"~a = {}\", not_a)\n",
"\n",
"\n",
"operator_demo()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/elementwise_add.ipynb
+1226
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/cute/notebooks/hello_world.ipynb
+181
View File
@@ -0,0 +1,181 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Your First Program with CuTe DSL\n",
"\n",
"## Introduction\n",
"\n",
"Welcome! In this tutorial, we'll write a simple \"Hello World\" program that runs on your GPU using CuTe DSL. This will help you understand the basics of GPU programming with our framework.\n",
"\n",
"### What You'll Learn\n",
"\n",
"- How to write code that runs on both CPU (host) and GPU (device),\n",
"- How to launch a GPU kernel (a function that runs on the GPU),\n",
"- Basic CUDA concepts like threads and thread blocks,\n",
"\n",
"### Step 1: Import Required Libraries\n",
"\n",
"First, let's import the libraries we need:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"### Step 2: Write Our GPU Kernel\n",
"A GPU kernel is a function that runs on the GPU. Here's a simple kernel that prints \"Hello World\".\n",
"Key concepts:\n",
"- `@cute.kernel`: This decorator tells CUTLASS that this function should run on the GPU\n",
"- `cute.arch.thread_idx()`: Gets the ID of the current GPU thread (like a worker's ID number)\n",
"- We only want one thread to print the message (thread 0) to avoid multiple prints"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def kernel():\n",
" # Get the x component of the thread index (y and z components are unused)\n",
" tidx, _, _ = cute.arch.thread_idx()\n",
" # Only the first thread (thread 0) prints the message\n",
" if cutlass.dynamic_expr(tidx == 0):\n",
" cute.printf(\"Hello world\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Step 3: Write Our Host Function\n",
"\n",
"Now we need a function that sets up the GPU and launches our kernel.\n",
"Key concepts:\n",
"- `@cute.jit`: This decorator is for functions that run on the CPU but can launch GPU code\n",
"- We need to initialize CUDA before using the GPU\n",
"- `.launch()` tells CUDA how many blocks, threads, shared memory, etc. to use"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def hello_world():\n",
" # Print hello world from host code\n",
" cute.printf(\"hello world\")\n",
"\n",
" # Launch kernel\n",
" kernel().launch(\n",
" grid=(1, 1, 1), # Single thread block\n",
" block=(32, 1, 1), # One warp (32 threads) per thread block\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Step 4: Run Our Program\n",
"\n",
"There are 2 ways we can run our program:\n",
"\n",
"1. compile and run immediately\n",
"2. separate compilation which allows us to compile the code once and run multiple times\n",
" \n",
"Please note the `Compiling...` for Method 2 prints before the \"Hello world\" of the first kernel. This shows the asynchronous behavior between CPU and GPU prints. "
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Running hello_world()...\n",
"Compiling...\n",
"hello world\n",
"Hello world\n",
"Compiling with PTX/CUBIN dumped...\n",
"Running compiled version...\n",
"hello world\n",
"Hello world\n"
]
}
],
"source": [
"# Initialize CUDA context for launching a kernel with error checking\n",
"# We make context initialization explicit to allow users to control the context creation\n",
"# and avoid potential issues with multiple contexts\n",
"cutlass.cuda.initialize_cuda_context()\n",
"\n",
"# Method 1: Just-In-Time (JIT) compilation - compiles and runs the code immediately\n",
"print(\"Running hello_world()...\")\n",
"hello_world()\n",
"\n",
"# Method 2: Compile first (useful if you want to run the same code multiple times)\n",
"print(\"Compiling...\")\n",
"hello_world_compiled = cute.compile(hello_world)\n",
"\n",
"# Dump PTX/CUBIN files while compiling\n",
"from cutlass.cute import KeepPTX, KeepCUBIN\n",
"\n",
"print(\"Compiling with PTX/CUBIN dumped...\")\n",
"hello_world_compiled_ptx_on = cute.compile[KeepPTX, KeepCUBIN](hello_world)\n",
"\n",
"# Run the pre-compiled version\n",
"print(\"Running compiled version...\")\n",
"hello_world_compiled()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/images/blocked_gemm.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 9.8 KiB

examples/python/CuTeDSL/cute/notebooks/images/cuda_graphs_image.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 8.4 KiB

examples/python/CuTeDSL/cute/notebooks/images/software_pipelining_ab_stages_minus_1.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 22 KiB

examples/python/CuTeDSL/cute/notebooks/images/software_pipelining_ab_stages_minus_2.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 20 KiB

examples/python/CuTeDSL/cute/notebooks/print.ipynb
+500
View File
@@ -0,0 +1,500 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Printing with CuTe DSL\n",
"\n",
"This notebook demonstrates the different ways to print values in CuTe and explains the important distinction between static (compile-time) and dynamic (runtime) values.\n",
"\n",
"## Key Concepts\n",
"- Static values: Known at compile time\n",
"- Dynamic values: Only known at runtime\n",
"- Different printing methods for different scenarios\n",
"- Layout representation in CuTe\n",
"- Tensor visualization and formatting"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Print Example Function\n",
"\n",
"The `print_example` function demonstrates several important concepts:\n",
"\n",
"### 1. Python's `print` vs CuTe's `cute.printf`\n",
"- `print`: Can only show static values at compile time\n",
"- `cute.printf`: Can display both static and dynamic values at runtime\n",
"\n",
"### 2. Value Types\n",
"- `a`: Dynamic `Int32` value (runtime)\n",
"- `b`: Static `Constexpr[int]` value (compile-time)\n",
"\n",
"### 3. Layout Printing\n",
"Shows how layouts are represented differently in static vs dynamic contexts:\n",
"- Static context: Unknown values shown as `?`\n",
"- Dynamic context: Actual values displayed"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def print_example(a: cutlass.Int32, b: cutlass.Constexpr[int]):\n",
" \"\"\"\n",
" Demonstrates different printing methods in CuTe and how they handle static vs dynamic values.\n",
"\n",
" This example shows:\n",
" 1. How Python's `print` function works with static values at compile time but can't show dynamic values\n",
" 2. How `cute.printf` can display both static and dynamic values at runtime\n",
" 3. The difference between types in static vs dynamic contexts\n",
" 4. How layouts are represented in both printing methods\n",
"\n",
" Args:\n",
" a: A dynamic Int32 value that will be determined at runtime\n",
" b: A static (compile-time constant) integer value\n",
" \"\"\"\n",
" # Use Python `print` to print static information\n",
" print(\">>>\", b) # => 2\n",
" # `a` is dynamic value\n",
" print(\">>>\", a) # => ?\n",
"\n",
" # Use `cute.printf` to print dynamic information\n",
" cute.printf(\">?? {}\", a) # => 8\n",
" cute.printf(\">?? {}\", b) # => 2\n",
"\n",
" print(\">>>\", type(a)) # => <class 'cutlass.Int32'>\n",
" print(\">>>\", type(b)) # => <class 'int'>\n",
"\n",
" layout = cute.make_layout((a, b))\n",
" print(\">>>\", layout) # => (?,2):(1,?)\n",
" cute.printf(\">?? {}\", layout) # => (8,2):(1,8)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Compile and Run\n",
"\n",
"**Direct Compilation and Run**\n",
" - `print_example(cutlass.Int32(8), 2)`\n",
" - Compiles and runs in one step will execute both static and dynamic print\n",
" * `>>>` stands for static print\n",
" * `>??` stands for dynamic print"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
">>> 2\n",
">>> ?\n",
">>> Int32\n",
">>> <class 'int'>\n",
">>> (?,2):(1,?)\n",
">?? 8\n",
">?? 2\n",
">?? (8,2):(1,8)\n"
]
}
],
"source": [
"print_example(cutlass.Int32(8), 2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Compile Function\n",
"\n",
"When compiles the function with `cute.compile(print_example, cutlass.Int32(8), 2)`, Python interpreter \n",
"traces code and only evaluate static expression and print static information."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
">>> 2\n",
">>> ?\n",
">>> Int32\n",
">>> <class 'int'>\n",
">>> (?,2):(1,?)\n"
]
}
],
"source": [
"print_example_compiled = cute.compile(print_example, cutlass.Int32(8), 2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Call compiled function\n",
"\n",
"Only print out runtime information"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
">?? 8\n",
">?? 2\n",
">?? (8,2):(1,8)\n"
]
}
],
"source": [
"print_example_compiled(cutlass.Int32(8))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Format String Example\n",
"\n",
"The `format_string_example` function shows an important limitation:\n",
"- F-strings in CuTe are evaluated at compile time\n",
"- This means dynamic values won't show their runtime values in f-strings\n",
"- Use `cute.printf` when you need to see runtime values"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Direct run output:\n",
"a: ?, b: 2\n",
"layout: (?,2):(1,?)\n"
]
}
],
"source": [
"@cute.jit\n",
"def format_string_example(a: cutlass.Int32, b: cutlass.Constexpr[int]):\n",
" \"\"\"\n",
" Format string is evaluated at compile time.\n",
" \"\"\"\n",
" print(f\"a: {a}, b: {b}\")\n",
"\n",
" layout = cute.make_layout((a, b))\n",
" print(f\"layout: {layout}\")\n",
"\n",
"\n",
"print(\"Direct run output:\")\n",
"format_string_example(cutlass.Int32(8), 2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Printing Tensor Examples\n",
"\n",
"CuTe provides specialized functionality for printing tensors through the `print_tensor` operation. The `cute.print_tensor` takes the following parameter:\n",
"- `Tensor` (required): A CuTe tensor object that you want to print. The tensor must support load and store operations\n",
"- `verbose` (optional, default=False): A boolean flag that controls the level of detail in the output. When set to True, it will print indices details for each element in the tensor.\n",
"\n",
"Below example code shows the difference between verbose ON and OFF, and how to print a sub range of the given tensor."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_basic(x: cute.Tensor):\n",
" # Print the tensor\n",
" print(\"Basic output:\")\n",
" cute.print_tensor(x)\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_verbose(x: cute.Tensor):\n",
" # Print the tensor with verbose mode\n",
" print(\"Verbose output:\")\n",
" cute.print_tensor(x, verbose=True)\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_slice(x: cute.Tensor, coord: tuple):\n",
" # slice a 2D tensor from the 3D tensor\n",
" sliced_data = cute.slice_(x, coord)\n",
" y = cute.make_rmem_tensor(sliced_data.layout, sliced_data.element_type)\n",
" # Convert to TensorSSA format by loading the sliced data into the fragment\n",
" y.store(sliced_data.load())\n",
" print(\"Slice output:\")\n",
" cute.print_tensor(y)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The default `cute.print_tensor` will output CuTe tensor with datatype, storage space, CuTe layout information, and print data in torch-style format."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Basic output:\n",
"tensor(raw_ptr(0x000000000a5f1d50: f32, generic, align<4>) o (4,3,2):(6,2,1), data=\n",
" [[[ 0.000000, 2.000000, 4.000000, ],\n",
" [ 6.000000, 8.000000, 10.000000, ],\n",
" [ 12.000000, 14.000000, 16.000000, ],\n",
" [ 18.000000, 20.000000, 22.000000, ]],\n",
"\n",
" [[ 1.000000, 3.000000, 5.000000, ],\n",
" [ 7.000000, 9.000000, 11.000000, ],\n",
" [ 13.000000, 15.000000, 17.000000, ],\n",
" [ 19.000000, 21.000000, 23.000000, ]]])\n"
]
}
],
"source": [
"def tensor_print_example1():\n",
" shape = (4, 3, 2)\n",
"\n",
" # Creates [0,...,23] and reshape to (4, 3, 2)\n",
" data = np.arange(24, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_basic(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example1()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The verbosed print will show coodination details of each element in the tensor. The below example shows how we index element in a 2D 4x3 tensor space."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Verbose output:\n",
"tensor(raw_ptr(0x000000000a814cc0: f32, generic, align<4>) o (4,3):(3,1), data= (\n",
"\t(0,0)= 0.000000\n",
"\t(0,1)= 1.000000\n",
"\t(0,2)= 2.000000\n",
"\t(1,0)= 3.000000\n",
"\t(1,1)= 4.000000\n",
"\t(1,2)= 5.000000\n",
"\t(2,0)= 6.000000\n",
"\t(2,1)= 7.000000\n",
"\t(2,2)= 8.000000\n",
"\t(3,0)= 9.000000\n",
"\t(3,1)= 10.000000\n",
"\t(3,2)= 11.000000\n",
")\n"
]
}
],
"source": [
"def tensor_print_example2():\n",
" shape = (4, 3)\n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_verbose(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example2()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To print a subset elements in the given Tensor, we can use cute.slice_ to select a range of the given tensor, load them into register and then print the values with `cute.print_tensor`."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Slice output:\n",
"tensor(raw_ptr(0x00007ffeeae1fc60: f32, rmem, align<32>) o (4):(3), data=\n",
" [ 0.000000, ],\n",
" [ 3.000000, ],\n",
" [Slice output:\n",
" 6.000000, ],\n",
" [ 9.000000, ])\n",
"tensor(raw_ptr(0x00007ffeeae1fc60: f32, rmem, align<32>) o (3):(1), data=\n",
" [ 3.000000, ],\n",
" [ 4.000000, ],\n",
" [ 5.000000, ])\n"
]
}
],
"source": [
"def tensor_print_example3():\n",
" shape = (4, 3)\n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_slice(from_dlpack(data), (None, 0))\n",
" print_tensor_slice(from_dlpack(data), (1, None))\n",
"\n",
"\n",
"tensor_print_example3()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To print the tensor in device memory, you can use `cute.print_tensor` within CuTe JIT kernels."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def print_tensor_gpu(src: cute.Tensor):\n",
" print(src)\n",
" cute.print_tensor(src)\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_host(src: cute.Tensor):\n",
" print_tensor_gpu(src).launch(grid=(1, 1, 1), block=(1, 1, 1))"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<ptr<f32, gmem> o (4,3):(3,1)>\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x00007f5f81200400: f32, gmem, align<4>) o (4,3):(3,1), data=\n",
" [[-0.690547, -0.274619, -1.659539, ],\n",
" [-1.843524, -1.648711, 1.163431, ],\n",
" [-0.716668, -1.900705, 0.592515, ],\n",
" [ 0.711333, -0.552422, 0.860237, ]])\n"
]
}
],
"source": [
"import torch\n",
"\n",
"\n",
"def tensor_print_example4():\n",
" a = torch.randn(4, 3, device=\"cuda\")\n",
" cutlass.cuda.initialize_cuda_context()\n",
" print_tensor_host(from_dlpack(a))\n",
"\n",
"\n",
"tensor_print_example4()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Currently, `cute.print_tensor` only supports tensor with integer data types and `Float16`/`Float32`/`Float64` floating point data types. We will support more data types in the future."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/tensor.ipynb
+330
View File
@@ -0,0 +1,330 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Tensor\n",
"\n",
"A tensor in CuTe is created through the composition of two key components:\n",
"\n",
"1. An **Engine** (E) - A random-access, pointer-like object that supports:\n",
" - Offset operation: `e + d → e` (offset engine by elements of a layout's codomain)\n",
" - Dereference operation: `*e → v` (dereference engine to produce value)\n",
"\n",
"2. A **Layout** (L) - Defines the mapping from coordinates to offsets\n",
"\n",
"A tensor is formally defined as the composition of an engine E with a layout L, expressed as `T = E ∘ L`. When evaluating a tensor at coordinate c, it:\n",
"\n",
"1. Maps the coordinate c to the codomain using the layout\n",
"2. Offsets the engine accordingly\n",
"3. Dereferences the result to obtain the tensor's value\n",
"\n",
"This can be expressed mathematically as:\n",
"\n",
"```\n",
"T(c) = (E ∘ L)(c) = *(E + L(c))\n",
"```\n",
"\n",
"## Example Usage\n",
"\n",
"Here's a simple example of creating a tensor using pointer and layout `(8,5):(5,1)` and fill with ones:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def create_tensor_from_ptr(ptr: cute.Pointer):\n",
" layout = cute.make_layout((8, 5), stride=(5, 1))\n",
" tensor = cute.make_tensor(ptr, layout)\n",
" tensor.fill(1)\n",
" cute.print_tensor(tensor)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This creates a tensor where:\n",
"- The engine is a pointer\n",
"- The layout with shape `(8, 5)` and stride `(5, 1)`\n",
"- The resulting tensor can be evaluated using coordinates defined by the layout\n",
"\n",
"We can test this by allocating buffer with torch and run test with pointer to torch tensor"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"from cutlass.torch import dtype as torch_dtype\n",
"import cutlass.cute.runtime as cute_rt\n",
"\n",
"a = torch.randn(8, 5, dtype=torch_dtype(cutlass.Float32))\n",
"ptr_a = cute_rt.make_ptr(cutlass.Float32, a.data_ptr())\n",
"\n",
"create_tensor_from_ptr(ptr_a)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## DLPACK support \n",
"\n",
"CuTe DSL is designed to support dlpack protocol natively. This offers easy integration with frameworks \n",
"supporting DLPack, e.g. torch, numpy, jax, tensorflow, etc.\n",
"\n",
"For more information, please refer to DLPACK project: https://github.com/dmlc/dlpack\n",
"\n",
"Calling `from_dlpack` can convert any tensor or ndarray object supporting `__dlpack__` and `__dlpack_device__`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_dlpack(src: cute.Tensor):\n",
" print(src)\n",
" cute.print_tensor(src)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"a = torch.randn(8, 5, dtype=torch_dtype(cutlass.Float32))\n",
"\n",
"print_tensor_dlpack(from_dlpack(a))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"a = np.random.randn(8, 8).astype(np.float32)\n",
"\n",
"print_tensor_dlpack(from_dlpack(a))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Tensor Evaluation Methods\n",
"\n",
"Tensors support two primary methods of evaluation:\n",
"\n",
"### 1. Full Evaluation\n",
"When applying the tensor evaluation with a complete coordinate c, it computes the offset, applies it to the engine, \n",
"and dereferences it to return the stored value. This is the straightforward case where you want to access \n",
"a specific element of the tensor.\n",
"\n",
"### 2. Partial Evaluation (Slicing)\n",
"When evaluating with an incomplete coordinate c = c' ⊕ c* (where c* represents the unspecified portion), \n",
"the result is a new tensor which is a slice of the original tensor with its engine offset to account for \n",
"the coordinates that were provided. This operation can be expressed as:\n",
"\n",
"```\n",
"T(c) = (E ∘ L)(c) = (E + L(c')) ∘ L(c*) = T'(c*)\n",
"```\n",
"\n",
"Slicing effectively reduces the dimensionality of the tensor, creating a sub-tensor that can be \n",
"further evaluated or manipulated."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def tensor_access_item(a: cute.Tensor):\n",
" # access data using linear index\n",
" cute.printf(\n",
" \"a[2] = {} (equivalent to a[{}])\",\n",
" a[2],\n",
" cute.make_identity_tensor(a.layout.shape)[2],\n",
" )\n",
" cute.printf(\n",
" \"a[9] = {} (equivalent to a[{}])\",\n",
" a[9],\n",
" cute.make_identity_tensor(a.layout.shape)[9],\n",
" )\n",
"\n",
" # access data using n-d coordinates, following two are equivalent\n",
" cute.printf(\"a[2,0] = {}\", a[2, 0])\n",
" cute.printf(\"a[2,4] = {}\", a[2, 4])\n",
" cute.printf(\"a[(2,4)] = {}\", a[2, 4])\n",
"\n",
" # assign value to tensor@(2,4)\n",
" a[2, 3] = 100.0\n",
" a[2, 4] = 101.0\n",
" cute.printf(\"a[2,3] = {}\", a[2, 3])\n",
" cute.printf(\"a[(2,4)] = {}\", a[(2, 4)])\n",
"\n",
"\n",
"# Create a tensor with sequential data using torch\n",
"data = torch.arange(0, 8 * 5, dtype=torch.float32).reshape(8, 5)\n",
"tensor_access_item(from_dlpack(data))\n",
"\n",
"print(data)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Tensor as memory view\n",
"\n",
"In CUDA programming, different memory spaces have different characteristics in terms of access speed, scope, and lifetime:\n",
"\n",
"- **generic**: Default memory space that can refer to any other memory space.\n",
"- **global memory (gmem)**: Accessible by all threads across all blocks, but has higher latency.\n",
"- **shared memory (smem)**: Accessible by all threads within a block, with much lower latency than global memory.\n",
"- **register memory (rmem)**: Thread-private memory with the lowest latency, but limited capacity.\n",
"- **tensor memory (tmem)**: Specialized memory introduced in NVIDIA Blackwell architecture for tensor operations.\n",
"\n",
"When creating tensors in CuTe, you can specify the memory space to optimize performance based on your access patterns.\n",
"\n",
"For more information on CUDA memory spaces, see the [CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#memory-hierarchy).\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Coordinate Tensors\n",
"\n",
"### Definition and Properties\n",
"\n",
"A coordinate tensor $T: Z^n → Z^m$ is a mathematical structure that establishes a mapping between coordinate spaces. Unlike standard tensors that map coordinates to scalar values, coordinate tensors map coordinates to other coordinates, forming a fundamental building block for tensor operations and transformations.\n",
"\n",
"### Examples\n",
"\n",
"Consider a `(4,4)` coordinate tensor:\n",
"\n",
"**Row-Major Layout (C-style):**\n",
"\\begin{bmatrix} \n",
"(0,0) & (0,1) & (0,2) & (0,3) \\\\\n",
"(1,0) & (1,1) & (1,2) & (1,3) \\\\\n",
"(2,0) & (2,1) & (2,2) & (2,3) \\\\\n",
"(3,0) & (3,1) & (3,2) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"**Column-Major Layout (Fortran-style):**\n",
"\\begin{bmatrix}\n",
"(0,0) & (1,0) & (2,0) & (3,0) \\\\\n",
"(0,1) & (1,1) & (2,1) & (3,1) \\\\\n",
"(0,2) & (1,2) & (2,2) & (3,2) \\\\\n",
"(0,3) & (1,3) & (2,3) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"### Identity Tensor\n",
"\n",
"An identity tensor $I$ is a special case of a coordinate tensor that implements the identity mapping function:\n",
"\n",
"**Definition:**\n",
"For a given shape $S = (s_1, s_2, ..., s_n)$, the identity tensor $I$ satisfies: $I(c) = c, \\forall c \\in \\prod_{i=1}^n [0, s_i)$\n",
"\n",
"**Properties:**\n",
"1. **Bijective Mapping**: The identity tensor establishes a one-to-one correspondence between coordinates.\n",
"2. **Layout Invariance**: The logical structure remains constant regardless of the underlying memory layout.\n",
"3. **Coordinate Preservation**: For any coordinate c, I(c) = c.\n",
"\n",
"\n",
"CuTe establishes an isomorphism between 1-D indices and N-D coordinates through lexicographical ordering. For a coordinate c = (c₁, c₂, ..., cₙ) in an identity tensor with shape S = (s₁, s₂, ..., sₙ):\n",
"\n",
"**Linear Index Formula:**\n",
"$\\text{idx} = c_1 + \\sum_{i=2}^{n} \\left(c_i \\prod_{j=1}^{i-1} s_j\\right)$\n",
"\n",
"**Example:**\n",
"```python\n",
"# Create an identity tensor from a given shape\n",
"coord_tensor = make_identity_tensor(layout.shape())\n",
"\n",
"# Access coordinate using linear index\n",
"coord = coord_tensor[linear_idx] # Returns the N-D coordinate\n",
"```\n",
"\n",
"This bidirectional mapping enables efficient conversion from linear indices to N-dimensional coordinates, facilitating tensor operations and memory access patterns."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def print_tensor_coord(a: cute.Tensor):\n",
" coord_tensor = cute.make_identity_tensor(a.layout.shape)\n",
" print(coord_tensor)\n",
" cute.print_tensor(coord_tensor)\n",
"\n",
"\n",
"a = torch.randn(8, 4, dtype=torch_dtype(cutlass.Float32))\n",
"print_tensor_coord(from_dlpack(a))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/tensorssa.ipynb
+495
View File
@@ -0,0 +1,495 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Introduction to the TensorSSA in CuTe DSL\n",
"\n",
"This tutorial introduces what is the `TensorSSA` and why we need it. We also give some examples to show how to use `TensorSSA`.\n",
"\n",
"## What is TensorSSA\n",
"\n",
"`TensorSSA` is a Python class that represents a tensor value in Static Single Assignment (SSA) form within the CuTe DSL. You can think of it as a tensor residing in a (simulated) register.\n",
"\n",
"## Why TensorSSA\n",
"\n",
"`TensorSSA` encapsulates the underlying MLIR tensor value into an object that's easier to manipulate in Python. By overloading numerous Python operators (like `+`, `-`, `*`, `/`, `[]`, etc.), it allows users to express tensor computations (primarily element-wise operations and reductions) in a more Pythonic way. These element-wise operations are then translated into optimized vectorization instructions.\n",
"\n",
"It's part of the CuTe DSL, serving as a bridge between the user-described computational logic and the lower-level MLIR IR, particularly for representing and manipulating register-level data.\n",
"\n",
"## When to use TensorSSA\n",
"\n",
"`TensorSSA` is primarily used in the following scenarios:\n",
"\n",
"### Load from memory and store to memory"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def load_and_store(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
" \"\"\"\n",
" Load data from memory and store the result to memory.\n",
"\n",
" :param res: The destination tensor to store the result.\n",
" :param a: The source tensor to be loaded.\n",
" :param b: The source tensor to be loaded.\n",
" \"\"\"\n",
" a_vec = a.load()\n",
" print(f\"a_vec: {a_vec}\") # prints `a_vec: vector<12xf32> o (3, 4)`\n",
" b_vec = b.load()\n",
" print(f\"b_vec: {b_vec}\") # prints `b_vec: vector<12xf32> o (3, 4)`\n",
" res.store(a_vec + b_vec)\n",
" cute.print_tensor(res)\n",
"\n",
"\n",
"a = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
"b = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
"c = np.zeros(12).reshape((3, 4)).astype(np.float32)\n",
"load_and_store(from_dlpack(c), from_dlpack(a), from_dlpack(b))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Register-Level Tensor Operations\n",
"\n",
"When writing kernel logic, various computations, transformations, slicing, etc., are performed on data loaded into registers."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def apply_slice(src: cute.Tensor, dst: cute.Tensor, indices: cutlass.Constexpr):\n",
" \"\"\"\n",
" Apply slice operation on the src tensor and store the result to the dst tensor.\n",
"\n",
" :param src: The source tensor to be sliced.\n",
" :param dst: The destination tensor to store the result.\n",
" :param indices: The indices to slice the source tensor.\n",
" \"\"\"\n",
" src_vec = src.load()\n",
" dst_vec = src_vec[indices]\n",
" print(f\"{src_vec} -> {dst_vec}\")\n",
" if cutlass.const_expr(isinstance(dst_vec, cute.TensorSSA)):\n",
" dst.store(dst_vec)\n",
" cute.print_tensor(dst)\n",
" else:\n",
" dst[0] = dst_vec\n",
" cute.print_tensor(dst)\n",
"\n",
"\n",
"def slice_1():\n",
" src_shape = (4, 2, 3)\n",
" dst_shape = (4, 3)\n",
" indices = (None, 1, None)\n",
"\n",
" \"\"\"\n",
" a:\n",
" [[[ 0. 1. 2.]\n",
" [ 3. 4. 5.]]\n",
"\n",
" [[ 6. 7. 8.]\n",
" [ 9. 10. 11.]]\n",
"\n",
" [[12. 13. 14.]\n",
" [15. 16. 17.]]\n",
"\n",
" [[18. 19. 20.]\n",
" [21. 22. 23.]]]\n",
" \"\"\"\n",
" a = np.arange(np.prod(src_shape)).reshape(*src_shape).astype(np.float32)\n",
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_1()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def slice_2():\n",
" src_shape = (4, 2, 3)\n",
" dst_shape = (1,)\n",
" indices = 10\n",
" a = np.arange(np.prod(src_shape)).reshape(*src_shape).astype(np.float32)\n",
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_2()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Arithmetic Operations\n",
"\n",
"As we mentioned earlier, there're many tensor operations whose operands are `TensorSSA`. And they are all element-wise operations. We give some examples below.\n",
"\n",
"### Binary Operations\n",
"\n",
"For binary operations, the LHS operand is `TensorSSA` and the RHS operand can be either `TensorSSA` or `Numeric`. When the RHS is `Numeric`, it will be broadcast to a `TensorSSA`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_1(a: cute.Tensor, b: cute.Tensor):\n",
" a_vec = a.load()\n",
" b_vec = b.load()\n",
"\n",
" add_res = a_vec + b_vec\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - b_vec\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * b_vec\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / b_vec\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // b_vec\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % b_vec\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n",
"b = np.empty((3,), dtype=np.float32)\n",
"b.fill(2.0)\n",
"binary_op_1(from_dlpack(a), from_dlpack(b))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_2(a: cute.Tensor, c: cutlass.Constexpr):\n",
" a_vec = a.load()\n",
"\n",
" add_res = a_vec + c\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - c\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * c\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / c\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // c\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % c\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n",
"c = 2.0\n",
"binary_op_2(from_dlpack(a), c)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_3(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
" a_vec = a.load()\n",
" b_vec = b.load()\n",
"\n",
" gt_res = a_vec > b_vec\n",
" res.store(gt_res)\n",
"\n",
" \"\"\"\n",
" ge_res = a_ >= b_ # [False, True, False]\n",
" lt_res = a_ < b_ # [True, False, True]\n",
" le_res = a_ <= b_ # [True, False, True]\n",
" eq_res = a_ == b_ # [False, False, False]\n",
" \"\"\"\n",
"\n",
"\n",
"a = np.array([1, 2, 3], dtype=np.float32)\n",
"b = np.array([2, 1, 4], dtype=np.float32)\n",
"res = np.empty((3,), dtype=np.bool_)\n",
"binary_op_3(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
"print(res) # prints [False, True, False]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_4(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
" a_vec = a.load()\n",
" b_vec = b.load()\n",
"\n",
" xor_res = a_vec ^ b_vec\n",
" res.store(xor_res)\n",
"\n",
" # or_res = a_vec | b_vec\n",
" # res.store(or_res) # prints [3, 2, 7]\n",
"\n",
" # and_res = a_vec & b_vec\n",
" # res.store(and_res) # prints [0, 2, 0]\n",
"\n",
"\n",
"a = np.array([1, 2, 3], dtype=np.int32)\n",
"b = np.array([2, 2, 4], dtype=np.int32)\n",
"res = np.empty((3,), dtype=np.int32)\n",
"binary_op_4(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
"print(res) # prints [3, 0, 7]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Unary Operations"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def unary_op_1(res: cute.Tensor, a: cute.Tensor):\n",
" a_vec = a.load()\n",
"\n",
" sqrt_res = cute.math.sqrt(a_vec)\n",
" cute.print_tensor(sqrt_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" sin_res = cute.math.sin(a_vec)\n",
" res.store(sin_res)\n",
" cute.print_tensor(sin_res) # prints [-0.756802, -0.756802, -0.756802]\n",
"\n",
" exp2_res = cute.math.exp2(a_vec)\n",
" cute.print_tensor(exp2_res) # prints [16.000000, 16.000000, 16.000000]\n",
"\n",
"\n",
"a = np.array([4.0, 4.0, 4.0], dtype=np.float32)\n",
"res = np.empty((3,), dtype=np.float32)\n",
"unary_op_1(from_dlpack(res), from_dlpack(a))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Reduction Operation\n",
"\n",
"The `TensorSSA`'s `reduce` method applies a specified reduction operation (`ReductionOp.ADD`, \n",
"`ReductionOp.MUL`, `ReductionOp.MAX`, `ReductionOp.MIN`) starting with an initial value, and \n",
"performs this reduction along the dimensions specified by the `reduction_profile`. The result \n",
"is typically a new `TensorSSA` with reduced dimensions or a scalar value if it reduces across \n",
"all axes."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@cute.jit\n",
"def reduction_op(a: cute.Tensor):\n",
" \"\"\"\n",
" Apply reduction operation on the src tensor.\n",
"\n",
" :param src: The source tensor to be reduced.\n",
" \"\"\"\n",
" a_vec = a.load()\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=0)\n",
" cute.printf(red_res) # prints 21.000000\n",
"\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=(None, 1))\n",
" cute.print_tensor(red_res) # prints [6.000000, 15.000000]\n",
"\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 1.0, reduction_profile=(1, None))\n",
" cute.print_tensor(red_res) # prints [6.000000, 8.000000, 10.000000]\n",
"\n",
"\n",
"a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)\n",
"reduction_op(from_dlpack(a))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Broadcast\n",
"\n",
"`TensorSSA` supports broadcasting operations following NumPy's broadcasting rules. Broadcasting \n",
"allows you to perform operations on arrays of different shapes when certain conditions are met. \n",
"The key rules are:\n",
"\n",
"1. Source shape is padded with 1's to match the rank of target shape\n",
"2. The size in each mode of source shape must either be 1 or equal to target shape\n",
"3. After broadcasting, all modes should match target shape\n",
"\n",
"Let's look at some examples of broadcasting in action:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"\n",
"\n",
"@cute.jit\n",
"def broadcast_examples():\n",
" a = cute.make_rmem_tensor((1, 3), dtype=cutlass.Float32)\n",
" a[0] = 0.0\n",
" a[1] = 1.0\n",
" a[2] = 2.0\n",
" a_val = a.load()\n",
" cute.print_tensor(a_val.broadcast_to((4, 3)))\n",
" # tensor(raw_ptr(0x00007ffe26625740: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
" # [[ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ]])\n",
"\n",
" c = cute.make_rmem_tensor((4, 1), dtype=cutlass.Float32)\n",
" c[0] = 0.0\n",
" c[1] = 1.0\n",
" c[2] = 2.0\n",
" c[3] = 3.0\n",
" cute.print_tensor(a.load() + c.load())\n",
" # tensor(raw_ptr(0x00007ffe26625780: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
" # [[ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 1.000000, 2.000000, 3.000000, ],\n",
" # [ 2.000000, 3.000000, 4.000000, ],\n",
" # [ 3.000000, 4.000000, 5.000000, ]])\n",
"\n",
"\n",
"broadcast_examples()"
]
},
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "raw"
}
},
"source": [
"The examples above demonstrate two key broadcasting scenarios:\n",
"\n",
"1. **Row Vector Broadcasting**: In the first example, we create a row vector `a` with shape \n",
" (1, 3) containing values [0.0, 1.0, 2.0]. When we broadcast it to shape (4, 3), the values \n",
" are repeated across the first dimension, resulting in:\n",
" ```\n",
" [[0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0],\n",
" [0.0, 1.0, 2.0]]\n",
" ```\n",
" This demonstrates how a row vector can be broadcast to create multiple identical rows.\n",
"\n",
"2. **Column Vector and Row Vector Addition**: In the second example, we have:\n",
" - A row vector `a` with shape (1, 3) containing [0.0, 1.0, 2.0]\n",
" - A column vector `c` with shape (4, 1) containing [0.0, 1.0, 2.0, 3.0]\n",
" \n",
" When we add these together, both vectors are broadcast to shape (4, 3):\n",
" - The row vector is broadcast vertically (4 times)\n",
" - The column vector is broadcast horizontally (3 times)\n",
" \n",
" The result is:\n",
" ```\n",
" [[0.0 + 0.0, 1.0 + 0.0, 2.0 + 0.0],\n",
" [0.0 + 1.0, 1.0 + 1.0, 2.0 + 1.0],\n",
" [0.0 + 2.0, 1.0 + 2.0, 2.0 + 2.0],\n",
" [0.0 + 3.0, 1.0 + 3.0, 2.0 + 3.0]]\n",
" ```\n",
" =\n",
" ```\n",
" [[0.0, 1.0, 2.0],\n",
" [1.0, 2.0, 3.0],\n",
" [2.0, 3.0, 4.0],\n",
" [3.0, 4.0, 5.0]]\n",
" ```\n",
"\n",
"This demonstrates how `TensorSSA` can automatically handle broadcasting of both row and column \n",
"vectors in arithmetic operations, following the broadcasting rules where each dimension must \n",
"either be 1 or match the target size. The broadcasting is handled implicitly during operations, \n",
"making it easy to work with tensors of different shapes.\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv3_12",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/cute/notebooks/tour_to_sol_gemm.ipynb
+1023
View File
File diff suppressed because it is too large Load Diff