v4.4 tag release update. (#3032)
This commit is contained in:
Junkai-Wu
@@ -29,7 +29,6 @@
|
||||
import sys
|
||||
import os
|
||||
from typing import Tuple
|
||||
import torch
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
@@ -125,6 +124,8 @@ def tensor_op_gemm_wrapper(
|
||||
|
||||
|
||||
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
|
||||
import torch
|
||||
|
||||
print("\nRunning TensorOpGemm test with:")
|
||||
print(f"Tensor dimensions: {mnkl}")
|
||||
|
||||
|
||||
@@ -60,11 +60,8 @@ import os
|
||||
import sys
|
||||
from typing import Type, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass.torch import dtype as torch_dtype
|
||||
from cutlass.cute.runtime import make_ptr
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -205,6 +202,9 @@ def tensor_op_gemm_wrapper(
|
||||
|
||||
|
||||
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
|
||||
import torch
|
||||
from cutlass.torch import dtype as torch_dtype
|
||||
|
||||
print("\nRunning TensorOpGemm test with:")
|
||||
print(f"Tensor dimensions: {mnkl}")
|
||||
|
||||
|
||||
627
examples/python/CuTeDSL/ampere/cooperative_launch.py
Normal file
627
examples/python/CuTeDSL/ampere/cooperative_launch.py
Normal file
@@ -0,0 +1,627 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Cooperative Launch Example:
|
||||
|
||||
This module demonstrates CUDA Cooperative Launch functionality. It implements a
|
||||
global barrier that synchronizes ALL threads across the entire GPU grid.
|
||||
|
||||
In traditional CUDA kernel launches, there is no guarantee that all thread blocks
|
||||
will be resident on the GPU simultaneously. This means that thread blocks may
|
||||
execute in waves (some finish before others start) and attempting to synchronize
|
||||
across blocks can cause deadlock.
|
||||
|
||||
**Cooperative Launch** solves this by guaranteeing that all thread blocks launch
|
||||
atomically and simultaneously.
|
||||
|
||||
For more details, see the CUDA Programming Guide official documentation:
|
||||
https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cooperative-groups.html#when-to-use-cudalaunchcooperativekernel
|
||||
|
||||
Cooperative Launch Limitations:
|
||||
|
||||
Cooperative launch has strict grid size constraints.
|
||||
If you exceed this limit, cudaLaunchCooperativeKernel returns
|
||||
cudaErrorCooperativeLaunchTooLarge.
|
||||
|
||||
This example demonstrates both a successful cooperative launch with a small grid
|
||||
and an expected failure when exceeding the grid size limit.
|
||||
|
||||
Usage:
|
||||
|
||||
Run directly:
|
||||
$ python cooperative_launch.py
|
||||
|
||||
This will:
|
||||
1. Demonstrate expected failure with too many thread blocks
|
||||
2. Successfully run a cooperative kernel with grid-wide barrier
|
||||
3. Print confirmation that all threads synchronized successfully
|
||||
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass._mlir import ir
|
||||
from cutlass.cutlass_dsl import (
|
||||
dsl_user_op, # Decorator for user-defined device operations
|
||||
DSLCudaRuntimeError, # Exception type for CUDA runtime errors
|
||||
extract_mlir_values, # Extract MLIR values from the object
|
||||
new_from_mlir_values, # Create a new instance from MLIR values
|
||||
)
|
||||
|
||||
# Function to check cuda errors
|
||||
from cutlass.base_dsl.runtime.cuda import checkCudaErrors
|
||||
|
||||
# LLVM dialect for inline PTX assembly generation
|
||||
from cutlass._mlir.dialects import llvm
|
||||
|
||||
# CUDA Python bindings for runtime API (memory allocation, synchronization, etc.)
|
||||
import cuda.bindings.runtime as cuda_runtime
|
||||
|
||||
|
||||
class GlobalBarrier:
|
||||
"""
|
||||
A grid-wide barrier for synchronizing ALL thread blocks on the GPU.
|
||||
|
||||
This class implements a cooperative barrier that enables grid-wide
|
||||
synchronization. It requires cooperative launch to function correctly.
|
||||
|
||||
Design Overview:
|
||||
|
||||
The barrier uses a single 32-bit integer in global memory with the
|
||||
following bit layout:
|
||||
|
||||
┌──────────────────────────────────────────────────────────────────┐
|
||||
│ Bit 31 │ Bits 30-0 │
|
||||
│ ────────── │ ───────────────────────────────────────────────────│
|
||||
│ Phase Bit │ Arrival Counter (supports up to 2^31 - 1 blocks) │
|
||||
└──────────────────────────────────────────────────────────────────┘
|
||||
|
||||
Capacity:
|
||||
|
||||
- Maximum thread blocks: 2^31 - 1 = 2,147,483,647 blocks
|
||||
|
||||
Memory Ordering:
|
||||
|
||||
The barrier uses specific memory ordering semantics:
|
||||
|
||||
- Release semantics on arrival (atom.add.release.gpu)
|
||||
|
||||
- Acquire semantics on wait (ld.global.acquire.gpu)
|
||||
|
||||
Usage Example:
|
||||
|
||||
Host-side setup:
|
||||
|
||||
>>> barrier_ptr = GlobalBarrier.allocate() # Allocate barrier memory
|
||||
|
||||
Device-side usage (inside a kernel):
|
||||
|
||||
>>> barrier = GlobalBarrier(barrier_ptr)
|
||||
>>>
|
||||
>>> # Do some work...
|
||||
>>>
|
||||
>>> barrier.arrive_and_wait() # Synchronize all blocks
|
||||
>>>
|
||||
>>> # All blocks proceed together after this point
|
||||
|
||||
Warning:
|
||||
|
||||
This barrier requires cooperative launch! Using it with a regular launch
|
||||
can result in a deadlock because not all thread blocks may be resident
|
||||
simultaneously.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def allocate() -> cute.runtime.Pointer:
|
||||
"""
|
||||
Allocate and initialize barrier memory on the GPU.
|
||||
|
||||
This function allocates device memory for the barrier.
|
||||
It must be called before launching any kernel that uses the barrier.
|
||||
"""
|
||||
ptr = checkCudaErrors(cuda_runtime.cudaMalloc(4))
|
||||
|
||||
# This sets all 32 bits to 0:
|
||||
# - Phase bit (bit 31) = 0
|
||||
# - Counter (bits 30-0) = 0
|
||||
checkCudaErrors(cuda_runtime.cudaMemset(ptr, 0, 4))
|
||||
|
||||
# Create a pointer with the following properties:
|
||||
# - Type: Uint32 (32-bit unsigned integer)
|
||||
# - Address: the allocated device pointer
|
||||
# - Address Space: gmem (global memory)
|
||||
barrier_ptr = cute.runtime.make_ptr(
|
||||
cutlass.Uint32, # Element type
|
||||
ptr, # Raw CUDA pointer
|
||||
cute.AddressSpace.gmem, # Memory address space
|
||||
)
|
||||
|
||||
return barrier_ptr
|
||||
|
||||
@staticmethod
|
||||
def free(barrier_ptr: cute.Pointer):
|
||||
"""
|
||||
Free the barrier memory on the GPU.
|
||||
|
||||
This function frees the device memory for the barrier.
|
||||
It must be called after the barrier is no longer needed.
|
||||
"""
|
||||
checkCudaErrors(cuda_runtime.cudaFree(barrier_ptr._pointer))
|
||||
|
||||
@dsl_user_op
|
||||
def __init__(
|
||||
self,
|
||||
barrier_ptr: cute.Pointer,
|
||||
*,
|
||||
phase: Optional[cutlass.Uint32] = None,
|
||||
is_leader: Optional[cutlass.Boolean] = None,
|
||||
number_of_thread_blocks: Optional[cutlass.Uint32] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""
|
||||
Initialize a GlobalBarrier instance on the device.
|
||||
|
||||
This constructor is called by each thread when the kernel
|
||||
starts. It sets up the barrier state for this thread's participation
|
||||
in grid-wide synchronization.
|
||||
|
||||
Each thread stores the following:
|
||||
|
||||
- A reference to the shared barrier memory
|
||||
- Whether it's the leader thread of its block
|
||||
- The current phase for barrier tracking
|
||||
- The total number of thread blocks in the grid
|
||||
"""
|
||||
# The barrier is shared across ALL thread blocks, so it must be in
|
||||
# global memory. Shared memory (smem) is block-local and wouldn't work.
|
||||
if barrier_ptr.memspace != cute.AddressSpace.gmem:
|
||||
raise ValueError(
|
||||
"GlobalBarrier requires barrier_ptr to be in global memory (gmem)"
|
||||
)
|
||||
|
||||
# Store barrier pointer reference
|
||||
self.barrier_ptr = barrier_ptr
|
||||
|
||||
# Initialize phase tracking
|
||||
# Phase starts at 0 for the first barrier, then alternates:
|
||||
# First barrier: wait for phase 1
|
||||
# Second barrier: wait for phase 0
|
||||
# Third barrier: wait for phase 1
|
||||
# ... and so on
|
||||
if phase is not None:
|
||||
self.phase = phase
|
||||
else:
|
||||
self.phase = cutlass.Uint32(0)
|
||||
|
||||
if is_leader is not None:
|
||||
self.is_leader = is_leader
|
||||
else:
|
||||
# Determine if this thread is the block leader
|
||||
# Get this thread's position within its block
|
||||
tidx, tidy, tidz = cute.arch.thread_idx()
|
||||
|
||||
# Leader is the thread at position (0, 0, 0) in the block
|
||||
# We use bitwise AND to combine the three conditions efficiently
|
||||
self.is_leader = (
|
||||
cutlass.Boolean(tidx == 0) # First in X dimension
|
||||
& cutlass.Boolean(tidy == 0) # First in Y dimension
|
||||
& cutlass.Boolean(tidz == 0) # First in Z dimension
|
||||
)
|
||||
|
||||
if number_of_thread_blocks is not None:
|
||||
self.number_of_thread_blocks = number_of_thread_blocks
|
||||
else:
|
||||
# Calculate total number of thread blocks in the grid
|
||||
# Get grid dimensions (how many blocks in each dimension)
|
||||
gidx, gidy, gidz = cute.arch.grid_dim()
|
||||
|
||||
# Total blocks = gridDim.x × gridDim.y × gridDim.z
|
||||
# This is needed to know when ALL blocks have arrived
|
||||
self.number_of_thread_blocks = cutlass.Uint32(gidx * gidy * gidz)
|
||||
|
||||
@dsl_user_op
|
||||
@cute.jit
|
||||
def arrive(self, *, loc=None, ip=None):
|
||||
"""
|
||||
Arrive at the barrier without waiting.
|
||||
|
||||
This signals that the calling thread block has reached the barrier
|
||||
point, but does not wait for other blocks. Use this when you want
|
||||
to overlap computation with barrier synchronization.
|
||||
|
||||
This method must be called by ALL threads in the block,
|
||||
not just the leader. The internal block-level sync ensures all
|
||||
threads in the block agree before the leader signals arrival.
|
||||
"""
|
||||
# Ensure ALL threads in this block have reached this point before
|
||||
# the leader signals arrival. This is critical for correctness!
|
||||
cute.arch.sync_threads(loc=loc, ip=ip)
|
||||
|
||||
# Only the leader thread performs atomic operations to minimize
|
||||
# contention on the barrier memory location
|
||||
if self.is_leader:
|
||||
# Atomically increment the arrival counter by 1
|
||||
# The atomic add returns the value before the add, so we add 1
|
||||
# to get the current value after our arrival
|
||||
barrier_value = (
|
||||
self._increment_barrier(cutlass.Uint32(1), loc=loc, ip=ip) + 1
|
||||
)
|
||||
|
||||
# Check if we're the last block to arrive
|
||||
# Mask out the phase bit (bit 31) to get just the counter value
|
||||
# Compare against total number of thread blocks
|
||||
if (barrier_value & ~(1 << 31)) == self.number_of_thread_blocks:
|
||||
# Flip phase and reset counter
|
||||
# We add a value that simultaneously:
|
||||
# 1. Flips bit 31 (adds 2^31)
|
||||
# 2. Resets counter to 0 (subtracts N, where N was the count)
|
||||
#
|
||||
# Example with 8 blocks:
|
||||
# Current: 0x00000008 (phase=0, counter=8)
|
||||
# Add: 0x80000000 - 8 = 0x7FFFFFF8
|
||||
# Result: 0x80000000 (phase=1, counter=0) ✓
|
||||
#
|
||||
# This works because we're doing modular arithmetic and the
|
||||
# counter wraps correctly
|
||||
self._increment_barrier(
|
||||
cutlass.Uint32((1 << 31) - self.number_of_thread_blocks),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
def _read_barrier(self, *, loc=None, ip=None) -> cutlass.Uint32:
|
||||
"""
|
||||
Read the barrier value with acquire memory semantics.
|
||||
|
||||
This is an internal method that reads the 32-bit barrier value from
|
||||
global memory using GPU-scope acquire semantics.
|
||||
|
||||
Notes
|
||||
-----
|
||||
PTX Instruction:
|
||||
Uses ld.global.acquire.gpu.b32 which is a:
|
||||
|
||||
- Global memory load (ld.global)
|
||||
- With acquire semantics (.acquire)
|
||||
- At GPU scope (.gpu) - visible across all thread blocks
|
||||
- For 32-bit data (.b32)
|
||||
|
||||
Inline Assembly:
|
||||
We use LLVM inline assembly because CuTe DSL may not have a direct
|
||||
high-level API for acquire loads. The assembly string format:
|
||||
|
||||
- $0: Output operand (the loaded value)
|
||||
- $1: Input operand (the address to load from)
|
||||
|
||||
"""
|
||||
# Use inline PTX assembly for the acquire-semantics load
|
||||
return cutlass.Uint32(
|
||||
llvm.inline_asm(
|
||||
# Return type: 32-bit unsigned integer
|
||||
cutlass.Uint32.mlir_type,
|
||||
# Input arguments: barrier pointer address
|
||||
# We convert the pointer to an integer (64-bit address)
|
||||
[self.barrier_ptr.toint().ir_value(loc=loc, ip=ip)],
|
||||
# PTX instruction
|
||||
"ld.global.acquire.gpu.b32 $0, [$1];",
|
||||
# Constraint string
|
||||
# "=r" : Output is a 32-bit register (write-only)
|
||||
# "l" : Input is a 64-bit register (pointer address)
|
||||
"=r,l",
|
||||
# Assembly attributes
|
||||
# Mark as having side effects
|
||||
has_side_effects=True,
|
||||
# No special stack alignment needed
|
||||
is_align_stack=False,
|
||||
# Use AT&T syntax (required for LLVM inline asm)
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
# MLIR location and insertion point
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
def _increment_barrier(
|
||||
self, value: cutlass.Uint32, *, loc=None, ip=None
|
||||
) -> cutlass.Uint32:
|
||||
"""
|
||||
Atomically increment the barrier with release memory semantics.
|
||||
|
||||
This is an internal method that performs an atomic add on the barrier
|
||||
value using GPU-scope release semantics.
|
||||
|
||||
Notes
|
||||
-----
|
||||
PTX Instruction:
|
||||
Uses atom.add.release.gpu.u32 which is a:
|
||||
|
||||
- Atomic operation (atom)
|
||||
- Addition (.add)
|
||||
- With release semantics (.release)
|
||||
- At GPU scope (.gpu)
|
||||
- For unsigned 32-bit integers (.u32)
|
||||
|
||||
Atomicity:
|
||||
The atomic add is guaranteed to be indivisible - no other thread
|
||||
can see a partial update or interleave with this operation.
|
||||
|
||||
Return Value:
|
||||
Atomic operations return the OLD value, not the new value.
|
||||
This is why the caller adds 1 to get the current count.
|
||||
"""
|
||||
# Atomic add using inline PTX assembly with release semantics
|
||||
return cutlass.Uint32(
|
||||
llvm.inline_asm(
|
||||
# Return type: 32-bit unsigned integer (the old value)
|
||||
cutlass.Uint32.mlir_type,
|
||||
# Input arguments: (barrier address, value to add)
|
||||
[
|
||||
self.barrier_ptr.toint().ir_value(
|
||||
loc=loc, ip=ip
|
||||
), # Barrier address
|
||||
value.ir_value(loc=loc, ip=ip), # Value to add
|
||||
],
|
||||
# PTX instruction
|
||||
"atom.add.release.gpu.u32 $0, [$1], $2;",
|
||||
# Constraint string
|
||||
# "=r" : Output is a 32-bit register
|
||||
# "l" : First input is 64-bit (pointer)
|
||||
# "r" : Second input is 32-bit (value)
|
||||
"=r,l,r",
|
||||
# Assembly attributes
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
# MLIR metadata
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
@cute.jit
|
||||
def wait(self, *, loc=None, ip=None):
|
||||
"""
|
||||
Wait for all thread blocks to arrive at the barrier.
|
||||
|
||||
This method blocks (spins) until all thread blocks have called
|
||||
arrive() on the barrier. It does NOT signal arrival itself - use
|
||||
arrive_and_wait() if you need to both arrive and wait.
|
||||
|
||||
IMPORTANT: This method MUST be called by ALL threads in the block.
|
||||
The internal sync_threads ensures all threads proceed together.
|
||||
|
||||
Algorithm:
|
||||
- Leader thread spins, reading barrier with acquire semantics
|
||||
- Waits until phase bit matches expected value
|
||||
- Block-level sync ensures all threads proceed together
|
||||
- Update local phase tracking for next barrier
|
||||
|
||||
"""
|
||||
# Leader thread: spin-wait for phase flip
|
||||
if self.is_leader:
|
||||
# Calculate expected phase (opposite of current phase)
|
||||
# XOR with 1 flips: 0→1, 1→0
|
||||
expected = self.phase ^ 1
|
||||
|
||||
# Initial read of barrier value
|
||||
barrier_value = self._read_barrier(loc=loc, ip=ip)
|
||||
|
||||
# Spin loop: wait until phase matches expected
|
||||
# Extract phase bit (bit 31)
|
||||
# Compare against expected phase value
|
||||
while (barrier_value >> 31) != expected:
|
||||
# Keep reading barrier until phase flips
|
||||
# The acquire semantics ensure memory ordering
|
||||
barrier_value = self._read_barrier(loc=loc, ip=ip)
|
||||
|
||||
# Block-level synchronization
|
||||
# Ensure all threads in the block wait for the leader to see the
|
||||
# phase flip before any thread proceeds
|
||||
cute.arch.sync_threads(loc=loc, ip=ip)
|
||||
|
||||
# Update phase for next barrier
|
||||
# Flip local phase: 0→1 or 1→0
|
||||
# This prepares for the next barrier synchronization
|
||||
self.phase = self.phase ^ 1
|
||||
|
||||
@dsl_user_op
|
||||
def arrive_and_wait(self, *, loc=None, ip=None):
|
||||
"""
|
||||
Arrive at the barrier AND wait for all other thread blocks.
|
||||
|
||||
This is the most common barrier operation - it combines arrive()
|
||||
and wait() into a single call. All thread blocks will be synchronized
|
||||
after this call returns.
|
||||
|
||||
IMPORTANT: This method MUST be called by ALL threads in the block.
|
||||
|
||||
Semantics
|
||||
---------
|
||||
|
||||
Logically equivalent to:
|
||||
|
||||
>>> barrier.arrive() # Signal we've reached this point
|
||||
>>> barrier.wait() # Wait for everyone else
|
||||
"""
|
||||
# Execute both phases: arrive then wait
|
||||
self.arrive(loc=loc, ip=ip)
|
||||
self.wait(loc=loc, ip=ip)
|
||||
|
||||
def __extract_mlir_values__(self) -> List[ir.Value]:
|
||||
"""
|
||||
Extract MLIR values from the GlobalBarrier instance.
|
||||
"""
|
||||
|
||||
assert len(extract_mlir_values(self.barrier_ptr)) == 1
|
||||
assert len(extract_mlir_values(self.is_leader)) == 1
|
||||
assert len(extract_mlir_values(self.phase)) == 1
|
||||
assert len(extract_mlir_values(self.number_of_thread_blocks)) == 1
|
||||
|
||||
return (
|
||||
extract_mlir_values(self.barrier_ptr)
|
||||
+ extract_mlir_values(self.is_leader)
|
||||
+ extract_mlir_values(self.phase)
|
||||
+ extract_mlir_values(self.number_of_thread_blocks)
|
||||
)
|
||||
|
||||
def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GlobalBarrier":
|
||||
"""
|
||||
Create a new GlobalBarrier instance from MLIR values.
|
||||
"""
|
||||
assert len(values) == 4, f"Expected 4 IR values, but got {len(values)}"
|
||||
return GlobalBarrier(
|
||||
barrier_ptr=new_from_mlir_values(self.barrier_ptr, [values[0]]),
|
||||
is_leader=new_from_mlir_values(self.is_leader, [values[1]]),
|
||||
phase=new_from_mlir_values(self.phase, [values[2]]),
|
||||
number_of_thread_blocks=new_from_mlir_values(
|
||||
self.number_of_thread_blocks, [values[3]]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def cooperative_kernel(barrier_ptr: cute.Pointer):
|
||||
"""
|
||||
Example kernel demonstrating cooperative launch with grid-wide barrier.
|
||||
|
||||
This kernel shows how to use the GlobalBarrier class to synchronize all
|
||||
thread blocks in a grid. It performs 10 iterations, with a barrier
|
||||
synchronization after each iteration.
|
||||
|
||||
Launch Requirements: This kernel MUST be launched with cooperative=True.
|
||||
|
||||
"""
|
||||
# Initialize the barrier for this thread
|
||||
# Each thread creates its own GlobalBarrier instance, all sharing the
|
||||
# same underlying barrier_ptr in global memory
|
||||
barrier = GlobalBarrier(barrier_ptr=barrier_ptr)
|
||||
|
||||
for i in range(10):
|
||||
# Synchronize all thread blocks across the entire grid
|
||||
# After this call, ALL blocks have completed iterations 0..i
|
||||
barrier.arrive_and_wait()
|
||||
|
||||
# Get block and thread indices
|
||||
bidx, bidy, bidz = cute.arch.block_idx()
|
||||
tidx, tidy, tidz = cute.arch.thread_idx()
|
||||
|
||||
# Check if this is the leader block (first block in the grid)
|
||||
leader_cluster = bidx == 0 and bidy == 0 and bidz == 0
|
||||
|
||||
# Check if this is the leader thread (first thread in the block)
|
||||
leader_thread = tidx == 0 and tidy == 0 and tidz == 0
|
||||
|
||||
# Only the single leader thread of the leader block prints
|
||||
if leader_cluster and leader_thread:
|
||||
cute.printf("All threads arrived at barrier for the %dth iteration\n", i)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# KERNEL LAUNCH WRAPPERS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@cute.jit
|
||||
def run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer):
|
||||
"""
|
||||
Launch the cooperative kernel with a reasonable grid size.
|
||||
|
||||
This wrapper launches the cooperative_kernel with a grid of 8 thread blocks
|
||||
(2×2×2), where each block contains 128 threads (32×2×2).
|
||||
|
||||
Notes: The cooperative=True flag is ESSENTIAL. It tells CUDA to:
|
||||
- Verify the grid fits within hardware limits
|
||||
- Launch all blocks atomically
|
||||
"""
|
||||
cooperative_kernel(barrier_ptr).launch(
|
||||
grid=(2, 2, 2), # 8 thread blocks
|
||||
block=(32, 2, 2), # 128 threads per block
|
||||
cooperative=True, # Enable cooperative launch semantics
|
||||
)
|
||||
|
||||
|
||||
@cute.jit
|
||||
def xfail_run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer):
|
||||
"""
|
||||
Demonstrate cooperative launch failure with an oversized grid.
|
||||
|
||||
This wrapper intentionally launches with a grid that exceeds
|
||||
the limits, demonstrating how cooperative launch fails.
|
||||
|
||||
This launch is expected to fail with cudaErrorCooperativeLaunchTooLarge.
|
||||
|
||||
This demonstrates proper error handling for cooperative launch.
|
||||
|
||||
See Also
|
||||
--------
|
||||
The main() function shows how to properly catch and handle this error.
|
||||
"""
|
||||
# Attempt to launch with way too many blocks
|
||||
cooperative_kernel(barrier_ptr).launch(
|
||||
grid=(10000, 1, 1), # 10,000 blocks
|
||||
block=(1024, 1, 1), # 1,024 threads per block
|
||||
cooperative=True, # Cooperative launch will reject this
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Initialize CUDA context
|
||||
cutlass.cuda.initialize_cuda_context()
|
||||
|
||||
# Allocate barrier memory
|
||||
# Allocate 4 bytes in device global memory for the barrier state
|
||||
barrier_ptr = GlobalBarrier.allocate()
|
||||
|
||||
# Demonstrate expected failure (grid too large)
|
||||
expectedly_failed = False
|
||||
try:
|
||||
# Attempt to launch with 10,000 blocks - this WILL fail
|
||||
xfail_run_cooperative_kernel(barrier_ptr)
|
||||
except DSLCudaRuntimeError as e:
|
||||
# Verify we got the expected error code
|
||||
assert (
|
||||
e.error_code == cuda_runtime.cudaError_t.cudaErrorCooperativeLaunchTooLarge
|
||||
)
|
||||
expectedly_failed = True
|
||||
finally:
|
||||
# Ensure the failure actually happened (test validation)
|
||||
assert expectedly_failed
|
||||
|
||||
# Run successful cooperative kernel
|
||||
# Launch with a reasonable grid size that fits hardware constraints
|
||||
run_cooperative_kernel(barrier_ptr)
|
||||
|
||||
# Synchronize and clean up
|
||||
checkCudaErrors(cuda_runtime.cudaDeviceSynchronize())
|
||||
|
||||
# Free the barrier memory
|
||||
GlobalBarrier.free(barrier_ptr)
|
||||
@@ -28,7 +28,6 @@
|
||||
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import time
|
||||
from typing import Type
|
||||
|
||||
@@ -36,7 +35,6 @@ from typing import Type
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.cute.testing as testing
|
||||
import cutlass.torch as cutlass_torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
"""
|
||||
@@ -252,6 +250,8 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128):
|
||||
cC = cute.zipped_divide(idC, tiler=tiler_mn)
|
||||
print(f"[DSL INFO] coord tensor = {cC.type}")
|
||||
|
||||
kernel_name = f"cutlass_dsl_elementwise_add_kernel"
|
||||
elementwise_add_kernel.set_name_prefix(kernel_name)
|
||||
elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(
|
||||
grid=[cute.size(gC, mode=[1]), 1, 1],
|
||||
block=[cute.size(tv_layout, mode=[0]), 1, 1],
|
||||
@@ -270,6 +270,12 @@ def run_elementwise_add(
|
||||
warmup_iterations=2,
|
||||
iterations=200,
|
||||
):
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("Ampere GPU is required to run this example!")
|
||||
|
||||
print("\nRunning Elementwise Add test with:")
|
||||
print(f"Tensor dimensions: [{M}, {N}]")
|
||||
print(f"Input and Output Data type: {dtype}")
|
||||
@@ -304,6 +310,8 @@ def run_elementwise_add(
|
||||
else:
|
||||
c_tensor = c
|
||||
|
||||
elementwise_add.set_name_prefix("host_prefix")
|
||||
|
||||
print("Compiling kernel with cute.compile ...")
|
||||
start_time = time.time()
|
||||
compiled_func = cute.compile(
|
||||
@@ -386,9 +394,6 @@ if __name__ == "__main__":
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("Ampere GPU is required to run this example!")
|
||||
|
||||
run_elementwise_add(
|
||||
args.M,
|
||||
args.N,
|
||||
|
||||
372
examples/python/CuTeDSL/ampere/elementwise_add_autotune.py
Normal file
372
examples/python/CuTeDSL/ampere/elementwise_add_autotune.py
Normal file
@@ -0,0 +1,372 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import argparse
|
||||
from typing import Any, Callable, Type
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.cute.testing as testing
|
||||
|
||||
"""
|
||||
In this example we revisit the elementwise add example and use the autotune_jit decorator to
|
||||
autotune the kernel.
|
||||
|
||||
To run this example:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python examples/ampere/elementwise_add_autotune.py --M 3 --N 12
|
||||
python examples/ampere/elementwise_add_autotune.py --M 1024 --N 512
|
||||
python examples/ampere/elementwise_add_autotune.py --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 1000
|
||||
|
||||
"""
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def elementwise_add_kernel(
|
||||
gA: cute.Tensor,
|
||||
gB: cute.Tensor,
|
||||
gC: cute.Tensor,
|
||||
cC: cute.Tensor, # coordinate tensor
|
||||
shape: cute.Shape,
|
||||
thr_layout: cute.Layout,
|
||||
val_layout: cute.Layout,
|
||||
):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
bidx, _, _ = cute.arch.block_idx()
|
||||
|
||||
# slice for CTAs
|
||||
# logical id -> address
|
||||
blk_coord = ((None, None), bidx)
|
||||
blkA = gA[blk_coord] # (TileM,TileN)
|
||||
blkB = gB[blk_coord] # (TileM,TileN)
|
||||
blkC = gC[blk_coord] # (TileM,TileN)
|
||||
blkCrd = cC[blk_coord] # (TileM, TileN)
|
||||
|
||||
# # declare the atoms which will be used later for memory copy
|
||||
copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)
|
||||
copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type)
|
||||
|
||||
tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
|
||||
tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
|
||||
tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout)
|
||||
|
||||
thr_copy_A = tiled_copy_A.get_slice(tidx)
|
||||
thr_copy_B = tiled_copy_B.get_slice(tidx)
|
||||
thr_copy_C = tiled_copy_C.get_slice(tidx)
|
||||
|
||||
thrA = thr_copy_A.partition_S(blkA)
|
||||
thrB = thr_copy_B.partition_S(blkB)
|
||||
thrC = thr_copy_C.partition_S(blkC)
|
||||
|
||||
# allocate fragments for gmem->rmem
|
||||
frgA = cute.make_rmem_tensor_like(thrA)
|
||||
frgB = cute.make_rmem_tensor_like(thrB)
|
||||
frgC = cute.make_rmem_tensor_like(thrC)
|
||||
|
||||
thrCrd = thr_copy_C.partition_S(blkCrd)
|
||||
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
|
||||
|
||||
for i in range(0, cute.size(frgPred), 1):
|
||||
val = cute.elem_less(thrCrd[i], shape)
|
||||
frgPred[i] = val
|
||||
|
||||
# Print per thread predicate mask
|
||||
# if tidx == 0 and bidx == 0:
|
||||
# cute.printf("block_dim = {}", cute.arch.grid_dim())
|
||||
# cute.printf("shape = {}", shape)
|
||||
# cute.print_tensor(thrA)
|
||||
# cute.print_tensor(thrB)
|
||||
# cute.print_tensor(frgPred)
|
||||
|
||||
##########################################################
|
||||
# Move data to reg address space
|
||||
##########################################################
|
||||
|
||||
cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
|
||||
cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
|
||||
|
||||
# if tidx == 0 and bidx == 0:
|
||||
# cute.print_tensor(frgA)
|
||||
# cute.print_tensor(frgB)
|
||||
|
||||
# Load data before use. The compiler will optimize the copy and load
|
||||
# operations to convert some memory ld/st into register uses.
|
||||
result = frgA.load() + frgB.load()
|
||||
|
||||
# Save the results back to registers. Here we reuse b's registers.
|
||||
frgC.store(result)
|
||||
|
||||
# Copy the results back to c
|
||||
cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)
|
||||
|
||||
|
||||
@testing.autotune_jit(
|
||||
params_dict={"copy_bits": [64, 128]},
|
||||
update_on_change=["M", "N"],
|
||||
warmup_iterations=100,
|
||||
iterations=100,
|
||||
)
|
||||
@cute.jit
|
||||
def elementwise_add_autotune(mA, mB, mC, M, N, copy_bits: cutlass.Constexpr = 128):
|
||||
dtype = mA.element_type
|
||||
vector_size = copy_bits // dtype.width
|
||||
|
||||
thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
|
||||
val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
|
||||
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
|
||||
|
||||
gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
idC = cute.make_identity_tensor(mC.shape)
|
||||
cC = cute.zipped_divide(idC, tiler=tiler_mn)
|
||||
|
||||
elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(
|
||||
grid=[cute.size(gC, mode=[1]), 1, 1],
|
||||
block=[cute.size(tv_layout, mode=[0]), 1, 1],
|
||||
)
|
||||
|
||||
|
||||
class ElementwiseAddWrapper:
|
||||
"""
|
||||
This class mimics more advanced kernel development, where a class encapsulates
|
||||
pieces of the kernel implementation.
|
||||
|
||||
The can_implement method can be used to check if the kernel can be implemented
|
||||
for the given arguments.
|
||||
|
||||
The __call__ method is the actual cute.jit function.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, copy_bits: cutlass.Constexpr = 128):
|
||||
self.copy_bits = copy_bits
|
||||
|
||||
def can_implement(self, mA, mB, mC, M, N):
|
||||
return self.copy_bits in [64, 128]
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, mA, mB, mC, M, N):
|
||||
dtype = mA.element_type
|
||||
vector_size = self.copy_bits // dtype.width
|
||||
|
||||
thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
|
||||
val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
|
||||
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
|
||||
|
||||
gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))
|
||||
idC = cute.make_identity_tensor(mC.shape)
|
||||
cC = cute.zipped_divide(idC, tiler=tiler_mn)
|
||||
|
||||
elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(
|
||||
grid=[cute.size(gC, mode=[1]), 1, 1],
|
||||
block=[cute.size(tv_layout, mode=[0]), 1, 1],
|
||||
)
|
||||
|
||||
|
||||
def tune_class(mA, mB, mC, M, N):
|
||||
"""
|
||||
This function is used to autotune the elementwise add kernel which is wrapped in a class.
|
||||
An internal function is defined to compile the class with the given arguments.
|
||||
The internal function is then passed to the benchmarking.tune function to autotune.
|
||||
The best parameters are then used to instantiate the class.
|
||||
|
||||
:param mA: Input tensor A
|
||||
:type mA: cute.Tensor
|
||||
:param mB: Input tensor B
|
||||
:type mB: cute.Tensor
|
||||
:param mC: Output tensor C
|
||||
:type mC: cute.Tensor
|
||||
:param M: Number of rows in the input tensors
|
||||
:type M: int
|
||||
:param N: Number of columns in the input tensors
|
||||
:type N: int
|
||||
:return: An instance of the ElementwiseAddWrapper class with the best parameters
|
||||
:rtype: ElementwiseAddWrapper
|
||||
"""
|
||||
|
||||
def compile_class(a, b, c, M, N, copy_bits=128) -> Callable[[], Any]:
|
||||
kernel = ElementwiseAddWrapper(copy_bits)
|
||||
if not kernel.can_implement(a, b, c, M, N):
|
||||
raise ValueError(f"Cannot implement kernel for copy_bits={copy_bits}")
|
||||
compiled_kernel = cute.compile(kernel, a, b, c, M, N)
|
||||
return lambda: compiled_kernel(a, b, c, M, N)
|
||||
|
||||
params = testing.tune(
|
||||
compile_class,
|
||||
params_dict={"copy_bits": [1, 64, 128]},
|
||||
kernel_arguments=testing.JitArguments(mA, mB, mC, M, N),
|
||||
)
|
||||
return ElementwiseAddWrapper(**params)
|
||||
|
||||
|
||||
def run_elementwise_add(
|
||||
M_start,
|
||||
M_range,
|
||||
M_step,
|
||||
N_start,
|
||||
N_range,
|
||||
N_step,
|
||||
dtype: Type[cutlass.Numeric],
|
||||
skip_ref_check=False,
|
||||
warmup_iterations=2,
|
||||
iterations=200,
|
||||
):
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("Ampere GPU is required to run this example!")
|
||||
|
||||
for M in range(M_start, M_start + M_range + 1, M_step):
|
||||
for N in range(N_start, N_start + N_range + 1, N_step):
|
||||
print("\nRunning Elementwise Add test with:")
|
||||
print(f"Tensor dimensions: [{M}, {N}]")
|
||||
print(f"Input and Output Data type: {dtype}")
|
||||
|
||||
torch_dtype = cutlass_torch.dtype(dtype)
|
||||
if dtype.is_integer:
|
||||
a = torch.randint(
|
||||
0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
b = torch.randint(
|
||||
0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
else:
|
||||
a = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
|
||||
b = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
|
||||
|
||||
c = torch.zeros_like(a)
|
||||
|
||||
print("Input tensor shapes:")
|
||||
print(f"a: {a.shape}, dtype: {a.dtype}")
|
||||
print(f"b: {b.shape}, dtype: {b.dtype}")
|
||||
print(f"c: {c.shape}, dtype: {c.dtype}\n")
|
||||
|
||||
elementwise_class = tune_class(a, b, c, M, N)
|
||||
|
||||
if not skip_ref_check:
|
||||
print("Verifying results for class ...")
|
||||
torch.testing.assert_close(a + b, c)
|
||||
print("Results verified successfully!")
|
||||
c = torch.zeros_like(a)
|
||||
|
||||
elementwise_add_autotune(a, b, c, M, N)
|
||||
|
||||
if not skip_ref_check:
|
||||
print("Verifying results for autotuned function ...")
|
||||
torch.testing.assert_close(a + b, c)
|
||||
print("Results verified successfully!")
|
||||
|
||||
def generate_kernel_arguments():
|
||||
if dtype.is_integer:
|
||||
a = torch.randint(
|
||||
0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
b = torch.randint(
|
||||
0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
else:
|
||||
a = torch.randn(
|
||||
M, N, device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
b = torch.randn(
|
||||
M, N, device=torch.device("cuda"), dtype=torch_dtype
|
||||
)
|
||||
|
||||
c = torch.zeros_like(a)
|
||||
|
||||
return testing.JitArguments(a, b, c, M, N)
|
||||
|
||||
avg_time_us = testing.benchmark(
|
||||
elementwise_add_autotune,
|
||||
workspace_generator=generate_kernel_arguments,
|
||||
workspace_count=10,
|
||||
warmup_iterations=warmup_iterations,
|
||||
iterations=iterations,
|
||||
)
|
||||
|
||||
# Print execution results
|
||||
print(
|
||||
f"Kernel execution time for cute.jit kernel with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms"
|
||||
)
|
||||
print(
|
||||
f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
|
||||
)
|
||||
|
||||
compiled_class = cute.compile(elementwise_class, a, b, c, M, N)
|
||||
|
||||
avg_time_us = testing.benchmark(
|
||||
compiled_class,
|
||||
workspace_generator=generate_kernel_arguments,
|
||||
workspace_count=10,
|
||||
warmup_iterations=warmup_iterations,
|
||||
iterations=iterations,
|
||||
)
|
||||
|
||||
print(
|
||||
f"Kernel execution time for Class Wrapper with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms"
|
||||
)
|
||||
print(
|
||||
f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
|
||||
)
|
||||
parser.add_argument("--M", default=1024, type=int)
|
||||
parser.add_argument("--M_range", default=0, type=int)
|
||||
parser.add_argument("--M_step", default=1024, type=int)
|
||||
parser.add_argument("--N", default=1024, type=int)
|
||||
parser.add_argument("--N_range", default=0, type=int)
|
||||
parser.add_argument("--N_step", default=1024, type=int)
|
||||
parser.add_argument("--warmup_iterations", default=2, type=int)
|
||||
parser.add_argument("--iterations", default=100, type=int)
|
||||
parser.add_argument("--skip_ref_check", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
run_elementwise_add(
|
||||
args.M,
|
||||
args.M_range,
|
||||
args.M_step,
|
||||
args.N,
|
||||
args.N_range,
|
||||
args.N_step,
|
||||
dtype=cutlass.Float32,
|
||||
skip_ref_check=args.skip_ref_check,
|
||||
warmup_iterations=args.warmup_iterations,
|
||||
iterations=args.iterations,
|
||||
)
|
||||
print("\nPASS")
|
||||
@@ -36,8 +36,6 @@ from typing import List, Type
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass.cute as cute
|
||||
import cutlass.cute.testing as testing
|
||||
import cutlass.torch as cutlass_torch
|
||||
import torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
import cutlass
|
||||
@@ -274,6 +272,8 @@ def leaky_relu(x, alpha, *, loc=None, ip=None):
|
||||
|
||||
|
||||
def leaky_relu_ref(x, alpha):
|
||||
import torch
|
||||
|
||||
return torch.where(x > 0, x, alpha * x)
|
||||
|
||||
|
||||
@@ -287,6 +287,9 @@ def run_and_verify(
|
||||
warmup_iterations=2,
|
||||
iterations=100,
|
||||
):
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("NVIDIA GPU is required to run this example!")
|
||||
|
||||
|
||||
@@ -30,13 +30,11 @@ import argparse
|
||||
from types import SimpleNamespace
|
||||
from typing import Type, Callable
|
||||
|
||||
import torch
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass.cute.testing as testing
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.nvgpu import cpasync, warp
|
||||
import cutlass.torch as cutlass_torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
import cutlass.pipeline as pipeline
|
||||
import cutlass.utils as utils
|
||||
@@ -1162,6 +1160,9 @@ def run(
|
||||
use_cold_l2: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
# Skip unsupported testcase
|
||||
if not FlashAttentionForwardAmpere.can_implement(
|
||||
dtype,
|
||||
@@ -1237,8 +1238,12 @@ def run(
|
||||
torch_stream = torch.cuda.current_stream()
|
||||
# Get the raw stream pointer as a CUstream
|
||||
current_stream = cuda.CUstream(torch_stream.cuda_stream)
|
||||
# Pass compile options if needed
|
||||
compile_options = ""
|
||||
# compile the fa2 forward pass
|
||||
compiled_fa2_fwd = cute.compile(fa2_fwd, q, k, v, o, softmax_scale, current_stream)
|
||||
compiled_fa2_fwd = cute.compile(
|
||||
fa2_fwd, q, k, v, o, softmax_scale, current_stream, options=compile_options
|
||||
)
|
||||
|
||||
if not skip_ref_check:
|
||||
compiled_fa2_fwd(q, k, v, o, softmax_scale, current_stream)
|
||||
@@ -1317,7 +1322,6 @@ if __name__ == "__main__":
|
||||
default=False,
|
||||
help="Use circular buffer tensor sets to ensure L2 cold cache",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
run(
|
||||
args.dtype,
|
||||
|
||||
@@ -29,10 +29,8 @@
|
||||
from typing import Type
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass
|
||||
import cutlass.torch as cutlass_torch
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
from cutlass._mlir.dialects import llvm
|
||||
@@ -838,28 +836,25 @@ class HSTUAttentionForwardAmpere(object):
|
||||
|
||||
|
||||
def run_pytorch_hstu_test(
|
||||
dtype: torch.dtype,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
rab: torch.Tensor,
|
||||
dtype,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
rab,
|
||||
is_causal: bool,
|
||||
):
|
||||
"""Generate the reference output of the HSTU attention with Pytorch.
|
||||
|
||||
:param dtype: data type of the input tensors
|
||||
:type dtype: torch.dtype
|
||||
:param q: query tensor
|
||||
:type q: torch.Tensor
|
||||
:param k: key tensor
|
||||
:type k: torch.Tensor
|
||||
:param v: value tensor
|
||||
:type v: torch.Tensor
|
||||
:param rab: RAB tensor
|
||||
:type rab: torch.Tensor
|
||||
:param is_causal: whether to use causal masking
|
||||
:type is_causal: bool
|
||||
"""
|
||||
import torch
|
||||
|
||||
q = q.to(dtype)
|
||||
k = k.to(dtype)
|
||||
v = v.to(dtype)
|
||||
@@ -922,6 +917,9 @@ def run(
|
||||
"""
|
||||
assert dtype == cutlass.Float16 or dtype == cutlass.BFloat16
|
||||
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
torch_stream = torch.cuda.current_stream()
|
||||
stream = cuda.CUstream(torch_stream.cuda_stream)
|
||||
|
||||
|
||||
@@ -29,8 +29,6 @@
|
||||
from functools import partial
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
from cutlass._mlir.dialects import llvm
|
||||
@@ -49,7 +47,7 @@ Situations like:
|
||||
motivate developers to inline PTX themselves.
|
||||
|
||||
In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred,
|
||||
vote.sync.uni.pred, and use the corresponding ops in nvvm_wrappers.py for the test.
|
||||
vote.sync.uni.pred, and use the corresponding ops in nvvm dialect for the test.
|
||||
|
||||
You can refer to the documentation of `inline_asm op in llvm dialect <https://mlir.llvm.org/docs/Dialects/LLVM/#llvminline_asm-llvminlineasmop>`_
|
||||
and `vote.sync <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`_
|
||||
@@ -61,8 +59,8 @@ To run this example:
|
||||
|
||||
python examples/ampere/inline_ptx.py
|
||||
|
||||
The example will run the vote kernel with inline PTX and nvvm dialect separately.
|
||||
The results from inline PTX and nvvm dialect will be verified correspondingly.
|
||||
The example will run the vote kernel with inline ptx and nvvm dialect separately.
|
||||
The results from inline ptx and nvvm dialect will be verified correspondingly.
|
||||
|
||||
"""
|
||||
|
||||
@@ -184,6 +182,8 @@ def vote(
|
||||
|
||||
|
||||
def run():
|
||||
import torch
|
||||
|
||||
ballot_ptx = torch.randint(
|
||||
0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
|
||||
)
|
||||
@@ -230,14 +230,11 @@ def run():
|
||||
torch.testing.assert_close(ballot_ptx, ballot_nvvm)
|
||||
print("Verifying any results...")
|
||||
torch.testing.assert_close(any_ptx, any_nvvm)
|
||||
print(torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE))))
|
||||
assert torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE)))
|
||||
print("Verifying all results...")
|
||||
torch.testing.assert_close(all_ptx, all_nvvm)
|
||||
assert torch.all(all_ptx == all(i < 10 for i in range(WARP_SIZE)))
|
||||
print("Verifying uni results...")
|
||||
torch.testing.assert_close(uni_ptx, uni_nvvm)
|
||||
assert torch.all(uni_ptx == (len(set(i < 10 for i in range(WARP_SIZE))) == 1))
|
||||
|
||||
print("Results verified successfully!")
|
||||
|
||||
|
||||
|
||||
@@ -31,7 +31,6 @@ import time
|
||||
from typing import Tuple
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
import torch
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
@@ -643,6 +642,8 @@ def run(
|
||||
use_cold_l2: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
import torch
|
||||
|
||||
"""Execute SIMT GEMM operation and benchmark performance.
|
||||
|
||||
:param mnk: GEMM problem size (M, N, K, L)
|
||||
@@ -666,6 +667,7 @@ def run(
|
||||
:return: Execution time of the GEMM kernel in microseconds
|
||||
:rtype: float
|
||||
"""
|
||||
torch.manual_seed(1024)
|
||||
print("Running Ampere SIMT GEMM example:")
|
||||
print(f"mnk: {mnk}")
|
||||
print(f"A major: {a_major}, B major: {b_major}, C major: {c_major}")
|
||||
@@ -851,8 +853,6 @@ if __name__ == "__main__":
|
||||
args = parser.parse_args()
|
||||
print("Running SIMT GEMM example:")
|
||||
|
||||
torch.manual_seed(1024)
|
||||
|
||||
run(
|
||||
args.mnk,
|
||||
args.a_major,
|
||||
|
||||
@@ -28,7 +28,6 @@
|
||||
|
||||
import cutlass.cute as cute
|
||||
import cutlass
|
||||
import torch
|
||||
import numpy as np
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@@ -175,6 +174,8 @@ def host(
|
||||
|
||||
|
||||
def run_and_verify(const_a, const_b, const_c):
|
||||
import torch
|
||||
|
||||
dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda")
|
||||
dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda")
|
||||
dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda")
|
||||
|
||||
@@ -30,12 +30,9 @@ import argparse
|
||||
import math
|
||||
from typing import Tuple, Type
|
||||
|
||||
import torch
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.cute.testing as testing
|
||||
import cutlass.torch as cutlass_torch
|
||||
import cutlass.utils as utils
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@@ -849,6 +846,9 @@ def run(
|
||||
use_cold_l2: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
import torch
|
||||
import cutlass.torch as cutlass_torch
|
||||
|
||||
print("Running Ampere tensor core GEMM example:")
|
||||
print(f"mnkl: {mnkl}")
|
||||
print(
|
||||
|
||||
Reference in New Issue
Block a user