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

CUTLASS 2.10 (#615)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2022-09-03 18:48:46 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent ca23ff7924
commit b72cbf957d
289 changed files with 43708 additions and 2513 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/library/scripts/pycutlass/test/gemm/__init__.py
View File
tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py
+93
View File
@@ -0,0 +1,93 @@
import pycutlass
from pycutlass import *
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_testbed import test_all_gemm
class GemmBF16TensorOpSm80(unittest.TestCase):
def SM80_Device_Gemm_bf16n_bf16n_f32t_tensor_op_f32_64x128x64_32x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.bfloat16, element_b=cutlass.bfloat16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[64, 128, 64],
stages=4, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.bfloat16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.bfloat16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_bf16t_bf16t_bf16t_tensor_op_f32_128x256x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.bfloat16, element_b=cutlass.bfloat16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[64, 128, 32],
stages=6, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.bfloat16, layout=cutlass.RowMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.bfloat16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.bfloat16, layout=cutlass.RowMajor,
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "multistage"))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**24, 2**24)
unittest.main()
tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py
+425
View File
@@ -0,0 +1,425 @@
import pycutlass
from pycutlass import *
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_testbed import test_all_gemm
class GemmF16Sm80(unittest.TestCase):
def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.BatchedIdentitySwizzle
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor,
direct_store=True
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16n_f16n_f16t_tensor_op_f32_128x128x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 64],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16n_f16n_f32n_tensor_op_f32_128x256x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 256, 64],
stages=3, warp_count=[2, 4, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16n_f16n_f32t_tensor_op_f32_256x128x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[256, 128, 64],
stages=3, warp_count=[4, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16n_f16t_f16t_tensor_op_f16_sliced_k_128x64x64_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float16, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 64, 64],
stages=3, warp_count=[2, 1, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float16
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_GemmUniversal_f16n_f16t_f32t_tensor_op_f32_64x64x32_32x32x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float16, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[64, 64, 32],
stages=10, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float16
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16n_f16t_f32t_tensor_op_f32_256x128x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[256, 128, 64],
stages=3, warp_count=[4, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_test_SM80_Device_Gemm_f16t_f16n_f16t_tensor_op_f16_sliced_k_128x64x64_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 64, 64],
stages=3, warp_count=[2, 1, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16t_f16t_f32n_tensor_op_f32_128x256x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 256, 64],
stages=3, warp_count=[2, 4, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.RowMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f16t_f16t_f32t_tensor_op_f32_128x256x64_64x64x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16],
element_a=cutlass.float16, element_b=cutlass.float16,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 256, 64],
stages=3, warp_count=[2, 4, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**24, 2**24)
unittest.main()
tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py
+138
View File
@@ -0,0 +1,138 @@
import pycutlass
from pycutlass import *
from pycutlass.memory_manager import get_allocated_size
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_testbed import test_all_gemm
class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 8],
element_a=cutlass.float32, element_b=cutlass.float32,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add_fast_bf16
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
B = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f32n_f32n_f32t_tensor_op_f32_128x128x32_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 8],
element_a=cutlass.float32, element_b=cutlass.float32,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
B = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f32n_f32n_f32t_tensor_op_fast_accurate_f32_64x64x32_32x32x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 8],
element_a=cutlass.float32, element_b=cutlass.float32,
element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add_fast_f32
)
tile_description = TileDescription(
threadblock_shape=[64, 64, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
B = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**24, 2**24)
pycutlass.compiler.load_from_cache()
unittest.main()
tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py
+95
View File
@@ -0,0 +1,95 @@
import pycutlass
from pycutlass import *
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_testbed import test_all_gemm
class GemmF64TensorOpSm80(unittest.TestCase):
def test_SM80_Device_Gemm_f64n_f64t_f64t_tensor_op_f64_32x32x16_16x16x16(self):
math_inst = MathInstruction(
instruction_shape=[8, 8, 4],
element_a=cutlass.float64, element_b=cutlass.float64,
element_accumulator=cutlass.float64, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[32, 32, 16],
stages=4, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
# alignment 1 restricted for double
A = TensorDescription(
element=cutlass.float64, layout=cutlass.ColumnMajor,
alignment=1
)
B = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
C = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
def test_SM80_Device_Gemm_f64t_f64n_f64t_tensor_op_f64_64x64x16_32x32x16(self):
math_inst = MathInstruction(
instruction_shape=[8, 8, 4],
element_a=cutlass.float64, element_b=cutlass.float64,
element_accumulator=cutlass.float64, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[64, 64, 16],
stages=4, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
# alignment 1 restricted for double
A = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
B = TensorDescription(
element=cutlass.float64, layout=cutlass.ColumnMajor,
alignment=1
)
C = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "universal"))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**24, 2**24)
unittest.main()
tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py
+197
View File
@@ -0,0 +1,197 @@
import pycutlass
from pycutlass import *
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_grouped_testbed import TestbedGrouped
class GemmGroupedSm80(unittest.TestCase):
def test_SM80_Device_GemmGrouped_f16n_f16t_f32n_tensor_op_f32_128x128x32_64x64x32(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16], element_a=cutlass.float16,
element_b=cutlass.float16, element_accumulator=cutlass.float32,
opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
testbed = TestbedGrouped(operation=operation)
self.assertTrue(testbed.run(24))
def test_SM80_Device_GemmGrouped_f64t_f64t_f64n_tensor_op_f64_64x64x16_32x32x16(self):
math_inst = MathInstruction(
instruction_shape=[8, 8, 4], element_a=cutlass.float64,
element_b=cutlass.float64, element_accumulator=cutlass.float64,
opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[64, 64, 16],
stages=4, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
B = TensorDescription(
element=cutlass.float64, layout=cutlass.RowMajor,
alignment=1
)
C = TensorDescription(
element=cutlass.float64, layout=cutlass.ColumnMajor,
alignment=1
)
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
testbed = TestbedGrouped(operation=operation)
self.assertTrue(testbed.run(24))
def test_SM80_Device_GemmGrouped_f32t_f32t_f32t_simt_f32_128x64x8_64x32x1(self):
math_inst = MathInstruction(
instruction_shape=[1, 1, 1], element_a=cutlass.float32,
element_b=cutlass.float32, element_accumulator=cutlass.float32,
opcode_class=cutlass.OpClass.Simt,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 64, 8],
stages=4, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=1
)
B = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=1
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.RowMajor,
alignment=1
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
testbed = TestbedGrouped(operation=operation)
self.assertTrue(testbed.run(27))
def test_SM80_Device_GemmGrouped_f16n_f16t_f32n_tensor_op_f32_128x128x32_64x64x32_cache(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 16], element_a=cutlass.float16,
element_b=cutlass.float16, element_accumulator=cutlass.float32,
opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 32],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
B = TensorDescription(
element=cutlass.float16, layout=cutlass.ColumnMajor,
alignment=8
)
C = TensorDescription(
element=cutlass.float32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
testbed = TestbedGrouped(operation=operation)
self.assertTrue(testbed.run(5))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**26, 2**26)
unittest.main()
tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py
+219
View File
@@ -0,0 +1,219 @@
import pycutlass
from pycutlass import *
from pycutlass.test import *
import unittest
from pycutlass.test.gemm_testbed import test_all_gemm
class GemmS8TensorOpF32Sm80(unittest.TestCase):
def test_SM80_Device_Gemm_s8t_s8n_s8t_tensor_op_s32_64x64x64_32x32x64(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 32],
element_a=cutlass.int8, element_b=cutlass.int8,
element_accumulator=cutlass.int32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add_saturate
)
tile_description = TileDescription(
threadblock_shape=[64, 64, 64],
stages=6, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajorInterleaved32,
alignment=16
)
B = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajorInterleaved32,
alignment=16
)
C = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajorInterleaved32,
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "interleaved"))
def test_SM80_Device_Gemm_s8t_s8n_s8t_tensor_op_s32_256x128x128_64x64x128(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 32],
element_a=cutlass.int8, element_b=cutlass.int8,
element_accumulator=cutlass.int32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 128],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajor,
alignment=16
)
B = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajor,
alignment=16
)
C = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajor,
alignment=16
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "multistage"))
def test_SM80_Device_Gemm_s8t_s8n_s8n_tensor_op_s32_128x128x128_64x64x128(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 32],
element_a=cutlass.int8, element_b=cutlass.int8,
element_accumulator=cutlass.int32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 128],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajor,
alignment=16
)
B = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajor,
alignment=16
)
C = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajor,
alignment=16
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "multistage"))
def test_SM80_Device_Gemm_s8t_s8n_s32n_tensor_op_s32_128x128x128_64x64x128(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 32],
element_a=cutlass.int8, element_b=cutlass.int8,
element_accumulator=cutlass.int32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 128],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajor,
alignment=16
)
B = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajor,
alignment=16
)
C = TensorDescription(
element=cutlass.int32, layout=cutlass.ColumnMajor,
alignment=4
)
element_epilogue = cutlass.int32
epilogue_functor = EpilogueFunctor.LinearCombinationClamp
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "multistage"))
def test_SM80_Device_Gemm_s8t_s8n_s32t_tensor_op_s32_128x128x128_64x64x128(self):
math_inst = MathInstruction(
instruction_shape=[16, 8, 32],
element_a=cutlass.int8, element_b=cutlass.int8,
element_accumulator=cutlass.int32, opcode_class=cutlass.OpClass.TensorOp,
math_operation=MathOperation.multiply_add
)
tile_description = TileDescription(
threadblock_shape=[128, 128, 128],
stages=3, warp_count=[2, 2, 1],
math_instruction=math_inst, min_compute=80, max_compute=80
)
A = TensorDescription(
element=cutlass.int8, layout=cutlass.RowMajor,
alignment=16
)
B = TensorDescription(
element=cutlass.int8, layout=cutlass.ColumnMajor,
alignment=16
)
C = TensorDescription(
element=cutlass.int32, layout=cutlass.RowMajor,
alignment=4
)
element_epilogue = cutlass.int32
epilogue_functor = EpilogueFunctor.LinearCombinationClamp
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
self.assertTrue(test_all_gemm(operation, "multistage"))
if __name__ == '__main__':
pycutlass.get_memory_pool(2**24, 2**24)
unittest.main()
tools/library/scripts/pycutlass/test/gemm/run_all_tests.py
+9
View File
@@ -0,0 +1,9 @@
import pycutlass
import unittest
if __name__ == '__main__':
pycutlass.get_memory_pool(2**26, 2**26)
loader = unittest.TestLoader()
tests = loader.discover('./', 'gemm_*.py')
testRunner = unittest.runner.TextTestRunner()
testRunner.run(tests)