v4.4 tag release update. (#3032)

examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py

@@ -30,15 +30,12 @@ import argparse
 from typing import Optional, Type, Tuple, Union
 import cuda.bindings.driver as cuda
 
-import torch
-
 import cutlass
 import cutlass.cute as cute
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
-import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.cute.testing as testing
 
@@ -767,6 +764,14 @@ class DenseGemmKernel:
         acc_pipeline.consumer_wait(acc_consumer_state)
 
         if cutlass.const_expr(self.use_tma_store):
+            # TODO: segment fault if I put here
+            # sC = smem.allocate_tensor(
+            #     element_type=self.c_dtype,
+            #     layout=c_smem_layout_staged.outer,
+            #     byte_alignment=128,
+            #     swizzle=c_smem_layout_staged.inner,
+            # )
+
             assert tma_atom_c is not None and sC is not None
             self.epilogue_tma_store(
                 tidx,
@@ -983,7 +988,10 @@ class DenseGemmKernel:
             c_buffer = subtile_idx % self.num_c_stage
             cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
             # Fence and barrier to make sure shared memory store is visible to TMA store
-            cute.arch.fence_proxy("async.shared", space="cta")
+            cute.arch.fence_proxy(
+                "async.shared",
+                space="cta",
+            )
             pipeline.sync(barrier_id=1)
 
             # TMA store C to global memory
@@ -1214,7 +1222,7 @@ class DenseGemmKernel:
         """
         acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2])
         tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
-        return sm100_utils.get_num_tmem_alloc_cols(tCtAcc_fake)
+        return utils.get_num_tmem_alloc_cols(tCtAcc_fake)
 
     def is_valid_dtypes(
         self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric]
@@ -1365,6 +1373,7 @@ class DenseGemmKernel:
         """
         is_valid = True
 
+        # TODO: move to utils
         def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape):
             major_mode_idx = 0 if is_mode0_major else 1
             num_major_elements = tensor_shape[major_mode_idx]
@@ -1446,6 +1455,9 @@ class DenseGemmKernel:
 
 
 def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype):
+    import torch
+    import cutlass.torch as cutlass_torch
+
     torch.manual_seed(1111)
 
     a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
@@ -1474,6 +1486,9 @@ def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype):
 
 
 def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance):
+    import torch
+    import cutlass.torch as cutlass_torch
+
     # Copy gpu result back
     kernel_result = c_torch_gpu.cpu()
 
@@ -1563,6 +1578,7 @@ def run(
     print(f"Iterations: {iterations}")
     print(f"Skip reference checking: {skip_ref_check}")
     print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
+    import torch
 
     # Unpack parameters
     m, n, k, l = mnkl
@@ -1600,6 +1616,8 @@ def run(
         compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance)
 
     def generate_tensors():
+        import cutlass.torch as cutlass_torch
+
         a_tensor, _ = cutlass_torch.cute_tensor_like(
             a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
         )