CUTLASS 3.5.0 (#1411)

This commit is contained in:
Vijay Thakkar
2024-03-19 17:51:04 -04:00
committed by GitHub
parent ffa34e7075
commit 629f4653c3
468 changed files with 48730 additions and 7253 deletions

View File

@@ -121,7 +121,7 @@ def get_option_registry():
this._option_registry = OptionRegistry(device_cc())
return this._option_registry
this.__version__ = '3.4.1'
this.__version__ = '3.5.0'
from cutlass.backend import create_memory_pool
from cutlass.emit.pytorch import pytorch

View File

@@ -244,7 +244,7 @@ def get_gemm_arguments_3x(mainloop_arguments, epilogue_functor, scheduler_args,
class _HardwareInfo(ctypes.Structure):
_fields_ = [
("device_id", ctypes.c_int),
("sm_count", ctypes.c_int)
("sm_count", ctypes.c_int),
]
class _GemmArguments(ctypes.Structure):

View File

@@ -122,7 +122,7 @@ class LinearCombination(EpilogueFunctorBase):
:param element_output: data type used to load and store tensors
:param epilogue_vector_length: number of elements computed per operation.
Usually it is 128/sizeof_bits<ElementOutput_>, but we use 64 and 32 sometimes
Usually it is 128/sizeof_bits_v<ElementOutput_>, but we use 64 and 32 sometimes
when there are not enough data to store
:param element_accumulator: Accumulator data type
@@ -207,7 +207,7 @@ class LinearCombinationClamp(LinearCombination):
:param element_output: data type used to load and store tensors
:param epilogue_vector_length: number of elements computed per operation.
Usually it is 128/sizeof_bits<ElementOutput_>, but we use 64 and 32 sometimes
Usually it is 128/sizeof_bits_v<ElementOutput_>, but we use 64 and 32 sometimes
when there are not enough data to store
:param element_accumulator: Accumulator data type
@@ -260,7 +260,7 @@ class FastLinearCombinationClamp(EpilogueFunctorBase):
:param element_output: data type used to load and store tensors
:param epilogue_vector_length: number of elements computed per operation.
Usually it is 128/sizeof_bits<ElementOutput_>, but we use 64 and 32 sometimes
Usually it is 128/sizeof_bits_v<ElementOutput_>, but we use 64 and 32 sometimes
when there are not enough data to store
"""
@@ -310,7 +310,7 @@ class LinearCombinationGeneric(LinearCombination):
:param element_output: data type used to load and store tensors
:param epilogue_vector_length: number of elements computed per operation.
Usually it is 128/sizeof_bits<ElementOutput_>, but we use 64 and 32 sometimes
Usually it is 128/sizeof_bits_v<ElementOutput_>, but we use 64 and 32 sometimes
when there are not enough data to store
:param element_accumulator: Accumulator data type

View File

@@ -299,7 +299,7 @@ class Sm90ColumnReductionImpl(ColumnReductionImpl):
self._type_decl = f"""
using {self.name_camel} = cutlass::epilogue::fusion::Sm90ColReduction<
{op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0,
{op_tag(self.reg_reduce_fn)}, {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0,
typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]},
{DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]},
{self.stride_mnl}
@@ -321,7 +321,7 @@ class Sm90RowReductionImpl(RowReductionImpl):
self._type_decl = f"""
using {self.name_camel} = cutlass::epilogue::fusion::Sm90RowReduction<
{op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0 /* Stages */,
{op_tag(self.reg_reduce_fn)}, {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0 /* Stages */,
typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]},
{DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]},
{self.stride_mnl}

View File

@@ -565,7 +565,9 @@ class GemmArguments3x(GemmArguments2x):
)
# Set hardware info
hw_info_ = hw_info(0, device_sm_count())
hw_info_ = hw_info(
0, device_sm_count(),
)
self.arguments = argument_type(
int(self.gemm_mode),
@@ -1300,7 +1302,7 @@ using DeviceKernel = cutlass::gemm::device::GemmUniversalAdapter<${operation_nam
# Support built-in epilogue functors or user-defined functions
if operation.tile_description.stages is None or operation.tile_description.stages == 0:
stage_count_type = "cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>"
stage_count_type = "cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>"
else:
stage_count_type = "_" + str(operation.tile_description.stages)