CUTLASS 3.2.1 (#1113)
* Updates for 3.2.1 release. * Minor fix in gemm op profiler for raster order. * Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
387
test/python/cutlass/gemm/gemm_testbed.py
Normal file
387
test/python/cutlass/gemm/gemm_testbed.py
Normal file
@@ -0,0 +1,387 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
from math import prod
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
|
||||
from cutlass import (
|
||||
DataType,
|
||||
DataTypeSize,
|
||||
GemmUniversalMode,
|
||||
LayoutType,
|
||||
OpcodeClass,
|
||||
ShortDataTypeNames,
|
||||
SwizzlingFunctor
|
||||
)
|
||||
|
||||
from cutlass.backend import compiler
|
||||
from cutlass.backend.gemm_operation import GemmArguments, GemmOperationUniversal
|
||||
from cutlass.backend.memory_manager import get_allocated_size
|
||||
from cutlass.backend.reduction_operation import ReductionArguments, ReductionOperation
|
||||
from cutlass.shape import GemmCoord, MatrixCoord
|
||||
from cutlass.utils.datatypes import torch_type
|
||||
|
||||
|
||||
class GemmUniversalLauncher:
|
||||
def __init__(
|
||||
self,
|
||||
operation,
|
||||
seed=2080,
|
||||
verification=True,
|
||||
iterations=500,
|
||||
compiler_mode= "nvcc",
|
||||
**kwargs,
|
||||
) -> None:
|
||||
# Create the reduction kernel, if needed
|
||||
self.reduction_operation: ReductionOperation = ReductionOperation(
|
||||
shape=MatrixCoord(4, 32 * operation.C.alignment),
|
||||
C=operation.C,
|
||||
element_accumulator=operation.tile_description.math_instruction.element_accumulator,
|
||||
element_compute=operation.epilogue_functor.element_epilogue,
|
||||
epilogue_functor=operation.epilogue_functor,
|
||||
count=operation.C.alignment,
|
||||
)
|
||||
|
||||
self.math_operation = operation.tile_description.math_instruction.math_operation
|
||||
self.verification = verification
|
||||
|
||||
if compiler_mode == "nvcc":
|
||||
compiler.nvcc()
|
||||
elif compiler_mode == "nvrtc":
|
||||
compiler.nvrtc()
|
||||
else:
|
||||
raise Exception(f"Unexpected compiler string {compiler_mode}")
|
||||
|
||||
op_list = [operation]
|
||||
if operation.arch < 90:
|
||||
# Split K via Python is currently only supported for pre-SM90 kernels
|
||||
op_list.append(self.reduction_operation)
|
||||
|
||||
compiler.add_module(op_list, bypass_cache=False)
|
||||
|
||||
self.operation = operation
|
||||
|
||||
self.dtype_A = torch_type(operation.A.element)
|
||||
self.dtype_B = torch_type(operation.B.element)
|
||||
self.dtype_C = torch_type(operation.C.element)
|
||||
self.dtype_D = torch_type(operation.C.element)
|
||||
|
||||
accumulator_size = DataTypeSize[operation.tile_description.math_instruction.element_accumulator]
|
||||
element_size = DataTypeSize[operation.A.element]
|
||||
|
||||
if element_size == 1:
|
||||
self.rand_max = 1
|
||||
self.rand_min = 0
|
||||
elif element_size <= 8:
|
||||
self.rand_max = 1
|
||||
self.rand_min = -1
|
||||
elif element_size == 16:
|
||||
self.rand_max = 4
|
||||
self.rand_min = -4
|
||||
else:
|
||||
self.rand_max = 8
|
||||
self.rand_min = -8
|
||||
|
||||
self.seed = seed
|
||||
|
||||
self.compute_type = operation.epilogue_functor.element_epilogue
|
||||
self.accumulator_type = operation.tile_description.math_instruction.element_accumulator
|
||||
|
||||
def print_problem_size(self, p, mode, batch_count):
|
||||
if mode == GemmUniversalMode.Gemm:
|
||||
mode = "Gemm"
|
||||
elif mode == GemmUniversalMode.Batched:
|
||||
mode = "GemmBatched"
|
||||
elif mode == GemmUniversalMode.GemmSplitKParallel:
|
||||
mode = "GemmSplitKParallel"
|
||||
print(f"problem: {p.m}, {p.n}, {p.k}\n batch_count: {batch_count}\n mode: {mode}")
|
||||
|
||||
def uniform_init(self, shape, dtype, layout):
|
||||
size = prod(shape)
|
||||
if dtype.is_floating_point:
|
||||
data = torch.ceil(torch.empty(size=(size,), dtype=dtype, device="cuda").uniform_(self.rand_min - 0.5, self.rand_max - 0.5))
|
||||
else:
|
||||
# PyTorch does not currently support integer-typed matrix multiplications on GPU.
|
||||
# Fall back to CPU for integer type references.
|
||||
data = torch.empty(size=(size,), dtype=dtype, device="cpu").random_(self.rand_min, self.rand_max + 1)
|
||||
|
||||
if dtype == torch.float64 or dtype == torch.float32:
|
||||
data = data.to("cpu")
|
||||
|
||||
data_ref = data.reshape(shape)
|
||||
|
||||
if layout == LayoutType.RowMajor:
|
||||
data_cutlass = data_ref
|
||||
else:
|
||||
data_cutlass = data_ref.transpose(-1, -2).contiguous()
|
||||
|
||||
data_cutlass = data_cutlass.to("cuda")
|
||||
return data_cutlass, data_ref
|
||||
|
||||
def reference(self, problem_size, tensor_A, tensor_B, tensor_C, alpha, beta):
|
||||
# If any tensor is on CPU, place all tensors on CPU unless only
|
||||
# tensor C is on CPU
|
||||
devices = [x.device.type for x in [tensor_A, tensor_B, tensor_C]]
|
||||
if "cpu" in devices and devices != ["cuda", "cuda", "cpu"]:
|
||||
device = torch.device("cpu")
|
||||
else:
|
||||
device = tensor_A.device
|
||||
|
||||
tensor_A = tensor_A.to(device)
|
||||
tensor_B = tensor_B.to(device)
|
||||
tensor_C = tensor_C.to(device)
|
||||
|
||||
dtype = torch_type(self.compute_type)
|
||||
alpha_torch = torch.tensor([alpha], device=device).to(dtype)
|
||||
beta_torch = torch.tensor([beta], device=device).to(dtype)
|
||||
|
||||
tmp = tensor_A @ tensor_B
|
||||
tensor_D_ref = (alpha_torch * tmp) + (tensor_C * beta_torch)
|
||||
return tensor_D_ref.to(self.dtype_D)
|
||||
|
||||
def run(self, mode, problem_size, batch_count=1, split_k_slices=1, alpha=1.0, beta=0.0):
|
||||
torch.random.manual_seed(self.seed)
|
||||
|
||||
# Assign an actual batch count in cases where we are not running in batched mode.
|
||||
# This is to differentiate between the number of split K slices and the batch count,
|
||||
# which are overloaded within the single `batch_count` variable.
|
||||
if mode == GemmUniversalMode.Batched:
|
||||
true_batch_count = batch_count
|
||||
else:
|
||||
true_batch_count = 1
|
||||
|
||||
def transpose(layout):
|
||||
if layout == LayoutType.RowMajor:
|
||||
return LayoutType.ColumnMajor
|
||||
else:
|
||||
return LayoutType.RowMajor
|
||||
|
||||
tensor_A, tensor_A_ref = self.uniform_init(
|
||||
(true_batch_count, problem_size.m, problem_size.k),
|
||||
self.dtype_A,
|
||||
self.operation.A.layout if not self.operation.switched else transpose(self.operation.B.layout),
|
||||
)
|
||||
tensor_B, tensor_B_ref = self.uniform_init(
|
||||
(true_batch_count, problem_size.k, problem_size.n),
|
||||
self.dtype_B,
|
||||
self.operation.B.layout if not self.operation.switched else transpose(self.operation.A.layout),
|
||||
)
|
||||
tensor_C, tensor_C_ref = self.uniform_init(
|
||||
(true_batch_count, problem_size.m, problem_size.n),
|
||||
self.dtype_C,
|
||||
self.operation.C.layout if not self.operation.switched else transpose(self.operation.C.layout),
|
||||
)
|
||||
tensor_D = torch.zeros_like(tensor_C)
|
||||
|
||||
if self.compute_type in [DataType.s8, DataType.s32, DataType.u8, DataType.u32]:
|
||||
alpha = int(alpha)
|
||||
beta = int(beta)
|
||||
|
||||
#
|
||||
# Launch kernel
|
||||
#
|
||||
|
||||
arguments = GemmArguments(
|
||||
operation=self.operation,
|
||||
problem_size=problem_size,
|
||||
A=tensor_A,
|
||||
B=tensor_B,
|
||||
C=tensor_C,
|
||||
D=tensor_D,
|
||||
output_op=self.operation.epilogue_type(alpha, beta),
|
||||
gemm_mode=mode,
|
||||
split_k_slices=split_k_slices,
|
||||
batch=batch_count,
|
||||
)
|
||||
|
||||
if mode == GemmUniversalMode.GemmSplitKParallel:
|
||||
reduction_arguments = ReductionArguments(
|
||||
self.reduction_operation,
|
||||
problem_size=[problem_size.m, problem_size.n],
|
||||
partitions=split_k_slices,
|
||||
workspace=arguments.ptr_D,
|
||||
destination=tensor_D,
|
||||
source=tensor_C,
|
||||
output_op=self.reduction_operation.epilogue_type(alpha, beta),
|
||||
)
|
||||
|
||||
self.operation.run(arguments)
|
||||
|
||||
if mode == GemmUniversalMode.GemmSplitKParallel:
|
||||
self.reduction_operation.run(reduction_arguments)
|
||||
|
||||
passed = True
|
||||
|
||||
if self.verification:
|
||||
if mode == GemmUniversalMode.GemmSplitKParallel:
|
||||
reduction_arguments.sync()
|
||||
else:
|
||||
arguments.sync()
|
||||
tensor_D_ref = self.reference(
|
||||
problem_size,
|
||||
tensor_A_ref,
|
||||
tensor_B_ref,
|
||||
tensor_C_ref,
|
||||
alpha,
|
||||
beta,
|
||||
)
|
||||
|
||||
tensor_D_ref = tensor_D_ref.to('cuda')
|
||||
|
||||
if self.operation.switched or self.operation.C.layout == LayoutType.ColumnMajor:
|
||||
tensor_D = tensor_D.transpose(-1, -2).contiguous()
|
||||
|
||||
passed = tensor_D.equal(tensor_D_ref)
|
||||
|
||||
try:
|
||||
assert passed
|
||||
except AssertionError:
|
||||
self.print_problem_size(problem_size, mode, batch_count)
|
||||
del arguments
|
||||
if mode == GemmUniversalMode.GemmSplitKParallel:
|
||||
del reduction_arguments
|
||||
|
||||
cur_size = get_allocated_size()
|
||||
assert cur_size == 0, f"{cur_size} B of memory were not released after this run"
|
||||
|
||||
return passed
|
||||
|
||||
|
||||
def test_all_gemm(operation: "GemmOperationUniversal", testcase="universal", compilation_mode="nvcc"):
|
||||
passed = True
|
||||
|
||||
minimum_operand_element_size = min(
|
||||
DataTypeSize[operation.A.element], DataTypeSize[operation.B.element]
|
||||
)
|
||||
opcode_class = operation.tile_description.math_instruction.opcode_class
|
||||
|
||||
if opcode_class == OpcodeClass.Simt:
|
||||
alignment = 1
|
||||
else:
|
||||
alignment = 128 // minimum_operand_element_size
|
||||
|
||||
alignment_m = alignment
|
||||
alignment_n = alignment
|
||||
alignment_k = alignment
|
||||
|
||||
# INT8 alignment constraints
|
||||
if opcode_class == OpcodeClass.Simt:
|
||||
A_is_s8 = operation.A.element == DataType.s8
|
||||
B_is_s8 = operation.B.element == DataType.s8
|
||||
|
||||
if A_is_s8 and operation.A.layout == LayoutType.ColumnMajor:
|
||||
alignment_m = 4
|
||||
if B_is_s8 == DataType.s8 and operation.A.layout == LayoutType.RowMajor:
|
||||
alignment_n = 4
|
||||
if A_is_s8 and B_is_s8 and (operation.A.layout == LayoutType.RowMajor or operation.B.layout == LayoutType.ColumnMajor):
|
||||
alignment_k = 4
|
||||
|
||||
threadblock_k = operation.tile_description.threadblock_shape[2]
|
||||
|
||||
assert testcase != "interleaved"
|
||||
|
||||
supports_split_k = operation.arch < 90 and not operation.swizzling_functor == SwizzlingFunctor.StreamK
|
||||
|
||||
if testcase == "multistage":
|
||||
modes = [GemmUniversalMode.Gemm]
|
||||
problem_size_m = [16, 528]
|
||||
problem_size_n = [16, 528]
|
||||
problem_size_k = [
|
||||
threadblock_k,
|
||||
threadblock_k * operation.tile_description.stages
|
||||
+ operation.tile_description.math_instruction.instruction_shape[2],
|
||||
]
|
||||
problem_alpha = [1.0]
|
||||
problem_beta = [0.0]
|
||||
batch_counts = [1]
|
||||
else:
|
||||
modes = [GemmUniversalMode.Gemm]
|
||||
batch_counts = [1, 2, 3, 5, 7]
|
||||
if supports_split_k:
|
||||
modes.append(GemmUniversalMode.GemmSplitKParallel)
|
||||
|
||||
problem_size_m = [alignment_m, 512 - 3 * alignment_m]
|
||||
problem_size_n = [alignment_n, 512 - 2 * alignment_n]
|
||||
if operation.tile_description.stages is None:
|
||||
stages_for_k_calc = 7
|
||||
else:
|
||||
stages_for_k_calc = operation.tile_description.stages
|
||||
problem_size_k = [
|
||||
alignment_k,
|
||||
threadblock_k * stages_for_k_calc - alignment_k,
|
||||
threadblock_k * stages_for_k_calc * 3 - alignment_k,
|
||||
]
|
||||
problem_alpha = [1.0]
|
||||
problem_beta = [2.0]
|
||||
|
||||
testbed = GemmUniversalLauncher(operation, compiler_mode=compilation_mode)
|
||||
|
||||
for mode in modes:
|
||||
for m in problem_size_m:
|
||||
for n in problem_size_n:
|
||||
for k in problem_size_k:
|
||||
for batch_count in batch_counts:
|
||||
for alpha in problem_alpha:
|
||||
for beta in problem_beta:
|
||||
# skip very small K problems
|
||||
if testcase == "universal":
|
||||
if k // batch_count < 2 * threadblock_k:
|
||||
continue
|
||||
|
||||
problem_size = GemmCoord(m, n, k)
|
||||
|
||||
if supports_split_k:
|
||||
split_k_slices = batch_count
|
||||
else:
|
||||
split_k_slices = 1
|
||||
|
||||
overridden_mode = mode
|
||||
if mode == GemmUniversalMode.Gemm and batch_count > 1:
|
||||
overridden_mode = GemmUniversalMode.Batched
|
||||
|
||||
passed = testbed.run(
|
||||
overridden_mode,
|
||||
problem_size,
|
||||
batch_count,
|
||||
split_k_slices,
|
||||
alpha,
|
||||
beta,
|
||||
)
|
||||
|
||||
if not passed:
|
||||
return False
|
||||
|
||||
return passed
|
||||
Reference in New Issue
Block a user