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

CUTLASS 2.10 updates (#622)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2022-09-12 21:26:30 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent beae168f90
commit e773429f7e
96 changed files with 8365 additions and 1667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py
+10 -10
View File
@@ -17,7 +17,7 @@ class GemmBF16TensorOpSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -33,15 +33,15 @@ class GemmBF16TensorOpSm80(unittest.TestCase):
alignment=4
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, cutlass.float32)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -58,7 +58,7 @@ class GemmBF16TensorOpSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -74,15 +74,15 @@ class GemmBF16TensorOpSm80(unittest.TestCase):
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, cutlass.float32)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py
+50 -30
View File
@@ -18,7 +18,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -36,13 +36,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.BatchedIdentitySwizzle
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor,
direct_store=True
)
@@ -60,7 +62,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -78,13 +80,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -101,7 +105,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -119,13 +123,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -142,7 +148,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -160,13 +166,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -183,7 +191,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -201,13 +209,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float16
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -224,7 +234,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -242,13 +252,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float16
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -265,7 +277,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -283,13 +295,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -306,7 +320,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -324,13 +338,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -347,7 +363,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -365,13 +381,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -388,7 +406,7 @@ class GemmF16Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -406,13 +424,15 @@ class GemmF16Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py
+15 -9
View File
@@ -19,7 +19,7 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -37,13 +37,15 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -61,7 +63,7 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -79,13 +81,15 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -102,7 +106,7 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -120,13 +124,15 @@ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py
+10 -6
View File
@@ -17,7 +17,7 @@ class GemmF64TensorOpSm80(unittest.TestCase):
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
math_instruction=math_inst
)
# alignment 1 restricted for double
@@ -36,13 +36,15 @@ class GemmF64TensorOpSm80(unittest.TestCase):
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -59,7 +61,7 @@ class GemmF64TensorOpSm80(unittest.TestCase):
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
math_instruction=math_inst
)
# alignment 1 restricted for double
@@ -78,13 +80,15 @@ class GemmF64TensorOpSm80(unittest.TestCase):
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py
+20 -16
View File
@@ -18,7 +18,7 @@ class GemmGroupedSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -37,14 +37,15 @@ class GemmGroupedSm80(unittest.TestCase):
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
80,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
@@ -64,7 +65,7 @@ class GemmGroupedSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -83,14 +84,15 @@ class GemmGroupedSm80(unittest.TestCase):
)
element_epilogue = cutlass.float64
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
80,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
@@ -110,7 +112,7 @@ class GemmGroupedSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -129,14 +131,15 @@ class GemmGroupedSm80(unittest.TestCase):
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
80,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
@@ -156,7 +159,7 @@ class GemmGroupedSm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -175,14 +178,15 @@ class GemmGroupedSm80(unittest.TestCase):
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.LinearCombination
epilogue_functor = LinearCombination(
C.element, C.alignment,
math_inst.element_accumulator, element_epilogue)
swizzling_functor = cutlass.BatchedIdentitySwizzle
for precompute_mode in [SchedulerMode.Device, SchedulerMode.Host]:
operation = GemmOperationGrouped(
tile_description.minimum_compute_capability,
80,
tile_description, A, B, C,
element_epilogue,
epilogue_functor, swizzling_functor,
precompute_mode=precompute_mode
)
tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py
+28 -21
View File
@@ -1,5 +1,6 @@
import pycutlass
from pycutlass import *
from pycutlass.epilogue import LinearCombinationClamp
from pycutlass.test import *
import unittest
@@ -17,7 +18,7 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -33,15 +34,15 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
alignment=8
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
epilogue_functor = FastLinearCombinationClamp(
C.element, C.alignment
)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -58,7 +59,7 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -74,15 +75,15 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
alignment=16
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
epilogue_functor = FastLinearCombinationClamp(
C.element, C.alignment
)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -99,7 +100,7 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -115,15 +116,15 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
alignment=16
)
element_epilogue = cutlass.float32
epilogue_functor = EpilogueFunctor.FastLinearCombinationClamp
epilogue_functor = FastLinearCombinationClamp(
C.element, C.alignment
)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -140,7 +141,7 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -158,13 +159,16 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
element_epilogue = cutlass.int32
epilogue_functor = EpilogueFunctor.LinearCombinationClamp
epilogue_functor = LinearCombinationClamp(
C.element, C.alignment, math_inst.element_accumulator,
element_epilogue
)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)
@@ -181,7 +185,7 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
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
math_instruction=math_inst
)
A = TensorDescription(
@@ -199,13 +203,16 @@ class GemmS8TensorOpF32Sm80(unittest.TestCase):
element_epilogue = cutlass.int32
epilogue_functor = EpilogueFunctor.LinearCombinationClamp
epilogue_functor = LinearCombinationClamp(
C.element, C.alignment, math_inst.element_accumulator,
element_epilogue
)
swizzling_functor = cutlass.IdentitySwizzle1
operation = GemmOperationUniversal(
arch=80, tile_description=tile_description,
A=A, B=B, C=C, element_epilogue=element_epilogue,
A=A, B=B, C=C,
epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
)