ANIKET SHIVAM
co-authored by
Aniket Shivam
Aniket Shivam
parent
9b8166e3f0
commit
d572cc1aab
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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()],
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user