Rename python/cutlass to python/cutlass_cppgen (#2652)
This commit is contained in:
Jack Kosaian
committed by
Haicheng Wu
Haicheng Wu
parent
4260d4aef9
commit
177a82e251
700
python/cutlass_cppgen/backend/conv2d_operation.py
Normal file
700
python/cutlass_cppgen/backend/conv2d_operation.py
Normal file
@ -0,0 +1,700 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2025 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 __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
from typing import Union
|
||||
|
||||
from cutlass_cppgen.utils.lazy_import import lazy_import
|
||||
cuda = lazy_import("cuda.cuda")
|
||||
from cutlass_library import SubstituteTemplate
|
||||
import numpy as np
|
||||
|
||||
from cutlass_library import (
|
||||
ConvKindNames,
|
||||
ConvKindTag,
|
||||
DataTypeNames,
|
||||
DataTypeSize,
|
||||
DataTypeTag,
|
||||
IteratorAlgorithmNames,
|
||||
IteratorAlgorithmTag,
|
||||
LayoutTag,
|
||||
LayoutType,
|
||||
MathOperation,
|
||||
MathOperationTag,
|
||||
OpcodeClass,
|
||||
OpcodeClassNames,
|
||||
OpcodeClassTag,
|
||||
OperationKind,
|
||||
ShortDataTypeNames,
|
||||
ShortLayoutTypeNames,
|
||||
SplitKMode,
|
||||
StrideSupport,
|
||||
StrideSupportTag,
|
||||
SwizzlingFunctor,
|
||||
SwizzlingFunctorTag,
|
||||
get_complex_from_real,
|
||||
)
|
||||
|
||||
from cutlass_cppgen.backend.arguments import ArgumentBase
|
||||
from cutlass_cppgen.backend.c_types import dim3_, get_conv2d_arguments
|
||||
from cutlass_cppgen.backend.library import (
|
||||
EmissionType,
|
||||
TensorDescription,
|
||||
TileDescription,
|
||||
)
|
||||
from cutlass_cppgen.backend.memory_manager import device_mem_alloc
|
||||
from cutlass_cppgen.backend.operation import ExecutableOperation, LaunchConfiguration
|
||||
from cutlass_cppgen.backend.utils.device import to_device_ptr
|
||||
from cutlass_cppgen.shape import GemmCoord
|
||||
|
||||
|
||||
class Conv2dArguments(ArgumentBase):
|
||||
"""
|
||||
Argument wrapper for Conv2d. It encodes problem information and
|
||||
user-provide tensors into the kernel's argument.
|
||||
|
||||
:param operation: the Conv2d operation to take the argument
|
||||
:type operation: :class:`cutlass_cppgen.backend.Conv2dOperation`
|
||||
:param problem_size: the Conv2d problem size
|
||||
:type problem_size: :class:`cutlass_cppgen.shape.Conv2dProblemSize`
|
||||
:param A: tensor A
|
||||
:type A: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray
|
||||
:param B: tensor B
|
||||
:type B: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray
|
||||
:param C: tensor C
|
||||
:type C: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray
|
||||
:param D: tensor D
|
||||
:type D: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray
|
||||
:param split_k_mode: conv2d split K mode, defaults to cutlass_library.library.SplitKMode.Serial
|
||||
:type split_k_mode: cutlass_library.library.SplitKMode, optional
|
||||
:param output_op: output operator, optional
|
||||
:type output_op: :class:`cutlass_cppgen.backend.LinearCombinationFunctorArguments`
|
||||
:param stream: cuda stream, defaults to cuda.cuda.CUstream(0)
|
||||
:type stream: :class:`cuda.cuda.CUstream`
|
||||
"""
|
||||
|
||||
def __init__(self, operation, problem_size, A, B, C, D,
|
||||
split_k_mode=SplitKMode.Serial, **kwargs, ) -> None:
|
||||
self.operation = operation
|
||||
self.conv_kind = operation.conv_kind
|
||||
self.layout_A = operation.A.layout
|
||||
self.layout_B = operation.B.layout
|
||||
self.layout_C = operation.C.layout
|
||||
|
||||
self.element_A = operation.A.element
|
||||
self.element_B = operation.B.element
|
||||
self.element_C = operation.C.element
|
||||
|
||||
if self.layout_C == LayoutType.TensorNC32HW32:
|
||||
raise Exception("Layout type TensorNC32HW32 is not currently supported")
|
||||
|
||||
super().__init__(A, B, C, D, **kwargs)
|
||||
|
||||
if "split_k_slices" in kwargs.keys() and kwargs["split_k_slices"] > 1:
|
||||
self.split_k_mode = split_k_mode
|
||||
self.split_k_slices = kwargs["split_k_slices"]
|
||||
else:
|
||||
self.split_k_mode = SplitKMode.Serial
|
||||
self.split_k_slices = 1
|
||||
|
||||
if "output_op" in kwargs.keys() and self.split_k_mode != SplitKMode.Parallel:
|
||||
self.output_op = kwargs["output_op"]
|
||||
else:
|
||||
self.output_op = self.operation.epilogue_type(1.0, 0.0)
|
||||
|
||||
self.problem_size = problem_size
|
||||
self.problem_size.split_k_slices = self.split_k_slices
|
||||
|
||||
self.initialize()
|
||||
|
||||
def get_arguments(self):
|
||||
tc_numel = -1
|
||||
if hasattr(self, "tensor_c_numel"):
|
||||
tc_numel = self.tensor_c_numel
|
||||
|
||||
self.c_arguments = self.operation.argument_type(
|
||||
int(self.conv_kind),
|
||||
self.problem_size.ctype,
|
||||
int(to_device_ptr(self.ptr_A)),
|
||||
int(to_device_ptr(self.ptr_B)),
|
||||
int(to_device_ptr(self.ptr_C)),
|
||||
int(to_device_ptr(self.ptr_D)),
|
||||
tc_numel,
|
||||
self.output_op,
|
||||
int(self.split_k_mode)
|
||||
)
|
||||
|
||||
def initialize(self):
|
||||
self.launch_config = self.operation.rt_module.plan(self)
|
||||
|
||||
self.get_arguments()
|
||||
|
||||
# Allocate and initialize device workspace
|
||||
device_workspace_size = self.operation.rt_module.get_workspace_size(self.c_arguments)
|
||||
if device_workspace_size > 0:
|
||||
self.workspace_buffer = device_mem_alloc(device_workspace_size)
|
||||
workspace_ptr = self.workspace_buffer.ptr
|
||||
err, = cuda.cuMemsetD32(
|
||||
workspace_ptr, 0, device_workspace_size // 4)
|
||||
else:
|
||||
workspace_ptr = None
|
||||
|
||||
self.semaphore = 0
|
||||
if workspace_ptr is not None and self.split_k_mode == SplitKMode.Parallel:
|
||||
self.ptr_D = workspace_ptr
|
||||
# Reset arguments now that ptr_D has been updated
|
||||
self.get_arguments()
|
||||
elif workspace_ptr is not None and self.split_k_mode == SplitKMode.Serial:
|
||||
self.semaphore = workspace_ptr
|
||||
|
||||
params_ = self.operation.rt_module.get_args(
|
||||
self.c_arguments, ctypes.c_void_p(int(self.semaphore)))
|
||||
self.host_workspace = bytearray(params_.contents)
|
||||
self.device_workspace = None
|
||||
|
||||
def sync(self):
|
||||
"""
|
||||
Synchronize the arguments. If the input tensor is in host,
|
||||
copy it from device to host.
|
||||
"""
|
||||
return super().sync()
|
||||
|
||||
|
||||
class Conv2dRT(ExecutableOperation):
|
||||
"""
|
||||
Conv2dRT manages the CUTLASS runtime components
|
||||
"""
|
||||
|
||||
KernelTemplate = r"""
|
||||
extern "C"
|
||||
__global__ void
|
||||
${operation_name}(${operation_name}${operation_suffix}::Params params) {
|
||||
|
||||
// Dynamic shared memory base pointer
|
||||
extern __shared__ int SharedStorageBase[];
|
||||
|
||||
// Declare pointer to dynamic shared memory.
|
||||
${operation_name}${operation_suffix}::SharedStorage *shared_storage =
|
||||
reinterpret_cast<${operation_name}${operation_suffix}::SharedStorage *>(SharedStorageBase);
|
||||
|
||||
${operation_name}${operation_suffix} op;
|
||||
|
||||
op(params, *shared_storage);
|
||||
}
|
||||
"""
|
||||
|
||||
HostTemplate = r"""
|
||||
extern "C" {
|
||||
// Get the size of params in bytes
|
||||
int ${operation_name}_get_param_size(){
|
||||
return sizeof(${operation_name}${operation_suffix}::Params);
|
||||
}
|
||||
|
||||
// Get the size of dynamic shared memory in bytes
|
||||
int ${operation_name}_shared_memory_size() {
|
||||
return int(sizeof(${operation_name}${operation_suffix}::SharedStorage));
|
||||
}
|
||||
|
||||
using ElementA = typename ${operation_name}_base::ElementA;
|
||||
using ElementB = typename ${operation_name}_base::ElementB;
|
||||
using ElementC = typename ${operation_name}_base::ElementC;
|
||||
using LayoutA = typename ${operation_name}_base::LayoutA;
|
||||
using LayoutB = typename ${operation_name}_base::LayoutB;
|
||||
using LayoutC = typename ${operation_name}_base::LayoutC;
|
||||
using EpilogueOutputOp = typename ${operation_name}_base::EpilogueOutputOp;
|
||||
|
||||
struct ${operation_name}_TemporaryArgs {
|
||||
int conv_kind;
|
||||
cutlass::conv::Conv2dProblemSize problem_size;
|
||||
ElementA* ptr_A;
|
||||
ElementB* ptr_B;
|
||||
ElementC* ptr_C;
|
||||
ElementC* ptr_D;
|
||||
int tensor_c_numel;
|
||||
typename EpilogueOutputOp::Params epilogue_params;
|
||||
int split_k_mode;
|
||||
};
|
||||
|
||||
typename ${operation_name}${operation_suffix}::Arguments
|
||||
construct_arguments(${operation_name}_TemporaryArgs args) {
|
||||
cutlass::conv::Operator conv_operator = static_cast<cutlass::conv::Operator>(args.conv_kind);
|
||||
auto tc_A = cutlass::conv::implicit_gemm_tensor_a_extent(conv_operator, args.problem_size);
|
||||
auto tc_B = cutlass::conv::implicit_gemm_tensor_b_extent(conv_operator, args.problem_size);
|
||||
auto tc_C = cutlass::conv::implicit_gemm_tensor_c_extent(conv_operator, args.problem_size);
|
||||
auto tc_D = cutlass::conv::implicit_gemm_tensor_c_extent(conv_operator, args.problem_size);
|
||||
|
||||
auto size_C = tc_C.at(0) * tc_C.at(1) * tc_C.at(2) * tc_C.at(3);
|
||||
if (args.tensor_c_numel >= 0 && args.tensor_c_numel == tc_C.at(3) && args.tensor_c_numel < size_C) {
|
||||
// C is interpreted as bias
|
||||
tc_C = {0, 0, 0, 0};
|
||||
}
|
||||
|
||||
cutlass::TensorRef<ElementA, LayoutA> tref_A(args.ptr_A, LayoutA::packed(tc_A));
|
||||
cutlass::TensorRef<ElementB, LayoutA> tref_B(args.ptr_B, LayoutB::packed(tc_B));
|
||||
cutlass::TensorRef<ElementC, LayoutA> tref_C(args.ptr_C, LayoutC::packed(tc_C));
|
||||
cutlass::TensorRef<ElementC, LayoutA> tref_D(args.ptr_D, LayoutC::packed(tc_D));
|
||||
|
||||
return {
|
||||
args.problem_size,
|
||||
tref_A,
|
||||
tref_B,
|
||||
tref_C,
|
||||
tref_D,
|
||||
args.epilogue_params,
|
||||
static_cast<cutlass::conv::SplitKMode>(args.split_k_mode)
|
||||
};
|
||||
}
|
||||
|
||||
// Get the params as byte array
|
||||
char* ${operation_name}_get_params(${operation_name}_TemporaryArgs args, int *semaphore=nullptr) {
|
||||
auto arguments = construct_arguments(args);
|
||||
typename ${operation_name}${operation_suffix}::Params* params;
|
||||
params = new ${operation_name}${operation_suffix}::Params(arguments, semaphore);
|
||||
|
||||
char *bytes = ((char*)(params));
|
||||
char *output = new char[sizeof(${operation_name}${operation_suffix}::Params)];
|
||||
for (unsigned int i = 0; i < sizeof(${operation_name}${operation_suffix}::Params); i ++)
|
||||
output[i] = bytes[i];
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
dim3 ${operation_name}_get_grid_shape(
|
||||
int conv_kind,
|
||||
cutlass::conv::Conv2dProblemSize problem_size,
|
||||
cutlass::gemm::GemmCoord tile_size,
|
||||
int split_k_slices
|
||||
) {
|
||||
|
||||
using Swizzle = typename ${operation_name}_base::ThreadblockSwizzle;
|
||||
auto tiled_shape = Swizzle::get_tiled_shape(
|
||||
static_cast<cutlass::conv::Operator>(conv_kind),
|
||||
problem_size,
|
||||
tile_size,
|
||||
split_k_slices);
|
||||
|
||||
return Swizzle::get_grid_shape(tiled_shape);
|
||||
}
|
||||
|
||||
size_t ${operation_name}_get_workspace_size(${operation_name}_TemporaryArgs args) {
|
||||
auto arguments = construct_arguments(args);
|
||||
|
||||
// Temporarily define device::-level Conv2d so that we can call get_workspace_size
|
||||
using DeviceConv = cutlass::conv::device::ImplicitGemmConvolution<${operation_name}_base>;
|
||||
return DeviceConv::get_workspace_size(arguments);
|
||||
}
|
||||
}
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, operation: "Conv2dOperation"):
|
||||
super().__init__(operation)
|
||||
self.extra_funcs = {
|
||||
"get_grid_shape": dim3_,
|
||||
"get_workspace_size": ctypes.c_uint64
|
||||
}
|
||||
self.argument_type, self.epilogue_type = get_conv2d_arguments(operation.epilogue_functor)
|
||||
self.argtype = [ctypes.POINTER(self.argument_type), ctypes.c_void_p]
|
||||
self.conv_kind = operation.conv_kind
|
||||
|
||||
self.operation: Conv2dOperation = operation
|
||||
|
||||
self.emitter = EmitConv2dInstance("_type")
|
||||
|
||||
self.threads = operation.tile_description.num_threads
|
||||
|
||||
self.swizzle_functor = operation.swizzling_functor
|
||||
|
||||
def emit(self):
|
||||
return self.emitter.emit(self.operation)
|
||||
|
||||
def plan(self, arguments: Conv2dArguments):
|
||||
tile_size = GemmCoord(
|
||||
self.operation.tile_description.threadblock_shape[0],
|
||||
self.operation.tile_description.threadblock_shape[1],
|
||||
self.operation.tile_description.threadblock_shape[2],
|
||||
)
|
||||
|
||||
grid = self.get_grid_shape(
|
||||
int(self.conv_kind),
|
||||
arguments.problem_size.ctype,
|
||||
tile_size.ctype,
|
||||
arguments.split_k_slices
|
||||
)
|
||||
|
||||
return LaunchConfiguration(
|
||||
[grid.x, grid.y, grid.z], [self.threads, 1, 1],
|
||||
self.shared_memory_capacity)
|
||||
|
||||
def initialize(self):
|
||||
err, = cuda.cuFuncSetAttribute(
|
||||
self.kernel,
|
||||
attrib=cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
|
||||
value=self.shared_memory_capacity)
|
||||
if err != cuda.CUresult.CUDA_SUCCESS:
|
||||
raise RuntimeError(f"CUDA Error: {err}")
|
||||
|
||||
|
||||
class Conv2dOperation:
|
||||
"""
|
||||
CUTLASS Conv2d operation description.
|
||||
|
||||
:param conv_kind: convolution operator
|
||||
:type conv_kind: :class:`cutlass_library.library.ConvKind`
|
||||
|
||||
:param iterator_algorithm: Selects among several implementation
|
||||
variants trading off performance with simplicity
|
||||
:type iterator_algorithm: :class:`cutlass_library.library.IteratorAlgorithm`
|
||||
|
||||
:param arch: GPU compute capability (sm_xx)
|
||||
:type arch: int
|
||||
|
||||
:param tile_description: tile description
|
||||
:type tile_description: :class:`cutlass_cppgen.backend.TileDescription`
|
||||
|
||||
:param A: tensor A description
|
||||
:type A: :class:`cutlass_cppgen.backend.TensorDescription`
|
||||
|
||||
:param B: tensor B description
|
||||
:type B: :class:`cutlass_cppgen.backend.TensorDescription`
|
||||
|
||||
:param C: tensor C description
|
||||
:type C: :class:`cutlass_cppgen.backend.TensorDescription`
|
||||
|
||||
:param D: tensor D description
|
||||
:type D: :class:`cutlass_cppgen.backend.TensorDescription`
|
||||
|
||||
:param element_epilogue: element type for computation in epilogue \
|
||||
:type element_epilogue: cutlass_library.library.DataType
|
||||
|
||||
:param stride_support: distinguish among partial specializations that \
|
||||
accelerate certain problems where convolution stride is unit \
|
||||
:type stride_support: :class:`cutlass_library.library.StrideSupport`
|
||||
|
||||
:param epilogue_functor: convolution epilogue functor
|
||||
:type epilogue_functor: :class:`EpilogueFunctor`
|
||||
|
||||
:param swizzling_functor: threadblock swizzling functor
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
conv_kind,
|
||||
iterator_algorithm,
|
||||
arch: int,
|
||||
tile_description: TileDescription,
|
||||
A: TensorDescription,
|
||||
B: TensorDescription,
|
||||
C: TensorDescription,
|
||||
stride_support,
|
||||
epilogue_functor,
|
||||
swizzling_functor=SwizzlingFunctor.Identity1,
|
||||
emission_type=EmissionType.Kernel,
|
||||
**kwargs
|
||||
):
|
||||
self.operation_kind: OperationKind = OperationKind.Conv2d
|
||||
self.arch: int = arch
|
||||
self.tile_description: TileDescription = tile_description
|
||||
self.conv_kind = conv_kind
|
||||
self.A: TensorDescription = A
|
||||
self.B: TensorDescription = B
|
||||
self.C: TensorDescription = C
|
||||
self.epilogue_functor = epilogue_functor
|
||||
self.iterator_algorithm = iterator_algorithm
|
||||
self.stride_support = stride_support
|
||||
self.swizzling_functor = swizzling_functor
|
||||
|
||||
self.emission_type = emission_type
|
||||
|
||||
self.rt_module: Conv2dRT = Conv2dRT(self)
|
||||
self.argument_type = self.rt_module.argument_type
|
||||
self.epilogue_type = self.rt_module.epilogue_type
|
||||
|
||||
def run(self, arguments: Conv2dArguments) -> cuda.CUresult:
|
||||
"""
|
||||
Launch the cuda kernel with input arguments
|
||||
|
||||
:param arguments: conv2d arguments
|
||||
:type arguments: :class:`cutlass_cppgen.backend.Conv2dArguments`
|
||||
"""
|
||||
|
||||
# launch the kernel
|
||||
err = self.rt_module.run(
|
||||
arguments.host_workspace,
|
||||
arguments.device_workspace,
|
||||
arguments.launch_config,
|
||||
arguments.stream
|
||||
)
|
||||
|
||||
if err != cuda.CUresult.CUDA_SUCCESS:
|
||||
raise RuntimeError(f"CUDA Error {err}")
|
||||
|
||||
return err
|
||||
|
||||
#
|
||||
# Get function name
|
||||
#
|
||||
|
||||
def procedural_name(self):
|
||||
"""The full procedural name indicates architecture, extended name, tile size, and layout."""
|
||||
return self.configuration_name()
|
||||
|
||||
def configuration_name(self):
|
||||
"""The full procedural name indicates architecture, extended name, tile size, and layout."""
|
||||
|
||||
opcode_class_name = OpcodeClassNames[
|
||||
self.tile_description.math_instruction.opcode_class
|
||||
]
|
||||
|
||||
threadblock = "%dx%d_%dx%d" % (
|
||||
self.tile_description.threadblock_shape[0],
|
||||
self.tile_description.threadblock_shape[1],
|
||||
self.tile_description.threadblock_shape[2],
|
||||
self.tile_description.stages,
|
||||
)
|
||||
|
||||
if self.stride_support == StrideSupport.Unity:
|
||||
configuration_name = "cutlass_sm${arch}_${opcode_class}_${extended_name}_${threadblock}_${layout}_unity_stride_align${alignment}"
|
||||
else:
|
||||
configuration_name = "cutlass_sm${arch}_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}"
|
||||
|
||||
return SubstituteTemplate(
|
||||
configuration_name,
|
||||
{
|
||||
"arch": str(self.arch),
|
||||
"opcode_class": opcode_class_name,
|
||||
"extended_name": self.extended_name(),
|
||||
"threadblock": threadblock,
|
||||
"layout": self.layout_name(),
|
||||
"alignment": "%d" % self.A.alignment
|
||||
},
|
||||
)
|
||||
|
||||
def extended_name(self):
|
||||
"""Append data types if they differ from compute type."""
|
||||
if self.C.element != self.tile_description.math_instruction.element_accumulator and \
|
||||
self.A.element != self.tile_description.math_instruction.element_accumulator:
|
||||
extended_name = "${element_c}_${core_name}_${element_a}"
|
||||
elif self.C.element == self.tile_description.math_instruction.element_accumulator and \
|
||||
self.A.element != self.tile_description.math_instruction.element_accumulator:
|
||||
extended_name = "${core_name}_${element_a}"
|
||||
else:
|
||||
extended_name = "${core_name}"
|
||||
|
||||
extended_name = SubstituteTemplate(extended_name, {
|
||||
"element_a": DataTypeNames[self.A.element],
|
||||
"element_c": DataTypeNames[self.C.element],
|
||||
"core_name": self.core_name(),
|
||||
})
|
||||
|
||||
return extended_name
|
||||
|
||||
def layout_name(self):
|
||||
return "%s" % (ShortLayoutTypeNames[self.A.layout])
|
||||
|
||||
def core_name(self):
|
||||
"""The basic operation kind is prefixed with a letter indicating the accumulation type."""
|
||||
|
||||
intermediate_type = ""
|
||||
|
||||
if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp:
|
||||
inst_shape = "%dx%dx%d" % tuple(
|
||||
self.tile_description.math_instruction.instruction_shape)
|
||||
if self.tile_description.math_instruction.element_a != self.A.element and \
|
||||
self.tile_description.math_instruction.element_a != self.accumulator_type():
|
||||
intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
|
||||
else:
|
||||
inst_shape = ""
|
||||
|
||||
return "%s%s%s%s_%s" % (
|
||||
ShortDataTypeNames[self.accumulator_type()],
|
||||
inst_shape,
|
||||
intermediate_type,
|
||||
ConvKindNames[self.conv_kind],
|
||||
IteratorAlgorithmNames[self.iterator_algorithm]
|
||||
)
|
||||
|
||||
def is_complex(self):
|
||||
complex_operators = [
|
||||
MathOperation.multiply_add_complex,
|
||||
MathOperation.multiply_add_complex_gaussian,
|
||||
]
|
||||
return self.tile_description.math_instruction.math_operation in complex_operators
|
||||
|
||||
def accumulator_type(self):
|
||||
accum = self.tile_description.math_instruction.element_accumulator
|
||||
|
||||
if self.is_complex():
|
||||
return get_complex_from_real(accum)
|
||||
|
||||
return accum
|
||||
|
||||
def device_op(self):
|
||||
"""
|
||||
Returns a new Conv2dOperation object that is constructed with emission type
|
||||
``EmissionType.Device``.
|
||||
|
||||
:return: operation ready for device-level code emission
|
||||
:rtype: Conv2dOperation
|
||||
"""
|
||||
return Conv2dOperation(
|
||||
self.conv_kind, self.iterator_algorithm, self.arch, self.tile_description,
|
||||
self.A, self.B, self.C, self.stride_support, self.epilogue_functor, self.swizzling_functor,
|
||||
emission_type=EmissionType.Device)
|
||||
|
||||
|
||||
###################################################################################################
|
||||
#
|
||||
# Emits single instances of a CUTLASS device-wide operator
|
||||
#
|
||||
###################################################################################################
|
||||
|
||||
|
||||
class EmitConv2dInstance:
|
||||
def __init__(self, operation_suffix=""):
|
||||
self.operation_suffix = operation_suffix
|
||||
self.includes = [
|
||||
"cutlass/cutlass.h",
|
||||
"cutlass/conv/kernel/default_conv2d_fprop.h",
|
||||
"cutlass/conv/kernel/default_conv2d_dgrad.h",
|
||||
"cutlass/conv/kernel/default_conv2d_wgrad.h",
|
||||
"cutlass/conv/device/implicit_gemm_convolution.h"
|
||||
]
|
||||
self.template = """
|
||||
// Conv2d${conv_kind_name} ${iterator_algorithm_name} kernel instance "${operation_name}"
|
||||
using ${operation_name}_base =
|
||||
typename cutlass::conv::kernel::DefaultConv2d${conv_kind_name}<
|
||||
${element_a},
|
||||
${layout_a},
|
||||
${element_b},
|
||||
${layout_b},
|
||||
${element_c},
|
||||
${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k} >,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue_functor},
|
||||
${swizzling_functor},
|
||||
${stages},
|
||||
${math_operator},
|
||||
${iterator_algorithm},
|
||||
${stride_support},
|
||||
${align_a},
|
||||
${align_b}
|
||||
>::Kernel;
|
||||
|
||||
struct ${operation_name}${operation_suffix}:
|
||||
public ${operation_name}_base { };
|
||||
|
||||
"""
|
||||
|
||||
self.template_device = """
|
||||
// Conv2d operation ${operation_name}
|
||||
|
||||
using Conv2d${conv_kind_name}Kernel = typename cutlass::conv::kernel::DefaultConv2d${conv_kind_name}<
|
||||
${element_a},
|
||||
${layout_a},
|
||||
${element_b},
|
||||
${layout_b},
|
||||
${element_c},
|
||||
${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k} >,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue_functor},
|
||||
${swizzling_functor},
|
||||
${stages},
|
||||
${math_operator},
|
||||
${iterator_algorithm},
|
||||
${stride_support},
|
||||
${align_a},
|
||||
${align_b}
|
||||
>::Kernel;
|
||||
|
||||
using DeviceKernel =
|
||||
typename cutlass::conv::device::ImplicitGemmConvolution<Conv2d${conv_kind_name}Kernel>;
|
||||
"""
|
||||
|
||||
def emit(self, operation):
|
||||
warp_shape = [int(operation.tile_description.threadblock_shape[idx] /
|
||||
operation.tile_description.warp_count[idx]) for idx in range(3)]
|
||||
|
||||
epilogue_vector_length = int(min(
|
||||
operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element])
|
||||
|
||||
values = {
|
||||
"operation_name": operation.procedural_name(),
|
||||
"operation_suffix": self.operation_suffix,
|
||||
"conv_kind": ConvKindTag[operation.conv_kind],
|
||||
"conv_kind_name": ConvKindNames[operation.conv_kind].capitalize(),
|
||||
"element_a": DataTypeTag[operation.A.element],
|
||||
"layout_a": LayoutTag[operation.A.layout],
|
||||
"element_b": DataTypeTag[operation.B.element],
|
||||
"layout_b": LayoutTag[operation.B.layout],
|
||||
"element_c": DataTypeTag[operation.C.element],
|
||||
"layout_c": LayoutTag[operation.C.layout],
|
||||
"element_accumulator": DataTypeTag[operation.accumulator_type()],
|
||||
"opcode_class": OpcodeClassTag[operation.tile_description.math_instruction.opcode_class],
|
||||
"arch": "cutlass::arch::Sm%d" % operation.arch,
|
||||
"threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
|
||||
"threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
|
||||
"threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
|
||||
"warp_shape_m": str(warp_shape[0]),
|
||||
"warp_shape_n": str(warp_shape[1]),
|
||||
"warp_shape_k": str(warp_shape[2]),
|
||||
"instruction_shape_m": str(operation.tile_description.math_instruction.instruction_shape[0]),
|
||||
"instruction_shape_n": str(operation.tile_description.math_instruction.instruction_shape[1]),
|
||||
"instruction_shape_k": str(operation.tile_description.math_instruction.instruction_shape[2]),
|
||||
"epilogue_vector_length": str(epilogue_vector_length),
|
||||
"epilogue_functor": operation.epilogue_functor.emit(),
|
||||
"swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
|
||||
"stages": str(operation.tile_description.stages),
|
||||
"iterator_algorithm": IteratorAlgorithmTag[operation.iterator_algorithm],
|
||||
"iterator_algorithm_name": IteratorAlgorithmNames[operation.iterator_algorithm].capitalize(),
|
||||
"stride_support": StrideSupportTag[operation.stride_support],
|
||||
"math_operator": "cutlass::arch::OpMultiplyAddComplex" if operation.is_complex() else MathOperationTag[operation.tile_description.math_instruction.math_operation],
|
||||
"align_a": str(operation.A.alignment),
|
||||
"align_b": str(operation.B.alignment),
|
||||
}
|
||||
|
||||
if operation.emission_type == EmissionType.Kernel:
|
||||
conv2d_template = self.template
|
||||
else:
|
||||
conv2d_template = self.template_device
|
||||
|
||||
return SubstituteTemplate(conv2d_template, values)
|
||||
Reference in New Issue
Block a user