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

v4.3 update. (#2709)

Browse Source 
* v4.3 update.

* Update the cute_dsl_api changelog's doc link

* Update version to 4.3.0

* Update the example link

* Update doc to encourage user to install DSL from requirements.txt

---------

Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
2025-10-22 02:26:30 +08:00
committed by GitHub
parent e6e2cc29f5
commit b1d6e2c9b3
244 changed files with 59272 additions and 10455 deletions
Download Patch File Download Diff File Expand all files Collapse all files

599
examples/python/CuTeDSL/notebooks/async_pipeline.ipynb Normal file
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": 2,
"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), smem=SharedStorage.size_in_bytes()\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": 6,
"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), smem=SharedStorage.size_in_bytes()\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
}

460
examples/python/CuTeDSL/notebooks/benchmark_autotune.ipynb Normal file
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
}

225
examples/python/CuTeDSL/notebooks/composed_layout.ipynb Normal file
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
}

9
examples/python/CuTeDSL/notebooks/cuda_graphs.ipynb
View File

@@ -26,10 +26,11 @@
"source": [
"# import torch for CUDA graphs\n",
"import torch\n",
"import cutlass\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"
]
@@ -61,13 +62,15 @@
" \"\"\"\n",
" cute.printf(\"Hello world\")\n",
"\n",
"\n",
"@cute.jit\n",
"def hello_world(stream : CUstream):\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",
@@ -585,7 +588,7 @@
"\n",
"# Calculate the time spent when launching kernels in a stream\n",
"# Results are in ms\n",
"stream_time = start.elapsed_time(end) \n",
"stream_time = start.elapsed_time(end)\n",
"\n",
"# Warmup our GPU again\n",
"g.replay()\n",

151
examples/python/CuTeDSL/notebooks/cute_layout_algebra.ipynb
View File

@@ -90,7 +90,9 @@
" \"\"\"\n",
" Demonstrates coalesce operation flattening and combining modes\n",
" \"\"\"\n",
" layout = cute.make_layout((2, (1, 6)), stride=(1, (cutlass.Int32(6), 2))) # Dynamic stride\n",
" layout = cute.make_layout(\n",
" (2, (1, 6)), stride=(1, (cutlass.Int32(6), 2))\n",
" ) # Dynamic stride\n",
" result = cute.coalesce(layout)\n",
"\n",
" print(\">>> Original:\", layout)\n",
@@ -98,6 +100,7 @@
" print(\">>> Coalesced:\", result)\n",
" cute.printf(\">?? Coalesced: {}\", result)\n",
"\n",
"\n",
"coalesce_example()"
]
},
@@ -275,8 +278,7 @@
" 3. for all i, 0 <= i < size(@a layout), @a result(i) == @a layout(i)\n",
" \"\"\"\n",
" layout = cute.make_layout(\n",
" ((2, (3, 4)), (3, 2), 1),\n",
" stride=((4, (8, 24)), (2, 6), 12)\n",
" ((2, (3, 4)), (3, 2), 1), stride=((4, (8, 24)), (2, 6), 12)\n",
" )\n",
" result = cute.coalesce(layout)\n",
"\n",
@@ -288,21 +290,26 @@
" original_size = cute.size(layout)\n",
" coalesced_size = cute.size(result)\n",
" print(f\"Original size: {original_size}, Coalesced size: {coalesced_size}\")\n",
" assert coalesced_size == original_size, \\\n",
" f\"Size mismatch: original {original_size}, coalesced {coalesced_size}\"\n",
" \n",
" assert coalesced_size == original_size, (\n",
" f\"Size mismatch: original {original_size}, coalesced {coalesced_size}\"\n",
" )\n",
"\n",
" print(\">>> 2. Checking depth of coalesced layout <= 1:\")\n",
" depth = cute.depth(result)\n",
" print(f\"Depth of coalesced layout: {depth}\")\n",
" assert depth <= 1, f\"Depth of coalesced layout should be <= 1, got {depth}\"\n",
"\n",
" print(\">>> 3. Checking layout functionality remains the same after the coalesce operation:\")\n",
" print(\n",
" \">>> 3. Checking layout functionality remains the same after the coalesce operation:\"\n",
" )\n",
" for i in cutlass.range_constexpr(original_size):\n",
" original_value = layout(i)\n",
" coalesced_value = result(i)\n",
" print(f\"Index {i}: original {original_value}, coalesced {coalesced_value}\")\n",
" assert coalesced_value == original_value, \\\n",
" assert coalesced_value == original_value, (\n",
" f\"Value mismatch at index {i}: original {original_value}, coalesced {coalesced_value}\"\n",
" )\n",
"\n",
"\n",
"coalesce_post_conditions()"
]
@@ -338,11 +345,12 @@
"\n",
" # Coalesce with mode-wise profile (1,1) = coalesce both modes\n",
" result = cute.coalesce(layout, target_profile=(1, 1))\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Original: \", layout)\n",
" print(\">>> Coalesced Result: \", result)\n",
"\n",
"\n",
"bymode_coalesce_example()"
]
},
@@ -387,18 +395,19 @@
" \"\"\"\n",
" Demonstrates basic layout composition R = A ◦ B\n",
" \"\"\"\n",
" A = cute.make_layout((6, 2), stride=(cutlass.Int32(8), 2)) # Dynamic stride\n",
" A = cute.make_layout((6, 2), stride=(cutlass.Int32(8), 2)) # Dynamic stride\n",
" B = cute.make_layout((4, 3), stride=(3, 1))\n",
" R = cute.composition(A, B)\n",
"\n",
" # Print static and dynamic information\n",
" print(\">>> Layout A:\", A)\n",
" cute.printf(\">?? Layout A: {}\", A)\n",
" print(\">>> Layout B:\", B) \n",
" print(\">>> Layout B:\", B)\n",
" cute.printf(\">?? Layout B: {}\", B)\n",
" print(\">>> Composition R = A ◦ B:\", R)\n",
" cute.printf(\">?? Composition R: {}\", R)\n",
"\n",
"\n",
"composition_example()"
]
},
@@ -438,14 +447,8 @@
" Shows difference between static and dynamic composition results\n",
" \"\"\"\n",
" # Static version - using compile-time values\n",
" A_static = cute.make_layout(\n",
" (10, 2), \n",
" stride=(16, 4)\n",
" )\n",
" B_static = cute.make_layout(\n",
" (5, 4), \n",
" stride=(1, 5)\n",
" )\n",
" A_static = cute.make_layout((10, 2), stride=(16, 4))\n",
" B_static = cute.make_layout((5, 4), stride=(1, 5))\n",
" R_static = cute.composition(A_static, B_static)\n",
"\n",
" # Static print shows compile-time info\n",
@@ -457,20 +460,21 @@
" # Dynamic version - using runtime Int32 values\n",
" A_dynamic = cute.make_layout(\n",
" (cutlass.Int32(10), cutlass.Int32(2)),\n",
" stride=(cutlass.Int32(16), cutlass.Int32(4))\n",
" stride=(cutlass.Int32(16), cutlass.Int32(4)),\n",
" )\n",
" B_dynamic = cute.make_layout(\n",
" (cutlass.Int32(5), cutlass.Int32(4)),\n",
" stride=(cutlass.Int32(1), cutlass.Int32(5))\n",
" stride=(cutlass.Int32(1), cutlass.Int32(5)),\n",
" )\n",
" R_dynamic = cute.composition(A_dynamic, B_dynamic)\n",
" \n",
"\n",
" # Dynamic printf shows runtime values\n",
" cute.printf(\">?? Dynamic composition:\")\n",
" cute.printf(\">?? A_dynamic: {}\", A_dynamic)\n",
" cute.printf(\">?? B_dynamic: {}\", B_dynamic)\n",
" cute.printf(\">?? R_dynamic: {}\", R_dynamic)\n",
"\n",
"\n",
"composition_static_vs_dynamic_layout()"
]
},
@@ -511,12 +515,12 @@
" \"\"\"\n",
" # Define the original layout A\n",
" A = cute.make_layout(\n",
" (cutlass.Int32(12), (cutlass.Int32(4), cutlass.Int32(8))), \n",
" stride=(cutlass.Int32(59), (cutlass.Int32(13), cutlass.Int32(1)))\n",
" (cutlass.Int32(12), (cutlass.Int32(4), cutlass.Int32(8))),\n",
" stride=(cutlass.Int32(59), (cutlass.Int32(13), cutlass.Int32(1))),\n",
" )\n",
"\n",
" # Define the tiler for by-mode composition\n",
" tiler = (3, 8) # Apply 3:1 to mode-0 and 8:1 to mode-1\n",
" tiler = (3, 8) # Apply 3:1 to mode-0 and 8:1 to mode-1\n",
"\n",
" # Apply by-mode composition\n",
" result = cute.composition(A, tiler)\n",
@@ -529,6 +533,7 @@
" print(\">>> By-mode Composition Result:\", result)\n",
" cute.printf(\">?? By-mode Composition Result: {}\", result)\n",
"\n",
"\n",
"bymode_composition_example()"
]
},
@@ -571,19 +576,20 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((4, 2, 3), stride=(2, 1, 8)) # (4,2,3):(2,1,8)\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout(4, stride=2) # Apply to layout 4:2\n",
" \n",
"\n",
" # Apply logical divide\n",
" result = cute.logical_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Divide Result:\", result)\n",
" cute.printf(\">?? Logical Divide Result: {}\", result)\n",
"\n",
"\n",
"logical_divide_1d_example()"
]
},
@@ -620,21 +626,26 @@
" Result Shape : ((TileM,RestM), (TileN,RestN), L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply logical divide\n",
" result = cute.logical_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Divide Result:\", result)\n",
" cute.printf(\">?? Logical Divide Result: {}\", result)\n",
"\n",
"\n",
"logical_divide_2d_example()"
]
},
@@ -673,21 +684,26 @@
" Result Shape : ((TileM,TileN), (RestM,RestN,L,...))\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply zipped divide\n",
" result = cute.zipped_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Zipped Divide Result:\", result)\n",
" cute.printf(\">?? Zipped Divide Result: {}\", result)\n",
"\n",
"\n",
"zipped_divide_example()"
]
},
@@ -724,21 +740,26 @@
" Result Shape : ((TileM,TileN), RestM, RestN, L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply tiled divide\n",
" result = cute.tiled_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Tiled Divide Result:\", result)\n",
" cute.printf(\">?? Tiled Divide Result: {}\", result)\n",
"\n",
"\n",
"tiled_divide_example()"
]
},
@@ -775,21 +796,26 @@
" Result Shape : (TileM, TileN, RestM, RestN, L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply flat divide\n",
" result = cute.flat_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Flat Divide Result:\", result)\n",
" cute.printf(\">?? Flat Divide Result: {}\", result)\n",
"\n",
"\n",
"flat_divide_example()"
]
},
@@ -834,19 +860,20 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 2), stride=(4, 1)) # (2,2):(4,1)\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout(6, stride=1) # Apply to layout 6:1\n",
" \n",
"\n",
" # Apply logical product\n",
" result = cute.logical_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Product Result:\", result)\n",
" cute.printf(\">?? Logical Product Result: {}\", result)\n",
"\n",
"\n",
"logical_product_1d_example()"
]
},
@@ -886,16 +913,16 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 5), stride=(5, 1))\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout((3, 4), stride=(1, 3))\n",
" \n",
"\n",
" # Apply blocked product\n",
" blocked_result = cute.blocked_product(layout, tiler=tiler)\n",
"\n",
" # Apply raked product\n",
" raked_result = cute.raked_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
@@ -904,6 +931,7 @@
" cute.printf(\">?? Blocked Product Result: {}\", blocked_result)\n",
" cute.printf(\">?? Raked Product Result: {}\", raked_result)\n",
"\n",
"\n",
"blocked_raked_product_example()"
]
},
@@ -950,16 +978,16 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 5), stride=(5, 1))\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout((3, 4), stride=(1, 3))\n",
"\n",
" # Apply zipped product\n",
" zipped_result = cute.zipped_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Apply tiled product\n",
" tiled_result = cute.tiled_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Apply flat product\n",
" flat_result = cute.flat_product(layout, tiler=tiler)\n",
"\n",
@@ -973,6 +1001,7 @@
" cute.printf(\">?? Tiled Product Result: {}\", tiled_result)\n",
" cute.printf(\">?? Flat Product Result: {}\", flat_result)\n",
"\n",
"\n",
"zipped_tiled_flat_product_example()"
]
}

15
examples/python/CuTeDSL/notebooks/data_types.ipynb
View File

@@ -6,8 +6,6 @@
"metadata": {},
"outputs": [],
"source": [
"from typing import List\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute"
]
@@ -83,12 +81,13 @@
"@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",
" 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",
" print(\"b(static) =\", b) # prints `b(static) = 5`\n",
" cute.printf(\"b(dynamic) = {}\", b) # prints `b(dynamic) = 5`\n",
"\n",
"\n",
"bar()"
]
@@ -154,6 +153,7 @@
" f = e.to(cutlass.Int8)\n",
" cute.printf(\"Int32({}) => Int8({}) (truncated due to range limitation)\", e, f)\n",
"\n",
"\n",
"type_conversion()"
]
},
@@ -241,7 +241,8 @@
" not_a = ~a\n",
" cute.printf(\"~a = {}\", not_a)\n",
"\n",
"operator_demo()\n"
"\n",
"operator_demo()"
]
}
],

1124
examples/python/CuTeDSL/notebooks/elementwise_add.ipynb
View File

File diff suppressed because it is too large Load Diff

27
examples/python/CuTeDSL/notebooks/hello_world.ipynb
View File

@@ -27,8 +27,8 @@
"metadata": {},
"outputs": [],
"source": [
"import cutlass \n",
"import cutlass.cute as cute "
"import cutlass\n",
"import cutlass.cute as cute"
]
},
{
@@ -80,14 +80,13 @@
"source": [
"@cute.jit\n",
"def hello_world():\n",
"\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",
" grid=(1, 1, 1), # Single thread block\n",
" block=(32, 1, 1), # One warp (32 threads) per thread block\n",
" )"
]
},
@@ -107,7 +106,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 5,
"metadata": {},
"outputs": [
{
@@ -115,17 +114,19 @@
"output_type": "stream",
"text": [
"Running hello_world()...\n",
"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",
"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",
"# 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",
@@ -137,6 +138,14 @@
"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",
"# Alternatively, compile with string based options like\n",
"# cute.compile(hello_world, options=\"--keep-ptx --keep-cubin\") would also work.\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()"

45
examples/python/CuTeDSL/notebooks/print.ipynb
View File

@@ -83,8 +83,8 @@
" 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)"
" print(\">>>\", layout) # => (?,2):(1,?)\n",
" cute.printf(\">?? {}\", layout) # => (8,2):(1,8)"
]
},
{
@@ -221,6 +221,7 @@
" 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)"
]
@@ -246,23 +247,26 @@
"source": [
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_basic(x : cute.Tensor):\n",
"def print_tensor_basic(x: cute.Tensor):\n",
" # Print the tensor\n",
" print(\"Basic output:\")\n",
" cute.print_tensor(x)\n",
" \n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_verbose(x : cute.Tensor):\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",
"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_fragment(sliced_data.layout, sliced_data.element_type)\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",
@@ -302,12 +306,13 @@
"source": [
"def tensor_print_example1():\n",
" shape = (4, 3, 2)\n",
" \n",
"\n",
" # Creates [0,...,23] and reshape to (4, 3, 2)\n",
" data = np.arange(24, dtype=np.float32).reshape(*shape) \n",
" \n",
" data = np.arange(24, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_basic(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example1()"
]
},
@@ -348,12 +353,13 @@
"source": [
"def tensor_print_example2():\n",
" shape = (4, 3)\n",
" \n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape) \n",
" \n",
" data = np.arange(12, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_verbose(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example2()"
]
},
@@ -390,13 +396,14 @@
"source": [
"def tensor_print_example3():\n",
" shape = (4, 3)\n",
" \n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape) \n",
" \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()"
]
},
@@ -418,9 +425,10 @@
" 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))"
" print_tensor_gpu(src).launch(grid=(1, 1, 1), block=(1, 1, 1))"
]
},
{
@@ -449,11 +457,14 @@
],
"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()"
]
},

176
examples/python/CuTeDSL/notebooks/tensor.ipynb
View File

@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -43,7 +43,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -69,24 +69,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x000000000736b0c0: f32, generic, align<4>) o (8,5):(5,1), data=\n",
" [[ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" ...\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ]])\n"
]
}
],
"outputs": [],
"source": [
"import torch\n",
"\n",
@@ -115,12 +100,13 @@
},
{
"cell_type": "code",
"execution_count": 4,
"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",
@@ -129,25 +115,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<ptr<f32, generic> o (8,5):(5,1)>\n",
"tensor(raw_ptr(0x0000000007559340: f32, generic, align<4>) o (8,5):(5,1), data=\n",
" [[-1.151769, 1.019397, -0.371175, -0.717776, 0.502176, ],\n",
" [ 0.114282, 0.900084, 0.320770, 1.564574, -0.632329, ],\n",
" [-0.570140, 0.178112, -0.423079, 1.936198, 0.003355, ],\n",
" ...\n",
" [-2.425393, -0.275528, 1.267157, -0.811101, -0.985456, ],\n",
" [ 0.777889, -2.114074, 0.357184, -0.321312, -0.938138, ],\n",
" [ 1.959564, 1.797602, 0.116901, 0.306198, -1.837295, ]])\n"
]
}
],
"outputs": [],
"source": [
"a = torch.randn(8, 5, dtype=torch_dtype(cutlass.Float32))\n",
"\n",
@@ -156,25 +126,9 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<ptr<f32, generic> o (8,8):(8,1)>\n",
"tensor(raw_ptr(0x0000000007979da0: f32, generic, align<4>) o (8,8):(8,1), data=\n",
" [[ 0.122739, -0.605744, -1.442022, ..., -0.356501, -0.993329, -0.091110, ],\n",
" [ 0.278448, 0.318482, -0.276867, ..., 1.542181, -1.701539, -0.309454, ],\n",
" [ 0.563565, -0.753936, 0.131214, ..., 0.437912, -0.482277, -0.051540, ],\n",
" ...\n",
" [-1.974096, -0.177881, 0.426807, ..., -1.579115, -0.304974, 0.451164, ],\n",
" [ 0.149851, -0.704689, -0.295063, ..., -0.653001, 0.008871, 0.903916, ],\n",
" [ 1.188619, 1.519662, 1.270734, ..., 0.404082, 0.173200, 0.093476, ]])\n"
]
}
],
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
@@ -211,39 +165,23 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"a[2] = 10.000000 (equivalent to a[(2,0)])\n",
"a[9] = 6.000000 (equivalent to a[(1,1)])\n",
"a[2,0] = 10.000000\n",
"a[2,4] = 14.000000\n",
"a[(2,4)] = 14.000000\n",
"a[2,3] = 100.000000\n",
"a[(2,4)] = 101.000000\n",
"tensor([[ 0., 1., 2., 3., 4.],\n",
" [ 5., 6., 7., 8., 9.],\n",
" [ 10., 11., 12., 100., 101.],\n",
" [ 15., 16., 17., 18., 19.],\n",
" [ 20., 21., 22., 23., 24.],\n",
" [ 25., 26., 27., 28., 29.],\n",
" [ 30., 31., 32., 33., 34.],\n",
" [ 35., 36., 37., 38., 39.]])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def tensor_access_item(a: cute.Tensor):\n",
" # access data using linear index\n",
" cute.printf(\"a[2] = {} (equivalent to a[{}])\", a[2],\n",
" cute.make_identity_tensor(a.layout.shape)[2])\n",
" cute.printf(\"a[9] = {} (equivalent to a[{}])\", a[9],\n",
" cute.make_identity_tensor(a.layout.shape)[9])\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",
@@ -251,14 +189,14 @@
" 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",
" 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",
"data = torch.arange(0, 8 * 5, dtype=torch.float32).reshape(8, 5)\n",
"tensor_access_item(from_dlpack(data))\n",
"\n",
"print(data)"
@@ -287,14 +225,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### Coordinate Tensor\n",
"## Coordinate Tensors\n",
"\n",
"A coordinate tensor is a special type of tensor that maps coordinates to coordinates rather than to values. \n",
"The key distinction is that while regular tensors map coordinates to some value type (like numbers), \n",
"coordinate tensors map coordinates to other coordinates.\n",
"### Definition and Properties\n",
"\n",
"For example, given a shape (4,4), a coordinate tensor using row-major layout would appear as:\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",
@@ -302,8 +243,7 @@
"(3,0) & (3,1) & (3,2) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"The same shape with a column-major layout would appear as:\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",
@@ -311,40 +251,50 @@
"(0,3) & (1,3) & (2,3) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"The key points about coordinate tensors are:\n",
"- Each element in the tensor is itself a coordinate tuple (i,j) rather than a scalar value\n",
"- The coordinates map to themselves - so position (1,2) contains the coordinate (1,2)\n",
"- The layout (row-major vs column-major) determines how these coordinate tuples are arranged in memory\n",
"### Identity Tensor\n",
"\n",
"For example, coordinate tensors can be created using the `make_identity_tensor` utility:\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 creates a tensor that maps each coordinate to itself, providing a reference point for understanding how other layouts transform these coordinates."
"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": 8,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<(0,0) o (8,4):(1@0,1@1)>\n"
]
}
],
"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",
"a = torch.randn(8,4, dtype=torch_dtype(cutlass.Float32))\n",
"\n",
"a = torch.randn(8, 4, dtype=torch_dtype(cutlass.Float32))\n",
"print_tensor_coord(from_dlpack(a))"
]
}

79
examples/python/CuTeDSL/notebooks/tensorssa.ipynb
View File

@@ -10,8 +10,7 @@
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"import numpy as np\n",
"import torch"
"import numpy as np"
]
},
{
@@ -55,12 +54,13 @@
" :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",
" 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",
" 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",
@@ -101,6 +101,7 @@
" 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",
@@ -124,6 +125,7 @@
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_1()"
]
},
@@ -141,6 +143,7 @@
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_2()"
]
},
@@ -169,22 +172,22 @@
" 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",
" 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",
" 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",
" 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",
" 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(res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(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",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
@@ -206,22 +209,23 @@
" 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",
" 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",
" 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",
" 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",
" 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",
" 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",
" 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",
@@ -251,11 +255,12 @@
" 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]\n"
"print(res) # prints [False, True, False]"
]
},
{
@@ -278,11 +283,12 @@
" # 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]"
"print(res) # prints [3, 0, 7]"
]
},
{
@@ -303,14 +309,15 @@
" 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",
" 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",
" 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",
" 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",
@@ -344,26 +351,14 @@
" :param src: The source tensor to be reduced.\n",
" \"\"\"\n",
" a_vec = a.load()\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 0.0,\n",
" reduction_profile=0\n",
" )\n",
" cute.printf(red_res) # prints 21.000000\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(\n",
" cute.ReductionOp.ADD,\n",
" 0.0,\n",
" reduction_profile=(None, 1)\n",
" )\n",
" cute.print_tensor(red_res) # prints [6.000000, 15.000000]\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(\n",
" cute.ReductionOp.ADD,\n",
" 1.0,\n",
" reduction_profile=(1, None)\n",
" )\n",
" cute.print_tensor(red_res) # prints [6.000000, 8.000000, 10.000000]\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",
@@ -399,7 +394,7 @@
"\n",
"@cute.jit\n",
"def broadcast_examples():\n",
" a = cute.make_fragment((1,3), dtype=cutlass.Float32)\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",
@@ -411,7 +406,7 @@
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ]])\n",
"\n",
" c = cute.make_fragment((4,1), dtype=cutlass.Float32)\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",
@@ -494,7 +489,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.10"
"version": "3.12.11"
}
},
"nbformat": 4,