CUTLASS 3.2.1 (#1113)
* Updates for 3.2.1 release. * Minor fix in gemm op profiler for raster order. * Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
100
test/python/cutlass/evt/evt_compute_sm80_90.py
Normal file
100
test/python/cutlass/evt/evt_compute_sm80_90.py
Normal file
@@ -0,0 +1,100 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
#
|
||||
################################################################################
|
||||
"""
|
||||
Unit test for compute node in SM90
|
||||
"""
|
||||
|
||||
import logging
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass.backend import *
|
||||
from cutlass.epilogue import *
|
||||
from cutlass import swizzle
|
||||
|
||||
from utils.evt_testbed import EVTTestBed, EVTTestCaseBase
|
||||
|
||||
cutlass.set_log_level(logging.WARNING)
|
||||
|
||||
|
||||
@unittest.skipIf(device_cc() not in [80, 90], "This unittest is for Sm80 and Sm90 only")
|
||||
class TestEVTComputeSM90(EVTTestCaseBase):
|
||||
|
||||
def test_arith(self):
|
||||
"""
|
||||
Test Arithmatic op
|
||||
"""
|
||||
def evt_arith_compute(accum, C, alpha, beta, gamma):
|
||||
D = ((accum + C) * alpha - gamma) / beta
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.5,
|
||||
"beta": 0.5,
|
||||
"gamma": 2.5,
|
||||
"D": self.fake_tensor(self.element, (l, m, n))
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_arith_compute, example_inputs)
|
||||
input_keys = ["C", "alpha", "beta", "gamma"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_func_call(self):
|
||||
"""
|
||||
Test Function call
|
||||
"""
|
||||
def evt_func_call(accum, C, alpha, beta, gamma):
|
||||
D = multiply_add(relu(accum + alpha) + C, beta, gamma)
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.5,
|
||||
"beta": 0.5,
|
||||
"gamma": 2.5,
|
||||
"D": self.fake_tensor(self.element, (l, m, n))
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_func_call, example_inputs)
|
||||
input_keys = ["C", "alpha", "beta", "gamma"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
173
test/python/cutlass/evt/evt_layout_sm80_90.py
Normal file
173
test/python/cutlass/evt/evt_layout_sm80_90.py
Normal file
@@ -0,0 +1,173 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
#
|
||||
################################################################################
|
||||
|
||||
"""
|
||||
Unit test for store nodes in SM90
|
||||
"""
|
||||
|
||||
import logging
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass.backend import *
|
||||
from cutlass.epilogue import *
|
||||
|
||||
from utils.evt_testbed import EVTTestBed, EVTTestCaseBase
|
||||
|
||||
cutlass.set_log_level(logging.WARNING)
|
||||
|
||||
|
||||
@unittest.skipIf(device_cc() not in [80, 90], "This unittest is for Sm80 and Sm90 only")
|
||||
class TestEVTLayoutSM90(EVTTestCaseBase):
|
||||
|
||||
def test_permute_1(self):
|
||||
"""
|
||||
Returning a tensor with shape [m, n]
|
||||
"""
|
||||
def evt_permute(accum, alpha, C):
|
||||
F = alpha * accum
|
||||
F_permute = permute(F, indices=(0, 2, 1))
|
||||
D_permute = F_permute + permute(C, indices=(0, 2, 1))
|
||||
D = permute(D_permute, indices=(0, 2, 1))
|
||||
return D, F
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 0.5,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_permute, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() == 80, "This unittest is for cc = Sm90 only")
|
||||
def test_permute_2(self):
|
||||
"""
|
||||
Returning a tensor with shape [m, n]
|
||||
"""
|
||||
def evt_permute(accum, alpha, C):
|
||||
F = alpha * accum
|
||||
F_permute = permute(F, indices=(0, 2, 1))
|
||||
D = F_permute + C
|
||||
return D, F
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 0.5,
|
||||
"C": self.fake_tensor(self.element, (l, n, m)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"D": self.fake_tensor(self.element, (l, n, m)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_permute, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() == 80, "This unittest is for cc = Sm90 only")
|
||||
def test_permute_3(self):
|
||||
"""
|
||||
Returning a tensor with shape [m, n]
|
||||
"""
|
||||
def evt_permute(accum, alpha, C):
|
||||
F = alpha * accum
|
||||
F_permute = permute(F, indices=(1, 0, 2))
|
||||
D = F_permute + C
|
||||
return D, F
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 0.5,
|
||||
"C": self.fake_tensor(self.element, (m, l, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"D": self.fake_tensor(self.element, (m, l, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_permute, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_reshape(self):
|
||||
"""
|
||||
Test reshape
|
||||
"""
|
||||
def evt_reshape(accum, alpha, TensorE):
|
||||
F = alpha * accum
|
||||
E_reshape = reshape(TensorE, new_shape=(512, 1))
|
||||
D = F + E_reshape
|
||||
return D
|
||||
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (self.l, self.m, self.n)),
|
||||
"alpha": 0.5,
|
||||
"TensorE": self.fake_tensor(self.element, (16, 32)),
|
||||
"D": self.fake_tensor(self.element, (self.l, self.m, self.n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_reshape, example_inputs)
|
||||
input_keys = ["alpha", "TensorE"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify(self.problem_size, input_keys, result_keys, self.l)
|
||||
|
||||
def test_reshape2(self):
|
||||
"""
|
||||
Test reshape
|
||||
"""
|
||||
def evt_reshape(accum, alpha, TensorE):
|
||||
F = alpha * accum
|
||||
F_reshape = reshape(F, new_shape=(2, 3, 512, 256))
|
||||
D = F_reshape + TensorE
|
||||
return D
|
||||
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (self.l, self.m, self.n)),
|
||||
"alpha": 0.5,
|
||||
"TensorE": self.fake_tensor(self.element, (2, 3, 1, self.n)),
|
||||
"D": self.fake_tensor(self.element, (2, 3, self.m, self.n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_reshape, example_inputs)
|
||||
input_keys = ["alpha", "TensorE"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify(self.problem_size, input_keys, result_keys, self.l)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
142
test/python/cutlass/evt/evt_load_sm80_90.py
Normal file
142
test/python/cutlass/evt/evt_load_sm80_90.py
Normal file
@@ -0,0 +1,142 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
#
|
||||
################################################################################
|
||||
|
||||
"""
|
||||
Unit test for load nodes in SM90
|
||||
"""
|
||||
|
||||
import logging
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass.backend import *
|
||||
from cutlass.epilogue import *
|
||||
|
||||
from utils.evt_testbed import EVTTestBed, EVTTestCaseBase
|
||||
|
||||
cutlass.set_log_level(logging.WARNING)
|
||||
|
||||
|
||||
@unittest.skipIf(device_cc() not in [80, 90], "This unittest is for Sm80 and Sm90 only")
|
||||
class TestEVTLoadSM90(EVTTestCaseBase):
|
||||
|
||||
def test_tensor_load(self):
|
||||
"""
|
||||
Load extra tensor with shape [m, n]
|
||||
"""
|
||||
def evt_tensor_load(accum, C, aux, aux_batch):
|
||||
D = accum + C + aux + aux_batch
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"aux": self.fake_tensor(self.element, (m, n)),
|
||||
"aux_batch": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_tensor_load, example_inputs)
|
||||
input_keys = ["C", "aux", "aux_batch"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_row_broadcast(self):
|
||||
"""
|
||||
Load extra tensor with shape [1, n]
|
||||
"""
|
||||
def evt_row_broadcast(accum, C, bias, bias_batch):
|
||||
D = accum + C + bias + bias_batch
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"bias": self.fake_tensor(self.element, (n,)),
|
||||
"bias_batch": self.fake_tensor(np.float32, (l, 1, n)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_row_broadcast, example_inputs)
|
||||
input_keys = ["C", "bias", "bias_batch"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_column_broadcast(self):
|
||||
"""
|
||||
Load extra tensor with shape [m, 1]
|
||||
"""
|
||||
def evt_column_broadcast(accum, C, bias, bias_batch):
|
||||
D = accum + C + bias + bias_batch
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"bias": self.fake_tensor(self.element, (m, 1)),
|
||||
"bias_batch": self.fake_tensor(np.float32, (l, m, 1)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_column_broadcast, example_inputs)
|
||||
input_keys = ["C", "bias", "bias_batch"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_scalar_broadcast(self):
|
||||
"""
|
||||
Load extra tensor with shape [1, 1]
|
||||
"""
|
||||
def evt_scalar_broadcast(accum, C, alpha, alpha_batch):
|
||||
D = accum + C + alpha + alpha_batch
|
||||
return D
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 0.5,
|
||||
"alpha_batch": self.fake_tensor(np.float32, (l, 1, 1)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_scalar_broadcast, example_inputs)
|
||||
input_keys = ["C", "alpha", "alpha_batch"]
|
||||
result_keys = ["D"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
274
test/python/cutlass/evt/evt_mixed_sm80_90.py
Normal file
274
test/python/cutlass/evt/evt_mixed_sm80_90.py
Normal file
@@ -0,0 +1,274 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 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.
|
||||
#
|
||||
################################################################################
|
||||
|
||||
"""
|
||||
Unittest for mixed types of nodes in SM90
|
||||
"""
|
||||
|
||||
import logging
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass.backend import *
|
||||
from cutlass.epilogue import *
|
||||
from cutlass.swizzle import ThreadblockSwizzleStreamK
|
||||
|
||||
from utils.evt_testbed import EVTTestBed, EVTTestCaseBase
|
||||
|
||||
cutlass.set_log_level(logging.WARNING)
|
||||
|
||||
|
||||
@unittest.skipIf(device_cc() not in [80, 90], "This unittest is for Sm80 and Sm90 only")
|
||||
class TestEVTMixedSM90(EVTTestCaseBase):
|
||||
def test_mixed_dag(self):
|
||||
def evt_mixed_dag(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
if device_cc() == 80:
|
||||
aligments = [2, 4, 8]
|
||||
else:
|
||||
# Sm90 EVT currently only supports 128-bit alignment
|
||||
aligments = [8,]
|
||||
for align in aligments:
|
||||
for m, n, k, l in self.get_problem_sizes(align):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (l, m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_mixed_dag, example_inputs)
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() != 80, "This unittest is for cc = Sm80 only")
|
||||
def test_mixed_dag_float(self):
|
||||
def evt_mixed_dag(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
for align in [3, 2, 4]:
|
||||
for m, n, k, l in self.get_problem_sizes(align):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"cbias": self.fake_tensor(np.float32, (m, 1)),
|
||||
"rbias": self.fake_tensor(np.float32, (n,)),
|
||||
"D": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"F": self.fake_tensor(np.float32, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(np.float32, (n,)),
|
||||
"E_col_max": self.fake_tensor(np.float32, (m, 1))
|
||||
}
|
||||
launcher = EVTTestBed(DataType.f32, evt_mixed_dag, example_inputs)
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() != 80, "This unittest is for cc = Sm80 only")
|
||||
def test_mixed_dag_stage2(self):
|
||||
def evt_mixed_dag(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (l, m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_mixed_dag, example_inputs, epilogue_stages=2)
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() != 80, "This unittest is for cc = Sm80 only")
|
||||
def test_mixed_dag_partition_k(self):
|
||||
def evt_mixed_dag(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (l, m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
|
||||
tile_description = {
|
||||
"threadblock_shape": [128, 128, 64],
|
||||
"warp_count": [2, 2, 2]
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_mixed_dag, example_inputs, tile_description=tile_description, epilogue_stages=2)
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
@unittest.skipIf(device_cc() != 80, "This unittest is for cc = Sm80 only")
|
||||
def test_mixed_dag_stream_k(self):
|
||||
def evt_mixed_dag(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
# High per-sm occupancy tile_description
|
||||
tile_description = {
|
||||
"threadblock_shape": [128, 128, 32],
|
||||
"warp_count": [2, 2, 1],
|
||||
"stages": 3
|
||||
}
|
||||
tds = [None, tile_description]
|
||||
for td in tds:
|
||||
for m, n, k, l in self.get_problem_sizes(8, k=960, batch_count=[1, 3]):
|
||||
if l == 1:
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (m, n)),
|
||||
"F": self.fake_tensor(self.element, (m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
else:
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (l, m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
|
||||
if td is not None:
|
||||
launcher = EVTTestBed(
|
||||
self.element, evt_mixed_dag, example_inputs,
|
||||
tile_description=td,
|
||||
swizzling_functor=ThreadblockSwizzleStreamK, backend="torch")
|
||||
else:
|
||||
launcher = EVTTestBed(
|
||||
self.element, evt_mixed_dag, example_inputs,
|
||||
swizzling_functor=ThreadblockSwizzleStreamK, backend="torch")
|
||||
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_mixed_dag_no_batch(self):
|
||||
def evt_mixed_dag_no_batch(accum, alpha, C, beta, aux, cbias, rbias):
|
||||
F = alpha * accum + (beta * C + aux)
|
||||
F_row_max = max(F, dim=[0, 1])
|
||||
E = relu(F + 1) + cbias + rbias
|
||||
E_col_max = max(E, dim=[0, 2])
|
||||
D = E + F
|
||||
return D, F, F_row_max, E_col_max
|
||||
|
||||
for m, n, k, _ in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (m, n)),
|
||||
"alpha": 1.0,
|
||||
"C": self.fake_tensor(self.element, (m, n)),
|
||||
"beta": 1.0,
|
||||
"aux": self.fake_tensor(self.element, (m, n)),
|
||||
"cbias": self.fake_tensor(self.element, (m, 1)),
|
||||
"rbias": self.fake_tensor(self.element, (n,)),
|
||||
"D": self.fake_tensor(self.element, (m, n)),
|
||||
"F": self.fake_tensor(self.element, (m, n)),
|
||||
"F_row_max": self.fake_tensor(DataType.f32, (n,)),
|
||||
"E_col_max": self.fake_tensor(DataType.f32, (m, 1))
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_mixed_dag_no_batch, example_inputs)
|
||||
input_keys = ["alpha", "C", "beta", "aux", "cbias", "rbias"]
|
||||
result_keys = ["D", "F", "F_row_max", "E_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, 1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
155
test/python/cutlass/evt/evt_store_sm80_90.py
Normal file
155
test/python/cutlass/evt/evt_store_sm80_90.py
Normal file
@@ -0,0 +1,155 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
#
|
||||
################################################################################
|
||||
|
||||
"""
|
||||
Unit test for store nodes in SM90
|
||||
"""
|
||||
|
||||
import logging
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass.backend import *
|
||||
from cutlass.epilogue import *
|
||||
|
||||
from utils.evt_testbed import EVTTestBed, EVTTestCaseBase
|
||||
|
||||
cutlass.set_log_level(logging.WARNING)
|
||||
|
||||
|
||||
@unittest.skipIf(device_cc() not in [80, 90], "This unittest is for Sm80 and Sm90 only")
|
||||
class TestEVTStoreSM90(EVTTestCaseBase):
|
||||
|
||||
def test_aux_store(self):
|
||||
"""
|
||||
Returning a tensor with shape [m, n]
|
||||
"""
|
||||
def evt_aux_store(accum, alpha, C):
|
||||
F = alpha * accum
|
||||
D = F + C
|
||||
return D, F
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 0.5,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F": self.fake_tensor(self.element, (l, m, n)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_aux_store, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_col_reduce(self):
|
||||
"""
|
||||
Reduction [m, n] -> [m, 1]
|
||||
"""
|
||||
def evt_row_reduce(accum, alpha, C):
|
||||
acc_row_max = max(accum, dim=[2,])
|
||||
F = alpha * accum
|
||||
F_row_max = max(F, dim=[0, 2])
|
||||
D = F + C
|
||||
return D, F_row_max, acc_row_max
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 2.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_row_max": self.fake_tensor(np.float32, (m, 1)),
|
||||
"acc_row_max": self.fake_tensor(np.float32, (l, m, 1)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_row_reduce, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F_row_max", "acc_row_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_row_reduce(self):
|
||||
"""
|
||||
Reduction [m, n] -> [n]
|
||||
"""
|
||||
def evt_col_reduce(accum, alpha, C):
|
||||
acc_col_max = max(accum, dim=[1,])
|
||||
F = alpha * accum
|
||||
F_col_max = max(F, dim=[0, 1])
|
||||
D = F + C
|
||||
return D, F_col_max, acc_col_max
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 2.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"F_col_max": self.fake_tensor(np.float32, (n,)),
|
||||
"acc_col_max": self.fake_tensor(np.float32, (l, 1, n)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_col_reduce, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F_col_max", "acc_col_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
def test_scalar_reduce(self):
|
||||
"""
|
||||
Reduction [m, n] -> [1,]
|
||||
"""
|
||||
def evt_scalar_reduce(accum, alpha, C):
|
||||
acc_max = max(accum, dim=[1, 2])
|
||||
F = alpha * accum
|
||||
F_max = max(F, dim=[0, 1, 2])
|
||||
D = F + C
|
||||
return D, F_max, acc_max
|
||||
|
||||
for m, n, k, l in self.get_problem_sizes(8):
|
||||
example_inputs = {
|
||||
"accum": self.fake_tensor(self.element, (l, m, n)),
|
||||
"alpha": 2.0,
|
||||
"C": self.fake_tensor(self.element, (l, m, n)),
|
||||
"acc_max": self.fake_tensor(np.float32, (l, 1, 1)),
|
||||
"F_max": self.fake_tensor(np.float32, (1,)),
|
||||
"D": self.fake_tensor(self.element, (l, m, n)),
|
||||
}
|
||||
|
||||
launcher = EVTTestBed(self.element, evt_scalar_reduce, example_inputs)
|
||||
input_keys = ["C", "alpha"]
|
||||
result_keys = ["D", "F_max", "acc_max"]
|
||||
launcher.verify((m, n, k), input_keys, result_keys, l)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
44
test/python/cutlass/evt/run_all_tests.py
Normal file
44
test/python/cutlass/evt/run_all_tests.py
Normal file
@@ -0,0 +1,44 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import pathlib
|
||||
import unittest
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
loader = unittest.TestLoader()
|
||||
script_dir = str(pathlib.Path(__file__).parent.resolve()) + '/'
|
||||
tests = loader.discover(script_dir, 'evt_*.py')
|
||||
testRunner = unittest.runner.TextTestRunner()
|
||||
results = testRunner.run(tests)
|
||||
if not results.wasSuccessful():
|
||||
raise Exception('Test cases failed')
|
||||
230
test/python/cutlass/evt/utils/evt_testbed.py
Normal file
230
test/python/cutlass/evt/utils/evt_testbed.py
Normal file
@@ -0,0 +1,230 @@
|
||||
################################################################################
|
||||
#
|
||||
# Copyright (c) 20123 - 2023 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.
|
||||
#
|
||||
################################################################################
|
||||
|
||||
"""
|
||||
Testbed classes of EVT
|
||||
"""
|
||||
|
||||
import torch
|
||||
import unittest
|
||||
|
||||
import cutlass
|
||||
from cutlass import Tensor
|
||||
import cutlass.backend.evt
|
||||
from cutlass.profiler import CUDAEventProfiler
|
||||
from cutlass.shape import GemmCoord
|
||||
from cutlass.utils.datatypes import torch_type
|
||||
|
||||
|
||||
class EVTReferenceModule:
|
||||
def __init__(self, layout_A, layout_B, layout_C, epilogue_visitor):
|
||||
self.layout_A = layout_A
|
||||
self.layout_B = layout_B
|
||||
self.layout_C = layout_C
|
||||
self.epilogue_visitor = epilogue_visitor
|
||||
|
||||
def run(self, A, B, C, problem_size, alpha, beta, batch=1):
|
||||
if self.layout_A == cutlass.LayoutType.RowMajor:
|
||||
A_row = A.view((batch, problem_size.m, problem_size.k))
|
||||
else:
|
||||
A_col = A.view((batch, problem_size.k, problem_size.m))
|
||||
A_row = torch.permute(A_col, (0, 2, 1))
|
||||
|
||||
if self.layout_B == cutlass.LayoutType.RowMajor:
|
||||
B_row = B.view((batch, problem_size.k, problem_size.n))
|
||||
else:
|
||||
B_col = B.view((batch, problem_size.n, problem_size.k))
|
||||
B_row = torch.permute(B_col, (0, 2, 1))
|
||||
|
||||
if self.layout_C == cutlass.LayoutType.RowMajor:
|
||||
C_row = C.view((batch, problem_size.m, problem_size.n))
|
||||
else:
|
||||
C_col = C.view((batch, problem_size.n, problem_size.m))
|
||||
C_row = torch.permute(C_col, (0, 2, 1))
|
||||
|
||||
out_row = torch.matmul(A_row, B_row) * alpha + C_row * beta
|
||||
|
||||
if self.layout_C == cutlass.LayoutType.ColumnMajor:
|
||||
out = torch.permute(out_row, (0, 2, 1))
|
||||
else:
|
||||
out = out_row
|
||||
|
||||
return torch.flatten(out)
|
||||
|
||||
def __call__(self, A, B, C, problem_size, batch=1, epilogue_args=None):
|
||||
# Running the mainloop
|
||||
accum = self.run(
|
||||
A, B, C, problem_size, 1.0, 0.0, batch=batch
|
||||
).reshape(batch, problem_size.m, problem_size.n)
|
||||
|
||||
# Running the epilogue
|
||||
epilogue_args["accum"] = accum
|
||||
references = self.epilogue_visitor(**epilogue_args)
|
||||
|
||||
# Return the results
|
||||
if not isinstance(references, tuple):
|
||||
references = (references,)
|
||||
return references
|
||||
|
||||
|
||||
class EVTTestBed:
|
||||
"""
|
||||
Epilogue Visitor Testbed
|
||||
"""
|
||||
def __init__(self, element, evt_fn, example_inputs, profile=False, **kwargs) -> None:
|
||||
self.element = element
|
||||
layout = cutlass.LayoutType.RowMajor
|
||||
self.example_inputs = example_inputs
|
||||
|
||||
# Create the Gemm plan
|
||||
self.plan = cutlass.op.Gemm(element=element, layout=layout, element_accumulator=torch.float32)
|
||||
|
||||
if "tile_description" in kwargs:
|
||||
self.plan.tile_description = kwargs["tile_description"]
|
||||
|
||||
if "swizzling_functor" in kwargs:
|
||||
self.plan.swizzling_functor = kwargs["swizzling_functor"]
|
||||
|
||||
# Compile the epilogue visitor
|
||||
epilogue_visitor = cutlass.epilogue.trace(evt_fn, example_inputs)
|
||||
if "epilogue_stages" in kwargs:
|
||||
epilogue_visitor.epilogue_stages = kwargs["epilogue_stages"]
|
||||
self.plan.epilogue_visitor = epilogue_visitor
|
||||
|
||||
# Reference model
|
||||
self.reference_fn = EVTReferenceModule(layout, layout, layout, epilogue_visitor)
|
||||
|
||||
self.profile = profile
|
||||
|
||||
def get_torch_tensor(self, shape, dtype=None, fill=None):
|
||||
if dtype is None:
|
||||
dtype = self.element
|
||||
|
||||
dtype = torch_type(dtype)
|
||||
if fill is None:
|
||||
return torch.ceil(
|
||||
torch.empty(size=shape, dtype=dtype, device="cuda").uniform_(-4.5, 3.5)
|
||||
)
|
||||
else:
|
||||
return torch.full(shape, fill, dtype=dtype, device="cuda")
|
||||
|
||||
def verify(self, problem_size, input_keys, result_keys, batch_count=1):
|
||||
"""
|
||||
Verify the results
|
||||
"""
|
||||
problem_size = GemmCoord(*problem_size)
|
||||
|
||||
# Initiate the GEMM arguments
|
||||
tensor_A = self.get_torch_tensor((batch_count, problem_size.m, problem_size.k))
|
||||
tensor_B = self.get_torch_tensor((batch_count, problem_size.k, problem_size.n))
|
||||
|
||||
# Initialize the epilogue args
|
||||
epilogue_args = {}
|
||||
for key in self.example_inputs.keys():
|
||||
if key in input_keys:
|
||||
tensor = self.example_inputs[key]
|
||||
if isinstance(tensor, Tensor):
|
||||
epilogue_args[key] = self.get_torch_tensor(tensor.shape, tensor.element)
|
||||
else:
|
||||
epilogue_args[key] = tensor
|
||||
elif key in result_keys:
|
||||
tensor = self.example_inputs[key]
|
||||
if isinstance(tensor, Tensor):
|
||||
if "max" in key:
|
||||
fill = -1000
|
||||
else:
|
||||
fill = 0
|
||||
epilogue_args[key] = self.get_torch_tensor(tensor.shape, tensor.element, fill=fill)
|
||||
else:
|
||||
epilogue_args[key] = tensor
|
||||
|
||||
tensor_D = epilogue_args["D"]
|
||||
if "C" in epilogue_args:
|
||||
tensor_C = epilogue_args["C"]
|
||||
else:
|
||||
tensor_C = tensor_D
|
||||
# Run the device kernel
|
||||
self.plan.run(tensor_A, tensor_B, tensor_C, tensor_D, visitor_args=epilogue_args)
|
||||
|
||||
# Run the host reference
|
||||
evt_args_inputs = {}
|
||||
for key in input_keys:
|
||||
evt_args_inputs[key] = epilogue_args[key]
|
||||
|
||||
reference_results = self.reference_fn(
|
||||
tensor_A, tensor_B, tensor_C, problem_size, batch_count, evt_args_inputs)
|
||||
|
||||
# Compare the results
|
||||
for result, ref in zip(result_keys, reference_results):
|
||||
assert torch.equal(epilogue_args[result].flatten(), ref.flatten())
|
||||
|
||||
# Run profile
|
||||
if self.profile:
|
||||
profiler = CUDAEventProfiler(
|
||||
self.plan, 100, 100, tensor_A, tensor_B, tensor_C, tensor_D,
|
||||
visitor_args = epilogue_args
|
||||
)
|
||||
print(f"Cutlass Python Duration: {profiler()}")
|
||||
|
||||
|
||||
class EVTTestCaseBase(unittest.TestCase):
|
||||
"""
|
||||
Base class for EVT Unittest
|
||||
"""
|
||||
def __init__(self, methodName: str = "runTest", lmnk=(6, 512, 256, 128)) -> None:
|
||||
super().__init__(methodName)
|
||||
|
||||
self.element = cutlass.DataType.f16
|
||||
self.l, self.m, self.n, self.k = lmnk
|
||||
|
||||
self.problem_size = (self.m, self.n, self.k)
|
||||
|
||||
torch.random.manual_seed(42)
|
||||
|
||||
def fake_tensor(self, element, shape):
|
||||
return Tensor(element=element, shape=shape, layout_tag=cutlass.LayoutType.RowMajor)
|
||||
|
||||
def get_problem_sizes(self, alignment, k=None, batch_count=[3,]):
|
||||
k = k if k else self.k
|
||||
problem_size_m = [alignment, 512 - 3 * alignment]
|
||||
problem_size_n = [alignment, 512 - alignment]
|
||||
if alignment % 8 == 0:
|
||||
problem_size_m.append(768)
|
||||
problem_size_n.append(768)
|
||||
problem_size_l = batch_count
|
||||
problem_sizes = []
|
||||
for m in problem_size_m:
|
||||
for n in problem_size_n:
|
||||
for l in problem_size_l:
|
||||
problem_sizes.append((m, n, k, l))
|
||||
|
||||
return problem_sizes
|
||||
Reference in New Issue
Block a user