warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py

@@ -0,0 +1,96 @@
+# test/unit/conv/device/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.cu
+import pycutlass
+from pycutlass.conv2d_operation import *
+from pycutlass import *
+from pycutlass.test import *
+from pycutlass.utils.device import device_cc
+import unittest
+
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
+class Conv2dWgradImplicitGemmF32nhwcF32nhwcF32nhwcSimtF32SM80(unittest.TestCase):
+    def test_SM80_Device_Conv2d_Wgrad_Analytic_ImplicitGemm_f32nhwc_f32nhwc_f32nhwc_simt_f32(self):
+        math_inst = MathInstruction(
+            instruction_shape=[1, 1, 1],
+            element_a=cutlass.float32, element_b=cutlass.float32,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.Simt,
+            math_operation=MathOperation.multiply_add
+        )
+
+        A = TensorDescription(
+            element=math_inst.element_a, 
+            layout=cutlass.TensorNHWC,
+            alignment=4)
+        B = TensorDescription(
+            element=math_inst.element_b, 
+            layout=cutlass.TensorNHWC, 
+            alignment=4)
+        C = TensorDescription(
+            element=cutlass.float32,
+            layout=cutlass.TensorNHWC, 
+            alignment=1)
+
+        tile_description = TileDescription(
+            threadblock_shape=[128, 128, 8], stages=4, 
+            warp_count=[2, 4, 1],
+            math_instruction=math_inst
+        )
+
+        epilogue_functor = LinearCombination(
+            C.element, C.alignment, 
+            math_inst.element_accumulator, cutlass.float32)
+
+        operation = Conv2dOperation(
+            conv_kind=cutlass.conv.Operator.wgrad, iterator_algorithm=cutlass.conv.IteratorAlgorithm.analytic,
+            arch=80, tile_description=tile_description, A=A, B=B, C=C, 
+            stride_support=StrideSupport.Strided,
+            epilogue_functor=epilogue_functor,
+            swizzling_functor=cutlass.IdentitySwizzle1
+        )
+        
+        self.assertTrue(test_all_conv2d(operation))
+
+    def test_SM80_Device_Conv2d_Wgrad_Optimized_ImplicitGemm_f32nhwc_f32nhwc_f32nhwc_simt_f32(self):
+        math_inst = MathInstruction(
+            instruction_shape=[1, 1, 1],
+            element_a=cutlass.float32, element_b=cutlass.float32,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.Simt,
+            math_operation=MathOperation.multiply_add
+        )
+
+        A = TensorDescription(
+            element=math_inst.element_a, 
+            layout=cutlass.TensorNHWC,
+            alignment=4)
+        B = TensorDescription(
+            element=math_inst.element_b, 
+            layout=cutlass.TensorNHWC, 
+            alignment=4)
+        C = TensorDescription(
+            element=cutlass.float32,
+            layout=cutlass.TensorNHWC, 
+            alignment=1)
+
+        tile_description = TileDescription(
+            threadblock_shape=[128, 128, 8], stages=4, 
+            warp_count=[2, 4, 1],
+            math_instruction=math_inst
+        )
+
+        epilogue_functor = LinearCombination(
+            C.element, C.alignment, 
+            math_inst.element_accumulator, cutlass.float32)
+
+        operation = Conv2dOperation(
+            conv_kind=cutlass.conv.Operator.wgrad, iterator_algorithm=cutlass.conv.IteratorAlgorithm.optimized,
+            arch=80, tile_description=tile_description, A=A, B=B, C=C, 
+            stride_support=StrideSupport.Strided,
+            epilogue_functor=epilogue_functor,
+            swizzling_functor=cutlass.IdentitySwizzle1
+        )
+        
+        self.assertTrue(test_all_conv2d(operation))
+
+if __name__ == '__main__':
+    pycutlass.get_memory_pool(2**26, 2**26)
+    unittest.main()

warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -0,0 +1,107 @@
+# test/unit/conv/device/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.cu
+import pycutlass
+from pycutlass import *
+from pycutlass.test import *
+from pycutlass.utils.device import device_cc
+import unittest
+
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
+class Conv2dWgradImplicitGemmTF32nhwcTF32nhwcTF32nhwcTensorOpF32SM80(unittest.TestCase):
+    def test_SM80_Device_Conv2d_Wgrad_Optimized_ImplicitGemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32(self):
+        math_inst = MathInstruction(
+            instruction_shape=[16, 8, 8],
+            element_a=cutlass.float32, element_b=cutlass.float32,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
+            math_operation=MathOperation.multiply_add
+        )
+
+        A = TensorDescription(
+            element=math_inst.element_a, 
+            layout=cutlass.TensorNHWC,
+            alignment=4)
+        B = TensorDescription(
+            element=math_inst.element_b, 
+            layout=cutlass.TensorNHWC, 
+            alignment=4)
+        C = TensorDescription(
+            element=cutlass.float32,
+            layout=cutlass.TensorNHWC, 
+            alignment=8)
+
+        tile_description = TileDescription(
+            threadblock_shape=[128, 128, 16], stages=3, 
+            warp_count=[2, 2, 1],
+            math_instruction=math_inst
+        )
+
+        epilogue_functor = LinearCombination(
+            C.element, C.alignment, 
+            math_inst.element_accumulator, cutlass.float32)
+
+        operation = Conv2dOperation(
+            conv_kind=cutlass.conv.Operator.wgrad, iterator_algorithm=cutlass.conv.IteratorAlgorithm.optimized,
+            arch=80, tile_description=tile_description, A=A, B=B, C=C, 
+            stride_support=StrideSupport.Strided,
+            epilogue_functor=epilogue_functor,
+            swizzling_functor=cutlass.IdentitySwizzle1
+        )
+        
+        self.assertTrue(test_all_conv2d(operation))
+    
+    def test_SM80_Device_Conv2d_Wgrad_Optimized_ImplicitGemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_align1(self):
+        math_inst = MathInstruction(
+            instruction_shape=[16, 8, 8],
+            element_a=cutlass.float32, element_b=cutlass.float32,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
+            math_operation=MathOperation.multiply_add
+        )
+
+        A = TensorDescription(
+            element=math_inst.element_a, 
+            layout=cutlass.TensorNHWC,
+            alignment=1)
+        B = TensorDescription(
+            element=math_inst.element_b, 
+            layout=cutlass.TensorNHWC, 
+            alignment=1)
+        C = TensorDescription(
+            element=cutlass.float32,
+            layout=cutlass.TensorNHWC, 
+            alignment=4)
+
+        tile_description = TileDescription(
+            threadblock_shape=[128, 128, 32], stages=3, 
+            warp_count=[2, 2, 1],
+            math_instruction=math_inst
+        )
+
+        epilogue_functor = LinearCombination(
+            C.element, C.alignment, 
+            math_inst.element_accumulator, cutlass.float32)
+
+        operation = Conv2dOperation(
+            conv_kind=cutlass.conv.Operator.wgrad, iterator_algorithm=cutlass.conv.IteratorAlgorithm.optimized,
+            arch=80, tile_description=tile_description, A=A, B=B, C=C, 
+            stride_support=StrideSupport.Strided,
+            epilogue_functor=epilogue_functor,
+            swizzling_functor=cutlass.IdentitySwizzle1
+        )
+        
+        problem_sizes = [
+            cutlass.conv.Conv2dProblemSize(
+                cutlass.Tensor4DCoord(1, 8, 8, 1),
+                cutlass.Tensor4DCoord(1, 3, 3, 1),
+                cutlass.Tensor4DCoord(1, 1, 1, 1),
+                cutlass.MatrixCoord(1, 1),
+                cutlass.MatrixCoord(1, 1),
+                cutlass.conv.Mode.cross_correlation,
+                1, 1
+            ),
+        ]
+        
+        self.assertTrue(test_all_conv2d(operation, problem_sizes))
+    
+if __name__ == '__main__':
+    pycutlass.get_memory_pool(2**26, 2**26)
+    unittest.main()

warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/run_all_tests.py

@@ -0,0 +1,10 @@
+import pycutlass
+import unittest
+from pycutlass.memory_manager import *
+
+if __name__ == '__main__':
+    pycutlass.get_memory_pool(2**32, 2**32)
+    loader = unittest.TestLoader()
+    tests = loader.discover('./', 'conv2d_*.py')
+    testRunner = unittest.runner.TextTestRunner()
+    testRunner.run(tests)

warp/native/cutlass/tools/library/scripts/pycutlass/test/frontend/test_frontend.py

@@ -0,0 +1,146 @@
+#################################################################################################
+#
+# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# 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.
+#
+#################################################################################################
+## Test case for Pytorch
+import pycutlass
+import unittest
+from pycutlass import *
+from pycutlass.utils.device import device_cc
+import torch
+import cupy as cp
+
+
+class Test_Frontend(unittest.TestCase):
+    def setUp(self) -> None:
+        #
+        # define the cutlass operator
+        #
+        cc = device_cc()
+        math_inst = MathInstruction(
+            [1, 1, 1], cutlass.float32, cutlass.float32, cutlass.float32,
+            cutlass.OpClass.Simt, MathOperation.multiply_add
+        )
+
+        # Stages > 2 is supported only for compute capability 80 and beyond
+        stages = 4 if cc >= 80 else 2
+
+
+        tile_description = TileDescription(
+            [128, 128, 8], stages, [2, 4, 1],
+            math_inst
+        )
+
+        A = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        B = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        C = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        epilogue_functor = LinearCombination(
+            C.element, C.alignment, 
+            math_inst.element_accumulator, cutlass.float32)
+
+        self.operation = GemmOperationUniversal(
+            arch=cc, tile_description=tile_description,
+            A=A, B=B, C=C, 
+            epilogue_functor=epilogue_functor, 
+            swizzling_functor=cutlass.IdentitySwizzle1
+        )
+
+        pycutlass.compiler.add_module([self.operation,])
+
+
+    def test_torch_frontend(self):
+        problem_size = cutlass.gemm.GemmCoord(512, 256, 128)
+
+        tensor_A = torch.ceil(torch.empty(size=(problem_size.m(), problem_size.k()), dtype=torch.float32, device="cuda").uniform_(-8.5, 7.5))
+        tensor_B = torch.ceil(torch.empty(size=(problem_size.k(), problem_size.n()), dtype=torch.float32, device="cuda").uniform_(-8.5, 7.5))
+        tensor_C = torch.ceil(torch.empty(size=(problem_size.m(), problem_size.n()), dtype=torch.float32, device="cuda").uniform_(-8.5, 7.5))
+        tensor_D = torch.empty_like(tensor_C)
+        
+
+        alpha = 1.0
+        beta = 0.0
+
+        arguments = GemmArguments(
+            operation=self.operation, problem_size=problem_size,
+            A=tensor_A, B=tensor_B, C=tensor_C, D=tensor_D,
+            output_op=self.operation.epilogue_type(alpha, beta),
+            gemm_mode=cutlass.gemm.Mode.Gemm, split_k_splices=1
+        )
+
+        self.operation.run(arguments)
+
+        arguments.sync()
+
+        tensor_D_ref = alpha * tensor_A @ tensor_B + beta * tensor_C
+
+        self.assertTrue(torch.equal(tensor_D, tensor_D_ref))
+    
+    def test_cupy_frontend(self):
+        cp.cuda.set_allocator(rmm.rmm_cupy_allocator)
+
+        problem_size = cutlass.gemm.GemmCoord(512, 256, 128)
+
+        tensor_A = cp.ceil(cp.random.uniform(low=-8.5, high=7.5, size=(problem_size.m(), problem_size.k()), dtype=cp.float32))
+        tensor_B = cp.ceil(cp.random.uniform(low=-8.5, high=7.5, size=(problem_size.k(), problem_size.n()), dtype=cp.float32))
+        tensor_C = cp.ceil(cp.random.uniform(low=-8.5, high=7.5, size=(problem_size.m(), problem_size.n()), dtype=cp.float32))
+        tensor_D = cp.ones_like(tensor_C)
+
+        alpha = 1.0
+        beta = 1.0
+
+        tensor_D_ref = alpha * tensor_A @ tensor_B + beta * tensor_C
+
+        arguments = GemmArguments(
+            operation=self.operation, problem_size=problem_size,
+            A=tensor_A, B=tensor_B, C=tensor_C, D=tensor_D,
+            output_op=self.operation.epilogue_type(alpha, beta),
+            gemm_mode=cutlass.gemm.Mode.Gemm, split_k_splices=1
+        )
+
+        self.operation.run(arguments)
+
+        arguments.sync()
+
+        self.assertTrue(cp.array_equal(tensor_D, tensor_D_ref))
+
+
+
+if __name__ == '__main__':
+    pycutlass.get_memory_pool(2**32, 2**32)
+    unittest.main()

warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py

@@ -0,0 +1,96 @@
+import pycutlass
+from pycutlass import *
+from pycutlass.test import *
+import unittest
+
+from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
+
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
+class GemmBF16TensorOpSm80(unittest.TestCase):
+    def SM80_Device_Gemm_bf16n_bf16n_f32t_tensor_op_f32_64x128x64_32x64x64(self):
+        math_inst = MathInstruction(
+            instruction_shape=[16, 8, 16],
+            element_a=cutlass.bfloat16, element_b=cutlass.bfloat16,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
+            math_operation=MathOperation.multiply_add
+        )
+
+        tile_description = TileDescription(
+            threadblock_shape=[64, 128, 64],
+            stages=4, warp_count=[2, 2, 1],
+            math_instruction=math_inst
+        )
+
+        A = TensorDescription(
+            element=cutlass.bfloat16, layout=cutlass.ColumnMajor,
+            alignment=8
+        )
+        B = TensorDescription(
+            element=cutlass.bfloat16, layout=cutlass.ColumnMajor,
+            alignment=8
+        )
+        C = TensorDescription(
+            element=cutlass.float32, layout=cutlass.RowMajor,
+            alignment=4
+        )
+
+        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, 
+            epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
+        )
+
+        self.assertTrue(test_all_gemm(operation, "universal"))
+    
+    def test_SM80_Device_Gemm_bf16t_bf16t_bf16t_tensor_op_f32_128x256x64_64x64x64(self):
+        math_inst = MathInstruction(
+            instruction_shape=[16, 8, 16],
+            element_a=cutlass.bfloat16, element_b=cutlass.bfloat16,
+            element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
+            math_operation=MathOperation.multiply_add
+        )
+
+        tile_description = TileDescription(
+            threadblock_shape=[64, 128, 32],
+            stages=6, warp_count=[2, 2, 1],
+            math_instruction=math_inst
+        )
+
+        A = TensorDescription(
+            element=cutlass.bfloat16, layout=cutlass.RowMajor,
+            alignment=8
+        )
+        B = TensorDescription(
+            element=cutlass.bfloat16, layout=cutlass.RowMajor,
+            alignment=8
+        )
+        C = TensorDescription(
+            element=cutlass.bfloat16, layout=cutlass.RowMajor,
+            alignment=8
+        )
+
+        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, 
+            epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
+        )
+
+        self.assertTrue(test_all_gemm(operation, "multistage"))
+
+if __name__ == '__main__':
+    pycutlass.get_memory_pool(2**24, 2**24)
+    unittest.main()