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

CUTLASS 3.1 (#915)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 20:19:34 -07:00
committed by GitHub
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37184 additions and 16419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
examples/09_turing_tensorop_conv2dfprop/turing_tensorop_conv2dfprop.cu
View File

@@ -555,6 +555,7 @@ Result profile_convolution(Options const &options) {
LayoutOutput,
ElementComputeEpilogue,
ElementAccumulator,
ElementOutput,
cutlass::NumericConverterClamp<ElementOutput, ElementComputeEpilogue>
>(
problem_size,

328
examples/16_ampere_tensorop_conv2dfprop/ampere_tensorop_conv2dfprop.cu
View File

@@ -31,83 +31,181 @@
/**
This example shows how to run convolution kernels using functions and data structures
provided by CUTLASS using tensor cores; which we run on a NVIDIA Ampere GPU.
This example shows how to run CUTLASS's convolution kernels
based on the Implicit GEMM algorithm, that use the Tensor Cores
on an NVIDIA Ampere GPU.
Writing a single high performance convolution kernel is hard but do-able. Whereas writing
high performance kernels at scale which works for multiple problem sizes with good abstractions is
really hard. CUTLASS solves this problem by providing simplified abstractions to compose
multiple sections of implicit gemm kernel. When used properly, the kernels can hit peak performance
of GPU easily.
Writing a single high-performance convolution kernel is hard enough,
let alone writing kernels that perform well for multiple problem sizes
and use good software abstractions.
CUTLASS provides simplified abstractions
to compose multiple sections of a convolution kernel.
When used properly, the kernels can reach peak GPU performance.
CUTLASS divides a kernel into hierarchical composable sections. Which means, at each thread, warp
and thread-block level, they compute on their own tile-size with higher level of tile sizes being
composed from lower level ones. Multiple thread-tiles (tile size each thread computes) can be used
to form warp-tiles (tile size each warp computes) and multiple warp tiles can be used to compute
threadblock-tile (tile size computed by a threadblock).
CUTLASS divides a kernel into hierarchical composable sections
for each level of the GPU hardware hierarchy:
thread, warp, and threadblock.
Each section computes on its own tile shape,
with each higher level's tile shape
being composed from lower-level tile shapes.
Multiple thread tiles (the tile shape each thread computes)
can be used to form warp tiles (the tile shape each warp computes),
and multiple warp tiles can be used to compute threadblock tiles
(the tile shape computed by a threadblock).
In thie example, we split variable initialization into
1. Setting up data properties : describes how tensors are laid out in the memory and how the kernel
can view them (logical to physical mapping)
2. Setting up computation properties : describes how the above set tensors will be used to compute
output of convolution.
In thie example, we split variable initialization into two parts.
First, we setup the data types of the input tensor A, weights' tensor B and output tensor C along
with alpha, beta as the equation for convolution is C = alpha * Conv2dFprop(A, B) + beta * C. In CUTLASS,
the kernels first compute Conv2dFprop(A, B) and leave the rest of the computation to end of the kernel as
alpha * X + beta * C is a simple element-wise operation on X (Conv2dFprop(A, B)) and C. We call this as
epilogue of kernel. Hence, we setup data types for alpha and beta to be equal to
ElementComputeEpilogue = float. We use the data type for elements in input tensor A and B as
cutlass::half_t. We convey this to CUTLASS kernel by initializing template variables ElementAccumulator (float),
ElementComputeEpilogue (float), ElementInputA (cutlass::half_t), ElementInputB (cutlass::half_t),
ElementOutput (float). Communicating just the data type is not enough. As the data is laid out
linearly in memory, we have to convey the layout of tensors. We do that by initializing template
variables LayoutInputA, LayoutInputB and LayoutOutput to TensorNHWC cutlass variable. Next, we setup
rules to comptue alpha * X + beta * C which is called epilogue of the kernel. We initialize template
variable EpilogueOp, which takes the data type of output ElementOutput (float), the number of
elements per vector memory access (8), data type of accumulator (float) and data type of
computation of linear combination (alpha * X + beta * C).
1. Setting up data properties: describes how tensors are laid out in the memory
and how the kernel can view them (logical to physical mapping)
Now that we setup the properties of data, we have to setup properties of computation.
2. Setting up computation properties: describes how the above tensors
will be used to compute the output of convolution
Second, we create template variables of tile sizes for thread-block, warp and mma-op to 128x128x64,
64x64x64, 16x8x16 (MxNxK) respectively. When passed to instantiate CUTLASS Implicit GEMM kernel, it
internally deduces the amount of threads needed per thread-block, amount of shared memory, storing
data in bank-conflict free manner, and ton of other variables required to compose, initialize and
launch a high performance Implicit GEMM kernel. This is the beauty of CUTLASS, it relieves developer
from understanding and coding complicated hardware optimizations which can easily go wrong.
We begin by setting up the data types
of all the input and output elements of a convolution.
A convolution computes
C = alpha * Conv2dFprop(A, B) + beta * C,
so we set up data types for the input tensor A,
weights tensor B, output tensor C,
and the scaling factors alpha and beta.
CUTLASS divides the convolution into two parts:
the "mainloop" that computes X = Conv2dFprop(A, B),
and the "epilogue" that computes C = alpha * X + beta * C.
The epilogue is an element-wise operation on X and C.
In this case, it is a linear combination,
but other epilogues are possible.
CUTLASS also supports multiple MMA pipelines in a threadblock. What are MMA pipelines? MMA pipelines
constitute the whole process of loading input data from global memory to shared memory, loading data
from shared memory to registers, doing matrix multiplication, store to global memory. The below flow
sequence shows a typical mma multistage pipeline.
(see include/cutlass/conv/threadblock/implicit_gemm_multistage.h)
In this example, we want
tensor in global memory --cp_async--> tile in shared memory --smem loads--> registers
--mma--> registers --global stores--> output to global memory
* the scaling factors alpha and beta to be float,
NVIDIA Ampere uses `cp_async` to build multistage software pipeline to better hide latencies.
* the elements of A and B to be cutlass::half_t
(a 16-bit floating-point type),
* the elements of C to be float, and
There are few more template variables initialized such as, which threadblock tile of output matrix
is done which threadblock launched on an SM, CUDA SM architecture of GPU you want to run on.
* intermediate sums to be accumulated in float.
These are all put together to create a template variable which describes CUTLASS Implicit GEMM
kernel using cutlass::conv::device::ImplicitGemm template.
We convey this to the CUTLASS kernel
by setting the following template parameters.
The next step is to initialize physical data, instantiate and initialize CUTLASS kernel and run it.
We use CUTLASS utilities to initialize, fill, compare tensors as they are simple and doesn't come
in the way of learning CUTLASS.
* alpha and beta: ElementComputeEpilogue = float
Once all the tensors are initialized and filled with data, create arguments tuple to launch CUTLASS
kernel which takes problem size (N = 1, H = 64, W = 64, C = 128), filter size (K = 64,
R = 3, S = 3, C = 128 ), padding, strides, dilation, tensors, alpha, beta and the
important one, split k-dimension factor. Along with that, we query CUTLASS if any scratch-space
memory required by the kernel we instantiated. If yes, we create it and pass it along with other
arguments created to initialize CUTLASS kernel then, the kernel is launched.
* Elements of input tensor A: ElementInputA = cutlass::half_t
In this example, we later on launch a reference convolution kernel (from CUTLASS utilities) to
compare if the output from CUTLASS kernel is same as the reference implicit GEMM kernel.
* Elements of input tensor B: ElementInputB = cutlass::half_t
* Elements of output tensor C: ElementOutput = float
* Accumulation type: ElementAccumulator = float
Next, we describe the layout of the input and output tensors.
We convey this to the CUTLASS kernel
by setting the following template parameters.
* Layout of input tensor A: LayoutInputA = TensorNHWC
* Layout of input tensor B: LayoutInputB = TensorNHWC
* Layout of output tensor C: LayoutOutput = TensorNHWC
After that, we set up rules to compute the epilogue.
The epilogue in this case is a simple linear combination
C = alpha * X + beta * C.
Thus, we set the kernel's template parameter EpilogueOp
to LinearCombination. LinearCombination itself
has template parameters:
* the element type of the output tensor (ElementOutput),
* the number of elements per vector memory access (8),
* the data type of the accumulator (ElementAccumulator),
* and the data type used to compute the linear combination
(ElementComputeEpilogue).
We then define the tile shapes
that each level of the computation uses.
We define these as types that encode the tile shapes
as compile-time integer values.
Each shape expresses the dimensions M x N x K.
Here, the letters refer to the dimensions
of a matrix-matrix multiply.
* ThreadblockShape defines the threadblock tile shape
as 128 x 128 x 64.
* WarpShape defines the warp tile shape as 64 x 64 x 64.
* InstructionShape defines the MMA
(matrix multiply-accumulate) operation shape
as 16 x 8 x 16.
These types become template arguments
of the kernel properties type
cutlass::conv::kernel::DefaultConv2dFprop.
The kernel uses these shapes to deduce
the number of threads needed per threadblock,
the required amount of shared memory,
the internal layouts needed to access
shared memory without bank conflicts,
and many other properties that the kernel needs
for good performance.
CUTLASS deduces all these properties automatically,
so that users don't have to.
DefaultConv2dFprop accepts other template parameters
that describe things like the target CUDA SM architecture.
CUTLASS also supports multiple MMA pipelines in a threadblock.
An MMA pipeline constitutes the whole process
of loading input data from global memory to shared memory,
loading data from shared memory to registers,
doing matrix multiplication,
and storing the result to global memory.
The below flow sequence shows a typical MMA multistage pipeline
(see include/cutlass/conv/threadblock/implicit_gemm_multistage.h).
tensor in global memory
--cp_async-->
tile in shared memory
--smem loads-->
registers
--mma-->
registers
--global stores-->
output to global memory
On NVIDIA Ampere, the kernel uses `cp_async`
to build a multistage software pipeline.
This helps it better hide latency.
At this point, we can define the actual CUTLASS kernel type
as the alias ImplicitGemm, a specialization of
cutlass::conv::device::ImplicitGemmConvolution.
The latter accepts the kernel properties type alias
Conv2dFpropKernel as its one template argument.
This example then sets up a test problem
and arguments to the kernel.
We use CUTLASS utilities to allocate
the input and output tensors
and fill them with sample input data.
We then create the kernel arguments
as an instance of ImplicitGemm::Arguments.
The arguments include
the problem size (N = 1, H = 64, W = 64, C = 128),
filter size (K = 64, R = 3, S = 3, C = 128),
padding, strides, dilation, tensors, alpha, beta,
and the split k-dimension factor.
We also query CUTLASS if the kernel we instantiated
requires any memory for scratch space.
If yes, we reserve scratch space and pass it along
with other arguments to initialize the CUTLASS kernel.
After lauching the CUTLASS kernel, this example runs
a reference convolution kernel (from CUTLASS utilities)
to check correctness.
*/
#include <iostream>
@@ -131,8 +229,8 @@ compare if the output from CUTLASS kernel is same as the reference implicit GEMM
#include "helper.h"
// The code section below describes datatype for input, output tensors and computation between
// elements
// Data types for input and output tensors
// and computation between elements
using ElementAccumulator = float; // Data type of accumulator
using ElementComputeEpilogue = float; // Data type of epilogue computation (alpha, beta)
using ElementInputA = cutlass::half_t; // Data type of elements in input tensor
@@ -143,39 +241,40 @@ using LayoutInputA = cutlass::layout::TensorNHWC;
using LayoutInputB = cutlass::layout::TensorNHWC;
using LayoutOutput = cutlass::layout::TensorNHWC;
// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM
// Whether to use tensor cores or regular SIMT cores on GPU SM
using MMAOp = cutlass::arch::OpClassTensorOp;
// This code section describes CUDA SM architecture number
// SM architecture number
using SmArch = cutlass::arch::Sm80;
// This code section describes the tile size a thread block will compute
using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 64>; // Threadblock tile shape
// Threadblock tile shape
using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 64>;
// This code section describes tile size a warp will compute
using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>; // Warp tile shape
// Warp tile shape
using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
// This code section describes the size of MMA op
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; // TensorCore instruction shape
// MMA (Tensor Core instruction, in this case) tile shape
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;
// This code section describes how threadblocks are scheduled on GPU
// How the kernel schedules threadblocks
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// Number of pipelines you want to use
// Number of pipeline stages to use
constexpr int NumStages = 3;
// This code section describe iterator algorithm selected is Analytic or Optimized
// Which iterator algorithm to use: Analytic or Optimized
static cutlass::conv::IteratorAlgorithm const IteratorAlgorithm = cutlass::conv::IteratorAlgorithm::kOptimized;
// This code section describes the epilogue part of the kernel, we use default value
// The epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
ElementOutput, // Data type of output matrix.
128 / cutlass::sizeof_bits<ElementOutput>::value, // The number of elements per vectorized.
128 / cutlass::sizeof_bits<ElementOutput>::value, // The number of elements per vectorized
// memory access. This becomes the vector width of
// math instructions in the epilogue too.
ElementAccumulator, // Data type of accumulator
ElementComputeEpilogue>; // Data type for alpha/beta in linear combination
// Kernel properties type
using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFprop<
ElementInputA, LayoutInputA,
ElementInputB, LayoutInputB,
@@ -193,6 +292,7 @@ using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFprop<
IteratorAlgorithm
>::Kernel;
// Type of the actual kernel
using ImplicitGemm = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -230,7 +330,7 @@ struct Options {
beta(0),
benchmark(false) { }
// Verify the problem size is compatible with the CUTLASS Convolution implementation.
// Verify that the problem size is compatible with CUTLASS's convolution implementation
bool valid() {
//
@@ -256,7 +356,7 @@ struct Options {
return true;
}
/// Updates input and filter sizes
/// Update input and filter sizes
void update(
cutlass::Tensor4DCoord input_size,
cutlass::Tensor4DCoord filter_size) {
@@ -270,7 +370,7 @@ struct Options {
padding.c() = filter_size.w() / 2;
}
// Parses the command line
// Parse command-line arguments
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
@@ -302,11 +402,11 @@ struct Options {
cmd.get_cmd_line_argument("k", filter_size.n());
cmd.get_cmd_line_argument("r", filter_size.h());
cmd.get_cmd_line_argument("s", filter_size.w());
filter_size.c() = input_size.c();
filter_size.c() = input_size.c();
cmd.get_cmd_line_argument("alpha", alpha);
cmd.get_cmd_line_argument("beta", beta);
cmd.get_cmd_line_argument("iterations", iterations);
cmd.get_cmd_line_argument("tag", tag);
@@ -320,12 +420,12 @@ struct Options {
}
}
/// Prints the usage statement.
/// Print an explanation of the command-line arguments
std::ostream & print_usage(std::ostream &out) const {
out << "16_ampere_tensorop_conv2dfprop example\n\n"
<< " This example uses Ampere's Tensor Core operators on F16 data types to compute\n"
<< " forward convolution on tensors of layout NHWC.\n\n"
<< " This example uses Ampere's Tensor Core operators on F16 data types\n"
<< " to compute forward convolution on tensors of layout NHWC.\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement.\n\n"
<< " --n=<int> Input tensor extent N\n"
@@ -350,7 +450,7 @@ struct Options {
return out;
}
/// Computes the output tensor size (NPQK)
cutlass::Tensor4DCoord output_size() const {
return cutlass::Tensor4DCoord(
@@ -360,19 +460,20 @@ struct Options {
filter_size.n());
}
/// Compute performance in GFLOP/s
/// Compute performance in Gflop/s
///
/// Gflop/s stands for billions (10^9) of
/// floating-point operations per second (Gflop/s).
double gflops(double runtime_s) const {
// Number of multiply-adds = NPQK * CRS
int64_t fmas = output_size().product() * int64_t(filter_size.h() * filter_size.w() * filter_size.c());
// Two flops per multiply-add
return 2.0 * double(fmas) / double(1.0e9) / runtime_s;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
struct Result {
double runtime_ms;
double gflops;
@@ -380,14 +481,14 @@ struct Result {
cutlass::Status reference_check;
cudaError_t error;
Result():
runtime_ms(0),
Result():
runtime_ms(0),
gflops(0),
status(cutlass::Status::kSuccess),
reference_check(cutlass::Status::kInvalid),
error(cudaSuccess) { }
static std::ostream & print_header(std::ostream &out, Options const &options) {
static std::ostream& print_header(std::ostream &out, Options const &options) {
if (!options.tag.empty()) {
out << "Name,";
@@ -404,7 +505,7 @@ struct Result {
out << options.tag << ",";
}
out
out
<< "conv_" << idx << ","
<< options.input_size.n() << ","
<< options.input_size.h() << ","
@@ -420,8 +521,6 @@ struct Result {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Runs one benchmark
Result profile_convolution(Options const &options) {
@@ -441,7 +540,7 @@ Result profile_convolution(Options const &options) {
// Initialize tensors
//
// Fill tensor A on host with uniform-distribution random data
// Fill tensor A on host with uniformly distributed random data
cutlass::reference::host::TensorFillRandomUniform(
tensor_a.host_view(),
1,
@@ -449,7 +548,7 @@ Result profile_convolution(Options const &options) {
ElementInputA(-8),
0);
// Fill tensor B on host with uniform-distribution random data
// Fill tensor B on host with uniformly distributed random data
cutlass::reference::host::TensorFillRandomUniform(
tensor_b.host_view(),
1,
@@ -457,7 +556,7 @@ Result profile_convolution(Options const &options) {
ElementInputB(-8),
0);
// Fill tensor C on host with uniform-distribution random data
// Fill tensor C on host with uniformly distributed random data
cutlass::reference::host::TensorFillRandomUniform(
tensor_c.host_view(),
1,
@@ -490,7 +589,7 @@ Result profile_convolution(Options const &options) {
int split_k_slices = 1;
// Construct Conv2dProblemSize with user defined output size
cutlass::conv::Conv2dProblemSize problem_size(
cutlass::conv::Conv2dProblemSize problem_size(
options.input_size,
options.filter_size,
options.padding,
@@ -501,7 +600,7 @@ Result profile_convolution(Options const &options) {
split_k_slices
);
// Construct ImplicitGemm::Argument structure with conv2d
// Construct ImplicitGemm::Argument structure with conv2d
// problem size, data pointers, and epilogue values
typename ImplicitGemm::Arguments arguments{
problem_size,
@@ -539,7 +638,7 @@ Result profile_convolution(Options const &options) {
//
// Optional reference check
//
if (options.reference_check) {
std::cout << "Verification on host...\n";
@@ -552,8 +651,7 @@ Result profile_convolution(Options const &options) {
ElementOutput,
LayoutOutput,
ElementComputeEpilogue,
ElementAccumulator,
cutlass::NumericConverter<ElementOutput, ElementComputeEpilogue>
ElementAccumulator
>(
problem_size,
tensor_a.host_ref(),
@@ -564,7 +662,7 @@ Result profile_convolution(Options const &options) {
options.beta
);
// Check if output from CUTLASS kernel and reference kernel are equal or not
// Check if CUTLASS kernel and reference kernel produced the same output
tensor_d.sync_host();
bool passed = cutlass::reference::host::TensorEquals(
@@ -589,14 +687,14 @@ Result profile_convolution(Options const &options) {
std::stringstream ss;
ss << "16_ampere_workspace_conv2dfprop_"
<< options.input_size.n() << "x" << options.input_size.h() << "x" << options.input_size.w() << "x" << options.input_size.c()
<< options.input_size.n() << "x" << options.input_size.h() << "x" << options.input_size.w() << "x" << options.input_size.c()
<< "_"
<< options.filter_size.n() << "x" << options.filter_size.h() << "x" << options.filter_size.w() << "x" << options.filter_size.c()
<< options.filter_size.n() << "x" << options.filter_size.h() << "x" << options.filter_size.w() << "x" << options.filter_size.c()
<< ".dat";
std::ofstream output_workspace(ss.str());
output_workspace
output_workspace
<< "Input = \n" << tensor_a.host_view() << "\n\n"
<< "Filters = \n" << tensor_b.host_view() << "\n\n";
@@ -616,7 +714,7 @@ Result profile_convolution(Options const &options) {
if (options.measure_performance) {
cudaEvent_t events[2];
for (auto & event : events) {
result.error = cudaEventCreate(&event);
if (result.error != cudaSuccess) {
@@ -632,7 +730,7 @@ Result profile_convolution(Options const &options) {
return result;
}
// Launch a sequence of implicit GEMM operations on the device
// Launch a sequence of implicit GEMM operations on the device.
for (int iteration = 0; iteration < options.iterations; ++iteration) {
result.status = implicit_gemm_op();
CUTLASS_CHECK(result.status);
@@ -652,7 +750,7 @@ Result profile_convolution(Options const &options) {
return result;
}
// Measure elapsed runtime
// Measure elapsed runtime.
float runtime_ms = 0;
result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]);
if (result.error != cudaSuccess) {
@@ -660,7 +758,7 @@ Result profile_convolution(Options const &options) {
return result;
}
// Print average runtime and GFLOPs.
// Print average run time and floating-point throughput (Gflop/s).
result.runtime_ms = double(runtime_ms) / double(options.iterations);
result.gflops = options.gflops(result.runtime_ms / 1000.0);
@@ -673,8 +771,6 @@ Result profile_convolution(Options const &options) {
return result;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char const **args) {
bool notSupported = false;
@@ -701,7 +797,7 @@ int main(int argc, char const **args) {
}
Options options;
options.parse(argc, args);
if (options.help) {
@@ -768,5 +864,3 @@ int main(int argc, char const **args) {
return 0;
}
/////////////////////////////////////////////////////////////////////////////////////////////////

3
examples/22_quaternion_conv/quaternion_conv.cu
View File

@@ -470,8 +470,7 @@ Result profile_convolution(Options const &options) {
ElementOutput,
LayoutOutput,
ElementComputeEpilogue,
ElementAccumulator,
cutlass::NumericConverter<ElementOutput, ElementComputeEpilogue>
ElementAccumulator
>(
problem_size,
tensor_a.host_ref(),

68
examples/24_gemm_grouped/gemm_grouped.cu
View File

@@ -37,7 +37,7 @@
leading dimensions and problem sizes are stored in arrays in GMEM.
This differs from "Batched Array" GEMM because the size of each GEMM problem in the Grouped GEMM
concept may be distinct.
concept may be distinct.
This benchmark program initializes a workspace with random problem sizes for a given number of
groups. Command line options enable overriding M, N, and/or K dimensions with uniform values to
@@ -186,7 +186,7 @@ struct Options {
//
// Methods
//
//
Options():
help(false),
@@ -216,7 +216,7 @@ struct Options {
cmd.get_cmd_line_argument("alignment", alignment, 8);
cmd.get_cmd_line_argument("groups", problem_count, 15);
cmd.get_cmd_line_argument("alpha", alpha, 1.0f);
cmd.get_cmd_line_argument("beta", beta, 0.0f);
cmd.get_cmd_line_argument("beta", beta, 0.0f);
cmd.get_cmd_line_argument("iterations", iterations, 20);
cmd.get_cmd_line_argument("streams", cuda_streams, 0);
cmd.get_cmd_line_argument("verbose", verbose, false);
@@ -455,13 +455,13 @@ struct Options {
/// Compute performance in GFLOP/s
double gflops(double runtime_s) const {
// Number of real-valued multiply-adds
// Number of real-valued multiply-adds
int64_t fmas = int64_t();
for (auto const & problem : problem_sizes) {
fmas += problem.product();
}
// Two flops per multiply-add
return 2.0 * double(fmas) / double(1.0e9) / runtime_s;
}
@@ -546,7 +546,7 @@ public:
template <typename Element>
void initialize_tensor(
Element *ptr,
size_t capacity,
size_t capacity,
cutlass::Distribution::Kind dist_kind,
uint32_t seed) {
@@ -578,7 +578,7 @@ public:
cutlass::reference::device::BlockFillRandomUniform(
ptr, capacity, seed, scope_max, scope_min, 0);
}
}
else if (dist_kind == cutlass::Distribution::Gaussian) {
cutlass::reference::device::BlockFillRandomGaussian(
@@ -589,7 +589,7 @@ public:
// Fill with increasing elements
cutlass::reference::device::BlockFillSequential(
ptr, capacity, Element(1), Element());
}
}
else {
// Fill with all 1s
@@ -674,13 +674,13 @@ public:
ptr_A.reset(problem_count());
ptr_A.copy_from_host(ptr_A_host.data());
ptr_B.reset(problem_count());
ptr_B.copy_from_host(ptr_B_host.data());
ptr_C.reset(problem_count());
ptr_C.copy_from_host(ptr_C_host.data());
ptr_D.reset(problem_count());
ptr_D.copy_from_host(ptr_D_host.data());
@@ -712,7 +712,7 @@ public:
MatrixCoord extent_A{problem.m(), problem.k()};
MatrixCoord extent_B{problem.k(), problem.n()};
MatrixCoord extent_C{problem.m(), problem.n()};
cutlass::TensorView<ElementA, LayoutA> view_A(block_A.get() + offset_A.at(i), layout_A, extent_A);
cutlass::TensorView<ElementB, LayoutB> view_B(block_B.get() + offset_B.at(i), layout_B, extent_B);
cutlass::TensorView<ElementC, LayoutC> view_C(block_C.get() + offset_C.at(i), layout_C, extent_C);
@@ -724,18 +724,18 @@ public:
cutlass::reference::device::GemmComplex<
ElementA, LayoutA,
ElementB, LayoutB,
ElementC, LayoutC,
ElementC, LayoutC,
ElementCompute, ElementAccumulator
>(
problem,
options.alpha,
options.alpha,
view_A,
Gemm::kTransformA,
view_B,
Gemm::kTransformB,
options.beta,
view_C,
view_Ref_device,
options.beta,
view_C,
view_Ref_device,
ElementAccumulator(0)
);
@@ -781,8 +781,8 @@ public:
std::cout << "Conventionally executed as " << this->options.problem_bins.size() << " batched GEMMs:\n";
for (auto const & bin : this->options.problem_bins) {
std::cout << " [" << bin_idx << "]: "
<< bin.first.m() << "-by-" << bin.first.n() << "-by-" << bin.first.k()
std::cout << " [" << bin_idx << "]: "
<< bin.first.m() << "-by-" << bin.first.n() << "-by-" << bin.first.k()
<< ", batch count: " << bin.second.size() << "\n";
++bin_idx;
@@ -832,7 +832,7 @@ public:
for (auto const & bin : this->options.problem_bins) {
int first_idx = bin.second.front();
bin_problem_sizes.push_back(this->options.problem_sizes.at(first_idx));
bin_count.push_back(int32_t(bin.second.size()));
@@ -974,7 +974,7 @@ public:
std::cerr << "CUTLASS error on line " << __LINE__ << std::endl;
return result;
}
}
//
@@ -1027,7 +1027,7 @@ public:
int last_stream_idx = 0;
for (int iter = 0; iter < this->options.iterations; ++iter) {
for (int bin_idx = 0; bin_idx < int32_t(bin_problem_sizes.size()); ++bin_idx) {
cutlass::gemm::GemmCoord const & problem = bin_problem_sizes[bin_idx];
@@ -1098,7 +1098,7 @@ public:
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl;
return result;
}
//
// Wait for work to be completed
//
@@ -1129,10 +1129,10 @@ public:
for (auto event : events) {
(void)cudaEventDestroy(event);
}
for (auto stream : cuda_streams) {
if (stream) {
(void)cudaStreamDestroy(stream);
(void)cudaStreamDestroy(stream);
}
}
@@ -1203,8 +1203,8 @@ public:
int tiles = Gemm::problem_tile_count(problem);
total_tiles += tiles;
std::cout << " [" << idx << "]: "
<< problem.m() << "-by-" << problem.n() << "-by-" << problem.k()
std::cout << " [" << idx << "]: "
<< problem.m() << "-by-" << problem.n() << "-by-" << problem.k()
<< " (" << tiles << " threadblock tiles)" << "\n";
++idx;
@@ -1442,12 +1442,12 @@ int main(int argc, char const **args) {
}
if (__CUDACC_VER_MAJOR__ < 11 || props.major < 8) {
//
// This example requires an NVIDIA Ampere-architecture GPU.
//
std::cout
std::cout
<< "CUTLASS's Grouped GEMM example requires a GPU of NVIDIA's Ampere Architecture or "
<< "later (compute capability 80 or greater).\n";
@@ -1497,9 +1497,9 @@ int main(int argc, char const **args) {
cutlass::gemm::GemmShape<64, 64, 32>,
cutlass::gemm::GemmShape<16, 8, 16>,
cutlass::epilogue::thread::LinearCombination<
ElementOutput,
ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
ElementAccumulator,
ElementAccumulator,
ElementAccumulator
>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,
@@ -1519,8 +1519,8 @@ int main(int argc, char const **args) {
cutlass::ComplexTransform::kNone,
8,
ElementOutput, LayoutC,
ElementAccumulator,
cutlass::arch::OpClassTensorOp,
ElementAccumulator,
cutlass::arch::OpClassTensorOp,
cutlass::arch::Sm80,
cutlass::gemm::GemmShape<128, 128, 32>,
cutlass::gemm::GemmShape<64, 64, 32>,
@@ -1531,7 +1531,7 @@ int main(int argc, char const **args) {
// NOTE: Threadblock swizzling is currently not supported by CUTLASS's grouped kernels.
// This parameter is passed in at present to match the APIs of other kernels. The parameter
// is unused within the kernel.
cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
4>::GemmKernel;
using GemmGrouped = cutlass::gemm::device::GemmGrouped<GemmKernel>;

2
examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu → examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm.cu
View File

@@ -181,7 +181,7 @@ struct Options {
<< " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n\n";
out << "\n\nExamples:\n\n"
<< "$ ./examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_complex_gemm --m=1024 --n=512 \\\n"
<< "$ ./examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm --m=1024 --n=512 \\\n"
<< " --alpha=2 --beta=0.707 \n\n";
return out;

6
examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/CMakeLists.txt
View File

@@ -27,9 +27,9 @@
# 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.
# Both filenames are shorter to avoid MAX_PATH issues on Windows.
cutlass_example_add_executable(
29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm
29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu
29_3xtf32_complex_gemm
29_3xtf32_complex_gemm.cu
)

1478
examples/39_gemm_permute/gemm_permute.cu
View File

File diff suppressed because it is too large Load Diff

510
examples/39_gemm_permute/layouts.h Normal file
View File

@@ -0,0 +1,510 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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.
*
**************************************************************************************************/
/*! \file
\brief Defines additional layout functions used in Permute GEMM example to simplify
computing reference permutations of 4/5D tensors when source data is column-major.
*/
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/cassert>
#else
#include "assert.h"
#endif
#include "cutlass/cutlass.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/coord.h"
#include "cutlass/tensor_coord.h"
namespace cutlass {
namespace layout {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Mapping function for 4-D CWHN tensors.
class TensorCWHN {
public:
/// Logical rank of tensor
static int const kRank = 4;
/// Rank of stride vector
static int const kStrideRank = 3;
/// Index type used for coordinates
using Index = int32_t;
/// Long index type used for offsets
using LongIndex = int64_t;
/// Logical coordinate (n, h, w, c)
using TensorCoord = Tensor4DCoord;
/// Stride vector
using Stride = Coord<kStrideRank>;
private:
//
// Data members
//
/// Stride data member - [n, hn, whn]
Stride stride_;
public:
//
// Methods
//
/// Constructor
CUTLASS_HOST_DEVICE
TensorCWHN(Stride const &stride = Stride(0)): stride_(stride) { }
/// Constructor
CUTLASS_HOST_DEVICE
TensorCWHN(
typename Stride::Index stride_h, ///< number of elements between adjacent N coordinates
typename Stride::Index stride_w, ///< number of elements between adjacent C coordinates
typename Stride::Index stride_c ///< number of elements between adjacent W coordinates
):
stride_(make_Coord(stride_h, stride_w, stride_c)) { }
/// Constructor
// Once convolutions implement 64b stride this ctor can be deleted
CUTLASS_HOST_DEVICE
TensorCWHN(Coord<kStrideRank, LongIndex> const &stride):
stride_(make_Coord(
static_cast<typename Stride::Index>(stride[0]),
static_cast<typename Stride::Index>(stride[1]),
static_cast<typename Stride::Index>(stride[2]))
) { }
/// Helper returns a layout to a tightly packed WCNH tensor.
CUTLASS_HOST_DEVICE
static TensorCWHN packed(TensorCoord const &extent) {
return TensorCWHN(
make_Coord(
extent.n(),
extent.h() * extent.n(),
extent.w() * extent.h() * extent.n()
)
);
}
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(TensorCoord const &coord) const {
return coord.n() +
LongIndex(stride_[0] * coord.h()) +
LongIndex(stride_[1] * coord.w()) +
LongIndex(stride_[2] * coord.c());
}
/// Returns the offset of a pitchlinear coordinate in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(PitchLinearCoord coord) const {
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
/// Compute the number of contiguous elements needed to store a tensor with the given size
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &extent) const {
// it does not make sense if the extent is larger than stride
// and we could not rely on the capacity calculation in such cases
// we could move this checkers to debug code only
if ((extent.n() > stride_[0])
|| (extent.h() * stride_[0] > stride_[1])
|| (extent.w() * stride_[1] > stride_[2])) {
assert(0);
}
return extent.c() * stride_[2];
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Mapping function for 4-D NHCW tensors.
class TensorNHCW {
public:
/// Logical rank of tensor
static int const kRank = 4;
/// Rank of stride vector
static int const kStrideRank = 3;
/// Index type used for coordinates
using Index = int32_t;
/// Long index type used for offsets
using LongIndex = int64_t;
/// Logical coordinate (n, h, w, c)
using TensorCoord = Tensor4DCoord;
/// Stride vector
using Stride = Coord<kStrideRank>;
private:
//
// Data members
//
/// Stride data member - [w, cw, hcw]
Stride stride_;
public:
//
// Methods
//
/// Constructor
CUTLASS_HOST_DEVICE
TensorNHCW(Stride const &stride = Stride(0)): stride_(stride) { }
/// Constructor
CUTLASS_HOST_DEVICE
TensorNHCW(
typename Stride::Index stride_c, ///< number of elements between adjacent C coordinates
typename Stride::Index stride_h, ///< number of elements between adjacent H coordinates
typename Stride::Index stride_n ///< number of elements between adjacent N coordinates
):
stride_(make_Coord(stride_c, stride_h, stride_n)) { }
/// Constructor
// Once convolutions implement 64b stride this ctor can be deleted
CUTLASS_HOST_DEVICE
TensorNHCW(Coord<kStrideRank, LongIndex> const &stride):
stride_(make_Coord(
static_cast<typename Stride::Index>(stride[0]),
static_cast<typename Stride::Index>(stride[1]),
static_cast<typename Stride::Index>(stride[2]))
) { }
/// Helper returns a layout to a tightly packed WCNH tensor.
CUTLASS_HOST_DEVICE
static TensorNHCW packed(TensorCoord const &extent) {
return TensorNHCW(
make_Coord(
extent.w(),
extent.c() * extent.w(),
extent.h() * extent.c() * extent.w()
)
);
}
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(TensorCoord const &coord) const {
return coord.w() +
LongIndex(stride_[0] * coord.c()) +
LongIndex(stride_[1] * coord.h()) +
LongIndex(stride_[2] * coord.n());
}
/// Returns the offset of a pitchlinear coordinate in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(PitchLinearCoord coord) const {
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
/// Compute the number of contiguous elements needed to store a tensor with the given size
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &extent) const {
// it does not make sense if the extent is larger than stride
// and we could not rely on the capacity calculation in such cases
// we could move this checkers to debug code only
if ((extent.w() > stride_[0])
|| (extent.c() * stride_[0] > stride_[1])
|| (extent.h() * stride_[1] > stride_[2])) {
assert(0);
}
return extent.n() * stride_[2];
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Mapping function for 4-D NHCW tensors.
class TensorNCWH {
public:
/// Logical rank of tensor
static int const kRank = 4;
/// Rank of stride vector
static int const kStrideRank = 3;
/// Index type used for coordinates
using Index = int32_t;
/// Long index type used for offsets
using LongIndex = int64_t;
/// Logical coordinate (n, h, w, c)
using TensorCoord = Tensor4DCoord;
/// Stride vector
using Stride = Coord<kStrideRank>;
private:
//
// Data members
//
/// Stride data member - [h, wh, cwh]
Stride stride_;
public:
//
// Methods
//
/// Constructor
CUTLASS_HOST_DEVICE
TensorNCWH(Stride const &stride = Stride(0)): stride_(stride) { }
/// Constructor
CUTLASS_HOST_DEVICE
TensorNCWH(
typename Stride::Index stride_w, ///< number of elements between adjacent C coordinates
typename Stride::Index stride_c, ///< number of elements between adjacent H coordinates
typename Stride::Index stride_n ///< number of elements between adjacent N coordinates
):
stride_(make_Coord(stride_w, stride_c, stride_n)) { }
/// Constructor
// Once convolutions implement 64b stride this ctor can be deleted
CUTLASS_HOST_DEVICE
TensorNCWH(Coord<kStrideRank, LongIndex> const &stride):
stride_(make_Coord(
static_cast<typename Stride::Index>(stride[0]),
static_cast<typename Stride::Index>(stride[1]),
static_cast<typename Stride::Index>(stride[2]))
) { }
/// Helper returns a layout to a tightly packed WCNH tensor.
CUTLASS_HOST_DEVICE
static TensorNCWH packed(TensorCoord const &extent) {
return TensorNCWH(
make_Coord(
extent.h(),
extent.w() * extent.h(),
extent.c() * extent.w() * extent.h()
)
);
}
/// Returns the offset of a coordinate (n, h, w, c) in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(TensorCoord const &coord) const {
return coord.h() +
LongIndex(stride_[0] * coord.w()) +
LongIndex(stride_[1] * coord.c()) +
LongIndex(stride_[2] * coord.n());
}
/// Returns the offset of a pitchlinear coordinate in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(PitchLinearCoord coord) const {
return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
/// Compute the number of contiguous elements needed to store a tensor with the given size
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &extent) const {
// it does not make sense if the extent is larger than stride
// and we could not rely on the capacity calculation in such cases
// we could move this checkers to debug code only
if ((extent.h() > stride_[0])
|| (extent.w() * stride_[0] > stride_[1])
|| (extent.c() * stride_[1] > stride_[2])) {
assert(0);
}
return extent.n() * stride_[2];
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Mapping function for 5-D CWHDN tensors.
class TensorCWHDN {
public:
/// Logical rank of tensor
static int const kRank = 5;
/// Rank of stride vector
static int const kStrideRank = 4;
/// Index type used for coordinates
using Index = int32_t;
/// Long index type used for offsets
using LongIndex = int64_t;
/// Logical coordinate (n, d, h, w, c)
using TensorCoord = Tensor5DCoord;
/// Stride vector
using Stride = Coord<kStrideRank>;
private:
//
// Data members
//
/// Stride data member - [n, dn, hdn, whdn]
Stride stride_;
public:
//
// Methods
//
/// Constructor
CUTLASS_HOST_DEVICE
TensorCWHDN(Stride const &stride = Stride(0)): stride_(stride) { }
/// Constructor
CUTLASS_HOST_DEVICE
TensorCWHDN(
typename Stride::Index n,
typename Stride::Index dn,
typename Stride::Index hdn,
typename Stride::Index whdn):
stride_(make_Coord(n, dn, hdn, whdn)) { }
/// Constructor
// Once convolutions implement 64b stride this ctor can be deleted
CUTLASS_HOST_DEVICE
TensorCWHDN(Coord<kStrideRank, LongIndex> const &stride):
stride_(make_Coord(
static_cast<typename Stride::Index>(stride[0]),
static_cast<typename Stride::Index>(stride[1]),
static_cast<typename Stride::Index>(stride[2]),
static_cast<typename Stride::Index>(stride[3]))
) { }
/// Helper returns a layout to a tightly packed CWHDN tensor.
CUTLASS_HOST_DEVICE
static TensorCWHDN packed(TensorCoord const &extent) {
return TensorCWHDN(
make_Coord(
extent.n(),
extent.d() * extent.n(),
extent.h() * extent.d() * extent.n(),
extent.w() * extent.h() * extent.d() * extent.n()
)
);
}
/// Returns the offset of a coordinate (n, d, h, w, c) in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(TensorCoord const &coord) const {
return coord.n() +
LongIndex(stride_[0] * coord.d()) +
LongIndex(stride_[1] * coord.h()) +
LongIndex(stride_[2] * coord.w()) +
LongIndex(stride_[3] * coord.c());
}
/// Returns the offset of a pitchlinear coordinate in linear memory.
CUTLASS_HOST_DEVICE
LongIndex operator()(PitchLinearCoord coord) const {
return coord.contiguous() + LongIndex(coord.strided() * stride_[3]);
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride stride() const {
return stride_;
}
/// Returns the stride of the layout
CUTLASS_HOST_DEVICE
Stride & stride() {
return stride_;
}
/// Compute the number of contiguous elements needed to store a tensor with the given size
CUTLASS_HOST_DEVICE
LongIndex capacity(TensorCoord const &extent) const {
// it does not make sense if the extent is larger than stride
// and we could not rely on the capacity calculation in such cases
// we could move this checkers to debug code only
if ((extent.n() > stride_[0])
|| (extent.d() * stride_[0] > stride_[1])
|| (extent.h() * stride_[1] > stride_[2])
|| (extent.w() * stride_[2] > stride_[3])) {
assert(0);
}
return extent.c() * stride_[3];
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace layout
} // namespace cutlass

344
examples/39_gemm_permute/permute_info.h Normal file
View File

@@ -0,0 +1,344 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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.
*
**************************************************************************************************/
/*! \file
\brief Contains additional metadata about layout permute functions used in the example.
*/
#include "cutlass/tensor_coord.h"
#include "cutlass/layout/permute.h"
/// Additional permutation metadata to facilitate testing/printing
template<typename PermuteLayout>
struct PermuteInfo;
/// Specialization for default case (no permute). Other specializations must follow this template.
template<>
struct PermuteInfo<cutlass::layout::NoPermute> {
/// Whether this is a BMM or GEMM permutation (NoPermute can actually be either)
static bool constexpr kBatched = false;
/// Minimal divisor for row extent
static int constexpr kRowFactor = 1;
/// Minimum divisor for column extent
static int constexpr kColumnFactor = 1;
/// Minimum divisor for batch size dimension
static int constexpr kBatchFactor = 1;
/// Tensor layout used in permutation operation
using Layout = cutlass::layout::PackedVectorLayout;
static std::string name() {
return "NoPermute";
}
/// User-friendly description of the permute operation
static std::string desc() {
return "no permutation";
}
/// Infer original higher-rank tensor shape from GEMM/BMM matrix extents.
/// For direct (output) permutations, must be a simple reshape of extent.
/// For inverse (input) permutations, must return shape *before* permute operation.
/// In case of NoPermute, simply use a linear (rank 1) view of the memory
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
return Layout::TensorCoord(extent.row() * extent.column() * batch_count);
}
/// Compute the permuted higher-rank tensor shape from the original shape.
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
return s;
}
};
template<int D1>
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>> {
static bool constexpr kBatched = true;
static int constexpr kRowFactor = 1;
static int constexpr kColumnFactor = 1;
static int constexpr kBatchFactor = D1;
using Layout = cutlass::layout::TensorNHWC;
static std::string name() {
return "Tensor4DPermuteBMM0213<" + std::to_string(D1) + ">";
}
static std::string desc() {
return "batched GEMM permutation [0, 2, 1, 3]";
}
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = batch_count / D1;
int D2 = extent.row();
int D3 = extent.column();
return {D0, D1, D2, D3};
}
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
return {s[0], s[2], s[1], s[3]};
}
};
template<int D1>
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajorInverse<D1>>
: public PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>> {
static bool constexpr kBatched = true;
static int constexpr kRowFactor = 1;
static int constexpr kColumnFactor = D1;
static int constexpr kBatchFactor = 1;
using Base = PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0213RowMajor<D1>>;
using Layout = typename Base::Layout;
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = batch_count;
int D2 = extent.row();
int D3 = extent.column() / D1;
return {D0, D1, D2, D3};
}
};
template<int D1>
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>> {
static bool constexpr kBatched = true;
static int constexpr kRowFactor = 1;
static int constexpr kColumnFactor = 1;
static int constexpr kBatchFactor = D1;
using Layout = cutlass::layout::TensorNHCW;
static std::string name() {
return "Tensor4DPermuteBMM0321<" + std::to_string(D1) + ">";
}
static std::string desc() {
return "batched GEMM permutation [0, 3, 2, 1]";
}
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = batch_count / D1;
int D2 = extent.row();
int D3 = extent.column();
return {D0, D1, D2, D3};
}
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
return {s[0], s[3], s[2], s[1]};
}
};
template<int D1>
struct PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajorInverse<D1>>
: public PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>> {
static bool constexpr kBatched = true;
static int constexpr kRowFactor = D1;
static int constexpr kColumnFactor = 1;
static int constexpr kBatchFactor = 1;
using Base = PermuteInfo<cutlass::layout::Tensor4DPermuteBMM0321ColumnMajor<D1>>;
using Layout = typename Base::Layout;
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = batch_count;
int D2 = extent.row() / D1;
int D3 = extent.column();
return {D0, D1, D2, D3};
}
};
template<int D1, int D2>
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = D1;
static int constexpr kColumnFactor = D2;
static int constexpr kBatchFactor = 1;
using Layout = cutlass::layout::TensorNHWC;
static std::string name() {
return "Tensor4DPermute0213<" + std::to_string(D1) + "," + std::to_string(D2) + ">";
}
static std::string desc() {
return "normal GEMM permutation [0, 2, 1, 3]";
}
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = extent.row() / D1;
int D3 = extent.column() / D2;
return {D0, D1, D2, D3};
}
static Layout::TensorCoord permute(Layout::TensorCoord const &s) {
return {s[0], s[2], s[1], s[3]};
}
};
template<int D1, int D2>
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajorInverse<D1, D2>>
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = D2;
static int constexpr kColumnFactor = D1;
static int constexpr kBatchFactor = 1;
using Base = PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>>;
using Layout = typename Base::Layout;
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int D0 = extent.row() / D2;
int D3 = extent.column() / D1;
return {D0, D1, D2, D3};
}
};
template<int D1, int D2>
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213ColumnMajor<D1, D2>>
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajor<D1, D2>> {
using Layout = cutlass::layout::TensorCWHN;
};
template<int D1, int D2>
struct PermuteInfo<cutlass::layout::Tensor4DPermute0213ColumnMajorInverse<D1, D2>>
: public PermuteInfo<cutlass::layout::Tensor4DPermute0213RowMajorInverse<D1, D2>> {
using Layout = cutlass::layout::TensorCWHN;
};
template<int T1, int T2, int T3>
struct PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = T1;
static int constexpr kColumnFactor = T2 * T3;
static int constexpr kBatchFactor = 1;
using Layout = cutlass::layout::TensorNDHWC;
static std::string name() {
return "Tensor5DPermute20314<" + std::to_string(T1) + "," + std::to_string(T2) + "," + std::to_string(T3) + ">";
}
static std::string desc() {
return "normal GEMM permutation [2, 0, 3, 1, 4]";
}
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count)
{
int const T0 = extent.row() / T1;
int const T4 = extent.column() / (T2 * T3);
return {T0, T1, T2, T3, T4};
}
static Layout::TensorCoord permute(Layout::TensorCoord const &s)
{
return {s[2], s[0], s[3], s[1], s[4]};
}
};
template<int T1, int T2, int T3>
struct PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajorInverse<T1, T2, T3>>
: public PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = T2;
static int constexpr kColumnFactor = T1 * T3;
static int constexpr kBatchFactor = 1;
using Base = PermuteInfo<cutlass::layout::Tensor5DPermute20314RowMajor<T1, T2, T3>>;
using Layout = typename Base::Layout;
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int const T0 = extent.row() / T2;
int const T4 = extent.column() / (T1 * T3);
return {T0, T1, T2, T3, T4};
}
};
template<int T1, int T2, int T3>
struct PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = T1;
static int constexpr kColumnFactor = T2 * T3;
static int constexpr kBatchFactor = 1;
using Layout = cutlass::layout::TensorCWHDN;
static std::string name() {
return "Tensor5DPermute02413<" + std::to_string(T1) + "," + std::to_string(T2) + "," + std::to_string(T3) + ">";
}
static std::string desc() {
return "normal GEMM permutation [0, 2, 4, 1, 3]";
}
using Coord = cutlass::Tensor5DCoord;
static Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count)
{
int const T0 = extent.row() / T1;
int const T4 = extent.column() / (T2 * T3);
return {T0, T1, T2, T3, T4};
}
static Layout::TensorCoord permute(Layout::TensorCoord const &s)
{
return {s[0], s[2], s[4], s[1], s[3]};
}
};
template<int T1, int T2, int T3>
struct PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajorInverse<T1, T2, T3>>
: public PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>> {
static bool constexpr kBatched = false;
static int constexpr kRowFactor = T2;
static int constexpr kColumnFactor = T1 * T3;
static int constexpr kBatchFactor = 1;
using Base = PermuteInfo<cutlass::layout::Tensor5DPermute02413ColumnMajor<T1, T2, T3>>;
using Layout = typename Base::Layout;
static typename Layout::TensorCoord original_shape(cutlass::MatrixCoord extent, int batch_count) {
int const T0 = extent.row() / T2;
int const T4 = extent.column() / (T1 * T3);
return {T0, T1, T2, T3, T4};
}
};

11
examples/40_cutlass_py/README.md
View File

@@ -1,10 +1,15 @@
# CUTLASS Python Interface Examples
This directory contains examples of using CUTLASS's Python interface. It consists of two types of examples:
# PyCUTLASS Examples
**NOTE:** This directory contains examples for PyCUTLASS, a Python library providing low-level
building blocks for emitting CUTLASS C++ kernels. For examples using CUTLASS's Pythonic interface,
see the [examples/python](/examples/python) directory.
Two types of examples are provided:
* _Basic examples_: minimal examples that illustrate how to set up GEMMs, convolutions, and grouped GEMM operations
* [_Customizable examples_](customizable): examples that allow one to specify a variety of template parameters for the given kernel
## Setting up the Python interface
Please follow the instructions [here](/tools/library/scripts/pycutlass/README.md#installation) to set up the Python API.
Please follow the instructions [here](/python/README.md#installation) to set up the PyCUTLASS.
## Running examples
Each of the basic examples can be run as follows:

47
examples/40_cutlass_py/conv2d.py
View File

@@ -38,10 +38,11 @@ import torch
import numpy as np
import sys
import cutlass
import pycutlass
from pycutlass import *
from pycutlass.utils.device import device_cc
import cutlass_bindings
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.reference_model import Conv2dReferenceModule
from cutlass.backend.utils.device import device_cc
parser = argparse.ArgumentParser(
@@ -76,11 +77,11 @@ pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
pycutlass.compiler.nvcc()
# Set up A, B, C and accumulator
A = TensorDescription(cutlass.float16, cutlass.TensorNHWC, alignment)
B = TensorDescription(cutlass.float16, cutlass.TensorNHWC, alignment)
C = TensorDescription(cutlass.float32, cutlass.TensorNHWC, alignment)
element_acc = cutlass.float32
element_epilogue = cutlass.float32
A = TensorDescription(cutlass_bindings.float16, cutlass_bindings.TensorNHWC, alignment)
B = TensorDescription(cutlass_bindings.float16, cutlass_bindings.TensorNHWC, alignment)
C = TensorDescription(cutlass_bindings.float32, cutlass_bindings.TensorNHWC, alignment)
element_acc = cutlass_bindings.float32
element_epilogue = cutlass_bindings.float32
# Select instruction shape based on the Tensor Core instructions supported
# by the device on which we are running
@@ -89,12 +90,14 @@ if cc == 70:
elif cc == 75:
instruction_shape = [16, 8, 8]
else:
# Use CUTLASS kernels for CC 80 by default (e.g., for cases in which SM86 is used)
cc = 80
instruction_shape = [16, 8, 16]
math_inst = MathInstruction(
instruction_shape,
A.element, B.element, element_acc,
cutlass.OpClass.TensorOp,
cutlass_bindings.OpClass.TensorOp,
MathOperation.multiply_add
)
@@ -108,8 +111,8 @@ tile_description = TileDescription(
epilogue_functor = pycutlass.LinearCombination(C.element, C.alignment, element_acc, element_epilogue)
operation = Conv2dOperation(
conv_kind=cutlass.conv.Operator.fprop,
iterator_algorithm=cutlass.conv.IteratorAlgorithm.optimized,
conv_kind=cutlass_bindings.conv.Operator.fprop,
iterator_algorithm=cutlass_bindings.conv.IteratorAlgorithm.optimized,
arch=cc, tile_description=tile_description,
A=A, B=B, C=C, stride_support=StrideSupport.Strided,
epilogue_functor=epilogue_functor
@@ -125,20 +128,20 @@ pycutlass.compiler.add_module(operations)
# Randomly initialize tensors
problem_size = cutlass.conv.Conv2dProblemSize(
cutlass.Tensor4DCoord(args.n, args.h, args.c, args.w),
cutlass.Tensor4DCoord(args.k, args.r, args.s, args.c),
cutlass.Tensor4DCoord(0, 0, 0, 0), # Padding
cutlass.MatrixCoord(1, 1), # Strides
cutlass.MatrixCoord(1, 1), # Dilation
cutlass.conv.Mode.cross_correlation,
problem_size = cutlass_bindings.conv.Conv2dProblemSize(
cutlass_bindings.Tensor4DCoord(args.n, args.h, args.c, args.w),
cutlass_bindings.Tensor4DCoord(args.k, args.r, args.s, args.c),
cutlass_bindings.Tensor4DCoord(0, 0, 0, 0), # Padding
cutlass_bindings.MatrixCoord(1, 1), # Strides
cutlass_bindings.MatrixCoord(1, 1), # Dilation
cutlass_bindings.conv.Mode.cross_correlation,
1, # Split k slices
1 # Groups
)
tensor_A_size = cutlass.conv.implicit_gemm_tensor_a_size(operation.conv_kind, problem_size)
tensor_B_size = cutlass.conv.implicit_gemm_tensor_b_size(operation.conv_kind, problem_size)
tensor_C_size = cutlass.conv.implicit_gemm_tensor_c_size(operation.conv_kind, problem_size)
tensor_A_size = cutlass_bindings.conv.implicit_gemm_tensor_a_size(operation.conv_kind, problem_size)
tensor_B_size = cutlass_bindings.conv.implicit_gemm_tensor_b_size(operation.conv_kind, problem_size)
tensor_C_size = cutlass_bindings.conv.implicit_gemm_tensor_c_size(operation.conv_kind, problem_size)
tensor_A = torch.ceil(torch.empty(size=(tensor_A_size,), dtype=torch.float16, device="cuda").uniform_(-8.5, 7.5))
tensor_B = torch.ceil(torch.empty(size=(tensor_B_size,), dtype=torch.float16, device="cuda").uniform_(-8.5, 7.5))

66
examples/40_cutlass_py/customizable/conv2d.py
View File

@@ -30,11 +30,11 @@
#
################################################################################
import numpy as np
import pycutlass
from pycutlass import *
from pycutlass.conv2d_operation import *
from pycutlass.utils import reference_model
from pycutlass.utils.device import device_cc
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.device import device_cc
from cutlass.backend.conv2d_operation import *
from cutlass.backend.utils.reference_model import Conv2dReferenceModule
import sys
import torch.nn.functional as F
@@ -62,7 +62,7 @@ parser.add_argument("-tacc", "--element_acc", default="float32", type=str,
help='Data type of accumulator')
parser.add_argument('-m', "--math", default="multiply_add",
type=str, choices=["multiply_add", "multiply_add_fast_bf16", "multiply_add_fast_f32"], help="math instruction")
parser.add_argument('-op', "--opcode", default="simt", type=str,
parser.add_argument('-op', "--opcode", default="Simt", type=str,
choices=["Simt", 'TensorOp'],
help='This option describes whether you want to use tensor \
cores (TensorOp) or regular SIMT cores (Simt) on GPU SM')
@@ -156,12 +156,12 @@ pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
np.random.seed(0)
element_a = getattr(cutlass, args.element_a)
element_b = getattr(cutlass, args.element_b)
element_c = getattr(cutlass, args.element_c)
element_acc = getattr(cutlass, args.element_acc)
element_a = getattr(cutlass_bindings, args.element_a)
element_b = getattr(cutlass_bindings, args.element_b)
element_c = getattr(cutlass_bindings, args.element_c)
element_acc = getattr(cutlass_bindings, args.element_acc)
math_operation = getattr(MathOperation, args.math)
opclass = getattr(cutlass.OpClass, args.opcode)
opclass = getattr(cutlass_bindings.OpClass, args.opcode)
math_inst = MathInstruction(
args.instruction_shape, element_a, element_b,
@@ -173,9 +173,9 @@ tile_description = TileDescription(
math_inst
)
layout_a = getattr(cutlass, args.layout_a)
layout_b = getattr(cutlass, args.layout_b)
layout_c = getattr(cutlass, args.layout_c)
layout_a = getattr(cutlass_bindings, args.layout_a)
layout_b = getattr(cutlass_bindings, args.layout_b)
layout_c = getattr(cutlass_bindings, args.layout_c)
A = TensorDescription(
element_a, layout_a, args.alignment_a
@@ -189,7 +189,7 @@ C = TensorDescription(
element_c, layout_c, args.alignment_c
)
element_epilogue = getattr(cutlass, args.element_epilogue)
element_epilogue = getattr(cutlass_bindings, args.element_epilogue)
if (args.activation_function == "identity"
or (args.split_k_mode == "Parallel" and args.split_k_slices > 1)):
#
@@ -200,10 +200,10 @@ else:
getattr(pycutlass, args.activation_function)(element_epilogue),
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
iterator_algorithm = getattr(cutlass.conv.IteratorAlgorithm, args.iterator_algorithm)
swizzling_functor = getattr(cutlass, args.swizzling_functor)
iterator_algorithm = getattr(cutlass_bindings.conv.IteratorAlgorithm, args.iterator_algorithm)
swizzling_functor = getattr(cutlass_bindings, args.swizzling_functor)
stride_support = getattr(StrideSupport, args.stride_support)
conv_kind = getattr(cutlass.conv.Operator, args.conv_kind)
conv_kind = getattr(cutlass_bindings.conv.Operator, args.conv_kind)
operation = Conv2dOperation(
conv_kind=conv_kind, iterator_algorithm=iterator_algorithm,
@@ -226,7 +226,7 @@ if args.split_k_mode == "Parallel" and args.split_k_slices > 1:
getattr(pycutlass, args.activation_function)(element_epilogue),
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
reduction_operation = ReductionOperation(
shape=cutlass.MatrixCoord(4, 32 * C.alignment),
shape=cutlass_bindings.MatrixCoord(4, 32 * C.alignment),
C=C, element_accumulator=element_acc,
element_compute=element_epilogue,
epilogue_functor=epilogue_functor_reduction,
@@ -236,34 +236,34 @@ if args.split_k_mode == "Parallel" and args.split_k_slices > 1:
pycutlass.compiler.add_module(operations)
problem_size = cutlass.conv.Conv2dProblemSize(
cutlass.Tensor4DCoord(args.nhwc[0], args.nhwc[1], args.nhwc[2], args.nhwc[3]),
cutlass.Tensor4DCoord(args.krsc[0], args.krsc[1], args.krsc[2], args.krsc[3]),
cutlass.Tensor4DCoord(args.pad[0], args.pad[1], args.pad[2], args.pad[3]),
cutlass.MatrixCoord(args.stride[0], args.stride[1]),
cutlass.MatrixCoord(args.dilation[0], args.dilation[1]),
cutlass.conv.Mode.cross_correlation,
problem_size = cutlass_bindings.conv.Conv2dProblemSize(
cutlass_bindings.Tensor4DCoord(args.nhwc[0], args.nhwc[1], args.nhwc[2], args.nhwc[3]),
cutlass_bindings.Tensor4DCoord(args.krsc[0], args.krsc[1], args.krsc[2], args.krsc[3]),
cutlass_bindings.Tensor4DCoord(args.pad[0], args.pad[1], args.pad[2], args.pad[3]),
cutlass_bindings.MatrixCoord(args.stride[0], args.stride[1]),
cutlass_bindings.MatrixCoord(args.dilation[0], args.dilation[1]),
cutlass_bindings.conv.Mode.cross_correlation,
args.split_k_slices, 1
)
# User-provide inputs
tensor_A_size = cutlass.conv.implicit_gemm_tensor_a_size(
tensor_A_size = cutlass_bindings.conv.implicit_gemm_tensor_a_size(
conv_kind, problem_size
)
tensor_B_size = cutlass.conv.implicit_gemm_tensor_b_size(
tensor_B_size = cutlass_bindings.conv.implicit_gemm_tensor_b_size(
conv_kind, problem_size
)
if args.bias:
tensor_C_size = cutlass.conv.implicit_gemm_tensor_c_extent(
tensor_C_size = cutlass_bindings.conv.implicit_gemm_tensor_c_extent(
conv_kind, problem_size
).at(3)
else:
tensor_C_size = cutlass.conv.implicit_gemm_tensor_c_size(
tensor_C_size = cutlass_bindings.conv.implicit_gemm_tensor_c_size(
conv_kind, problem_size
)
tensor_D_size = cutlass.conv.implicit_gemm_tensor_c_size(
tensor_D_size = cutlass_bindings.conv.implicit_gemm_tensor_c_size(
conv_kind, problem_size
)
@@ -288,12 +288,12 @@ arguments = Conv2dArguments(
operation=operation, problem_size=problem_size, A=tensor_A,
B=tensor_B, C=tensor_C, D=tensor_D,
output_op = operation.epilogue_type(*([args.alpha, args.beta] + args.activation_args)),
split_k_mode=getattr(cutlass.conv.SplitKMode, args.split_k_mode),
split_k_mode=getattr(cutlass_bindings.conv.SplitKMode, args.split_k_mode),
split_k_slices=problem_size.split_k_slices
)
if args.split_k_mode == "Parallel" and args.split_k_slices > 1:
implicit_gemm_size = cutlass.conv.implicit_gemm_problem_size(conv_kind, arguments.problem_size)
implicit_gemm_size = cutlass_bindings.conv.implicit_gemm_problem_size(conv_kind, arguments.problem_size)
reduction_arguments = ReductionArguments(
reduction_operation,
problem_size=[implicit_gemm_size.m(), implicit_gemm_size.n()],

36
examples/40_cutlass_py/customizable/gemm.py
View File

@@ -30,10 +30,10 @@
#
################################################################################
import numpy as np
import pycutlass
from pycutlass import *
from pycutlass.utils.device import device_cc
import cutlass
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.device import device_cc
import cutlass_bindings
from bfloat16 import bfloat16
import sys
@@ -62,7 +62,7 @@ parser.add_argument("-tacc", "--element_acc", default="float32", type=str,
help='Data type of accumulator')
parser.add_argument('-m', "--math", default="multiply_add",
type=str, choices=["multiply_add", "multiply_add_fast_bf16", "multiply_add_fast_f32"], help="math instruction")
parser.add_argument('-op', "--opcode", default="simt", type=str,
parser.add_argument('-op', "--opcode", default="Simt", type=str,
choices=["Simt", 'TensorOp'],
help="This option describes whether you want to use tensor \
cores (TensorOp) or regular SIMT cores (Simt) on GPU SM")
@@ -147,12 +147,12 @@ pycutlass.compiler.nvcc()
np.random.seed(0)
element_a = getattr(cutlass, args.element_a)
element_b = getattr(cutlass, args.element_b)
element_c = getattr(cutlass, args.element_c)
element_acc = getattr(cutlass, args.element_acc)
element_a = getattr(cutlass_bindings, args.element_a)
element_b = getattr(cutlass_bindings, args.element_b)
element_c = getattr(cutlass_bindings, args.element_c)
element_acc = getattr(cutlass_bindings, args.element_acc)
math_operation = getattr(MathOperation, args.math)
opclass = getattr(cutlass.OpClass, args.opcode)
opclass = getattr(cutlass_bindings.OpClass, args.opcode)
math_inst = MathInstruction(
args.instruction_shape, element_a, element_b,
@@ -164,9 +164,9 @@ tile_description = TileDescription(
math_inst
)
layout_a = getattr(cutlass, args.layout_a)
layout_b = getattr(cutlass, args.layout_b)
layout_c = getattr(cutlass, args.layout_c)
layout_a = getattr(cutlass_bindings, args.layout_a)
layout_b = getattr(cutlass_bindings, args.layout_b)
layout_c = getattr(cutlass_bindings, args.layout_c)
A = TensorDescription(
element_a, layout_a, args.alignment_a
@@ -180,7 +180,7 @@ C = TensorDescription(
element_c, layout_c, args.alignment_c
)
element_epilogue = getattr(cutlass, args.element_epilogue)
element_epilogue = getattr(cutlass_bindings, args.element_epilogue)
if (args.activation_function == "identity"
or (args.gemm_mode == "GemmSplitKParallel" and args.split_k_slices > 1)):
#
@@ -191,7 +191,7 @@ else:
getattr(pycutlass, args.activation_function)(element_epilogue),
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
swizzling_functor = getattr(cutlass, args.swizzling_functor)
swizzling_functor = getattr(cutlass_bindings, args.swizzling_functor)
visitor = args.epilogue_visitor is not None
@@ -275,7 +275,7 @@ if args.gemm_mode == "GemmSplitKParallel":
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
reduction_operation = ReductionOperation(
shape=cutlass.MatrixCoord(4, 32 * C.alignment),
shape=cutlass_bindings.MatrixCoord(4, 32 * C.alignment),
C=C, element_accumulator=element_acc,
element_compute=element_epilogue,
epilogue_functor=epilogue_functor_reduction,
@@ -287,7 +287,7 @@ pycutlass.compiler.add_module(operations)
# User-provide inputs
problem_size = cutlass.gemm.GemmCoord(
problem_size = cutlass_bindings.gemm.GemmCoord(
args.problem_size[0], args.problem_size[1], args.problem_size[2])
tensor_a_size = args.batch * problem_size.m() * problem_size.k()
@@ -384,7 +384,7 @@ arguments = GemmArguments(
operation=operation, problem_size=problem_size,
A=tensor_A, B=tensor_B, C=tensor_C, D=tensor_D,
output_op=output_op,
gemm_mode=getattr(cutlass.gemm.Mode, args.gemm_mode),
gemm_mode=getattr(cutlass_bindings.gemm.Mode, args.gemm_mode),
split_k_slices=args.split_k_slices, batch=args.batch
)

32
examples/40_cutlass_py/customizable/gemm_grouped.py
View File

@@ -30,9 +30,9 @@
#
################################################################################
import numpy as np
import pycutlass
from pycutlass import *
from pycutlass.utils.device import device_cc
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.device import device_cc
import csv
import sys
@@ -61,7 +61,7 @@ parser.add_argument("-tacc", "--element_acc", default="float32", type=str,
help='Data type of accumulator')
parser.add_argument('-m', "--math", default="multiply_add",
type=str, choices=["multiply_add", "multiply_add_fast_bf16", "multiply_add_fast_f32"], help="math instruction")
parser.add_argument('-op', "--opcode", default="simt", type=str,
parser.add_argument('-op', "--opcode", default="Simt", type=str,
choices=["Simt", 'TensorOp'], help='This option describes whether you want to use tensor \
cores (TensorOp) or regular SIMT cores (Simt) on GPU SM')
# tile description
@@ -111,7 +111,7 @@ parser.add_argument("-pm", "--precompute_mode",
default="Device", type=str, choices=["Host", "Device"],
help="Grouped Gemm Scheduing on device only (Device) or using host precompute (Host)")
# arguments
parser.add_argument("-p", "--problem_size_dir", type=str,
parser.add_argument("-p", "--problem_size_dir", type=str, default="grouped_gemm_problem_size.csv",
help="path to the csv file contains the problem sizes")
parser.add_argument("-alpha", "--alpha", default=1.0, type=float, help="alpha")
parser.add_argument("-beta", "--beta", default=0.0, type=float, help="beta")
@@ -139,12 +139,12 @@ pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
np.random.seed(0)
element_a = getattr(cutlass, args.element_a)
element_b = getattr(cutlass, args.element_b)
element_c = getattr(cutlass, args.element_c)
element_acc = getattr(cutlass, args.element_acc)
element_a = getattr(cutlass_bindings, args.element_a)
element_b = getattr(cutlass_bindings, args.element_b)
element_c = getattr(cutlass_bindings, args.element_c)
element_acc = getattr(cutlass_bindings, args.element_acc)
math_operation = getattr(MathOperation, args.math)
opclass = getattr(cutlass.OpClass, args.opcode)
opclass = getattr(cutlass_bindings.OpClass, args.opcode)
math_inst = MathInstruction(
args.instruction_shape, element_a, element_b,
@@ -156,9 +156,9 @@ tile_description = TileDescription(
math_inst
)
layout_a = getattr(cutlass, args.layout_a)
layout_b = getattr(cutlass, args.layout_b)
layout_c = getattr(cutlass, args.layout_c)
layout_a = getattr(cutlass_bindings, args.layout_a)
layout_b = getattr(cutlass_bindings, args.layout_b)
layout_c = getattr(cutlass_bindings, args.layout_c)
A = TensorDescription(
element_a, layout_a, args.alignment_a
@@ -172,7 +172,7 @@ C = TensorDescription(
element_c, layout_c, args.alignment_c
)
element_epilogue = getattr(cutlass, args.element_epilogue)
element_epilogue = getattr(cutlass_bindings, args.element_epilogue)
if args.activation_function == "identity":
epilogue_functor = getattr(pycutlass, args.epilogue_functor)(
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
@@ -180,7 +180,7 @@ else:
epilogue_functor = getattr(pycutlass, "LinearCombinationGeneric")(
getattr(pycutlass, args.activation_function)(element_epilogue),
C.element, C.alignment, math_inst.element_accumulator, element_epilogue)
swizzling_functor = getattr(cutlass, args.swizzling_functor)
swizzling_functor = getattr(cutlass_bindings, args.swizzling_functor)
precompute_mode = getattr(SchedulerMode, args.precompute_mode)
operation = GemmOperationGrouped(
@@ -203,7 +203,7 @@ with open(args.problem_size_dir) as csv_file:
reader = csv.reader(csv_file)
for row in reader:
problem_sizes.append(
cutlass.gemm.GemmCoord(int(row[0]), int(row[1]), int(row[2]))
cutlass_bindings.gemm.GemmCoord(int(row[0]), int(row[1]), int(row[2]))
)
problem_count = len(problem_sizes)

24
examples/40_cutlass_py/gemm.py
View File

@@ -37,10 +37,10 @@ import argparse
import numpy as np
import sys
import cutlass
import pycutlass
from pycutlass import *
from pycutlass.utils.device import device_cc
import cutlass_bindings
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.device import device_cc
parser = argparse.ArgumentParser(description="Launch a GEMM kernel from Python: 'D = alpha * A * B + beta * C'")
@@ -72,11 +72,11 @@ pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
pycutlass.compiler.nvcc()
# Set up A, B, C and accumulator
A = TensorDescription(cutlass.float16, cutlass.ColumnMajor, alignment)
B = TensorDescription(cutlass.float16, cutlass.RowMajor, alignment)
C = TensorDescription(cutlass.float32, cutlass.ColumnMajor, alignment)
element_acc = cutlass.float32
element_epilogue = cutlass.float32
A = TensorDescription(cutlass_bindings.float16, cutlass_bindings.ColumnMajor, alignment)
B = TensorDescription(cutlass_bindings.float16, cutlass_bindings.RowMajor, alignment)
C = TensorDescription(cutlass_bindings.float32, cutlass_bindings.ColumnMajor, alignment)
element_acc = cutlass_bindings.float32
element_epilogue = cutlass_bindings.float32
# Select instruction shape based on the Tensor Core instructions supported
# by the device on which we are running
@@ -85,12 +85,14 @@ if cc == 70:
elif cc == 75:
instruction_shape = [16, 8, 8]
else:
# Use CUTLASS kernels for CC 80 by default (e.g., for cases in which SM86 is used)
cc = 80
instruction_shape = [16, 8, 16]
math_inst = MathInstruction(
instruction_shape,
A.element, B.element, element_acc,
cutlass.OpClass.TensorOp,
cutlass_bindings.OpClass.TensorOp,
MathOperation.multiply_add
)
@@ -122,7 +124,7 @@ tensor_B = np.ceil(np.random.uniform(low=-8.5, high=7.5, size=(args.k * args.n,)
tensor_C = np.ceil(np.random.uniform(low=-8.5, high=7.5, size=(args.m * args.n,))).astype(np.float32)
tensor_D = np.zeros(shape=(args.m * args.n,)).astype(np.float32)
problem_size = cutlass.gemm.GemmCoord(args.m, args.n, args.k)
problem_size = cutlass_bindings.gemm.GemmCoord(args.m, args.n, args.k)
alpha = 1.
beta = 0.

26
examples/40_cutlass_py/gemm_grouped.py
View File

@@ -37,10 +37,10 @@ import argparse
import numpy as np
import sys
import cutlass
import pycutlass
from pycutlass import *
from pycutlass.utils.device import device_cc
import cutlass_bindings
import cutlass.backend as pycutlass
from cutlass.backend import *
from cutlass.backend.utils.device import device_cc
parser = argparse.ArgumentParser(description="Launch a grouped GEMM kernel from Python")
@@ -65,11 +65,11 @@ pycutlass.compiler.nvcc()
# Set up A, B, C and accumulator
alignment = 1
A = TensorDescription(cutlass.float16, cutlass.ColumnMajor, alignment)
B = TensorDescription(cutlass.float16, cutlass.RowMajor, alignment)
C = TensorDescription(cutlass.float32, cutlass.ColumnMajor, alignment)
element_acc = cutlass.float32
element_epilogue = cutlass.float32
A = TensorDescription(cutlass_bindings.float16, cutlass_bindings.ColumnMajor, alignment)
B = TensorDescription(cutlass_bindings.float16, cutlass_bindings.RowMajor, alignment)
C = TensorDescription(cutlass_bindings.float32, cutlass_bindings.ColumnMajor, alignment)
element_acc = cutlass_bindings.float32
element_epilogue = cutlass_bindings.float32
# Select instruction shape based on the Tensor Core instructions supported
# by the device on which we are running
@@ -78,12 +78,14 @@ if cc == 70:
elif cc == 75:
instruction_shape = [16, 8, 8]
else:
# Use CUTLASS kernels for CC 80 by default (e.g., for cases in which SM86 is used)
cc = 80
instruction_shape = [16, 8, 16]
math_inst = MathInstruction(
instruction_shape,
A.element, B.element, element_acc,
cutlass.OpClass.TensorOp,
cutlass_bindings.OpClass.TensorOp,
MathOperation.multiply_add
)
@@ -112,8 +114,8 @@ pycutlass.compiler.add_module(operations)
# Initialize tensors for each problem in the group
problem_sizes = [
cutlass.gemm.GemmCoord(128, 128, 64),
cutlass.gemm.GemmCoord(512, 256, 128)
cutlass_bindings.gemm.GemmCoord(128, 128, 64),
cutlass_bindings.gemm.GemmCoord(512, 256, 128)
]
problem_count = len(problem_sizes)

4
examples/45_dual_gemm/device/dual_gemm.h
View File

@@ -159,7 +159,7 @@ class DualGemm {
using Mma0 = typename cutlass::gemm::threadblock::DefaultMma<
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB0, kAlignmentB,
ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag,
ThreadblockShape, WarpShape,
ThreadblockShape, WarpShape,
InstructionShape, Stages, Operator>::ThreadblockMma;
using Mma1 = typename cutlass::gemm::threadblock::DefaultMma<
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB1, kAlignmentB,
@@ -348,7 +348,7 @@ public:
ThreadblockSwizzle threadblock_swizzle;
cutlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape(
args.problem_size,
args.problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
args.split_k_slices);

54
examples/45_dual_gemm/dual_gemm.cu
View File

@@ -167,10 +167,10 @@ bool run_nonfused_gemm_f16_sm80() {
std::cout << "Running Non-fused GEMMs FP16 TN GEMMs...\n";
bool pass = nonFusedGemm.run(
problem_size,
alpha0,
beta0,
alpha1,
problem_size,
alpha0,
beta0,
alpha1,
beta1,
true /* is_profiling */
);
@@ -248,10 +248,10 @@ bool run_fused_gemm_f16_sm80_shmem() {
std::cout << "Running Fused FP16 TN GEMMs + Epilogue2...\n";
bool passed = fusedGemm.run(
problem_size,
alpha0,
beta0,
alpha1,
problem_size,
alpha0,
beta0,
alpha1,
beta1
);
@@ -301,11 +301,11 @@ bool run_batched_fused_gemm_f16_sm80_shmem() {
std::cout << "Running Batched Fused FP16 TN GEMMs + Epilogue2...\n";
bool passed = fusedGemm.run(
batch_problem_size,
alpha0,
beta0,
alpha1,
beta1,
batch_problem_size,
alpha0,
beta0,
alpha1,
beta1,
kBatchCount,
false, /* broadcast_b1 */
false /* is_profiling */
@@ -358,11 +358,11 @@ bool run_broadcast_fused_gemm_f16_sm80_shmem() {
std::cout << "Running Broadcast Fused FP16 TN GEMMs + Epilogue2...\n";
bool passed = fusedGemm.run(
problem_size,
alpha0,
beta0,
alpha1,
beta1,
problem_size,
alpha0,
beta0,
alpha1,
beta1,
1, /* batch_count */
true, /* broadcast_b1 */
true /* is_profiling */
@@ -415,11 +415,11 @@ bool run_batched_broadcast_fused_gemm_f16_sm80_shmem() {
std::cout << "Running Batch Broadcast Fused FP16 TN GEMMs + Epilogue2...\n";
bool passed = fusedGemm.run(
batch_problem_size,
alpha0,
beta0,
alpha1,
beta1,
batch_problem_size,
alpha0,
beta0,
alpha1,
beta1,
kBatchCount,
true, /* broadcast_b1 */
false /* is_profiling */
@@ -444,11 +444,11 @@ int main() {
};
std::string test_name = (
"dual-gemm f16 bias=" +
std::to_string(kUseBias) +
" split_k_serial=" +
"dual-gemm f16 bias=" +
std::to_string(kUseBias) +
" split_k_serial=" +
std::to_string(kSplitKSerial) +
" batch_count=" +
" batch_count=" +
std::to_string(kBatchCount)
);

71
examples/45_dual_gemm/dual_gemm_run.h
View File

@@ -45,6 +45,7 @@
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_relu.h"
#include "cutlass/platform/platform.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/device/gemm_universal.h"
@@ -356,13 +357,13 @@ struct NonFusedDualGemmRun
for(int i = 0; i < runs; i++) {
status = gemm_op_0();
CUTLASS_CHECK(status);
}
cudaEventRecord(stop1);
for(int i = 0; i < runs; i++) {
status = gemm_op_1();
CUTLASS_CHECK(status);
}
@@ -564,22 +565,22 @@ struct DualFusedGemmRun
cutlass::HostTensor<
typename DualGemm::ElementA,
typename DualGemm::LayoutA> tensor_A0(
std::is_same<typename DualGemm::LayoutA, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.k()) :
cutlass::platform::is_same<typename DualGemm::LayoutA, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.k()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.k()));
cutlass::HostTensor<
typename DualGemm::ElementB,
typename DualGemm::LayoutB0> tensor_B0(
std::is_same<typename DualGemm::LayoutB0, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.k(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutB0, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.k(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.k(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> tensor_C0(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
@@ -589,22 +590,22 @@ struct DualFusedGemmRun
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> tensor_D0(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> reference_D0(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementB,
typename DualGemm::LayoutB1> tensor_B1(
std::is_same<typename DualGemm::LayoutB1, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.k(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutB1, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.k(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.k(), batch_count * problem_size.n()));
if (broadcast_b1) {
tensor_B1.resize({problem_size.k(), batch_count});
@@ -613,8 +614,8 @@ struct DualFusedGemmRun
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> tensor_C1(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
@@ -624,29 +625,29 @@ struct DualFusedGemmRun
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> tensor_D1(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> tensor_D2(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> reference_D1(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
cutlass::HostTensor<
typename DualGemm::ElementC,
typename DualGemm::LayoutC> reference_D2(
std::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::platform::is_same<typename DualGemm::LayoutC, cutlass::layout::RowMajor>::value ?
cutlass::MatrixCoord(batch_count * problem_size.m(), problem_size.n()) :
cutlass::MatrixCoord(problem_size.m(), batch_count * problem_size.n()));
CHECK_TRUE(initialize_tensor(tensor_A0.host_view(), init_A, seed + 2019));
@@ -712,16 +713,16 @@ struct DualFusedGemmRun
ref_B1 = {tensor_Bias1.device_data(), typename DualGemm::LayoutC::Stride(0)};
}
typename DualGemm::Arguments arguments{
(batch_count > 1 ?
cutlass::gemm::DualGemmMode::kBatched :
(batch_count > 1 ?
cutlass::gemm::DualGemmMode::kBatched :
cutlass::gemm::DualGemmMode::kGemm),
problem_size,
tensor_A0.device_ref(),
tensor_B0.device_ref(),
ref_B0,
DualGemm::kStoreD0 ? tensor_D0.device_ref() : nullptr_ref,
(broadcast_b1 ?
typename DualGemm::TensorRefB1(tensor_B1.device_data(), 0) :
(broadcast_b1 ?
typename DualGemm::TensorRefB1(tensor_B1.device_data(), 0) :
tensor_B1.device_ref()),
ref_B1,
DualGemm::kStoreD1 ? tensor_D1.device_ref() : nullptr_ref,
@@ -793,15 +794,15 @@ struct DualFusedGemmRun
using GemmUniversal0 = cutlass::gemm::device::GemmUniversal<
typename DualGemm::ElementA, typename DualGemm::LayoutA,
typename DualGemm::ElementB, typename DualGemm::LayoutB0,
typename DualGemm::ElementC, typename DualGemm::LayoutC,
typename DualGemm::ElementC, typename DualGemm::LayoutC,
ElementAccumulator
>;
GemmUniversal0 reference_gemm0;
typename GemmUniversal0::Arguments args0 {
(batch_count > 1 ?
cutlass::gemm::GemmUniversalMode::kBatched :
(batch_count > 1 ?
cutlass::gemm::GemmUniversalMode::kBatched :
cutlass::gemm::GemmUniversalMode::kGemm),
problem_size,
batch_count,
@@ -828,15 +829,15 @@ struct DualFusedGemmRun
using GemmUniversal1 = cutlass::gemm::device::GemmUniversal<
typename DualGemm::ElementA, typename DualGemm::LayoutA,
typename DualGemm::ElementB, typename DualGemm::LayoutB1,
typename DualGemm::ElementC, typename DualGemm::LayoutC,
typename DualGemm::ElementC, typename DualGemm::LayoutC,
ElementAccumulator
>;
GemmUniversal1 reference_gemm1;
typename GemmUniversal1::Arguments args1 {
(batch_count > 1 ?
cutlass::gemm::GemmUniversalMode::kBatched :
(batch_count > 1 ?
cutlass::gemm::GemmUniversalMode::kBatched :
cutlass::gemm::GemmUniversalMode::kGemm),
problem_size,
batch_count,
@@ -861,7 +862,7 @@ struct DualFusedGemmRun
CUTLASS_CHECK(status);
if(relu) {
cutlass::reference::device::TensorReLu(reference_D0.device_view());
cutlass::reference::device::TensorReLu(reference_D0.device_view());
cutlass::reference::device::TensorReLu(reference_D1.device_view());
}

16
examples/45_dual_gemm/kernel/dual_gemm.h
View File

@@ -300,7 +300,7 @@ struct DualGemm {
int offset_k = 0;
int problem_size_k = params.problem_size.k();
ElementA *ptr_A0 = static_cast<ElementA *>(params.ref_A0.data());
ElementA *ptr_A0 = static_cast<ElementA *>(params.ref_A0.data());
ElementB *ptr_B0 = static_cast<ElementB *>(params.ref_B0.data());
ElementB *ptr_B1 = static_cast<ElementB *>(params.ref_B1.data());
@@ -309,7 +309,7 @@ struct DualGemm {
//
if (params.mode == DualGemmMode::kGemm) {
if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
}
offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
@@ -413,11 +413,11 @@ struct DualGemm {
int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m();
ElementC *ptr_C0 = static_cast<ElementC *>(params.ref_C0.data());
ElementC *ptr_C1 = static_cast<ElementC *>(params.ref_C1.data());
ElementC *ptr_D0 = static_cast<ElementC *>(params.ref_D0.data());
ElementC *ptr_D1 = static_cast<ElementC *>(params.ref_D1.data());
ElementC *ptr_D2 = static_cast<ElementC *>(params.ref_D2.data());
ElementC *ptr_C0 = static_cast<ElementC *>(params.ref_C0.data());
ElementC *ptr_C1 = static_cast<ElementC *>(params.ref_C1.data());
ElementC *ptr_D0 = static_cast<ElementC *>(params.ref_D0.data());
ElementC *ptr_D1 = static_cast<ElementC *>(params.ref_D1.data());
ElementC *ptr_D2 = static_cast<ElementC *>(params.ref_D2.data());
// Construct the semaphore.
Semaphore semaphore(params.semaphore + block_idx, thread_idx);
@@ -425,7 +425,7 @@ struct DualGemm {
if (params.mode == DualGemmMode::kGemm) {
// If performing a reduction via split-K, fetch the initial synchronization
if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
// Fetch the synchronization lock initially but do not block.
semaphore.fetch();

30
examples/46_depthwise_simt_conv2dfprop/depthwise_simt_conv2dfprop.cu
View File

@@ -233,6 +233,17 @@ struct Options {
return false;
}
// Filter size passed through command line does not match filter size template parameter
if (filter_size.h() != FilterShape::kRow || filter_size.w() != FilterShape::kColumn) {
std::cerr << "Filter size passed in (" << filter_size.h() << "x" << filter_size.w() << ") "
<< "must match the FilterShape template parameter of the convolution "
<< "(" << FilterShape::kRow << "x" << FilterShape::kColumn << "). "
<< "To use the filter shape passed in, change the FilterShape template "
<< "parameter and recompile this example."
<< std::endl;
return false;
}
return true;
}
@@ -319,9 +330,9 @@ struct Options {
"table\n";
out << "\n\nExamples:\n\n"
<< "$ ./examples/45_depthwise_simt_conv2dfprop/45_depthwise_simt_conv2dfprop --n=32 "
<< "$ ./examples/46_depthwise_simt_conv2dfprop/46_depthwise_simt_conv2dfprop --n=32 "
"--h=224 --w=224 --c=128 --k=128 --g=128 --r=3 --s=3\n\n"
<< "$ ./examples/45_depthwise_simt_conv2dfprop/45_depthwise_simt_conv2dfprop --n=1 "
<< "$ ./examples/46_depthwise_simt_conv2dfprop/46_depthwise_simt_conv2dfprop --n=1 "
"--h=224 --w=224 --c=32 --k=32 --g=32 --r=3 --s=3 --splitk=10 --ref-check\n\n";
return out;
@@ -515,14 +526,13 @@ Result profile_convolution(Options const &options) {
ElementOutput,
LayoutOutput,
ElementComputeEpilogue,
ElementAccumulator,
cutlass::NumericConverter<ElementOutput, ElementComputeEpilogue> >(problem_size,
tensor_a.host_ref(),
tensor_b.host_ref(),
tensor_c.host_ref(),
tensor_ref_d.host_ref(),
options.alpha,
options.beta);
ElementAccumulator >(problem_size,
tensor_a.host_ref(),
tensor_b.host_ref(),
tensor_c.host_ref(),
tensor_ref_d.host_ref(),
options.alpha,
options.beta);
// Check if output from CUTLASS kernel and reference kernel are equal or not
tensor_d.sync_host();

26
examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm.cu
View File

@@ -60,6 +60,7 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
@@ -95,12 +96,13 @@ constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value; // M
// C/D matrix configuration
using ElementC = float; // Element type for C and D matrix operands
using LayoutC = cutlass::layout::ColumnMajor; // Layout type for C and D matrix operands
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
// Core kernel configurations
using ElementAccumulator = float; // Element type for internal accumulation
using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature
using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag
using TilesShape = Shape<_128,_128,_32>; // Threadblock-level tile size
using TileShape = Shape<_128,_128,_32>; // Threadblock-level tile size
using ClusterShape = Shape<_1,_2,_1>; // Shape of the threadblocks in a cluster
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; // Kernel to launch based on the default setting in the Collective Builder
@@ -110,15 +112,20 @@ using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
TilesShape, ClusterShape,
TileShape, ClusterShape,
cutlass::gemm::collective::StageCountAuto,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::epilogue::thread::LinearCombination<ElementC, 1, ElementAccumulator, ElementAccumulator>>;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
TileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutC, AlignmentC,
ElementC, LayoutC, AlignmentC,
cutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int>, // Indicates ProblemShape
@@ -308,11 +315,8 @@ typename Gemm::Arguments args_from_options(const Options &options)
typename Gemm::Arguments arguments{
cutlass::gemm::GemmUniversalMode::kGemm,
{options.m, options.n, options.k},
block_A.get(),
stride_A,
block_B.get(),
stride_B,
{block_C.get(), stride_C, block_D.get(), stride_D, {options.alpha, options.beta}}
{block_A.get(), stride_A, block_B.get(), stride_B},
{{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D}
};
return arguments;

135
examples/49_hopper_gemm_schedules_with_collective_builder/49_hopper_gemm_schedules_with_collective_builder.cu → examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu
View File

@@ -77,10 +77,27 @@
will fit in shared memory given the types of operands and the thread block shape, rather than simply using
a single default value.
Note that one does not need to use the CollectiveBuilder to declare CUTLASS 3 kernels; one can still provide
every template parameter to the gemm::collective::CollectiveMma. Specifying every template parameter in this
manner remains the primary API for using CUTLASS 3 kernels. The CollectiveBuilder is simply meant to be
a convenience interface.
CUTLASS 3.x provides builders for both collective mainloops and epilogues. The particular implementation of
the collective is specified via the schedule tags that corresond to the underlying collective's
dispatch policy. `gemm::collective::KernelScheduleAuto` and `epilogue::collective::EpilogueScheduleAuto`
are special cases of these schedules that allow the builder to also decide the dispatch policy for you,
therefore letting the builder pick the collective specialization.
CUTLASS builders make an attempt to pick the best schedule when `Auto` is provided such that the
assembled collctives have the best performance, but this is not a guarantee. A user relying on `Auto`
may get a free performance upgrade with newer CUTLASS releases in case we can provide more optimized
implementations that the builder can transparently assemble for `Auto`.
If a user decides to let the builders pick the collective specialization via `Auto` schedules,
they must be used for both mainloop and epilogue alike to ensure compatibility between the
chosen collectives. Additionally, if a user chooses to opt in to a specific schedule, non-`Auto`
schedules must be used for both mainloop and epilogue builder schedules, and these schedules
must be compatible.
One does not need to use the CollectiveBuilder to declare CUTLASS 3 kernels; one can still provide
every template parameter to the `gemm::collective::CollectiveMma`. Specifying every template parameter
in this manner remains the primary API for using CUTLASS 3 kernels. `CollectiveBuilder`s are
simply meant to be a convenience interface.
Note also that, while the selections made by CollectiveBuilder attempt to maximize performance, this is not
a guarantee. Furthermore, the behavior of the CollectiveBuilder when `Auto` parameters are provided is subject
@@ -94,7 +111,7 @@
extending the problem size with an additional tensor rank.
Example usage:
$ ./examples/49_hopper_gemm_schedules_with_collective_builder/49_hopper_gemm_schedules_with_collective_builder \
$ ./examples/49_hopper_with_collective_builder/49_collective_builder \
--m=2048 --n=2048 --k=2048 --l=2
*/
@@ -108,6 +125,7 @@
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
@@ -160,7 +178,7 @@ struct Options {
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "49_hopper_gemm_schedules_with_collective_builder\n\n"
out << "49_hopper_with_collective_builder\n\n"
<< " This example showcases the use of CUTLASS's collective operation builders to easily construct\n"
<< " performant kernels targeting NVIDIA's Hopper architecture.\n\n"
<< "Options:\n\n"
@@ -212,14 +230,24 @@ bool initialize_block(
// operation builders by specializing the GEMM only on the kernel schedule it will use and the
// number of pipeline stages.
//
// For either option, one can use a special `Auto` type that tells the CollectiveBuilder
// One can use a special `Auto` type that tells the CollectiveBuilder
// to select an appropriate value on its own. The CollectiveBuilder will attempt to select
// values that will result in the most-performant kernel, but this is not a guarantee. Furthermore,
// the behavior of the CollectiveBuilder with `Auto` types is subject to change in future releases
// configurations that will result in the most-performant kernel, but this is not a guarantee.
//
// If relying on 'Auto' schedules, all builders must use the 'Auto' schedule to ensure compatiblity.
// For example, if `KernelScheduleAuto` is used for the mainloop builder, `EpilogueScheduleAuto` must
// be used for the epilogue builder.
//
// Furthermore, if an override schedule is selected, both epilgoue and mainloop schedules must
// be specifically opt into a compatible selection.
//
// Behavior of the CollectiveBuilder with `Auto` types is subject to change in future releases
// -- do not rely on `Auto` if you require a specific scheduling policy.
template <
// Type of kernel schedule to generate
class KernelScheduleType = cutlass::gemm::collective::KernelScheduleAuto,
class MainloopScheduleType = cutlass::gemm::collective::KernelScheduleAuto,
// Type of epilogue schedule to generate
class EpilogueScheduleType = cutlass::epilogue::collective::EpilogueScheduleAuto,
// Number of pipeline stages to use
class StageCountType = cutlass::gemm::collective::StageCountAuto
>
@@ -230,22 +258,32 @@ struct ExampleRunner {
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
static constexpr int kAlignmentA = 8;
static constexpr int kAlignmentB = 8;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
static constexpr int AlignmentA = 8;
static constexpr int AlignmentB = 8;
static constexpr int AlignmentC = 8;
static constexpr int AlignmentD = 8;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
cutlass::half_t, LayoutA, kAlignmentA,
cutlass::half_t, LayoutB, kAlignmentB,
float,
Shape<_128,_128,_64>, Shape<_2,_1,_1>,
StageCountType,
KernelScheduleType
Shape<_128,_128,_64>, Shape<_1,_1,_1>,
cutlass::epilogue::collective::EpilogueTileAuto,
float, float,
cutlass::half_t, LayoutC, AlignmentC,
cutlass::half_t, LayoutD, AlignmentD,
EpilogueScheduleType
>::CollectiveOp;
using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::gemm::TagToStrideC_t<LayoutD>,
cutlass::epilogue::thread::LinearCombination<cutlass::half_t, 1, float, float>>;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
cutlass::half_t, LayoutA, AlignmentA,
cutlass::half_t, LayoutB, AlignmentB,
float,
Shape<_128,_128,_64>, Shape<_2,_1,_1>,
std::conditional_t<std::is_same_v<StageCountType, cutlass::gemm::collective::StageCountAuto>,
cutlass::gemm::collective::StageCountAutoCarveout<(int)sizeof(typename CollectiveEpilogue::SharedStorage)>,
StageCountType>,
MainloopScheduleType
>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>,
@@ -262,10 +300,10 @@ struct ExampleRunner {
using StrideC = typename Gemm::GemmKernel::StrideC;
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutTagA = decltype(cutlass::gemm::detail::stride_to_layout_tag_A<StrideA>());
using LayoutTagB = decltype(cutlass::gemm::detail::stride_to_layout_tag_B<StrideB>());
using LayoutTagC = decltype(cutlass::gemm::detail::stride_to_layout_tag_A<StrideC>());
using LayoutTagD = decltype(cutlass::gemm::detail::stride_to_layout_tag_A<StrideD>());
using LayoutTagA = cutlass::gemm::detail::StrideToLayoutTagA_t<StrideA>;
using LayoutTagB = cutlass::gemm::detail::StrideToLayoutTagB_t<StrideB>;
using LayoutTagC = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideC>;
using LayoutTagD = cutlass::gemm::detail::StrideToLayoutTagC_t<StrideD>;
//
// Data members
@@ -356,11 +394,8 @@ struct ExampleRunner {
typename Gemm::Arguments arguments{
cutlass::gemm::GemmUniversalMode::kGemm,
problem_size,
block_A.get(),
stride_A,
block_B.get(),
stride_B,
{block_C.get(), stride_C, block_D.get(), stride_D, {options.alpha, options.beta}},
{block_A.get(), stride_A, block_B.get(), stride_B},
{{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D},
hw_info
};
@@ -477,42 +512,48 @@ int main(int argc, char const **args) {
// selected and the maximum number of stages that can fit in shared memory will be selected.
//
// This example is equivalent to declaring
// ExampleRunner<cutlass::gemm::collective::KernelScheduleAuto, cutlass::gemm::collective::StageCountAuto>
// ExampleRunner<
// cutlass::gemm::collective::KernelScheduleAuto,
// cutlass::epilogue::collective::EpilogueScheduleAuto,
// cutlass::gemm::collective::StageCountAuto>
// Each of the `Auto` types indicate that the CollectiveBuilder should determine the scheduling policy and
// stage count. Note that the behavior of the CollectiveBuilder with `Auto` parameters is subject to change
// -- do not rely on `Auto` if you require a specific scheduling policy.
// If you opt in to a non-'Auto' schedule, make sure all collectives are built using specific, compatible schedules.
ExampleRunner<> auto_schedule_auto_stage_runner;
passed = auto_schedule_auto_stage_runner.run(options, hw_info);
print_result("Automatically-selected schedule and stage count", passed);
// One can override the stage count used in the GEMM by replacing cutlass::gemm::collective::StageCountAuto
// with the number of stages to use (5 in this case).
ExampleRunner<cutlass::gemm::collective::KernelScheduleAuto, _5> auto_schedule_5_stage_runner;
ExampleRunner<
cutlass::gemm::collective::KernelScheduleAuto,
cutlass::epilogue::collective::EpilogueScheduleAuto,
_5> auto_schedule_5_stage_runner;
passed = auto_schedule_5_stage_runner.run(options, hw_info);
print_result("Automatically-selected schedule with 5 stages", passed);
// One can also override the scheduling policy to use. In this case, use the KernelTma scheduling
// policy, which specifies that the Hopper TMA feature should be used.
ExampleRunner<cutlass::gemm::KernelTma> tma_schedule_auto_stage_runner;
// policy, which specifies that the Hopper TMA feature should be used, and we also use an epilgoue
// that does not use any shared memory.
ExampleRunner<cutlass::gemm::KernelTma, cutlass::epilogue::NoSmemWarpSpecialized> tma_schedule_auto_stage_runner;
passed = tma_schedule_auto_stage_runner.run(options, hw_info);
print_result("TMA schedule with automatically-selected stage count", passed);
// Here, we override the scheduling policy to use Hopper's TMA feature alongside the warp-specialized
// scheduling policy.
//
// Note that, as of the CUTLASS 3.0 release, this is the default scheduling policy
// used by the CollectiveBuilder, so this declaration is equivalent to ExampleRunner<> and
// ExampleRunner<cutlass::gemm::collective::KernelScheduleAuto>. However, this default is subject to
// change in future releases -- do not rely on `Auto` if you require a specific scheduling policy.
ExampleRunner<cutlass::gemm::KernelTmaWarpSpecialized> ws_schedule_auto_stage_runner;
// scheduling policy, and an epilgoue that does not use any shared memory.
ExampleRunner<cutlass::gemm::KernelTmaWarpSpecialized, cutlass::epilogue::NoSmemWarpSpecialized> ws_schedule_auto_stage_runner;
passed = ws_schedule_auto_stage_runner.run(options, hw_info);
print_result("Warp-specialized TMA schedule with automatically-selected stage count", passed);
// Finally, we override the scheduling policy to use Hopper's TMA feature, alongside the warp-specialized
// scheduling policy, leveraging persistent thread blocks.
ExampleRunner<cutlass::gemm::KernelTmaWarpSpecializedPersistent> ws_persistent_schedule_auto_stage_runner;
passed = ws_persistent_schedule_auto_stage_runner.run(options, hw_info);
print_result("Persistent warp-specialized TMA schedule with automatically-selected stage count", passed);
// scheduling policy, TMA-based epilogue, leveraging persistent thread blocks.
ExampleRunner<
cutlass::gemm::KernelTmaWarpSpecializedPingpong,
cutlass::epilogue::TmaWarpSpecialized> ws_pingpong_schedule_auto_stage_runner;
passed = ws_pingpong_schedule_auto_stage_runner.run(options, hw_info);
print_result("Ping-pong warp-specialized TMA schedule with automatically-selected stage count", passed);
#endif

7
examples/49_hopper_gemm_schedules_with_collective_builder/CMakeLists.txt → examples/49_hopper_gemm_with_collective_builder/CMakeLists.txt
View File

@@ -27,9 +27,8 @@
# 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.
# Both filenames are shorter to avoid MAX_PATH issues on Windows.
cutlass_example_add_executable(
49_hopper_gemm_schedules_with_collective_builder
49_hopper_gemm_schedules_with_collective_builder.cu
49_collective_builder
49_collective_builder.cu
)

22
examples/50_hopper_gemm_with_epilogue_swizzle/50_hopper_gemm_with_epilogue_swizzle.cu
View File

@@ -34,7 +34,7 @@
The following example shows how to assemble a custom GEMM kernel that spells out the Collectives
directly instead of using a builder and, in the process, instance a more efficient Epilogue
(from `cutlass/epilogue/collective/epilogue.hpp`) instead of using the default epilogue.
(from `cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp`) instead of using the default epilogue.
The GemmUniversal API takes 3 main template arguments:
(1) the problem shape / extents
@@ -65,7 +65,7 @@
#include "cute/tensor.hpp"
#include "cutlass/util/command_line.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/epilogue.hpp"
#include "cutlass/epilogue/collective/collective_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
@@ -122,7 +122,7 @@ struct Options {
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "50_hopper_gemm_with_vectorized_epilogue\n\n"
out << "50_hopper_gemm_with_epilogue_swizzle\n\n"
<< "Hopper GEMM Example with Epilogue Swizzle.\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
@@ -286,11 +286,8 @@ struct ExampleRunner {
typename Gemm::GemmKernel::Arguments arguments{
cutlass::gemm::GemmUniversalMode::kGemm,
problem_size,
block_A.get(),
stride_A,
block_B.get(),
stride_B,
{block_C.get(), stride_C, block_D.get(), stride_D, {options.alpha, options.beta}},
{block_A.get(), stride_A, block_B.get(), stride_B},
{{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D},
hw_info
};
@@ -443,11 +440,11 @@ int main(int argc, char const **args) {
cute::SM90_TMA_LOAD,
cute::SM90_TMA_LOAD_MULTICAST>::type;
using SmemLayoutAtomA = decltype(cute::GMMA::smem_selector<
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
GmmaMajorA, ElementA, decltype(cute::get<0>(TileShape{})), decltype(cute::get<2>(TileShape{}))
>());
using SmemLayoutAtomB = decltype(cute::GMMA::smem_selector<
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
GmmaMajorB, ElementB, decltype(cute::get<1>(TileShape{})), decltype(cute::get<2>(TileShape{}))
>());
@@ -494,14 +491,15 @@ int main(int argc, char const **args) {
Stride<_16,_1>>,
TileShapeS2R>;
using Epilogue = cutlass::epilogue::collective::Epilogue<
using Epilogue = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter<
cutlass::epilogue::collective::Epilogue<
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::gemm::TagToStrideC_t<LayoutD>,
cutlass::epilogue::thread::LinearCombination<int32_t, 1, int32_t, int32_t>,
SmemLayout,
Copy_Atom<DefaultCopy, ElementAcc>,
TiledCopyS2R,
Copy_Atom<DefaultCopy, ElementOutput>>;
Copy_Atom<DefaultCopy, ElementOutput>>>;
//
// Assembling the GemmKernel

15
examples/51_hopper_gett/gett_kernel.cuh
View File

@@ -37,7 +37,7 @@
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/collective/collective_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
namespace example {
@@ -88,10 +88,12 @@ gett_kernel(
cutlass::FloatRoundStyle::round_to_nearest, ElementC>;
// No changes are required to the default epilogue
using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
using CollectiveEpilogue = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter<
cutlass::epilogue::collective::DefaultEpilogue<
StrideC,
StrideD,
EpilogueThreadOp>;
EpilogueThreadOp,
cutlass::gemm::EpilogueDefault>>;
// CollectiveMma for GETTs can be built using the CollectiveBuilders
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
@@ -100,7 +102,7 @@ gett_kernel(
ElementB, StrideB, 128 / cutlass::sizeof_bits<ElementB>::value,
ElementAccumulator,
TileShape, Shape<_1,_2,_1>,
cutlass::gemm::collective::StageCountAuto,
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
@@ -115,9 +117,8 @@ gett_kernel(
typename GettOperator::Arguments args {
cutlass::gemm::GemmUniversalMode::kBatched,
problem_shape_mnkl,
ptr_A, stride_a_mkl,
ptr_B, stride_b_nkl,
{ ptr_C, stride_c_mnl, ptr_D, stride_d_mnl, {alpha, beta} }
{ ptr_A, stride_a_mkl, ptr_B, stride_b_nkl },
{ {alpha, beta}, ptr_C, stride_c_mnl, ptr_D, stride_d_mnl }
};
#if CUTLASS_DEBUG_TRACE_LEVEL > 0

2
examples/CMakeLists.txt
View File

@@ -129,7 +129,7 @@ foreach(EXAMPLE
46_depthwise_simt_conv2dfprop
47_ampere_gemm_universal_streamk
48_hopper_warp_specialized_gemm
49_hopper_gemm_schedules_with_collective_builder
49_hopper_gemm_with_collective_builder
50_hopper_gemm_with_epilogue_swizzle
51_hopper_gett
)

1
examples/common/helper.h
View File

@@ -31,6 +31,7 @@
#pragma once
#include "cuda_runtime.h"
#include <iostream>
/**
* Panic wrapper for unwinding CUTLASS errors

1
examples/cute/tutorial/CMakeLists.txt
View File

@@ -31,4 +31,3 @@ cutlass_example_add_executable(
sgemm_nt_1
sgemm_nt_1.cu
)

340
examples/python/00_basic_gemm.ipynb Normal file
View File

@@ -0,0 +1,340 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "1ef96b3f",
"metadata": {},
"source": [
"# Basic example of using the CUTLASS Python interface\n",
"This notebook walks through a basic example of using the CUTLASS Python interface to declare, compile, and run GEMMs.\n",
"\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cutlass/tree/master/examples/00_basic_gemm.ipynb)\n"
]
},
{
"cell_type": "markdown",
"id": "962324fd",
"metadata": {},
"source": [
"We first import various packages needed for the example and construct the input and output tensors that will be used in our example.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0e324219",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import random\n",
"\n",
"import cutlass\n",
"\n",
"# This controls whether ther C++ GEMM declaration will be printed at each step. Set to `false` to\n",
"# omit this information.\n",
"print_module = True\n",
"\n",
"m = 128\n",
"n = m\n",
"k = m\n",
"\n",
"dtype = np.float16\n",
"type_A = np.float16\n",
"type_B = np.float16\n",
"type_C = np.float16\n",
"type_D = np.float16\n",
"\n",
"np.random.seed(1234)\n",
"random.seed(1234)\n",
"scope_min = -4\n",
"scope_max = 4\n",
"tensor_A = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, k)).astype(type_A))\n",
"tensor_B = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(k, n)).astype(type_B))\n",
"tensor_C = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, n)).astype(type_C))\n",
"\n",
"alpha = np.float16(1.)\n",
"beta = np.float16(0.)\n",
"\n",
"tensor_D = np.zeros(tensor_C.shape).astype(type_D)"
]
},
{
"cell_type": "markdown",
"id": "f2c7bf48",
"metadata": {},
"source": [
"## Declaring and running a GEMM\n",
"To get started, one only needs to provide the tensors declared above to the `cutlass.op.Gemm` call.\n",
"This sets up a default GEMM operation for the given device on which you are running.\n",
"\n",
"Assuming that we are running on SM80, this default to using a GEMM that leverages FP16 Tensor Core operations.\n",
"\n",
"Calling `plan.run()` will generate the CUTLASS C++ kernel in question, compile it, and run it on the tensors we previously passed in. By setting `print_module` to `true`, the C++ code that is emitted is printed."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0dfd8975",
"metadata": {},
"outputs": [],
"source": [
"# We specify `element_accumulator` here so as to match the kernel run by NumPy below. However,\n",
"# specifying `element_accumulator` is not required if it is the same as `element`\n",
"plan = cutlass.Gemm(element=dtype, layout=cutlass.LayoutType.RowMajor, element_accumulator=np.float32)\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "4a5856de",
"metadata": {},
"source": [
"There are many other ways to construct a plan from `cutlass.op.Gemm` (e.g., by specifiying they types and layouts of each operand, by providing representative tensors as inputs). For more details on these, see the documentation in the `cutlass.op.Gemm` constructor."
]
},
{
"cell_type": "markdown",
"id": "945478ef",
"metadata": {},
"source": [
"We then compare the output to running the GEMM using NumPy."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b669de6",
"metadata": {},
"outputs": [],
"source": [
"tensor_D_numpy = (alpha * (tensor_A @ tensor_B)) + (beta * tensor_C)\n",
"np.testing.assert_array_equal(tensor_D, tensor_D_numpy)"
]
},
{
"cell_type": "markdown",
"id": "ee5cbbbe",
"metadata": {},
"source": [
"Note that one could use the same kernel just declared for tensors provided by other frameworks beyond NumPy, such as PyTorch or CuPy."
]
},
{
"cell_type": "markdown",
"id": "b6c86493",
"metadata": {},
"source": [
"## Changing operation modes\n",
"By default, the CUTLASS Python interface will try to use Tensor Core operations whenever possible. If the configuration provided to `cutlass.op.Gemm` is not supported on Tensor Cores, the interface will fall back to using a SIMT kernel.\n",
"\n",
"The operation mode currently in use can be returned via the `plan.opclass` property. In this case Tensor Core operations."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "529fda93",
"metadata": {},
"outputs": [],
"source": [
"print(plan.opclass)"
]
},
{
"cell_type": "markdown",
"id": "6d27c575",
"metadata": {},
"source": [
"Suppose that we don't want to use Tensor Cores for this GEMM. One can change to using CUTLASS's SIMT GEMMs by setting the plan's `opclass` field.\n",
"\n",
"As is shown in the printed output, the emitted kernel uses template parameters that fit CUTLASS's SIMT GEMMs.\n",
"\n",
"Also notice that, this time around, we provided tensor parameters to `plan.run()`. One is free to provide different parameters to `plan.run()` than were passed in at the initial call to `cutlass.op.Gemm`, provided that the passed-in tensors have the same data type and layout as those passed in on intialization."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6a44d35b",
"metadata": {},
"outputs": [],
"source": [
"tensor_D_simt = np.zeros(tensor_C.shape).astype(type_D)\n",
"plan.opclass = cutlass.OpcodeClass.Simt\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D_simt, alpha, beta, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "639dcb59",
"metadata": {},
"source": [
"If we compare the output of the Tensor Core and SIMT GEMMs we just ran we see that they are equal."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9b480853",
"metadata": {},
"outputs": [],
"source": [
"np.testing.assert_array_equal(tensor_D, tensor_D_simt)"
]
},
{
"cell_type": "markdown",
"id": "0cce1eae",
"metadata": {},
"source": [
"## Running cached kernels\n",
"You may have noticed that the `plan.run()` calls for the previous two kernels took some time to execute. This is because the kernel being emitted had not yet been compiled.\n",
"\n",
"CUTLASS caches compiled binaries so that recompilation isn't necessary every time a kernel is run. For example, if we change modes back to using Tensor Cores and call `plan.run()` again (with a different set of tensor parameters), you'll find the call to return much faster."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f8051e5e",
"metadata": {},
"outputs": [],
"source": [
"m = 2400\n",
"n = 3232\n",
"k = 4096\n",
"\n",
"tensor_A = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, k)).astype(type_A))\n",
"tensor_B = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(k, n)).astype(type_B))\n",
"tensor_C = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, n)).astype(type_C))\n",
"tensor_D = np.zeros(tensor_C.shape).astype(type_D)\n",
"\n",
"alpha = np.float16(1.)\n",
"beta = np.float16(2.)\n",
"\n",
"plan.opclass = cutlass.OpcodeClass.TensorOp\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "52a4e318",
"metadata": {},
"source": [
"## Running non-default GEMMs\n",
"The previous examples showed how it is simple to get started running a default GEMM kernel in CUTLASS. But, what do you do if you want a bit more control over the parameters to the GEMM?\n",
"\n",
"Under the hood, CUTLASS enumerates the different GEMM configuration parameters possible for this kernel from the CUTLASS profiler. The code below shows how one can access the tile descriptions for the kernels (e.g., cluster, threadblock, and warp shape)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1c593be1",
"metadata": {},
"outputs": [],
"source": [
"tiles = plan.tile_descriptions()\n",
"print('{} tile descriptions returned'.format(len(tiles)))\n",
"num_print = 10\n",
"print('First {} tile descriptions are:'.format(num_print))\n",
"for td in tiles[:num_print]:\n",
" print(td)"
]
},
{
"cell_type": "markdown",
"id": "dc3ad875",
"metadata": {},
"source": [
"Next, we'll pick one of these configurations at random and compile and run it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a8dc5287",
"metadata": {},
"outputs": [],
"source": [
"idx = random.randint(0, len(tiles)-1)\n",
"td = tiles[idx]\n",
"print('Tile description {} is: {}'.format(idx, td))\n",
"plan.compile(td)\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "c5a8b534",
"metadata": {},
"source": [
"One can also change the swizzling function used by the kernel. For example, one can modify the kernel to use the stream K feature of CUTLASS via:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e5e88d17",
"metadata": {},
"outputs": [],
"source": [
"# Stream K is only supported pre-SM90 (at least when this example was written)\n",
"if plan.cc != 90:\n",
" plan.swizzling_functor = cutlass.swizzle.ThreadblockSwizzleStreamK\n",
" plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "5a8ba2ba",
"metadata": {},
"source": [
"## Handling errors\n",
"The CUTLASS Python interface attempts to catch runtime and compilation errors in Python so as to provide more understandable error messages.\n",
"\n",
"Here's an example in which we try to use too many stages for a given GEMM kernel. Normally, this would result in a runtime error due to the GPU having insufficient shared memory to launch the kernel with 8 stages. The CUTLASS Python interface is able to detect this issue before compiling the kernel, and reports it back to the user."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fe7d0e42",
"metadata": {},
"outputs": [],
"source": [
"# td = tiles[0]\n",
"# td.stages = 8\n",
"# plan.compile(td)"
]
}
],
"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.8.10"
},
"vscode": {
"interpreter": {
"hash": "0466d96796c9cd8f7a1cad264ff326ececc950ba2420e0256d5105fc1a3c6e70"
}
}
},
"nbformat": 4,
"nbformat_minor": 5
}

202
examples/python/01_epilogue.ipynb Normal file
View File

@@ -0,0 +1,202 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "5d24a692",
"metadata": {},
"source": [
"# Example of using elementwise activation functions in the CUTLASS Python interface\n",
"This notebook walks through a basic example of using the CUTLASS Python interface to declare, compile, and run GEMMs with different epilogues.\n",
"\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cutlass/tree/master/examples/00_basic_gemm.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "3ca993fe",
"metadata": {},
"source": [
"We first import various packages needed for the example and construct the input and output tensors that will be used in our example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63a70a3c",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"import cutlass\n",
"\n",
"# This controls whether ther C++ GEMM declaration will be printed at each step. Set to `false` to\n",
"# omit this information.\n",
"print_module = True\n",
"\n",
"m = 256\n",
"n = m\n",
"k = m\n",
"\n",
"type_A = np.float16\n",
"type_B = np.float16\n",
"type_C = np.float16\n",
"type_D = np.float16\n",
"\n",
"np.random.seed(1234)\n",
"scope_min = -4\n",
"scope_max = 4\n",
"tensor_A = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, k)).astype(type_A))\n",
"tensor_B = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(k, n)).astype(type_B))\n",
"tensor_C = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, n)).astype(type_C))\n",
"\n",
"alpha = np.float16(1.)\n",
"beta = np.float16(0.)\n",
"\n",
"tensor_D = np.zeros(tensor_C.shape).astype(type_D)"
]
},
{
"cell_type": "markdown",
"id": "1eb0d95b",
"metadata": {},
"source": [
"## Run a GEMM with an identity activation function\n",
"To begin, we simply run a default GEMM with an identity activation function. This performs the well-known operation `D = alpha * (A @ B) + beta * C`. This is the default activation function used, and does not need to be specified."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8d257833",
"metadata": {},
"outputs": [],
"source": [
"plan = cutlass.op.Gemm(element=np.float16, layout=cutlass.LayoutType.RowMajor)\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "54961694",
"metadata": {},
"source": [
"## Run a GEMM with a ReLU element-wise activation function\n",
"CUTLASS makes it easy to support other element-wise activation functions. This results in performing an element-wise after the generic linear combination performed in a GEMM. If we call such an activation function `act`, the resulting formulation is:\n",
"```\n",
"D = alpha * (A @ B) + beta * C\n",
"D = act(D)\n",
"```\n",
"\n",
"Here, we will add a ReLU activation function. Given an input `x`, ReLU returns `max(x, 0)`.\n",
"\n",
"This is easy to do in CUTLASS. One only needs to set the plan's `activation` field."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5fe49443",
"metadata": {},
"outputs": [],
"source": [
"tensor_D_relu = np.zeros(tensor_C.shape).astype(type_D)\n",
"plan.activation = cutlass.epilogue.relu\n",
"plan.run(tensor_A, tensor_B, tensor_C, tensor_D_relu, print_module=print_module)"
]
},
{
"cell_type": "markdown",
"id": "455d0a37",
"metadata": {},
"source": [
"We can now verify that the result of the GEMM that used a ReLU activation function:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e32e7798",
"metadata": {},
"outputs": [],
"source": [
"relu_ref = (tensor_D >= 0).astype(type_D) * tensor_D\n",
"np.testing.assert_array_equal(relu_ref, tensor_D_relu)"
]
},
{
"cell_type": "markdown",
"id": "cf959171",
"metadata": {},
"source": [
"## Other element-wise activation functions\n",
"CUTLASS supports a variety of widely-used element-wise activation functions. We can obtain a list of these functions via the `get_activations()` method."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9e17d730",
"metadata": {},
"outputs": [],
"source": [
"activations = plan.activations()\n",
"for activation in activations:\n",
" print(activation)"
]
},
{
"cell_type": "markdown",
"id": "0e4599fa",
"metadata": {},
"source": [
"We can then run each of them:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9c3598c9",
"metadata": {},
"outputs": [],
"source": [
"for activation in activations:\n",
" print('=============================================================================================')\n",
" print(f'Compiling and running activation {activation}')\n",
" print('=============================================================================================')\n",
" plan.activation = activation\n",
" plan.run(tensor_A, tensor_B, tensor_C, tensor_D, print_module=print_module)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "751f8d92",
"metadata": {},
"outputs": [],
"source": []
}
],
"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.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

264
examples/python/02_pytorch_extension_grouped_gemm.ipynb Normal file
View File

@@ -0,0 +1,264 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "6acbea5d",
"metadata": {},
"source": [
"# Exporting a CUTLASS grouped GEMM kernel to a PyTorch CUDA extension\n",
"This notebook walks through a basic example of using the CUTLASS Python interface to declare\n",
"a grouped GEMM kernel and export it as a PyTorch CUDA extension.\n",
"\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cutlass/tree/master/examples/00_basic_gemm.ipynb)\n",
"\n",
"## Background on grouped GEMM\n",
"Grouped GEMM enables one to execute a set of GEMMs (each with potentially different sizes and strides)\n",
"in a single CUDA kernel. It can be thought of as a generalized version of a pointer-array GEMM,\n",
"without the requirement that the sizes and strides of each GEMM be the same.\n",
"\n",
"For example, if one has `p` GEMMs with sizes:\n",
"```text\n",
"M_1 x N_1 x K_1\n",
"M_2 x N_2 x K_2\n",
"...\n",
"M_p x N_p x K_p\n",
"```\n",
"CUTLASS's grouped GEMM will execute these in a single CUDA kernel.\n",
"\n",
"Grouped GEMM is particularly beneficial for saturating the GPU with many small problems that would\n",
"insufficiently utilize the device in isolation.\n",
"\n",
"## Declaring a grouped GEMM via the CUTLASS Python interface\n",
"A grouped GEMM operation is declared similarly to a GEMM operation in the CUTLASS Python interface: one\n",
"simply calls `cutlass.op.GroupedGemm`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fdcf21d8",
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import torch\n",
"\n",
"dtype = torch.float16\n",
"plan = cutlass.op.GroupedGemm(element=dtype, layout=cutlass.LayoutType.RowMajor)"
]
},
{
"cell_type": "markdown",
"id": "514f40a4",
"metadata": {},
"source": [
"We can then compile and run this operation on a group of GEMMs. We'll first set up some utility functions to initialize GEMMs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c2a7371e",
"metadata": {},
"outputs": [],
"source": [
"import random\n",
"random.seed(2023)\n",
"\n",
"# Utility function to initialize A, B, C, and D matrices corresponding to dimensions M, N, and K\n",
"def initialize(dtype, M, N, K):\n",
" sizes = [(M, K), (K, N), (M, N), (M, N)]\n",
" return [torch.randint(-3, 3, size, device='cuda').to(dtype) for size in sizes]\n",
"\n",
"# Utility function to generate `problems` GEMMs of random sizes\n",
"def generate_problems(problems):\n",
" valid_sizes = [128, 256, 512, 1024]\n",
" As, Bs, Cs, Ds = [], [], [], []\n",
" for _ in range(problems):\n",
" M, N, K = [random.choice(valid_sizes) for _ in range(3)]\n",
" A, B, C, D = initialize(dtype, M, N, K)\n",
" As.append(A)\n",
" Bs.append(B)\n",
" Cs.append(C)\n",
" Ds.append(D)\n",
" return As, Bs, Cs, Ds"
]
},
{
"cell_type": "markdown",
"id": "590a3bc5",
"metadata": {},
"source": [
"We'll next run a group of 50 GEMMs via the CUTLASS Python interface and via PyTorch."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "776c9233",
"metadata": {},
"outputs": [],
"source": [
"As, Bs, Cs, Ds, = generate_problems(50)\n",
"\n",
"plan.run(As, Bs, Cs, Ds, print_module=True)\n",
"Ds_torch = [a @ b for a, b in zip(As, Bs)]\n",
"\n",
"for d, d_torch in zip(Ds, Ds_torch):\n",
" assert torch.allclose(d, d_torch)"
]
},
{
"cell_type": "markdown",
"id": "766e4f03",
"metadata": {},
"source": [
"## Exporting the CUTLASS kernel to a PyTorch CUDA extension\n",
"The procedure above allows one to quickly experiment with using a CUTLASS kernels However, one might prefer to use the CUTLASS kernel via a [PyTorch CUDA extension](https://pytorch.org/tutorials/advanced/cpp_extension.html). This will avoids adding any runtime overheads associated with the Python portions of the CUTLASS Python interface.\n",
"\n",
"The CUTLASS Python interface provides simple solutions for creating PyTorch CUDA extensions for a CUTLASS kernel. These extensions can either be written out for a later \"ahead-of-time\" compilation, or be just-in-time compiled and returned to the user.\n",
"\n",
"To create a JIT-compiled module from the CUTLASS kernel we defined above, simply call the following:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3a98dee6",
"metadata": {},
"outputs": [],
"source": [
"op = plan.construct()\n",
"grouped_gemm = cutlass.emit.pytorch(op, name='grouped_gemm', cc=plan.cc, sourcedir='out', jit=True)"
]
},
{
"cell_type": "markdown",
"id": "c8ca3991",
"metadata": {},
"source": [
"The `cutlass.emit.pytorch` function emits:\n",
"* `out/grouped_gemm_kernel.cu`: This file contains the declaration of the CUTLASS kernel and a method to call it from PyTorch tensors\n",
"* `out/grouped_gemm.cpp`: This file contains a C++ wrapper around the aforementioned CUTLASS kernel\n",
"* `setup.py`: This file contains the `setuptools` script for building and installing the generated extension\n",
"\n",
"The extension can be build from within the `module_output` directory by running:\n",
"```bash\n",
"TORCH_CUDA_ARCH_LIST=\"8.0\" python setup.py install\n",
"```\n",
"Where `TORCH_ARCH_LIST` is set to the compute capability of the device on which the kernel will be run.\n",
"\n",
"See the PyTorch [\"Custom C++ and CUDA Extensions\"](https://pytorch.org/tutorials/advanced/cpp_extension.html) tutorial for more details on this.\n",
"\n",
"The PyTorch CUDA extension could be built for this module by running:\n",
"```bash\n",
"cd out\n",
"TORCH_CUDA_ARCH_LIST=\"8.0\" python setup.py\n",
"```\n",
"(assuming that one is building for SM80)\n",
"\n",
"One could then use the kernel in a later PyTorch module by running:\n",
"\n",
"```python\n",
"import torch\n",
"import grouped_gemm\n",
"\n",
"grouped_gemm.run(As, Bs)\n",
"```\n",
"\n",
"In this case, however, we set `jit=True`, which specifies that we would like to compile and load the PyTorch CUDA extension on the fly.\n",
"Under the hood, this leverages the [torch.utils.cpp_extension.load](https://pytorch.org/tutorials/advanced/cpp_extension.html) method\n",
"and returns back the loaded extension.\n",
"\n",
"We can then use the extension and compare its results to running the GEMMs via vanilla PyTorch GEMMs:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cecb26a4",
"metadata": {},
"outputs": [],
"source": [
"Ds = grouped_gemm.run(As, Bs)\n",
"Ds_torch = [a @ b for a, b in zip(As, Bs)]\n",
"for d, d_torch in zip(Ds, Ds_torch):\n",
" assert torch.allclose(d, d_torch)"
]
},
{
"cell_type": "markdown",
"id": "50db80e4",
"metadata": {},
"source": [
"Finally, we can profile our grouped GEMM extension:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b76805d3",
"metadata": {},
"outputs": [],
"source": [
"num_warmup = 20\n",
"num_profile = 100\n",
"\n",
"# Warmup iterations\n",
"for _ in range(num_warmup):\n",
" Ds = grouped_gemm.run(As, Bs)\n",
" Ds_torch = [a @ b for a, b in zip(As, Bs)]\n",
" torch.cuda.synchronize()\n",
"\n",
"# Timing iterations\n",
"import time\n",
"grouped = 0\n",
"nongrouped = 0\n",
"for _ in range(num_profile):\n",
" start = time.time()\n",
" Ds = grouped_gemm.run(As, Bs)\n",
" torch.cuda.synchronize()\n",
" grouped += time.time() - start\n",
"\n",
" start = time.time()\n",
" Ds_torch = [a @ b for a, b in zip(As, Bs)]\n",
" torch.cuda.synchronize()\n",
" nongrouped += time.time() - start\n",
"\n",
"print('Grouped: {:.3f} us'.format(grouped * 1e6/num_profile))\n",
"print('Non-Grouped: {:.3f} us'.format(nongrouped * 1e6/num_profile))\n",
"print('Speedup: {:.3f}'.format(nongrouped / grouped))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f22fc696",
"metadata": {},
"outputs": [],
"source": []
}
],
"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.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

14
examples/python/README.md Normal file
View File

@@ -0,0 +1,14 @@
# Examples of using the CUTLASS Python interface
* [00_basic_gemm](/examples/python/00_basic_gemm.ipynb)
Shows how declare, configure, compile, and run a CUTLASS GEMM using the Python interface
* [01_epilogue](/examples/python/01_epilogue.ipynb)
Shows how to fuse elementwise activation functions to GEMMs via the Python interface
* [02_pytorch_extension_grouped_gemm](/examples/python/02_pytorch_extension_grouped_gemm.ipynb)
Shows how to declare, compile, and run a grouped GEMM operation via the Python interface,
along with how the emitted kernel can be easily exported to a PyTorch CUDA extension.