CUTLASS 2.10 (#615)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
2022-09-03 18:48:46 -04:00
parent ca23ff7924
commit b72cbf957d
289 changed files with 43708 additions and 2513 deletions
@@ -0,0 +1 @@
+CUPY_CACHE_DIR=./ python test_frontend.py
@@ -0,0 +1,136 @@
+#################################################################################################
+#
+# 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 *
+import torch
+import cupy as cp
+
+
+class Test_Frontend(unittest.TestCase):
+    def setUp(self) -> None:
+        #
+        # define the cutlass operator
+        #
+        math_inst = MathInstruction(
+            [1, 1, 1], cutlass.float32, cutlass.float32, cutlass.float32,
+            cutlass.OpClass.Simt, MathOperation.multiply_add
+        )
+
+        tile_description = TileDescription(
+            [128, 128, 8], 4, [2, 4, 1],
+            math_inst, 80, 80
+        )
+
+        A = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        B = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        C = TensorDescription(
+            cutlass.float32, cutlass.RowMajor, 1
+        )
+
+        self.operation = GemmOperationUniversal(
+            arch=80, tile_description=tile_description,
+            A=A, B=B, C=C, element_epilogue=cutlass.float32,
+            epilogue_functor=EpilogueFunctor.LinearCombination, 
+            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=LinearCombinationFunctorArguments(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=LinearCombinationFunctorArguments(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()