v4.5 dev update. (#3153)

examples/python/CuTeDSL/ampere/dynamic_smem_size.py

@@ -26,9 +26,11 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+
 import cutlass.cute as cute
 import cutlass
 
+
 """
 Example of automatic shared memory size computation for configuring kernel launch
 
@@ -51,11 +53,15 @@ class SharedData:
 
 
 @cute.kernel
-def kernel():
+def kernel_static():
     """
     Example kernel that allocates shared memory.
     The total allocation will be automatically calculated when smem=None.
     """
+    tidx, _, _ = cute.arch.block_idx()
+    if tidx == 0:
+        cute.printf("Running kernel_static")
+
     allocator = cutlass.utils.SmemAllocator()
 
     # Allocate various types of shared memory
@@ -68,6 +74,8 @@ def kernel():
         byte_alignment=16,
         swizzle=None,
     )
+
+    cute.printf("Kernel launch smem size: {}", cute.arch.dynamic_smem_size())
     return
 
 
@@ -79,7 +87,7 @@ def kernel_no_smem():
     """
     tidx, _, _ = cute.arch.block_idx()
     if tidx == 0:
-        cute.printf("Hello world")
+        cute.printf("Running kernel_no_smem")
     return
 
 
@@ -89,26 +97,49 @@ if __name__ == "__main__":
 
     print("Launching kernel with auto smem size. (launch config `smem=None`)")
 
-    # Compile the example
+    # Compile the static example
     @cute.jit
-    def launch_kernel1():
-        k = kernel()
-        k.launch(
+    def launch_kernelno_smem():
+        kernel_no_smem().launch(
             grid=(1, 1, 1),
             block=(1, 1, 1),
         )
-        print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
+
+    # --------
+    print(f" > Run {kernel_no_smem.__name__}")
+    func = cute.compile(launch_kernelno_smem)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
 
     @cute.jit
-    def launch_kernel2():
-        k = kernel_no_smem()
-        k.launch(
+    def launch_kernel_static():
+        kernel_static().launch(
             grid=(1, 1, 1),
             block=(1, 1, 1),
+            # smem=None
+            # auto infer launch kernel static smem usage
         )
-        print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
 
-    cute.compile(launch_kernel1)
-    cute.compile(launch_kernel2)
+    # --------
+    print(f" > Run {kernel_static.__name__} with sufficient smem")
+    func = cute.compile(launch_kernel_static)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
+
+    @cute.jit
+    def launch_kernel_static_insufficient():
+        kernel_static().launch(
+            grid=(1, 1, 1),
+            block=(1, 1, 1),
+            # launch kernel with static smem usage exceeds cfg
+            # show warning
+            smem=16,
+        )
+
+    # --------
+    print(f" > Run {kernel_static.__name__} with insufficient smem, show warning:")
+    func = cute.compile(launch_kernel_static_insufficient)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
 
     print("PASS")

examples/python/CuTeDSL/ampere/hstu_attention.py

@@ -265,7 +265,6 @@ class HSTUAttentionForwardAmpere(object):
         ).launch(
             grid=grid_dim,
             block=[self._num_threads, 1, 1],
-            smem=SharedStorage.size_in_bytes(),
             stream=stream,
         )
 

examples/python/CuTeDSL/ampere/smem_allocator.py

@@ -129,8 +129,8 @@ def kernel(
     # ptr<i64, smem, align<128>>
     # ptr<f32, smem, align<8>>
     print(struct_in_smem.a.data_ptr())
-    print(struct_in_smem.b)
-    print(struct_in_smem.c.real)
+    print(struct_in_smem.b.ptr)
+    print(struct_in_smem.c.real.ptr)
     # ptr<i8, smem, align<512>>
     print(section_in_smem)
     # ptr<i64, smem, align<64>>
@@ -138,6 +138,17 @@ def kernel(
     # tensor<ptr<f16, smem, align<32>> o (16,4):(1,16)>
     print(tensor_in_smem)
 
+    # assign struct member array element
+    cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
+    struct_in_smem.a[0] = 2
+    cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
+
+    # assign struct member scalar
+    cute.printf("struct_in_smem.b.ptr = {}", struct_in_smem.b.ptr)
+    cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
+    struct_in_smem.b = 16
+    cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
+
     # fill MemRange tensor in struct and copy to dst
     a_tensor = struct_in_smem.a.get_tensor(cute.make_layout((8, 4)))
     a_tensor.fill(const_a)
@@ -169,7 +180,9 @@ def host(
 ):
     # Note: Shared Memory size is automatically calculated now
     kernel(const_a, dst_a, const_b, dst_b, const_c, dst_c).launch(
-        grid=(1, 1, 1), block=(1, 1, 1)
+        grid=(1, 1, 1),
+        block=(1, 1, 1),
+        # Automatically calculate the launch kernel shared memory usage when `smem=None`
     )
 
 

examples/python/CuTeDSL/ampere/tensorop_gemm.py

@@ -175,15 +175,6 @@ class TensorOpGemm:
             (self.cta_tiler[0], self.cta_tiler[1]),
         )
 
-        # Shared memory allocated for operations with A, B will be
-        # overwritten for operations on C. This is to improve performance
-        # by reducing the size of shared memory requested by each block
-        smem_size = max(
-            cute.size_in_bytes(mC.element_type, sC_layout),
-            cute.size_in_bytes(mA.element_type, sA_layout)
-            + cute.size_in_bytes(mB.element_type, sB_layout),
-        )
-
         # ///////////////////////////////////////////////////////////////////////////////
         # Tiled copy:
         # The majorness of tA/tB/tC follows the majorness of gA/gB/gC,
@@ -282,7 +273,6 @@ class TensorOpGemm:
         ).launch(
             grid=rasterization_remap_grid_dim,
             block=[self.num_threads, 1, 1],
-            smem=smem_size,
         )
 
     @cute.kernel
@@ -382,14 +372,36 @@ class TensorOpGemm:
             # tAgA: (CPY, CPY_M, CPY_K, k)     , tBgB: (CPY, CPY_N, CPY_K, k)
             # tAsA: (CPY, CPY_M, CPY_K, PIPE)  , tBsB: (CPY, CPY_N, CPY_K, PIPE)
             # ///////////////////////////////////////////////////////////////////////////////
+            @cute.struct
+            class SharedStorageAB:
+                a: cute.struct.Align[
+                    cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)],
+                    16,
+                ]
+                b: cute.struct.Align[
+                    cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)],
+                    16,
+                ]
+
+            @cute.struct
+            class SharedStorageC:
+                c: cute.struct.Align[
+                    cute.struct.MemRange[mC.element_type, cute.cosize(sC_layout)],
+                    16,
+                ]
+
             # Shared memory buffer
             smem = cutlass.utils.SmemAllocator()
-
-            sA = smem.allocate_tensor(mA.element_type, sA_layout, 16)
-            sB = smem.allocate_tensor(mB.element_type, sB_layout, 16)
-            sC = cute.make_tensor(
-                cute.recast_ptr(sA.iterator, dtype=self.c_dtype), sC_layout
+            # Shared memory allocated for operations with A, B will be
+            # overwritten for operations on C. This is to improve performance
+            # by reducing the size of shared memory requested by each block
+            storage = smem.allocate(
+                max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()),
+                byte_alignment=16,
             )
+            sA = SharedStorageAB(storage).a.get_tensor(sA_layout)
+            sB = SharedStorageAB(storage).b.get_tensor(sB_layout)
+            sC = SharedStorageC(storage).c.get_tensor(sC_layout)
 
             thr_copy_A = tiled_copy_A.get_slice(tidx)
             thr_copy_B = tiled_copy_B.get_slice(tidx)