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

CUTLASS 3.5.1 (#1623)

Browse Source 
* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/cutlass/__init__.py
View File

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

28
python/cutlass/backend/evt/backend/sm90_nodes.py
View File

@@ -154,20 +154,6 @@ using {self.name_camel} = cutlass::epilogue::fusion::Sm90ScalarBroadcast<
class Sm90RowBroadcastImpl(RowBroadcastImpl):
@property
def descriptor(self) -> str:
"""
Descriptor for Aux Load
"""
return f"{self.name_camel}Descriptor"
def decl_descriptor(self) -> str:
"""
Declare the descriptor type
"""
return f"\nusing {self.descriptor} = cutlass::epilogue::collective::detail::RowBroadcastDescriptor<EpilogueDescriptor, {DataTypeTag[self.element]}>;\n"
@property
def type_decl(self):
"""
@@ -176,22 +162,14 @@ class Sm90RowBroadcastImpl(RowBroadcastImpl):
if self._type_decl is not None:
return self._type_decl
self._type_decl = self.decl_descriptor()
self._type_decl += f"""
self._type_decl = f"""
using {self.name_camel} = cutlass::epilogue::fusion::Sm90RowBroadcast<
{self.descriptor}::Stages, typename EpilogueDescriptor::TileShape,
typename {self.descriptor}::Element, {self.stride_mnl}
0 /*Stages*/, typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]},
{self.stride_mnl}
>;
"""
return self._type_decl
def get_smem_size(self, cta_tile_mnk, epilogue_tile_mn, stages_c, stages_d, epi_tiles):
"""
Get the shared memory size based on epilogue_tile_mn, stages_c, and stages_d
"""
stages = (stages_c + epi_tiles - 1) // epi_tiles + 1
return (DataTypeSize[self.element] * cta_tile_mnk[1] * stages // 8, 16)
class Sm90ColumnBroadcastImpl(ColumnBroadcastImpl):