3.6.0 update (#2005)

* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>

python/cutlass/__init__.py

@@ -57,6 +57,19 @@ CUTLASS_PATH = os.getenv("CUTLASS_PATH", cutlass_library.source_path)
 # Alias CUTLASS_PATH as source_path
 source_path = CUTLASS_PATH
 
+_NVCC_VERSION = None
+def nvcc_version():
+    global _NVCC_VERSION
+    if _NVCC_VERSION is None:
+        import subprocess
+
+        # Attempt to get NVCC version
+        result = subprocess.run(['nvcc', '--version'], capture_output=True)
+        if result.returncode != 0:
+            raise Exception('Unable to run `nvcc --version')
+        _NVCC_VERSION = str(result.stdout).split(" release ")[-1].split(",")[0]
+    return _NVCC_VERSION
+
 _CUDA_INSTALL_PATH = None
 def cuda_install_path():
     """

python/cutlass/backend/c_types.py

@@ -139,7 +139,7 @@ def get_tile_scheduler_arguments_3x(
     splits: int = 1):
     max_swizzle_size = 1
     raster_order_option = 0 # Heuristic
-    if tile_scheduler == TileSchedulerType.Persistent:
+    if tile_scheduler in [TileSchedulerType.Default, TileSchedulerType.Persistent]:
         return _PersistentTileSchedulerArguments(
             max_swizzle_size,
             raster_order_option,

python/cutlass/backend/compiler.py

@@ -90,7 +90,7 @@ class CompilationOptions:
             opts.append(f"--include-path={incl}")
 
         arch_flag = f"-arch=sm_{self.arch}"
-        if self.arch == 90:
+        if self.arch == 90 and int(cutlass.nvcc_version().split('.')[0]) >= 12:
             arch_flag += "a"
         opts.append(arch_flag)
 
@@ -237,7 +237,7 @@ class ArtifactManager:
                 if incl not in includes:
                     includes.append(incl)
 
-        includes_host = ["builtin_types.h", "device_launch_parameters.h", "stddef.h"] + includes
+        includes_host = ["builtin_types.h", "device_launch_parameters.h", "cstddef"] + includes
         for incl in includes:
             source_buffer_device += SubstituteTemplate(
                 IncludeTemplate,

python/cutlass/backend/epilogue.py

@@ -44,6 +44,7 @@ from cutlass.utils.datatypes import is_numpy_tensor, is_torch_available, is_torc
 
 dtype2ctype = {
     DataType.f16: ctypes.c_uint16,
+    DataType.bf16: ctypes.c_uint16,
     DataType.f32: ctypes.c_float,
     DataType.f64: ctypes.c_double,
     DataType.s8: ctypes.c_int8,

python/cutlass/epilogue/evt_ops.py

@@ -59,18 +59,21 @@ def max(x, dim):
     elif is_torch_tensor(x):
         return torch.amax(x, dim)
 
+
 def maximum(x, y):
     if is_numpy_tensor(x):
         return np.maximum(x, y)
     elif is_torch_tensor(x):
         return torch.maximum(x, torch.tensor(y))
-    
+
+
 def minimum(x, y):
     if is_numpy_tensor(x):
         return np.minimum(x, y)
     elif is_torch_tensor(x):
         return torch.minimum(x, torch.tensor(y))
 
+
 ##############################################################################
 # Layout manipulate nodes
 ##############################################################################

python/cutlass/library_defaults.py

@@ -51,6 +51,20 @@ _generator_ccs = [50, 60, 61, 70, 75, 80, 90]
 # Strip any additional information from the CUDA version
 _cuda_version = __version__.split("rc")[0]
 
+# Check that Python CUDA version exceeds NVCC version
+_nvcc_version = cutlass.nvcc_version()
+_cuda_list = _cuda_version.split('.')
+_nvcc_list = _cuda_version.split('.')
+for val_cuda, val_nvcc in zip(_cuda_list, _nvcc_list):
+    if int(val_cuda) < int(val_nvcc):
+        raise Exception(f"Python CUDA version of {_cuda_version} must be greater than or equal to NVCC version of {_nvcc_version}")
+
+if len(_nvcc_list) > len(_cuda_list):
+    if len(_nvcc_list) != len(_cuda_list) + 1:
+        raise Exception(f"Malformatted NVCC version of {_nvcc_version}")
+    if _nvcc_list[:-1] == _cuda_list and int(_nvcc_list[-1]) != 0:
+        raise Exception(f"Python CUDA version of {_cuda_version} must be greater than or equal to NVCC version of {_nvcc_version}")
+
 
 class KernelsForDataType:
     """
@@ -278,7 +292,7 @@ class ArchOptions:
         ]
         manifest_args = cutlass_library.generator.define_parser().parse_args(args)
         manifest = cutlass_library.manifest.Manifest(manifest_args)
-        generate_function(manifest, _cuda_version)
+        generate_function(manifest, _nvcc_version)
 
         if operation_kind not in manifest.operations:
             # No kernels generated for this architecture, this could be because the CUDA