CUTLASS 3.1 (#915)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>

examples/40_cutlass_py/gemm.py

@@ -37,10 +37,10 @@ import argparse
 import numpy as np
 import sys
 
-import cutlass
-import pycutlass
-from pycutlass import *
-from pycutlass.utils.device import device_cc
+import cutlass_bindings
+import cutlass.backend as pycutlass
+from cutlass.backend import *
+from cutlass.backend.utils.device import device_cc
 
 
 parser = argparse.ArgumentParser(description="Launch a GEMM kernel from Python: 'D = alpha * A * B + beta * C'")
@@ -72,11 +72,11 @@ pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
 pycutlass.compiler.nvcc()
 
 # Set up A, B, C and accumulator
-A = TensorDescription(cutlass.float16, cutlass.ColumnMajor, alignment)
-B = TensorDescription(cutlass.float16, cutlass.RowMajor, alignment)
-C = TensorDescription(cutlass.float32, cutlass.ColumnMajor, alignment)
-element_acc = cutlass.float32
-element_epilogue = cutlass.float32
+A = TensorDescription(cutlass_bindings.float16, cutlass_bindings.ColumnMajor, alignment)
+B = TensorDescription(cutlass_bindings.float16, cutlass_bindings.RowMajor, alignment)
+C = TensorDescription(cutlass_bindings.float32, cutlass_bindings.ColumnMajor, alignment)
+element_acc = cutlass_bindings.float32
+element_epilogue = cutlass_bindings.float32
 
 # Select instruction shape based on the Tensor Core instructions supported
 # by the device on which we are running
@@ -85,12 +85,14 @@ if cc == 70:
 elif cc == 75:
     instruction_shape = [16, 8, 8]
 else:
+    # Use CUTLASS kernels for CC 80 by default (e.g., for cases in which SM86 is used)
+    cc = 80
     instruction_shape = [16, 8, 16]
 
 math_inst = MathInstruction(
     instruction_shape,
     A.element, B.element, element_acc,
-    cutlass.OpClass.TensorOp,
+    cutlass_bindings.OpClass.TensorOp,
     MathOperation.multiply_add
 )
 
@@ -122,7 +124,7 @@ tensor_B = np.ceil(np.random.uniform(low=-8.5, high=7.5, size=(args.k * args.n,)
 tensor_C = np.ceil(np.random.uniform(low=-8.5, high=7.5, size=(args.m * args.n,))).astype(np.float32)
 tensor_D = np.zeros(shape=(args.m * args.n,)).astype(np.float32)
 
-problem_size = cutlass.gemm.GemmCoord(args.m, args.n, args.k)
+problem_size = cutlass_bindings.gemm.GemmCoord(args.m, args.n, args.k)
 alpha = 1.
 beta = 0.