releaase 2.11 (#703)
This commit is contained in:
Aditya Atluri
@@ -0,0 +1,129 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import gen_turing_and_volta as api_generator
|
||||
import gen_sample as sample_creater
|
||||
import gen_cmake as cmake_creater
|
||||
import gen_verify as verify_creater
|
||||
import gen_device as b2b_fused_generator
|
||||
import replace_fix_impl_header
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Generates Fused Multi-GEMM CUTLASS Kernels")
|
||||
parser.add_argument("--config-file", default="config.json", help="JSON file containing configuration to generate")
|
||||
parser.add_argument("--gen-name", default="FusedMultiGemmForward", help="Specific the output name")
|
||||
parser.add_argument("--output-dir", default="", help="Specifies the output dir")
|
||||
parser.add_argument("--cutlass-dir", default="", help="Specifies the dependent CUTLASS repo dir")
|
||||
parser.add_argument("--gen-include-cutlass-dir", default="", help="Specifies the generated CUTLASS code include dir, if needed.")
|
||||
args = parser.parse_args()
|
||||
|
||||
gen_name = args.gen_name
|
||||
|
||||
cutlass_deps_dir = args.cutlass_dir
|
||||
|
||||
output_dir = args.output_dir
|
||||
output_dir += "/"
|
||||
|
||||
cutlass_deps_root = args.gen_include_cutlass_dir
|
||||
if cutlass_deps_root == '':
|
||||
cutlass_deps_root = cutlass_deps_dir + "/include/"
|
||||
cutlass_deps_root +='/'
|
||||
|
||||
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen")
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "fixed_impl"):
|
||||
os.mkdir(output_dir + "/" + "fixed_impl" )
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "sample"):
|
||||
os.mkdir(output_dir + "/" + "sample" )
|
||||
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "device"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "device")
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "kernel"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "kernel")
|
||||
if not os.path.exists(output_dir + "/" + "auto_gen" + "/" + "threadblock"):
|
||||
os.mkdir(output_dir + "/" + "auto_gen" + "/" + "threadblock")
|
||||
|
||||
with open(args.config_file, 'r') as infile:
|
||||
gemm_info_dict = json.load(infile)
|
||||
|
||||
keys = sorted(gemm_info_dict.keys())
|
||||
fuse_gemm_info = [gemm_info_dict[k] for k in keys]
|
||||
|
||||
|
||||
for_cutlass_gen_user_include_header_file = [
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination_leaky_relu.h",
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination.h",
|
||||
]
|
||||
|
||||
for_fused_wrapper = [
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination_leaky_relu.h",
|
||||
cutlass_deps_root + "cutlass/epilogue/thread/linear_combination.h",
|
||||
"auto_gen/device/" + gen_name + ".h",
|
||||
cutlass_deps_root + "cutlass/gemm/device/gemm_batched.h",
|
||||
cutlass_deps_root + "cutlass/cutlass.h",
|
||||
]
|
||||
|
||||
# Copy fixed implementation to the output directory
|
||||
fix_impl = replace_fix_impl_header.replace_fix_impl("../fixed_impl/", output_dir +"/fixed_impl/", cutlass_deps_root)
|
||||
fix_impl.gen_code()
|
||||
|
||||
auto_gen_output_dir = output_dir + "/auto_gen/"
|
||||
project_root = ""
|
||||
turing_plus = b2b_fused_generator.gen_device(fuse_gemm_info, gen_name, for_cutlass_gen_user_include_header_file, cutlass_deps_root, project_root, auto_gen_output_dir)
|
||||
turing_plus.gen_code(75, 'hmma1688', False)
|
||||
|
||||
api = api_generator.gen_one_API(fuse_gemm_info, gen_name, for_fused_wrapper, output_dir)
|
||||
api.gen_code()
|
||||
|
||||
# Generate C++ sample
|
||||
os.system("cp ../leaky_bias.h " + output_dir + "/sample/")
|
||||
os.system("cp ../utils.h " + output_dir + "/sample/")
|
||||
|
||||
sample_dir = output_dir + "/sample/"
|
||||
sample = sample_creater.gen_test(fuse_gemm_info, gen_name, for_cutlass_gen_user_include_header_file, sample_dir)
|
||||
sample.gen_cpp_sample()
|
||||
|
||||
cmake_gen = cmake_creater.gen_build_sys(cutlass_deps_dir, output_dir)
|
||||
cmake_gen.gen_code()
|
||||
|
||||
verify = verify_creater.gen_verify(fuse_gemm_info, gen_name, for_fused_wrapper, output_dir)
|
||||
verify.gen_code()
|
||||
@@ -0,0 +1,131 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
class gen_build_sys:
|
||||
def __init__(self, cutlass_deps_dir, output_dir = "../"):
|
||||
self.output_dir = output_dir
|
||||
self.cutlass_deps_dir = cutlass_deps_dir
|
||||
|
||||
def gen_top(self):
|
||||
code = ""
|
||||
code += '''\
|
||||
# Auto Generated code - Do not edit.
|
||||
|
||||
cmake_minimum_required(VERSION 3.8)
|
||||
project(CUTLASS_MULTI_GEMMS LANGUAGES CXX CUDA)
|
||||
find_package(CUDAToolkit)
|
||||
set(CUDA_PATH ${{CUDA_TOOLKIT_ROOT_DIR}})
|
||||
set(CUTLASS_PATH \"{cutlass_deps_dir}/include\")
|
||||
set(CUTLASS_UTIL_PATH \"{cutlass_deps_dir}/tools/util/include\")
|
||||
list(APPEND CMAKE_MODULE_PATH ${{CUDAToolkit_LIBRARY_DIR}})
|
||||
'''.format(cutlass_deps_dir=self.cutlass_deps_dir)
|
||||
|
||||
code += '''\
|
||||
set(GPU_ARCHS \"\" CACHE STRING
|
||||
\"List of GPU architectures (semicolon-separated) to be compiled for.\")
|
||||
|
||||
if(\"${GPU_ARCHS}\" STREQUAL \"\")
|
||||
set(GPU_ARCHS \"70\")
|
||||
endif()
|
||||
|
||||
foreach(arch ${GPU_ARCHS})
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -gencode arch=compute_${arch},code=sm_${arch}\")
|
||||
if(SM STREQUAL 70 OR SM STREQUAL 75)
|
||||
set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS} -DWMMA\")
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -DWMMA\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -DWMMA\")
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS}\")
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS}\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler -Wall\")
|
||||
|
||||
set(CMAKE_C_FLAGS_DEBUG \"${CMAKE_C_FLAGS_DEBUG} -Wall -O0\")
|
||||
set(CMAKE_CXX_FLAGS_DEBUG \"${CMAKE_CXX_FLAGS_DEBUG} -Wall -O0\")
|
||||
set(CMAKE_CUDA_FLAGS_DEBUG \"${CMAKE_CUDA_FLAGS_DEBUG} -O0 -G -Xcompiler -Wall\")
|
||||
|
||||
set(CMAKE_CXX_STANDARD 11)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
|
||||
if(CMAKE_CXX_STANDARD STREQUAL \"11\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} --expt-extended-lambda\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr\")
|
||||
endif()
|
||||
|
||||
set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -g -O3\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler -O3\")
|
||||
set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -Xcompiler=-fno-strict-aliasing\")
|
||||
|
||||
set(COMMON_HEADER_DIRS
|
||||
${PROJECT_SOURCE_DIR}
|
||||
${CUDAToolkit_INCLUDE_DIRS}
|
||||
)
|
||||
|
||||
set(COMMON_LIB_DIRS
|
||||
${CUDAToolkit_LIBRARY_DIR}
|
||||
)
|
||||
list(APPEND COMMON_HEADER_DIRS ${CUTLASS_PATH})
|
||||
list(APPEND COMMON_HEADER_DIRS ${CUTLASS_UTIL_PATH})
|
||||
'''
|
||||
code += '''\
|
||||
include_directories(
|
||||
${COMMON_HEADER_DIRS}
|
||||
)
|
||||
|
||||
link_directories(
|
||||
${COMMON_LIB_DIRS}
|
||||
)
|
||||
|
||||
add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
|
||||
add_definitions(-DGOOGLE_CUDA=1)
|
||||
|
||||
add_executable(sample
|
||||
sample/sample.cu
|
||||
one_api.cu
|
||||
)
|
||||
target_link_libraries(sample PRIVATE
|
||||
-lcudart
|
||||
-lnvToolsExt
|
||||
${CMAKE_THREAD_LIBS_INIT}
|
||||
)
|
||||
|
||||
if(NOT DEFINED LIB_INSTALL_PATH)
|
||||
set(LIB_INSTALL_PATH ${CMAKE_CURRENT_BINARY_DIR})
|
||||
endif()
|
||||
'''
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
top_code = self.gen_top()
|
||||
with open(self.output_dir + "CMakeLists.txt", "w") as f:
|
||||
f.write(top_code)
|
||||
@@ -0,0 +1,120 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import ast
|
||||
|
||||
fuse_gemm_info = [
|
||||
{
|
||||
'epilogue': {
|
||||
'tp': 'LeakyRelu', #'CustomizedLeaky_RELU'
|
||||
'bias': {'addbias': False, 'bias_tp': 'mat'},
|
||||
'args': [('float', 'leaky_alpha', 1.3), ],
|
||||
'func': '''
|
||||
y = max(leaky_alpha * x, x)
|
||||
y = y * x
|
||||
'''
|
||||
}
|
||||
},
|
||||
|
||||
]
|
||||
class AnalysisNodeVisitor(ast.NodeVisitor):
|
||||
def visit_Import(self,node):
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_ImportFrom(self,node):
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Assign(self,node):
|
||||
print('Node type: Assign and fields: ', node._fields)
|
||||
# print('Node type: Assign and targets value: ', node.targets, node.value)
|
||||
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_BinOp(self, node):
|
||||
print('Node type: BinOp and fields: ', node._fields)
|
||||
print('node op: ', type(node.op).__name__)
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Expr(self, node):
|
||||
print('Node type: Expr and fields: ', node._fields)
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Num(self,node):
|
||||
print('Node type: Num and fields: ', node._fields)
|
||||
print('Node type: Num: ', node.n)
|
||||
|
||||
def visit_Name(self,node):
|
||||
print('Node type: Name and fields: ', node._fields)
|
||||
print('Node type: Name and fields: ', type(node.ctx).__name__, node.id)
|
||||
|
||||
ast.NodeVisitor.generic_visit(self, node)
|
||||
|
||||
def visit_Str(self, node):
|
||||
print('Node type: Str and fields: ', node._fields)
|
||||
|
||||
class CodeVisitor(ast.NodeVisitor):
|
||||
def visit_BinOp(self, node):
|
||||
if isinstance(node.op, ast.Add):
|
||||
node.op = ast.Sub()
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Assign(self, node):
|
||||
print('Assign %s' % node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
print("Name:", node.id)
|
||||
self.generic_visit(node)
|
||||
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
print('Function Name:%s'% node.name.op)
|
||||
self.generic_visit(node)
|
||||
func_log_stmt = ast.Print(
|
||||
dest = None,
|
||||
values = [ast.Str(s = 'calling func: %s' % node.name, lineno = 0, col_offset = 0)],
|
||||
nl = True,
|
||||
lineno = 0,
|
||||
col_offset = 0,
|
||||
)
|
||||
node.body.insert(0, func_log_stmt)
|
||||
|
||||
visitor = AnalysisNodeVisitor()
|
||||
|
||||
code = \
|
||||
'''
|
||||
|
||||
a=max(leaky_alpha * x, x +1)
|
||||
|
||||
'''
|
||||
|
||||
visitor.visit(ast.parse(code))
|
||||
@@ -0,0 +1,477 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
from typing import *
|
||||
|
||||
import helper
|
||||
import gen_ir
|
||||
|
||||
import gen_kernel as gen_ker
|
||||
|
||||
|
||||
class gen_device:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, cutlass_deps_root, project_root, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.raw_gemm_info = fuse_gemm_info
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
self.user_header_file = user_header_file
|
||||
self.args = {}
|
||||
# device arg struct memebr
|
||||
self.arg_member = []
|
||||
self.gen_class_name = gen_class_name
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
self.__tempalate_arg_list = {'Stages': int, 'SplitKSerial': bool, 'IsBetaZero': bool, 'AlignmentA': int, 'AlignmentB': int}
|
||||
|
||||
self.file_name = output_dir + "/device/" +gen_class_name +".h"
|
||||
self.sample_dir = output_dir
|
||||
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
self.this_file_root = output_dir + "/device/"
|
||||
|
||||
self.first_use_1stage = False
|
||||
|
||||
## gen kernel
|
||||
self.gen_kernel = gen_ker.gen_kernel(self.tempalte_args, self.gen_class_name, self.b2b_num, output_dir, cutlass_deps_root, project_root)
|
||||
|
||||
|
||||
def __check_arg_type(self, temp_arg):
|
||||
if temp_arg in self.__tempalate_arg_list.keys():
|
||||
return self.__tempalate_arg_list[temp_arg]
|
||||
|
||||
find_sub = False
|
||||
for candidate_arg in self.__tempalate_arg_list.keys():
|
||||
if (temp_arg.find(candidate_arg) != -1):
|
||||
return self.__tempalate_arg_list[candidate_arg]
|
||||
|
||||
return 'typename'
|
||||
|
||||
# def gen_B2b2bGemm_class():
|
||||
def set_arch(self, sm_cap, mma_tp):
|
||||
if sm_cap == 75 or sm_cap == 80 or sm_cap == 86:
|
||||
self.arch = "cutlass::arch::Sm" + str(sm_cap)
|
||||
|
||||
if mma_tp is 'hmma1688':
|
||||
self.mma_shape = [16, 8, 8]
|
||||
self.mma_tp = 'hmma'
|
||||
elif mma_tp is 'imma8816':
|
||||
self.mma_tp = 'imma'
|
||||
self.mma_shape = [8, 8, 16]
|
||||
else:
|
||||
return 0
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''\
|
||||
/* Auto Generated code - Do not edit.*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_root}cutlass/cutlass.h\"
|
||||
#include \"{cutlass_root}cutlass/numeric_types.h\"
|
||||
#include \"{cutlass_root}cutlass/arch/arch.h\"
|
||||
#include \"{cutlass_root}cutlass/device_kernel.h\"
|
||||
|
||||
#include \"{cutlass_root}cutlass/gemm/threadblock/threadblock_swizzle.h\"
|
||||
|
||||
#include \"{cutlass_root}cutlass/gemm/device/default_gemm_configuration.h\"
|
||||
#include \"{cutlass_root}cutlass/epilogue/thread/linear_combination_relu.h\"
|
||||
#include \"{cutlass_root}cutlass/epilogue/thread/linear_combination.h\"
|
||||
|
||||
#include \"{project_root}../kernel/b2b_gemm.h\"
|
||||
#include \"{project_root}../kernel/default_b2b_gemm.h\"
|
||||
'''.format(cutlass_root=self.cutlass_deps_root, project_root=self.project_root, this_file_root=self.this_file_root)
|
||||
include_user_header = ""
|
||||
for header in self.user_header_file:
|
||||
include_user_header += "#include \"" + header + "\"\n"
|
||||
return code + include_user_header
|
||||
|
||||
def gen_code(self, sm_cap, mma_tp, ifprint = True):
|
||||
self.set_arch(sm_cap, mma_tp)
|
||||
|
||||
self.update_b2b_args()
|
||||
print(self.fuse_gemm_info)
|
||||
self.update_b2b_class_template_args()
|
||||
|
||||
func_code = self.gen_all_func()
|
||||
member_var_code = "private:\n typename B2bGemmKernel::Params params_;\n"
|
||||
|
||||
gen_code = gen_ir.gen_template_class(self.gen_class_name, self.tempalte_args, func_code + member_var_code)
|
||||
code = self.gen_include_header() + gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("device", gen_code)))
|
||||
|
||||
if ifprint:
|
||||
print(code)
|
||||
|
||||
print("[INFO]: Gen device code output Dir: is ", self.file_name)
|
||||
with open(self.file_name, 'w+') as f:
|
||||
f.write(code)
|
||||
|
||||
|
||||
gen_kernel = self.gen_kernel.gen_code(self.first_use_1stage)
|
||||
print(gen_kernel)
|
||||
|
||||
def update_b2b_class_template_args(self):
|
||||
for arg in self.args.keys():
|
||||
self.tempalte_args.append([self.__check_arg_type(arg), arg, self.args[arg]])
|
||||
|
||||
def update_b2b_args(self):
|
||||
|
||||
self.args['ElementA'] = helper.type_2_cutlass_type(self.fuse_gemm_info[0]['A_tp'])
|
||||
self.args['LayoutA'] = helper.type_2_cutlass_type(self.fuse_gemm_info[0]['A_format'])
|
||||
|
||||
cnt = 0
|
||||
|
||||
warp_M_tile = 32
|
||||
|
||||
# Determine maxmimum N_tile
|
||||
Max_Ntile = 0
|
||||
for layer in self.fuse_gemm_info:
|
||||
n_tile = layer['mnk'][1]
|
||||
if n_tile > Max_Ntile:
|
||||
Max_Ntile = n_tile
|
||||
if Max_Ntile >= 256:
|
||||
warp_M_tile = 16
|
||||
|
||||
stages_temp = []
|
||||
|
||||
for layer in self.fuse_gemm_info:
|
||||
cnt_str = str(cnt)
|
||||
B_tp_str= 'ElementB' + cnt_str
|
||||
B_format_str = 'LayoutB' + cnt_str
|
||||
C_tp_str= 'ElementC' + cnt_str
|
||||
C_format_str = 'LayoutC' + cnt_str
|
||||
Acc_str = 'ElementAccumulator' + cnt_str
|
||||
|
||||
self.args[B_tp_str] = helper.type_2_cutlass_type(layer['B_tp'])
|
||||
self.args[B_format_str] = helper.type_2_cutlass_type(layer['B_format'])
|
||||
self.args[C_tp_str] = helper.type_2_cutlass_type(layer['C_tp'])
|
||||
self.args[C_format_str] = helper.type_2_cutlass_type(layer['C_format'])
|
||||
self.args[Acc_str] = helper.type_2_cutlass_type(layer['Acc_tp'])
|
||||
|
||||
|
||||
mnk = layer['mnk'][:]
|
||||
|
||||
tile_mnk = mnk[:]
|
||||
|
||||
tile_mnk[2] = 32 # force the ktile is 32
|
||||
|
||||
#N tile gen
|
||||
if mnk[1] > 1024:
|
||||
assert(0)
|
||||
elif mnk[1] > 512:
|
||||
tile_mnk[1] = 1024
|
||||
elif mnk[1] > 256:
|
||||
tile_mnk[1] = 512
|
||||
elif mnk[1] > 128:
|
||||
tile_mnk[1] = 256
|
||||
elif mnk[1] > 64:
|
||||
tile_mnk[1] = 128
|
||||
elif mnk[1] > 32:
|
||||
tile_mnk[1] = 64
|
||||
else :
|
||||
tile_mnk[1] = 32
|
||||
|
||||
if tile_mnk[1] == 512:
|
||||
stages_temp.append(1)
|
||||
else:
|
||||
stages_temp.append(2)
|
||||
|
||||
tile_mnk[0] = 4 * warp_M_tile
|
||||
|
||||
|
||||
|
||||
epilogue_setted_type = helper.get_epilogue_tp(layer)
|
||||
cutlass_epilogue_name = "LinearCombinationRelu"
|
||||
if epilogue_setted_type.lower() == 'leakyrelu':
|
||||
cutlass_epilogue_name = "LinearCombinationLeakyRelu"
|
||||
elif epilogue_setted_type.lower() == 'identity':
|
||||
cutlass_epilogue_name = "LinearCombination"
|
||||
|
||||
epilogue_str = 'EpilogueOutputOp' + cnt_str
|
||||
if cnt != len(self.fuse_gemm_info) - 1:
|
||||
n = layer['mnk'][1]
|
||||
Fragments = tile_mnk[1] // 8 * 2
|
||||
self.args[epilogue_str] = "cutlass::epilogue::thread::" + cutlass_epilogue_name + "<ElementC0_, " + str(Fragments) +", ElementAccumulator0_, ElementAccumulator0_>"
|
||||
else:
|
||||
n = layer['mnk'][1]
|
||||
n_mod_8 = n % 4
|
||||
N_align_elements = 1
|
||||
if n_mod_8 == 0:
|
||||
N_align_elements = 8
|
||||
elif n_mod_8 == 4:
|
||||
N_align_elements = 4
|
||||
elif n_mod_8 == 2 or n_mod_8 == 6:
|
||||
N_align_elements = 2
|
||||
|
||||
self.args[epilogue_str] = "cutlass::epilogue::thread::" + cutlass_epilogue_name+ "<ElementC0_, " + str(N_align_elements) + ", ElementAccumulator0_, ElementAccumulator0_>"
|
||||
|
||||
|
||||
|
||||
ThreadBlockShape_str = 'ThreadblockShape' + cnt_str
|
||||
|
||||
self.args[ThreadBlockShape_str] = helper.cvt_2_cutlass_shape(tile_mnk)
|
||||
|
||||
WarpShape_str = 'WarpShape' + cnt_str
|
||||
tile_mnk[0] = warp_M_tile
|
||||
self.args[WarpShape_str] = helper.cvt_2_cutlass_shape(tile_mnk)
|
||||
cnt += 1
|
||||
|
||||
|
||||
self.args['ElementD'] = helper.type_2_cutlass_type(self.fuse_gemm_info[self.b2b_num - 1]['C_tp'])
|
||||
self.args['LayoutD'] = helper.type_2_cutlass_type(self.fuse_gemm_info[self.b2b_num - 1]['C_format'])
|
||||
|
||||
self.args['InstructionShape'] = helper.cvt_2_cutlass_shape(self.mma_shape)
|
||||
self.args['OperatorClass'] = 'arch::OpClassTensorOp'
|
||||
self.args['ArchTag'] = self.arch
|
||||
self.args['ThreadblockSwizzle'] = 'threadblock::GemmBatchedIdentityThreadblockSwizzle'
|
||||
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
self.args[helper.var_idx('Stages', i)] = "2"
|
||||
|
||||
self.args['AlignmentA'] = str(8)
|
||||
self.args['AlignmentB'] = str(8)
|
||||
self.args['SplitKSerial'] = 'false'
|
||||
self.args['Operator'] = 'typename DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB0_, ElementC0_, ElementAccumulator0_>::Operator'
|
||||
self.args['IsBetaZero'] = 'false'
|
||||
|
||||
|
||||
def gen_using_kernel(self):
|
||||
code = "using B2bGemmKernel = typename kernel::DefaultB2bGemm<\n"
|
||||
code += " " + "ElementA,\n"
|
||||
code += " " + "LayoutA,\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("ElementB", i) + ",\n"
|
||||
code += " " + helper.var_idx("LayoutB", i) + ",\n"
|
||||
code += " " + helper.var_idx("ElementC", i) + ",\n"
|
||||
code += " " + helper.var_idx("LayoutC", i) + ",\n"
|
||||
code += " " + helper.var_idx("ElementAccumulator", i) + ",\n"
|
||||
code += " " + helper.var_idx("EpilogueOutputOp", i) + ",\n"
|
||||
code += " " + helper.var_idx("ThreadblockShape", i) + ",\n"
|
||||
code += " " + helper.var_idx("WarpShape", i) + ",\n"
|
||||
|
||||
code += " " + "ElementD,\n"
|
||||
code += " " + "LayoutD,\n"
|
||||
code += " " + "InstructionShape,\n"
|
||||
code += " " + "OperatorClass,\n"
|
||||
code += " " + "ArchTag,\n"
|
||||
code += " " + "ThreadblockSwizzle,\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("Stages", i) + ",\n"
|
||||
|
||||
|
||||
code += " " + "AlignmentA,\n"
|
||||
code += " " + "AlignmentB,\n"
|
||||
code += " " + "SplitKSerial,\n"
|
||||
code += " " + "Operator,\n"
|
||||
code += " " + "IsBetaZero_\n"
|
||||
|
||||
code += ">::B2bGemmKernel;\n\n"
|
||||
|
||||
return code
|
||||
|
||||
def gen_args(self):
|
||||
|
||||
def gen_arg_member(b2b_num):
|
||||
data_members = []
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "GemmCoord"
|
||||
member_name = "problem_size_" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
member_type = "TensorRef<ElementA const, LayoutA>"
|
||||
member_name = "ref_A0"
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "TensorRef<ElementB" + str(i) + " const, LayoutB" + str(i) +">"
|
||||
member_name = "ref_B" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
member_type = "TensorRef<ElementC" + str(i) + " const, LayoutC" + str(i) +">"
|
||||
member_name = "ref_C" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
member_type = "TensorRef<ElementD, LayoutD>"
|
||||
member_name = helper.var_idx("ref_D", b2b_num - 1)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
for i in range(b2b_num):
|
||||
member_type = "typename EpilogueOutputOp" + str(i) + "::Params"
|
||||
member_name = "epilogue" + str(i)
|
||||
data_members.append((member_type, member_name))
|
||||
|
||||
data_members.append(('int', 'batch_count'))
|
||||
|
||||
return data_members
|
||||
|
||||
def gen_arg_struct_default_ctor(struct_name, data_members, inital_param_num, inital_value):
|
||||
constructs_code = gen_ir.indentation + "CUTLASS_HOST_DEVICE\n" + \
|
||||
gen_ir.indentation + struct_name + " (): "
|
||||
for i in range(inital_param_num):
|
||||
final_param = ','
|
||||
if i == inital_param_num - 1:
|
||||
final_param = '{ }'
|
||||
constructs_code += data_members[i][1] + inital_value + final_param
|
||||
|
||||
constructs_code += "\n"
|
||||
return constructs_code
|
||||
|
||||
def gen_arg_struct_ctor(struct_name, data_members):
|
||||
constructs_code = gen_ir.indentation + "CUTLASS_HOST_DEVICE\n" + \
|
||||
gen_ir.indentation + struct_name + " (\n"
|
||||
cnt = 0
|
||||
param_num = len(data_members)
|
||||
for param in data_members:
|
||||
final = ',\n'
|
||||
if cnt == param_num - 1:
|
||||
final = '\n):\n'
|
||||
constructs_code += gen_ir.indentation + param[0] + " " + param[1] + "_" + final
|
||||
cnt += 1
|
||||
|
||||
cnt = 0
|
||||
for param in data_members:
|
||||
final = '),\n'
|
||||
if cnt == param_num - 1:
|
||||
final = ") { }\n"
|
||||
constructs_code += gen_ir.indentation + param[1] + "(" + param[1] + "_" + final
|
||||
cnt += 1
|
||||
|
||||
constructs_code += "\n"
|
||||
return constructs_code
|
||||
|
||||
# (variable type, variable name)
|
||||
struct_member = gen_arg_member(self.b2b_num)
|
||||
self.arg_member = struct_member
|
||||
|
||||
codeBody = ""
|
||||
for each_member in struct_member:
|
||||
codeBody += gen_ir.indentation + each_member[0] + " " + each_member[1] + ";\n"
|
||||
|
||||
codeBody += gen_arg_struct_default_ctor("Arguments", struct_member, self.b2b_num, "(0,0,0)") + "\n"
|
||||
codeBody += gen_arg_struct_ctor("Arguments", struct_member) + "\n"
|
||||
struct_code = gen_ir.gen_struct("Arguments", codeBody)
|
||||
return struct_code
|
||||
|
||||
def gen_func_constructs(self):
|
||||
code = self.gen_class_name +"() {}"
|
||||
return code
|
||||
|
||||
def gen_func_initialize(self):
|
||||
code = "Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {\n" + \
|
||||
"// Determine grid shape\n" + \
|
||||
"ThreadblockSwizzle threadblock_swizzle;\n" + \
|
||||
"cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(\n" + \
|
||||
" args.problem_size_0, \n" + \
|
||||
" { ThreadblockShape0::kM, ThreadblockShape0::kN, ThreadblockShape0::kK },\n" + \
|
||||
" args.batch_count);\n" + \
|
||||
"// Initialize the Params structure\n" + \
|
||||
"params_ = typename B2bGemmKernel::Params{\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.problem_size_", i) + ",\n"
|
||||
code += " grid_shape,\n" + \
|
||||
" args.ref_A0.non_const_ref(),\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.ref_B", i) + ".non_const_ref(),\n"
|
||||
code += helper.var_idx(" args.ref_C", i) + ".non_const_ref(),\n"
|
||||
|
||||
code += helper.var_idx(" args.ref_D", self.b2b_num - 1) + ",\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx(" args.epilogue", i) + ",\n"
|
||||
|
||||
code += " args.batch_count\n"
|
||||
code += "};\n" + \
|
||||
"return Status::kSuccess;\n" + \
|
||||
"}\n"
|
||||
return code
|
||||
|
||||
def gen_func_run(self):
|
||||
code = "Status run(cudaStream_t stream = nullptr) {\n" + \
|
||||
"\n" + \
|
||||
" ThreadblockSwizzle threadblock_swizzle;\n" + \
|
||||
"\n" + \
|
||||
" dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);\n" + \
|
||||
" dim3 block(B2bGemmKernel::kThreadCount, 1, 1);\n" + \
|
||||
"\n" + \
|
||||
" cudaError_t result;\n" + \
|
||||
"\n" + \
|
||||
" int smem_size = int(sizeof(typename B2bGemmKernel::SharedStorage));\n" + \
|
||||
" if (smem_size >= (48 << 10)) {\n" + \
|
||||
" result = cudaFuncSetAttribute(Kernel<B2bGemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);\n" + \
|
||||
"\n" + \
|
||||
" if (result != cudaSuccess) {\n" + \
|
||||
" return Status::kErrorInternal;\n" + \
|
||||
" }\n" + \
|
||||
"\n" + \
|
||||
" result = cudaFuncSetAttribute(\n" + \
|
||||
" Kernel<B2bGemmKernel>,\n" + \
|
||||
" cudaFuncAttributePreferredSharedMemoryCarveout, 100);\n" + \
|
||||
"\n" + \
|
||||
" if (result != cudaSuccess) {\n" + \
|
||||
" return Status::kErrorInternal;\n" + \
|
||||
" }\n" + \
|
||||
" }\n" + \
|
||||
" cutlass::Kernel<B2bGemmKernel><<<grid, block, smem_size, stream>>>(params_);\n" + \
|
||||
" result = cudaGetLastError();\n" + \
|
||||
" return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;\n" + \
|
||||
" }\n"
|
||||
|
||||
return code
|
||||
def gen_func_operator(self):
|
||||
opeartor_with_arg_code = "Status operator()(\n" + \
|
||||
" Arguments const &args,\n" + \
|
||||
" void *workspace = nullptr,\n" + \
|
||||
" cudaStream_t stream = nullptr) {\n" + \
|
||||
" Status status = initialize(args, workspace);\n" + \
|
||||
" \n" + \
|
||||
" if (status == Status::kSuccess) {\n" + \
|
||||
" status = run(stream);\n" + \
|
||||
" }\n" + \
|
||||
" return status;\n" + \
|
||||
"}\n"
|
||||
operator_code = "Status operator()(\n" + \
|
||||
" cudaStream_t stream = nullptr) {\n" + \
|
||||
" Status status = run(stream);\n" + \
|
||||
" return status;\n" + \
|
||||
"}\n"
|
||||
return opeartor_with_arg_code + "\n" + operator_code
|
||||
|
||||
def gen_all_func(self):
|
||||
return self.gen_using_kernel() + "\n" + \
|
||||
self.gen_args() + "\n" + \
|
||||
self.gen_func_constructs() + "\n" + \
|
||||
self.gen_func_initialize() + "\n" + \
|
||||
self.gen_func_run() + "\n" + \
|
||||
self.gen_func_operator()
|
||||
@@ -0,0 +1,249 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
|
||||
|
||||
indentation = " "
|
||||
|
||||
|
||||
def append_word(word):
|
||||
code = ""
|
||||
code += word
|
||||
code += " "
|
||||
return code
|
||||
|
||||
|
||||
def gen_namespace(namespace, codeBody):
|
||||
code_gen = "namespace " + namespace + " {\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "} // namespace " + namespace + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_expression(type, lval, rval = None):
|
||||
code_gen = ""
|
||||
code_gen += append_word(type)
|
||||
code_gen += append_word(lval)
|
||||
if rval is not None:
|
||||
code_gen += append_word("=")
|
||||
code_gen += append_word(rval)
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_class(name, codeBody, inheritance_code = None):
|
||||
code_gen = ""
|
||||
if inheritance_code is None:
|
||||
code_gen = "class " + name + "{\n"
|
||||
else:
|
||||
code_gen = "class " + name + " : "+ inheritance_code + "{\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "}; // class " + name + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_struct(name, codeBody, specialized = None):
|
||||
specialized_code = ""
|
||||
if specialized is not None:
|
||||
specialized_code = "<" + specialized + ">"
|
||||
code_gen = "struct " + name + specialized_code + "{\n"
|
||||
code_gen += codeBody
|
||||
code_gen += "}; // struct " + name + "\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_arg(arg_type, arg_name, default_val = None):
|
||||
rval = None
|
||||
if default_val is not None:
|
||||
rval = str(default_val)
|
||||
|
||||
arg_typename = ""
|
||||
if arg_type is int:
|
||||
arg_typename = "int"
|
||||
elif arg_type is bool:
|
||||
arg_typename = "bool"
|
||||
else:
|
||||
arg_typename = "typename"
|
||||
|
||||
internal_arg_name = arg_name + "_"
|
||||
|
||||
code_gen = indentation
|
||||
code_gen += gen_expression(arg_typename, internal_arg_name, rval)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_args(args, set_default = True):
|
||||
arg_len = len(args)
|
||||
cnt = 1
|
||||
code_gen = ""
|
||||
for arg_tuple in args:
|
||||
arg_type = arg_tuple[0]
|
||||
arg_name = arg_tuple[1]
|
||||
arg_default_val = None
|
||||
if len(arg_tuple) == 3 and set_default:
|
||||
arg_default_val = arg_tuple[2]
|
||||
|
||||
code_gen += gen_template_arg(arg_type, arg_name, arg_default_val)
|
||||
if cnt != arg_len:
|
||||
code_gen += ",\n"
|
||||
cnt += 1
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_head(args, set_default = True):
|
||||
code_gen = "template <\n"
|
||||
code_gen += gen_template_args(args, set_default)
|
||||
code_gen += ">\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def export_template_args(args):
|
||||
code_gen = "public:\n"
|
||||
for arg_tuple in args:
|
||||
code_gen += indentation
|
||||
arg_type = arg_tuple[0]
|
||||
arg_name = arg_tuple[1]
|
||||
internal_arg_name = arg_name + "_"
|
||||
|
||||
typename = ""
|
||||
if arg_type is int:
|
||||
typename = "static int const"
|
||||
elif arg_type is bool:
|
||||
typename = "static bool const"
|
||||
else:
|
||||
typename = "using"
|
||||
|
||||
code_gen += gen_expression(typename, arg_name, internal_arg_name)
|
||||
code_gen += ";\n"
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_class(class_name, args, codeBody, set_default = True, inheritance_code = None):
|
||||
code_gen = ""
|
||||
|
||||
code_gen += gen_template_head(args, set_default)
|
||||
code_gen += gen_class(class_name, export_template_args(args) + codeBody, inheritance_code)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_template_struct(struct_name, args, codeBody, speicalized = None, set_default = True, export_args = True):
|
||||
code_gen = ""
|
||||
code_gen += gen_template_head(args, set_default)
|
||||
code = export_template_args(args) + codeBody
|
||||
if export_args is False:
|
||||
code = codeBody
|
||||
code_gen += gen_struct(struct_name, code , speicalized)
|
||||
|
||||
return code_gen
|
||||
|
||||
|
||||
def gen_declare_template_struct(name, *params):
|
||||
code = name + "<"
|
||||
cnt = 0
|
||||
param_num = len(params)
|
||||
for param in params:
|
||||
final = ", "
|
||||
if cnt == param_num - 1:
|
||||
final = ""
|
||||
code += param + final
|
||||
cnt += 1
|
||||
code += ">;\n"
|
||||
return code
|
||||
|
||||
|
||||
def filtered_param(params, name_and_value_pair, keep_ = False):
|
||||
rtn_template_args = []
|
||||
speicalized_template_args = []
|
||||
|
||||
for param in params:
|
||||
param_name = ""
|
||||
if len(param) >= 1:
|
||||
param_name = param[1]
|
||||
else:
|
||||
param_name = param[0]
|
||||
|
||||
hit_flag = False
|
||||
set_value = ""
|
||||
for n_v_pair in name_and_value_pair:
|
||||
|
||||
filter_name = n_v_pair[0]
|
||||
set_value = n_v_pair[1]
|
||||
|
||||
if param_name == (filter_name + "_") or param_name == filter_name :
|
||||
hit_flag = True
|
||||
break
|
||||
|
||||
|
||||
if hit_flag is False:
|
||||
rtn_template_args.append(param)
|
||||
|
||||
if hit_flag is True:
|
||||
speicalized_template_args.append(set_value)
|
||||
else:
|
||||
if keep_ is True:
|
||||
speicalized_template_args.append(param_name + "_")
|
||||
else:
|
||||
speicalized_template_args.append(param_name)
|
||||
|
||||
|
||||
specialized_template_arg_str = helper.list_2_string(speicalized_template_args)
|
||||
|
||||
return rtn_template_args, specialized_template_arg_str
|
||||
|
||||
|
||||
def gen_func(func_name, arg_lists, code_body, only_declare = False, with_cudaStream = True):
|
||||
code = "void " + func_name + "(\n"
|
||||
for arg in arg_lists:
|
||||
arg_tp = arg[0]
|
||||
arg_nm = arg[1]
|
||||
code += " " + arg_tp + " " + arg_nm + ",\n"
|
||||
code += "cudaStream_t stream)"
|
||||
if only_declare :
|
||||
return code
|
||||
code += "{\n"
|
||||
|
||||
code += code_body + "\n"
|
||||
code += "}\n"
|
||||
return code
|
||||
|
||||
|
||||
def indent_level(code, level = 0):
|
||||
rtn_code = ""
|
||||
for i in range(level):
|
||||
rtn_code += " "
|
||||
|
||||
rtn_code += code
|
||||
|
||||
return rtn_code
|
||||
@@ -0,0 +1,476 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import gen_ir
|
||||
import helper
|
||||
import gen_threadblock as gen_tb
|
||||
|
||||
|
||||
class gen_default_Gemm:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root):
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.template_param = template_param
|
||||
self.b2b_num = b2b_num
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
def gen_B2bMma(self, specialized_template_args):
|
||||
code = "using B2bMma = typename cutlass::gemm::threadblock::DefaultB2bMma<\n"
|
||||
code += specialized_template_args
|
||||
code += ">::ThreadblockB2bMma;\n"
|
||||
|
||||
# print(code)
|
||||
return code
|
||||
|
||||
def gen_epilogue(self):
|
||||
epilogue_code = ""
|
||||
epilogue_code += helper.var_idx("static const int kPartitionsK", self.b2b_num - 1) + helper.var_idx(" = ThreadblockShape", self.b2b_num - 1) + helper.var_idx("::kK / WarpShape", self.b2b_num - 1) + "::kK;\n"
|
||||
|
||||
epilogue_code += "using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<\n"
|
||||
epilogue_code += " " + helper.var_idx("ThreadblockShape", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("typename B2bMma::Operator", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("kPartitionsK", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("EpilogueOutputOp", self.b2b_num - 1) + ",\n"
|
||||
epilogue_code += " " + helper.var_idx("EpilogueOutputOp", self.b2b_num - 1) + "::kCount\n"
|
||||
epilogue_code += ">::Epilogue;\n"
|
||||
|
||||
epilogue_code += "using B2bGemmKernel = kernel::B2bGemm<B2bMma, Epilogue, ThreadblockSwizzle, SplitKSerial>;\n\n"
|
||||
|
||||
return epilogue_code
|
||||
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
/* Auto Generated code - Do not edit.*/
|
||||
|
||||
#pragma once
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/layout/matrix.h\"
|
||||
#include \"{cutlass_dir}cutlass/numeric_types.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/epilogue.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/thread/linear_combination.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/kernel/gemm_pipelined.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm75.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm70.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_sm80.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/default_mma_core_simt.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/threadblock/threadblock_swizzle.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_tensor_op.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h\"
|
||||
#include \"{cutlass_dir}cutlass/epilogue/threadblock/default_epilogue_simt.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/transform/threadblock/predicated_tile_iterator.h\"
|
||||
|
||||
#include \"../kernel/b2b_gemm.h\"
|
||||
#include \"../threadblock/default_b2b_mma.h\"
|
||||
'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
gen_using = ''
|
||||
# Generate default template struct
|
||||
gen_code = gen_ir.gen_template_struct("Default" + self.gen_class_name, self.template_param,"", speicalized = None, set_default=False)
|
||||
|
||||
|
||||
filter_list = []
|
||||
filter_list.append(('Stages', 2))
|
||||
filter_list.append(("OperatorClass", "arch::OpClassTensorOp"))
|
||||
filter_list.append(("ArchTag", "arch::Sm75"))
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
filter_list.append((helper.var_idx("LayoutC", i), "layout::RowMajor"))
|
||||
|
||||
|
||||
rtn_template_args, speicalized_template_args = gen_ir.filtered_param(self.template_param, filter_list, keep_= True)
|
||||
|
||||
|
||||
B2bMma_code = self.gen_B2bMma(speicalized_template_args)
|
||||
epilogue_and_rest_code = self.gen_epilogue()
|
||||
|
||||
gen_special_code = gen_ir.gen_template_struct("Default" + self.gen_class_name, rtn_template_args, B2bMma_code + epilogue_and_rest_code, speicalized = speicalized_template_args, set_default=False)
|
||||
|
||||
code = gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("kernel", gen_code + gen_special_code)))
|
||||
|
||||
return self.gen_include_header() + code
|
||||
|
||||
|
||||
class gen_Kernel:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root):
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.template_param = template_param
|
||||
self.b2bnum = b2b_num
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/matrix_coord.h\"\n'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
|
||||
def gen_Params(self):
|
||||
gen_param = ""
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + ";\n"
|
||||
gen_param += " " + "cutlass::gemm::GemmCoord grid_tiled_shape;\n"
|
||||
gen_param += " " + "typename B2bMma::IteratorA0::Params params_A0;\n"
|
||||
gen_param += " " + "typename B2bMma::IteratorA0::TensorRef ref_A0;\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::Params params_B", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::TensorRef ref_B", i) + ";\n"
|
||||
if i == self.b2bnum - 1:
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::Params params_C", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_C", i) + ";\n"
|
||||
|
||||
else:
|
||||
gen_param += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::Params params_C", i) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::TensorRef ref_C", i) + ";\n"
|
||||
|
||||
|
||||
|
||||
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::Params params_D", self.b2bnum - 1) + ";\n"
|
||||
gen_param += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_D", self.b2bnum - 1) + ";\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
gen_param += " " + helper.var_idx("typename OutputOp", i) + helper.var_idx("::Params output_op_", i) + ";\n"
|
||||
|
||||
gen_param += " " + 'int batch_count' + ";\n"
|
||||
gen_param += " " + 'int gemm_k_iterations_0' + ";\n"
|
||||
|
||||
|
||||
return gen_param
|
||||
|
||||
def gen_Memberfunc(self):
|
||||
code_default = "\nCUTLASS_HOST_DEVICE\n"
|
||||
code_default += "Params()"
|
||||
|
||||
code_default += " { } \n\n"
|
||||
|
||||
code_construct = "\nCUTLASS_HOST_DEVICE\n"
|
||||
code_construct += "Params(\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("cutlass::gemm::GemmCoord const & problem_size_", i) + ",\n"
|
||||
|
||||
code_construct += " " + "cutlass::gemm::GemmCoord const & grid_tiled_shape,\n"
|
||||
|
||||
code_construct += " " + "typename B2bMma::IteratorA0::TensorRef ref_A0,\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("typename B2bMma::IteratorB", i) + helper.var_idx("::TensorRef ref_B", i) + ",\n"
|
||||
if i == self.b2bnum - 1:
|
||||
code_construct += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_C", i) + ",\n"
|
||||
else:
|
||||
code_construct += " " + helper.var_idx("typename FusedAddBiasEpilogue", i) + helper.var_idx("::OutputTileIterator::TensorRef ref_C", i) + ",\n"
|
||||
|
||||
code_construct += " " + helper.var_idx("typename Epilogue::OutputTileIterator::TensorRef ref_D", self.b2bnum - 1) + ",\n"
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("typename OutputOp", i) + helper.var_idx("::Params output_op_", i) + helper.var_idx(" = typename OutputOp", i) + "::Params(),\n"
|
||||
|
||||
code_construct += " " + "int batch_count = 1\n"
|
||||
|
||||
code_construct += "):\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("problem_size_", i) + helper.var_idx("(problem_size_", i) + "),\n"
|
||||
|
||||
code_construct += " " + "grid_tiled_shape(grid_tiled_shape),\n"
|
||||
code_construct += " " + "params_A0(ref_A0.layout()),\n"
|
||||
code_construct += " " + "ref_A0(ref_A0),\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("params_B", i) + helper.var_idx("(ref_B", i) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_B", i) + helper.var_idx("(ref_B", i) + "),\n"
|
||||
code_construct += " " + helper.var_idx("params_C", i) + helper.var_idx("(ref_C", i) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_C", i) + helper.var_idx("(ref_C", i) + "),\n"
|
||||
|
||||
code_construct += " " + helper.var_idx("params_D", self.b2bnum - 1) + helper.var_idx("(ref_D", self.b2bnum - 1) + ".layout()),\n"
|
||||
code_construct += " " + helper.var_idx("ref_D", self.b2bnum - 1) + helper.var_idx("(ref_D", self.b2bnum - 1) + "),\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
code_construct += " " + helper.var_idx("output_op_", i) + helper.var_idx("(output_op_", i) + "), \n"
|
||||
|
||||
code_construct += " " + "batch_count(batch_count) {\n"
|
||||
code_construct += " " + helper.var_idx("gemm_k_iterations_", 0) + helper.var_idx(" = (problem_size_", 0) + helper.var_idx(".k() + B2bMma::Shape", 0) + helper.var_idx("::kK - 1) / B2bMma::Shape", 0) + "::kK;\n"
|
||||
|
||||
code_construct += "}\n"
|
||||
|
||||
return code_default + code_construct
|
||||
|
||||
def gen_using(self):
|
||||
code_using = ""
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
code_using += " " + helper.var_idx("using OutputOp", i) + helper.var_idx(" = typename B2bMma::OutputOp", i) + ";\n"
|
||||
|
||||
code_using += " " + helper.var_idx("using OutputOp", self.b2bnum - 1) + " = typename Epilogue::OutputOp;\n"
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
code_using += " " + helper.var_idx("using FusedAddBiasEpilogue", i) + helper.var_idx(" = typename B2bMma::FusedAddBiasEpilogue", i) +";\n"
|
||||
|
||||
|
||||
code_using += " " + "using WarpCount0 = typename B2bMma::WarpCount0;\n"
|
||||
code_using += " " + "static int const kThreadCount = 32 * WarpCount0::kCount;\n"
|
||||
|
||||
code_using += gen_ir.gen_struct("Params", self.gen_Params() + self.gen_Memberfunc())
|
||||
|
||||
code_using += "union SharedStorage {\n"
|
||||
code_using += " " + "typename B2bMma::B2bMmaSharedStorage main_loop;\n"
|
||||
code_using += " " + "typename Epilogue::SharedStorage epilogue;\n"
|
||||
code_using += "};\n"
|
||||
|
||||
return code_using
|
||||
|
||||
def gen_can_implement(self):
|
||||
gen_code = ""
|
||||
return gen_code
|
||||
|
||||
def gen_operator_and_constr(self):
|
||||
ctr_code = "CUTLASS_HOST_DEVICE\n"
|
||||
ctr_code += self.gen_class_name + "() { } \n\n"
|
||||
operator_code = "CUTLASS_DEVICE\n"
|
||||
operator_code += "void operator()(Params const ¶ms, SharedStorage &shared_storage) {\n"
|
||||
operator_code += " " + "ThreadblockSwizzle threadblock_swizzle;\n"
|
||||
operator_code += " " + "cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);\n"
|
||||
operator_code += " " + "int batch_idx = threadblock_tile_offset.k();\n"
|
||||
operator_code += " " + "if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||\n"
|
||||
operator_code += " " + "params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {\n"
|
||||
operator_code += " " + " " + "return;\n"
|
||||
operator_code += " " + "}\n"
|
||||
|
||||
operator_code += " " + "cutlass::MatrixCoord tb_offset_A0{\n"
|
||||
operator_code += " " + " " + "threadblock_tile_offset.m() * B2bMma::Shape0::kM,\n"
|
||||
operator_code += " " + " " + "0\n"
|
||||
operator_code += " " + "};\n"
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
operator_code += " " + helper.var_idx("cutlass::MatrixCoord tb_offset_B", i) + "{\n"
|
||||
operator_code += " " + " " + "0,\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.n() * B2bMma::Shape", i) + "::kN\n"
|
||||
operator_code += " " + "};\n"
|
||||
|
||||
operator_code += " " + "int thread_idx = threadIdx.x;\n\n"
|
||||
|
||||
operator_code += " " + "MatrixCoord threadblock_offset(\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.m() * B2bMma::Shape", self.b2bnum - 1) + "::kM,\n"
|
||||
operator_code += " " + " " + helper.var_idx("threadblock_tile_offset.n() * B2bMma::Shape", self.b2bnum - 1) + "::kN\n"
|
||||
operator_code += " " + ");\n"
|
||||
|
||||
operator_code += " " + "typename B2bMma::IteratorA0 iterator_A0(\n"
|
||||
operator_code += " " + " " + "params.params_A0,\n"
|
||||
operator_code += " " + " " + "params.ref_A0.data(),\n"
|
||||
operator_code += " " + " " + "params.problem_size_0.mk(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "tb_offset_A0);\n"
|
||||
|
||||
operator_code += " " + "iterator_A0.add_pointer_offset(batch_idx * params.problem_size_0.m() * params.problem_size_0.k());\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum):
|
||||
operator_code += " " + helper.var_idx("typename B2bMma::IteratorB", i ) + helper.var_idx(" iterator_B", i) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_B", i) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_B", i) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", i) + ".kn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + helper.var_idx("tb_offset_B", i) + ");\n"
|
||||
operator_code += " " + helper.var_idx("iterator_B", i) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", i) + helper.var_idx(".n() * params.problem_size_", i) + ".k());\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("typename FusedAddBiasEpilogue", i ) + helper.var_idx("::OutputTileIterator iterator_C", i) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_C", i) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_C", i) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_" , i) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset" + ");\n"
|
||||
operator_code += " " + helper.var_idx("int ref_C", i) + helper.var_idx("_stride = params.ref_C", i) + ".stride()[0];\n"
|
||||
operator_code += " " + helper.var_idx("iterator_C", i) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", i) + helper.var_idx(".n() * (ref_C", i) + helper.var_idx("_stride == 0 ? 1 : params.problem_size_", i) + ".m()));\n\n"
|
||||
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("FusedAddBiasEpilogue", i ) + helper.var_idx(" epilogue_", i ) + ";\n"
|
||||
|
||||
|
||||
operator_code += " " + "int warp_idx = __shfl_sync(0x1f, threadIdx.x / 32, 0);\n"
|
||||
operator_code += " " + "int lane_idx = threadIdx.x % 32;\n"
|
||||
|
||||
for i in range (self.b2bnum - 1):
|
||||
operator_code += " " + helper.var_idx("OutputOp", i) + helper.var_idx(" output_op_", i) + helper.var_idx("(params.output_op_", i) + ");\n"
|
||||
|
||||
operator_code += " " + "B2bMma b2bMma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);\n"
|
||||
|
||||
operator_code += " " + "typename B2bMma::FragmentC0 src_accum;\n"
|
||||
operator_code += " " + helper.var_idx("typename B2bMma::FragmentC", self.b2bnum - 1)+ " accumulators;\n"
|
||||
|
||||
operator_code += " " + "src_accum.clear();\n"
|
||||
operator_code += " " + "accumulators.clear();\n"
|
||||
operator_code += " " + "b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, "
|
||||
|
||||
for i in range(self.b2bnum):
|
||||
operator_code += helper.var_idx("iterator_B", i) + ", "
|
||||
|
||||
operator_code += "src_accum"
|
||||
if self.b2bnum != 1:
|
||||
operator_code += ", "
|
||||
for i in range(self.b2bnum - 1):
|
||||
operator_code += helper.var_idx("output_op_", i) + ", "
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
operator_code += helper.var_idx("epilogue_", i) + ", "
|
||||
|
||||
for i in range(self.b2bnum - 1):
|
||||
final = ", "
|
||||
if i == self.b2bnum - 2:
|
||||
final =""
|
||||
operator_code += helper.var_idx("iterator_C", i) + final
|
||||
operator_code += ");\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("OutputOp", self.b2bnum - 1) + helper.var_idx(" output_op_", self.b2bnum - 1) + helper.var_idx("(params.output_op_", self.b2bnum - 1) + ");\n"
|
||||
operator_code += " " + "threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);\n"
|
||||
|
||||
|
||||
|
||||
operator_code += " " + helper.var_idx("typename Epilogue::OutputTileIterator iterator_C", self.b2bnum - 1) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_C", self.b2bnum - 1) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_C", self.b2bnum - 1) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", self.b2bnum - 1) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset\n"
|
||||
operator_code += " " + ");\n"
|
||||
operator_code += " " + helper.var_idx("int ref_C", self.b2bnum - 1) + helper.var_idx("_stride = params.ref_C", self.b2bnum - 1) + ".stride()[0];\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("iterator_C", self.b2bnum - 1) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", self.b2bnum - 1) + helper.var_idx(".n() * (ref_C", self.b2bnum - 1) + helper.var_idx("_stride == 0 ? 1 : params.problem_size_", self.b2bnum - 1) + ".m()));\n\n"
|
||||
|
||||
operator_code += " " + helper.var_idx("typename Epilogue::OutputTileIterator iterator_D", self.b2bnum - 1) + "(\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.params_D", self.b2bnum - 1) + ",\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.ref_D", self.b2bnum - 1) + ".data(),\n"
|
||||
operator_code += " " + " " + helper.var_idx("params.problem_size_", self.b2bnum - 1) + ".mn(),\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "threadblock_offset\n"
|
||||
operator_code += " " + ");\n"
|
||||
operator_code += " " + helper.var_idx("iterator_D", self.b2bnum - 1) + helper.var_idx(".add_pointer_offset(batch_idx * params.problem_size_", self.b2bnum - 1) + helper.var_idx(".n() * params.problem_size_", self.b2bnum - 1) + ".m());\n\n"
|
||||
|
||||
|
||||
operator_code += " " + "Epilogue epilogue(\n"
|
||||
operator_code += " " + " " + "shared_storage.epilogue,\n"
|
||||
operator_code += " " + " " + "thread_idx,\n"
|
||||
operator_code += " " + " " + "warp_idx,\n"
|
||||
operator_code += " " + " " + "lane_idx\n"
|
||||
operator_code += " " + ");\n"
|
||||
|
||||
operator_code += " " + "epilogue("
|
||||
operator_code += helper.var_idx("output_op_", self.b2bnum - 1) + ", "
|
||||
operator_code += helper.var_idx("iterator_D", self.b2bnum - 1) + ", "
|
||||
operator_code += "accumulators, "
|
||||
operator_code += helper.var_idx("iterator_C", self.b2bnum - 1) + ");\n"
|
||||
operator_code += "}\n"
|
||||
|
||||
return ctr_code + operator_code
|
||||
|
||||
def gen_include_header(self):
|
||||
code = '''
|
||||
#pragma once
|
||||
|
||||
#include \"{cutlass_dir}cutlass/cutlass.h\"
|
||||
|
||||
#include \"{cutlass_dir}cutlass/gemm/gemm.h\"
|
||||
#include \"{cutlass_dir}cutlass/matrix_coord.h\"
|
||||
#include \"{cutlass_dir}cutlass/semaphore.h\"
|
||||
'''.format(cutlass_dir=self.cutlass_deps_root)
|
||||
return code
|
||||
def gen_code(self):
|
||||
|
||||
template_param = []
|
||||
template_param.append(("typename", "B2bMma"))
|
||||
template_param.append(("typename", "Epilogue"))
|
||||
template_param.append(("typename", "ThreadblockSwizzle"))
|
||||
template_param.append((bool, "SplitKSerial"))
|
||||
|
||||
code_body = ""
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_operator_and_constr()
|
||||
|
||||
struct_code = gen_ir.gen_template_struct(self.gen_class_name, template_param, code_body)
|
||||
code = self.gen_include_header()
|
||||
code += gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("kernel", struct_code)))
|
||||
|
||||
return self.gen_include_header() + code
|
||||
|
||||
|
||||
|
||||
class gen_kernel:
|
||||
def __init__(self, template_param, gen_class_name, b2b_num, output_dir, cutlass_deps_root, project_root):
|
||||
self.template_param = template_param
|
||||
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
self.gen_default_b2b_gemm = gen_default_Gemm(template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root)
|
||||
self.gen_Kerenl = gen_Kernel(template_param, gen_class_name, b2b_num, cutlass_deps_root, project_root)
|
||||
|
||||
# Include gen_threadBlock
|
||||
self.gen_threadBlock = gen_tb.gen_threadblock(template_param, gen_class_name, b2b_num, output_dir, cutlass_deps_root, project_root)
|
||||
|
||||
self.file_dir = output_dir + "/kernel/"
|
||||
|
||||
def gen_code(self, first_use_1stage):
|
||||
|
||||
default_b2b_gemm = self.gen_default_b2b_gemm.gen_code()
|
||||
|
||||
print("[INFO]: Gen kernel code [default_b2b_gemm.h]output Dir: is ", self.file_dir)
|
||||
|
||||
with open(self.file_dir + "default_b2b_gemm.h", "w+") as f:
|
||||
f.write(default_b2b_gemm)
|
||||
|
||||
kernel = self.gen_Kerenl.gen_code()
|
||||
print("[INFO]: Gen kernel code [b2b_gemm.h]output Dir: is ", self.file_dir)
|
||||
|
||||
with open(self.file_dir + "b2b_gemm.h", "w+") as f:
|
||||
f.write(kernel)
|
||||
|
||||
# Call code to gen threadblock
|
||||
self.gen_threadBlock.gen_code(first_use_1stage)
|
||||
@@ -0,0 +1,232 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
class gen_test:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name
|
||||
self.user_header_file = user_header_file
|
||||
self.sample_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
def gen_cpp_sample(self):
|
||||
code = "/* Auto Generated code - Do not edit.*/\n"
|
||||
code += "#include <stdio.h> \n"
|
||||
|
||||
code += "#include \"cutlass/gemm/device/gemm_batched.h\" \n"
|
||||
code += "#include \"cutlass/cutlass.h\" \n"
|
||||
|
||||
code += "#include \"../cutlass_irrelevant.h\" \n"
|
||||
code += "#include \"../cutlass_verify.h\" \n"
|
||||
|
||||
code += "#include \"leaky_bias.h\" \n"
|
||||
|
||||
code += "#include \"utils.h\" \n"
|
||||
|
||||
|
||||
|
||||
code += "int main(int args, char * argv[]) {\n"
|
||||
code += " " + "int M = atoi(argv[1]);\n"
|
||||
code += " " + "int K0 = " + str(self.fuse_gemm_info[0]['mnk'][0]) + ";\n"
|
||||
code += " " + "if(args == 3);\n"
|
||||
code += " " + " " + "K0 = atoi(argv[2]);\n"
|
||||
code += " " + "int B = 1;\n"
|
||||
code += " " + "if(args == 4);\n"
|
||||
code += " " + " " + "B = atoi(argv[3]);\n"
|
||||
|
||||
code += " " + "srand(1234UL);\n"
|
||||
code += " " + "int device_id = 0;\n"
|
||||
code += " " + "cudaGetDevice(&device_id);\n"
|
||||
code += " " + "cudaDeviceProp prop;\n"
|
||||
code += " " + "cudaGetDeviceProperties(&prop, device_id);\n"
|
||||
code += " " + "int sm = prop.major *10 + prop.minor;\n"
|
||||
code += "using ElementCompute = cutlass::half_t;\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("ElementCompute alpha", i) + " = ElementCompute(1);\n"
|
||||
addbias = helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i])
|
||||
if addbias:
|
||||
code += " " + helper.var_idx("ElementCompute beta", i) + " = ElementCompute(1);\n"
|
||||
else:
|
||||
code += " " + helper.var_idx("ElementCompute beta", i) + " = ElementCompute(0);\n"
|
||||
|
||||
code += " " + "size_t flops = 0;\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
m = self.fuse_gemm_info[i]['mnk'][0]
|
||||
n = self.fuse_gemm_info[i]['mnk'][1]
|
||||
k = self.fuse_gemm_info[i]['mnk'][2]
|
||||
|
||||
bias_shape = helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])
|
||||
|
||||
this_k = "K0"
|
||||
if (i > 0):
|
||||
this_k = str(k)
|
||||
|
||||
code += " " + "flops += size_t(2) * size_t(M) * size_t(B) * " + "size_t(" + str(n) + ") * size_t(" + this_k + ");\n"
|
||||
|
||||
code += " " + helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(" + "M" + ", " + str(n) + ", " + this_k + ");\n"
|
||||
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_A", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_", i) + ".k());\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_B", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".n() * problem_size_", i) + ".k());\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_C", i) + "(B * " + str(bias_shape[0]) + " * " + str(bias_shape[1]) + ");\n"
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_D_cutlass_ref", i) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_", i) + ".n());\n"
|
||||
|
||||
code += " " + helper.var_idx("Mat_A", i) + ".init();\n"
|
||||
code += " " + helper.var_idx("Mat_B", i) + ".init();\n"
|
||||
code += " " + helper.var_idx("Mat_C", i) + ".init();\n"
|
||||
|
||||
|
||||
|
||||
code += " " + helper.var_idx("memory_unit<cutlass::half_t> Mat_D", self.b2b_num - 1) + helper.var_idx("(B * problem_size_", i) + helper.var_idx(".m() * problem_size_",self.b2b_num - 1) + ".n());\n"
|
||||
|
||||
params = []
|
||||
params.append("M")
|
||||
params.append("B")
|
||||
|
||||
params.append("Mat_A0.device_ptr")
|
||||
for i in range(self.b2b_num):
|
||||
params.append(helper.var_idx("Mat_B", i) + ".device_ptr")
|
||||
params.append(helper.var_idx("Mat_C", i) + ".device_ptr")
|
||||
if i != self.b2b_num-1:
|
||||
params.append(helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr")
|
||||
params.append(helper.var_idx("Mat_D", self.b2b_num - 1) + ".device_ptr")
|
||||
|
||||
code += " " + "Param arguments = {\n"
|
||||
code += " " + " " + "M,\n"
|
||||
code += " " + " " + "K0,\n"
|
||||
code += " " + " " + "B,\n"
|
||||
|
||||
code += " " + " " + "reinterpret_cast<const void*>(Mat_A0.device_ptr),\n"
|
||||
cnt = 1
|
||||
for i in range(self.b2b_num):
|
||||
bias_flag = helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i])
|
||||
code += " " + " " + "reinterpret_cast<const void*>(" + helper.var_idx("Mat_B", i) + ".device_ptr" + "),\n"
|
||||
cnt += 1
|
||||
if bias_flag:
|
||||
code += " " + " " + "reinterpret_cast<const void*>(" + helper.var_idx("Mat_C", i) + ".device_ptr" + "),\n"
|
||||
cnt += 1
|
||||
else:
|
||||
code += " " + " " + "reinterpret_cast<const void*>(NULL),\n"
|
||||
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_value = str(arg[2])
|
||||
|
||||
code += " " + " " + helper.type_2_cutlass_type(acc_tp) + "(" + arg_value + "),\n"
|
||||
|
||||
if i != self.b2b_num - 1:
|
||||
code += " " + " " + "reinterpret_cast<void*>(" + helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr" + "),\n"
|
||||
else:
|
||||
code += " " + " " + "reinterpret_cast<void*>(" + helper.var_idx("Mat_D", i) + ".device_ptr" + ")};\n"
|
||||
|
||||
|
||||
|
||||
|
||||
code += " " + "TI(FUSED_CUTLASS);\n"
|
||||
code += " " + "for(int i = 0; i < 100; i++){\n"
|
||||
code += " " + " " + "one_api(arguments, sm, NULL);\n"
|
||||
|
||||
code += " " + "}\n"
|
||||
code += " " + "TO(FUSED_CUTLASS, \"FUSED_CUTLASS\", 100);\n"
|
||||
|
||||
code += "\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
N_str = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
code_this += " " + helper.var_idx("typename Gemm", i) + helper.var_idx("::Arguments arguments_", i) + "{\n"
|
||||
code_this += " " + " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
ldmA = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmA = "K0"
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmB = "K0"
|
||||
ldmC = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmBias = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
if self.fuse_gemm_info[i]['A_format'] is 'Col':
|
||||
ldmA = "M"
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
if self.fuse_gemm_info[i]['C_format'] is 'Col':
|
||||
ldmC = "M"
|
||||
|
||||
if i == 0:
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("Mat_A", i) + ".device_ptr), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
else:
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("Mat_D_cutlass_ref", i - 1) + ".device_ptr), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("Mat_B", i) + ".device_ptr), " + ldmB + "}, " + N_str + " * " + ldmB + ",\n"
|
||||
|
||||
M_bias = str(helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])[0])
|
||||
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("Mat_C", i) + ".device_ptr), " + ldmBias + "}, " + M_bias + " * " + N_str + ",\n"
|
||||
code_this += " " + " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("Mat_D_cutlass_ref", i) + ".device_ptr), " + ldmC + "}, " + "M * " + ldmC + ",\n"
|
||||
code_this += " " + " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_value = str(epilogue_arg[2])
|
||||
code_this += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_value) + ")"
|
||||
code_this += " " + " },\n"
|
||||
code_this += " " + " " + "B};\n"
|
||||
|
||||
code += code_this
|
||||
|
||||
|
||||
|
||||
code += " " + "TI(UNFUSED_CUTLASS);\n"
|
||||
code += " " + "for(int i = 0; i < 100; i++){\n"
|
||||
code += " " + " " + self.gen_class_name + "_verify(\n"
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + " " + " " + helper.var_idx("arguments_", i) + ",\n"
|
||||
code += " " + " " + " " + "NULL);\n"
|
||||
|
||||
code += " " + "}\n"
|
||||
code += " " + "TO(UNFUSED_CUTLASS, \"UNFUSED_CUTLASS\", 100);\n"
|
||||
|
||||
code += " " + helper.var_idx("Mat_D_cutlass_ref", self.b2b_num - 1) + ".d2h();\n"
|
||||
code += " " + helper.var_idx("Mat_D", self.b2b_num - 1) + ".d2h();\n"
|
||||
code += " " + helper.var_idx("check_result(Mat_D_cutlass_ref", self.b2b_num - 1) + helper.var_idx(".host_ptr, Mat_D", self.b2b_num - 1) \
|
||||
+ helper.var_idx(".host_ptr, Mat_D", self.b2b_num - 1) + ".elements);\n"
|
||||
|
||||
code += "\n\n}\n"
|
||||
|
||||
with open(self.sample_dir + "sample.cu", "w+") as f:
|
||||
f.write(code)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,456 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
class gen_turing_impl:
|
||||
def __init__(self,fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.class_name = gen_class_name
|
||||
self.gen_class_name = gen_class_name + "_turing_impl"
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
self.gen_turing_unfused = gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
def gen_using(self):
|
||||
code_using = "using b2b_gemm = typename cutlass::gemm::device::" + self.class_name + "<cutlass::half_t>;"
|
||||
|
||||
return code_using + "\n"
|
||||
|
||||
def gen_initialize(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " alpha", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(1);\n"
|
||||
beta = "(1)"
|
||||
|
||||
if helper.get_epilogue_add_bias_or_not(self.fuse_gemm_info[i]) is False:
|
||||
beta = "(0)"
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " beta", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + beta + ";\n"
|
||||
k_str = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
k_str = "K0"
|
||||
code_this += helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(M, " + str(self.fuse_gemm_info[i]['mnk'][1]) + ", " + k_str + ");\n"
|
||||
code += code_this
|
||||
code += "typename b2b_gemm::Arguments arguments{\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
|
||||
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("A", 0) + "), " + helper.var_idx("problem_size_", 0) + ".k()},\n"
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
ldmB = "K0"
|
||||
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmC = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("B", i) + "), " + ldmB + "},\n"
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("C", i) + "), " + ldmC + "},\n"
|
||||
code += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("D", self.b2b_num -1) + "), " + helper.var_idx("problem_size_", self.b2b_num - 1) + ".n()},\n"
|
||||
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
code += " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + epilogue_arg[1]
|
||||
code += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_name) + ")"
|
||||
code += "},\n"
|
||||
code += " " + "Batch};\n\n"
|
||||
|
||||
code += " " "b2b_gemm gemm_op;\n"
|
||||
code += " " + "gemm_op.initialize(arguments);\n"
|
||||
return code + "\n"
|
||||
|
||||
|
||||
|
||||
def gen_run(self):
|
||||
code = " " + "gemm_op(stream);\n"
|
||||
|
||||
return code
|
||||
|
||||
def gen_wrapper(self):
|
||||
code_body = ""
|
||||
|
||||
arg_lists = []
|
||||
arg_lists.append(["int", "M"])
|
||||
arg_lists.append(["int", "K0"])
|
||||
arg_lists.append(["int", "Batch"])
|
||||
arg_lists.append(["void*", helper.var_idx("A", 0)])
|
||||
for i in range(self.b2b_num):
|
||||
arg_lists.append(["void*", helper.var_idx("B", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("C", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("D", i)])
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
arg_lists.append([arg_tp, arg_name])
|
||||
|
||||
if self.b2b_num == 1:
|
||||
code_body += self.gen_turing_unfused.gen_using(False) #False -> Turing, True -> Volta
|
||||
code_body += self.gen_turing_unfused.gen_initialize()
|
||||
code_body += self.gen_turing_unfused.gen_run()
|
||||
else:
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_initialize()
|
||||
code_body += self.gen_run()
|
||||
|
||||
code = ir.gen_func(self.gen_class_name, arg_lists, code_body)
|
||||
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
|
||||
code = self.gen_wrapper()
|
||||
helper.write_2_headfile("turing_impl.h", self.output_dir, self.user_header_file + "\n" + code)
|
||||
|
||||
class gen_volta_turing_fuse_act_impl:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name + "_volta_impl"
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
def perf_tiling(self, layer_mnk):
|
||||
mnk = layer_mnk[:]
|
||||
block_tile = mnk[:]
|
||||
block_tile[2] = 32 # force the K tile to be 32
|
||||
|
||||
# M tile gen
|
||||
block_tile[0] = 32
|
||||
|
||||
# N tile gen
|
||||
if mnk[1] > 128:
|
||||
block_tile[1] = 256
|
||||
elif mnk[1] > 64:
|
||||
block_tile[1] = 128
|
||||
elif mnk[1] > 32:
|
||||
block_tile[1] = 64
|
||||
else :
|
||||
block_tile[1] = 32
|
||||
|
||||
warp_tile = block_tile[:]
|
||||
if block_tile[1] == 256:
|
||||
warp_tile[1] = 64
|
||||
elif block_tile[1] == 128:
|
||||
warp_tile[1] = 32
|
||||
elif block_tile[1] == 64:
|
||||
warp_tile[1] = 32
|
||||
else :
|
||||
warp_tile[1] = 32
|
||||
|
||||
warp_tile[0] = 32
|
||||
|
||||
return block_tile, warp_tile
|
||||
|
||||
|
||||
def process_epilogue(self, epilogue_tp, n, C_tp, Acc_tp):
|
||||
epilogue_setted_type = epilogue_tp
|
||||
cutlass_epilogue_name = "LinearCombinationRelu"
|
||||
if epilogue_setted_type.lower() == 'leakyrelu':
|
||||
cutlass_epilogue_name = "LinearCombinationLeakyRelu"
|
||||
elif epilogue_setted_type.lower() == 'identity':
|
||||
cutlass_epilogue_name = "LinearCombination"
|
||||
|
||||
|
||||
n_mod_8 = n % 4
|
||||
N_align_elements = 1
|
||||
if n_mod_8 == 0:
|
||||
N_align_elements = 8
|
||||
elif n_mod_8 == 4:
|
||||
N_align_elements = 4
|
||||
elif n_mod_8 == 2 or n_mod_8 == 6:
|
||||
N_align_elements = 2
|
||||
|
||||
epilogue_str = "cutlass::epilogue::thread::" + cutlass_epilogue_name+ "<" + C_tp + ", " + str(N_align_elements) + ", " + Acc_tp + ", " + Acc_tp + ">"
|
||||
|
||||
return epilogue_str
|
||||
|
||||
def gen_using(self, volta = True):
|
||||
code_using = ""
|
||||
volta_arch = "cutlass::arch::Sm70"
|
||||
volta_tc = "cutlass::gemm::GemmShape<8, 8, 4>"
|
||||
|
||||
turing_arch = "cutlass::arch::Sm75"
|
||||
turing_tc = "cutlass::gemm::GemmShape<16, 8, 8>"
|
||||
|
||||
arch = ""
|
||||
tc = ""
|
||||
if volta:
|
||||
arch = volta_arch
|
||||
tc = volta_tc
|
||||
else:
|
||||
arch = turing_arch
|
||||
tc = turing_tc
|
||||
|
||||
for i in range(self.b2b_num):
|
||||
|
||||
k = self.fuse_gemm_info[i]['mnk'][2]
|
||||
|
||||
k_mod_8 = k % 4
|
||||
ab_ldm = 1
|
||||
if k_mod_8 == 0:
|
||||
ab_ldm = 8
|
||||
elif k_mod_8 == 4:
|
||||
ab_ldm = 4
|
||||
elif k_mod_8 == 2 or k_mod_8 == 6:
|
||||
ab_ldm = 2
|
||||
|
||||
block_tile, warp_tile = self.perf_tiling(self.fuse_gemm_info[i]['mnk'])
|
||||
|
||||
this_gemm_config = helper.var_idx("using Gemm", i) + " = cutlass::gemm::device::GemmBatched<\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_format']) + ",\n"
|
||||
this_gemm_config += " " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + ",\n"
|
||||
this_gemm_config += " " + "cutlass::arch::OpClassTensorOp,\n"
|
||||
this_gemm_config += " " + arch + ",\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::GemmShape<" + str(block_tile[0]) + ", " + str(block_tile[1]) + ", " + str(block_tile[2]) + ">,\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::GemmShape<" + str(warp_tile[0]) + ", " + str(warp_tile[1]) + ", " + str(warp_tile[2]) + ">,\n"
|
||||
this_gemm_config += " " + tc + ",\n"
|
||||
this_gemm_config += " " + self.process_epilogue(helper.get_epilogue_tp(self.fuse_gemm_info[i]), self.fuse_gemm_info[i]['mnk'][1], helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']), helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp'])) + ",\n"
|
||||
this_gemm_config += " " + "cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,\n"
|
||||
this_gemm_config += " " + "2,\n"
|
||||
this_gemm_config += " " + str(ab_ldm) + ",\n"
|
||||
this_gemm_config += " " + str(ab_ldm) + ">;\n"
|
||||
|
||||
code_using += this_gemm_config + "\n"
|
||||
|
||||
return code_using + "\n"
|
||||
|
||||
def gen_initialize(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
|
||||
N_str = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " alpha", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(1);\n"
|
||||
beta = "(1)"
|
||||
if helper.get_epilogue_add_bias_or_not( self.fuse_gemm_info[i]) is False:
|
||||
beta = "(0)"
|
||||
code_this += helper.var_idx(helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + " beta", i) + " = " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + beta + ";\n"
|
||||
|
||||
k_str = str(self.fuse_gemm_info[i]['mnk'][2])
|
||||
if i == 0:
|
||||
k_str = "K0"
|
||||
code_this += helper.var_idx("cutlass::gemm::GemmCoord problem_size_", i) + "(M, " + str(self.fuse_gemm_info[i]['mnk'][1]) + ", " + k_str + ");\n"
|
||||
code_this += helper.var_idx("typename Gemm", i) + helper.var_idx("::Arguments arguments_", i) + "{\n"
|
||||
code_this += " " + helper.var_idx("problem_size_", i) + ",\n"
|
||||
ldmA = k_str
|
||||
ldmB = k_str
|
||||
ldmC = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
|
||||
ldmBias = str(helper.get_epilogue_bias_ldm(self.fuse_gemm_info[i]))
|
||||
|
||||
if self.fuse_gemm_info[i]['A_format'] is 'Col':
|
||||
ldmA = "M"
|
||||
if self.fuse_gemm_info[i]['B_format'] is 'Row':
|
||||
ldmB = str(self.fuse_gemm_info[i]['mnk'][1])
|
||||
if self.fuse_gemm_info[i]['C_format'] is 'Col':
|
||||
ldmC = "M"
|
||||
|
||||
if i == 0:
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("A", i) + "), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
else:
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['A_tp']) + "*>(" + helper.var_idx("D", i - 1) + "), " + ldmA + "}, " + "M * " + ldmA + ",\n"
|
||||
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['B_tp']) + "*>(" + helper.var_idx("B", i) + "), " + ldmB + "}, " + N_str + " * " + ldmB + ",\n"
|
||||
|
||||
M_bias = str(helper.get_epilogue_bias_shape(self.fuse_gemm_info[i])[0])
|
||||
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("C", i) + "), " + ldmBias + "}, " + M_bias + " * " + N_str + ",\n"
|
||||
code_this += " " + "{reinterpret_cast<" + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['C_tp']) + "*>(" + helper.var_idx("D", i) + "), " + ldmC + "}, " + "M * " + ldmC + ",\n"
|
||||
code_this += " " + "{ " + helper.var_idx("alpha", i) + ", " + helper.var_idx("beta", i)
|
||||
for epilogue_arg in helper.get_epilogue_args(self.fuse_gemm_info[i]):
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + epilogue_arg[1]
|
||||
code_this += ", " + helper.type_2_cutlass_type(self.fuse_gemm_info[i]['Acc_tp']) + "(" + str(arg_name) + ")"
|
||||
code_this += " },\n"
|
||||
code_this += " " + "Batch};\n"
|
||||
|
||||
code_this += " " + helper.var_idx("Gemm", i) + helper.var_idx(" gemm_op_", i) + ";\n"
|
||||
code_this += " " + helper.var_idx("gemm_op_", i) + helper.var_idx(".initialize(arguments_", i) + ", nullptr);\n"
|
||||
|
||||
code += code_this + "\n"
|
||||
return code + "\n"
|
||||
|
||||
|
||||
def gen_run(self):
|
||||
code = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_this = ""
|
||||
code_this += " " + helper.var_idx("gemm_op_", i) + "(stream);\n"
|
||||
|
||||
code += code_this
|
||||
return code
|
||||
|
||||
def gen_wrapper(self):
|
||||
code_body = ""
|
||||
|
||||
arg_lists = []
|
||||
arg_lists.append(["int", "M"])
|
||||
arg_lists.append(["int", "K0"])
|
||||
arg_lists.append(["int", "Batch"])
|
||||
arg_lists.append(["void*", helper.var_idx("A", 0)])
|
||||
for i in range(self.b2b_num):
|
||||
arg_lists.append(["void*", helper.var_idx("B", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("C", i)])
|
||||
arg_lists.append(["void*", helper.var_idx("D", i)])
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
arg_lists.append([arg_tp, arg_name])
|
||||
code_body += self.gen_using()
|
||||
code_body += self.gen_initialize()
|
||||
code_body += self.gen_run()
|
||||
|
||||
code = ir.gen_func(self.gen_class_name, arg_lists, code_body)
|
||||
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
code = self.gen_wrapper()
|
||||
helper.write_2_headfile("volta_impl.h", self.output_dir, self.user_header_file + "\n" + code)
|
||||
|
||||
class gen_one_API:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.gen_class_name = gen_class_name
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.output_dir = output_dir
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
|
||||
self.gen_volta = gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
self.gen_turing = gen_turing_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
|
||||
def gen_CUTLASS_irrelevant_API(self):
|
||||
code = ""
|
||||
code += "#include <cuda_runtime.h>\n"
|
||||
code += "#include <assert.h>\n"
|
||||
|
||||
param_name = "Fused" + str(self.b2b_num) + "xGemm_"
|
||||
for i in range(self.b2b_num):
|
||||
param_name += str(self.fuse_gemm_info[i]['mnk'][1]) + "_"
|
||||
param_name += "Params"
|
||||
params = ""
|
||||
params += " " + "int M;\n"
|
||||
params += " " + "int K0;\n"
|
||||
params += " " + "int Batch;\n"
|
||||
params += " " + "const void* A0;\n"
|
||||
for i in range(self.b2b_num):
|
||||
params += " " + "const void* " + helper.var_idx("B", i) + ";\n"
|
||||
params += " " + "const void* " + helper.var_idx("C", i) + ";\n"
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
acc_tp = helper.get_epilogue_compute_tp(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_tp = arg[0]
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
params += " " + arg_tp + " " + arg_name + ";\n"
|
||||
params += " " + "void* " + helper.var_idx("D", i) + ";\n"
|
||||
code += ir.gen_struct(param_name, params)
|
||||
code += "using Param = " + param_name + ";\n"
|
||||
code += "void one_api( const Param & param, int sm, cudaStream_t stream);\n"
|
||||
|
||||
|
||||
return code
|
||||
|
||||
def gen_one_api(self):
|
||||
code = ""
|
||||
code += "/* Auto Generated code - Do not edit.*/\n"
|
||||
code += "#include \"cutlass_irrelevant.h\"\n"
|
||||
code += "#include \"api.h\"\n"
|
||||
code += "void one_api( const Param & param, int sm, cudaStream_t stream) {\n"
|
||||
|
||||
code += " " + "if (sm == 70) \n"
|
||||
code += " " + " " + self.gen_class_name + "_volta_impl(param.M, param.K0, param.Batch, const_cast<void*>(param.A0), "
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx("const_cast<void*>(param.B", i) + "), "
|
||||
code += helper.var_idx("const_cast<void*>(param.C", i) + "), "
|
||||
code += helper.var_idx("param.D", i) + ", "
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
code += "param." + arg_name + ", "
|
||||
code += "stream);\n"
|
||||
code += " " + "else if(sm >= 75) \n"
|
||||
code += " " + " " + self.gen_class_name + "_turing_impl(param.M, param.K0, param.Batch, const_cast<void*>(param.A0), "
|
||||
for i in range(self.b2b_num):
|
||||
code += helper.var_idx("const_cast<void*>(param.B", i) + "), "
|
||||
code += helper.var_idx("const_cast<void*>(param.C", i) + "), "
|
||||
code += helper.var_idx("param.D", i) + ", "
|
||||
epilogue_args = helper.get_epilogue_args(self.fuse_gemm_info[i])
|
||||
for arg in epilogue_args:
|
||||
arg_name = helper.var_idx("Epilogue", i) + "_" + arg[1]
|
||||
code += "param." + arg_name + ", "
|
||||
code += "stream);\n"
|
||||
code += " " + "else assert(0);\n"
|
||||
code += "}\n"
|
||||
return code
|
||||
|
||||
def gen_code(self):
|
||||
|
||||
turing_code = self.gen_turing.gen_wrapper()
|
||||
volta_code = self.gen_volta.gen_wrapper()
|
||||
cutlass_irrelevant_code = self.gen_CUTLASS_irrelevant_API()
|
||||
|
||||
one_api_code = self.gen_one_api()
|
||||
with open(self.output_dir + "one_api.cu", "w+") as f:
|
||||
f.write(one_api_code)
|
||||
|
||||
helper.write_2_headfile("cutlass_irrelevant.h", self.output_dir, cutlass_irrelevant_code)
|
||||
|
||||
helper.write_2_headfile("api.h", self.output_dir, self.user_header_file + "\n" + turing_code + volta_code)
|
||||
@@ -0,0 +1,92 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import helper
|
||||
import gen_ir as ir
|
||||
|
||||
import gen_turing_and_volta as gen_basic
|
||||
|
||||
|
||||
class gen_verify:
|
||||
def __init__(self, fuse_gemm_info, gen_class_name, user_header_file, output_dir = "../"):
|
||||
self.fuse_gemm_info = fuse_gemm_info
|
||||
self.name = gen_class_name + "_verify"
|
||||
self.b2b_num = len(fuse_gemm_info)
|
||||
self.params = []
|
||||
self.user_header_file = ""
|
||||
for header in user_header_file:
|
||||
self.user_header_file += "#include \"" + header + "\"\n"
|
||||
self.seperate_cutlass = gen_basic.gen_volta_turing_fuse_act_impl(fuse_gemm_info, gen_class_name, user_header_file, output_dir)
|
||||
self.gen_params()
|
||||
self.output_dir = output_dir
|
||||
|
||||
|
||||
def gen_code(self):
|
||||
code = ""
|
||||
code += self.user_header_file
|
||||
code += self.seperate_cutlass.gen_using(False) #False -> Turing, True -> Volta
|
||||
|
||||
code_body = ""
|
||||
for i in range(self.b2b_num):
|
||||
code_body += " " + helper.var_idx("Gemm", i) + helper.var_idx(" gemm_op_", i) + ";\n"
|
||||
code_body += " " + helper.var_idx("gemm_op_", i) + helper.var_idx(".initialize(Arguments_", i) + ", nullptr);\n"
|
||||
|
||||
code_body += self.seperate_cutlass.gen_run()
|
||||
|
||||
code += ir.gen_func(self.name, self.params, code_body)
|
||||
helper.write_2_headfile("cutlass_verify.h", self.output_dir, code)
|
||||
|
||||
|
||||
def gen_params(self):
|
||||
for i in range(self.b2b_num):
|
||||
self.params.append(
|
||||
(
|
||||
helper.var_idx("typename Gemm", i)+ "::Arguments",
|
||||
helper.var_idx("Arguments_", i)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_params(self, declartion = True):
|
||||
code = ""
|
||||
if declartion:
|
||||
for param in self.params:
|
||||
code += param[0] + " " + param[1] + ";\n"
|
||||
|
||||
return code
|
||||
|
||||
|
||||
def gen_initialize():
|
||||
code = ""
|
||||
initialize_code = self.seperate_cutlass.gen_initialize()
|
||||
|
||||
code = ir.gen_func("initialize", [[]])
|
||||
@@ -0,0 +1,52 @@
|
||||
#!/bin/bash
|
||||
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
NUM_ARGS=3
|
||||
if [ $# -ne $NUM_ARGS ]; then
|
||||
echo "Usage: $0 <config_file> <output_directory> <cutlass_directory>"
|
||||
echo " config_file: JSON file containing configuration to run"
|
||||
echo " output_directory: directory to store results"
|
||||
echo " cutlass_directory: directory containing cutlass source"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
config_file=$1
|
||||
output_dir=$2
|
||||
cutlass_dir=$3
|
||||
|
||||
python3 gen_all_code.py \
|
||||
--config-file $config_file \
|
||||
--gen-name FusedMultiGemmForward \
|
||||
--output-dir $output_dir \
|
||||
--cutlass-dir $cutlass_dir
|
||||
@@ -0,0 +1,135 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
def type_2_cutlass_type(input_type = "fp16"):
|
||||
# float point type
|
||||
if input_type == "fp32":
|
||||
return "float"
|
||||
if input_type == "bf16":
|
||||
return "cutlass::bfloat16_t"
|
||||
if input_type == "fp16":
|
||||
return "cutlass::half_t"
|
||||
|
||||
# integer type
|
||||
if(input_type == "int32"):
|
||||
return "int32_t"
|
||||
if(input_type == "int8"):
|
||||
return "int8_t"
|
||||
|
||||
if input_type == 'Row':
|
||||
return 'cutlass::layout::RowMajor'
|
||||
if input_type == 'Col':
|
||||
return 'cutlass::layout::ColumnMajor'
|
||||
|
||||
def cvt_2_cutlass_shape(gemm_shape):
|
||||
# gemm shape
|
||||
if len(gemm_shape) == 3:
|
||||
val = "cutlass::gemm::GemmShape<" \
|
||||
+ str(gemm_shape[0]) + ", " \
|
||||
+ str(gemm_shape[1]) + ", " \
|
||||
+ str(gemm_shape[2]) + ">"
|
||||
return val
|
||||
|
||||
|
||||
def write_2_headfile(filename, file_dir, string):
|
||||
with open(file_dir + filename, 'w') as f:
|
||||
f.write("/* Auto Generated code - Do not edit.*/\n\n\n#pragma once\n" + string)
|
||||
|
||||
def var_idx(varaiable, index):
|
||||
return varaiable + str(index)
|
||||
|
||||
|
||||
def list_2_string(input_list, ):
|
||||
rtn_string = ""
|
||||
|
||||
cnt = 0
|
||||
|
||||
for element in input_list:
|
||||
final = ", \n"
|
||||
if cnt == len(input_list) - 1:
|
||||
final = "\n"
|
||||
cnt += 1
|
||||
rtn_string += str(element) + final
|
||||
|
||||
return rtn_string
|
||||
|
||||
|
||||
def get_epilouge_info(layer_info):
|
||||
return layer_info['epilogue']
|
||||
|
||||
def get_epilogue_tp(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['tp']
|
||||
|
||||
def get_epilogue_add_bias_or_not(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['bias']['addbias']
|
||||
|
||||
def get_epilogue_add_bias_tp(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['bias']['bias_tp']
|
||||
|
||||
def get_epilogue_args(layer_info):
|
||||
epilogue_info = get_epilouge_info(layer_info)
|
||||
return epilogue_info['args']
|
||||
|
||||
def get_epilogue_bias_shape(layer_info):
|
||||
bias_tp = get_epilogue_add_bias_tp(layer_info).lower()
|
||||
mn_shape = layer_info['mnk'][:-1]
|
||||
|
||||
if bias_tp == 'mat':
|
||||
mn_shape[0] = 'M'
|
||||
return mn_shape
|
||||
elif bias_tp == 'vec':
|
||||
mn_shape[0] = 1
|
||||
return mn_shape
|
||||
else:
|
||||
assert(0)
|
||||
|
||||
def get_epilogue_bias_ldm(layer_info):
|
||||
bias_tp = get_epilogue_add_bias_tp(layer_info).lower()
|
||||
mn_shape = layer_info['mnk'][:-1]
|
||||
|
||||
c_layout = layer_info['C_format'].lower()
|
||||
|
||||
if c_layout != 'row':
|
||||
assert(0)
|
||||
|
||||
if bias_tp == 'mat':
|
||||
return mn_shape[1]
|
||||
elif bias_tp == 'vec':
|
||||
return 0
|
||||
else:
|
||||
assert(0)
|
||||
|
||||
def get_epilogue_compute_tp(layer_info):
|
||||
return layer_info['Acc_tp']
|
||||
@@ -0,0 +1,67 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
import os
|
||||
|
||||
class replace_fix_impl:
|
||||
def __init__(self, src_dir, dst_dir, cutlass_deps_root):
|
||||
self.src_dir = src_dir
|
||||
self.dst_dir = dst_dir
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
|
||||
|
||||
|
||||
def gen_code(self):
|
||||
for sub_dir in os.walk(self.src_dir):
|
||||
files_in_sub_dir = sub_dir[2]
|
||||
|
||||
src_dirs = sub_dir[0]
|
||||
output_dirs = self.dst_dir + sub_dir[0][len(self.src_dir):]
|
||||
|
||||
if not os.path.exists(output_dirs):
|
||||
os.mkdir(output_dirs)
|
||||
|
||||
for f in files_in_sub_dir:
|
||||
with open(src_dirs +"/" + f, 'r') as current_file:
|
||||
output_lines = []
|
||||
lines = current_file.readlines()
|
||||
|
||||
for line in lines:
|
||||
if(len(line) >= len("#include \"cutlass") and line[:len("#include \"cutlass")] == "#include \"cutlass"):
|
||||
new_line = "#include \"" + self.cutlass_deps_root + line[len("#include \""):]
|
||||
# print(new_line)
|
||||
output_lines.append(new_line)
|
||||
else:
|
||||
output_lines.append(line)
|
||||
|
||||
with open(output_dirs + "/" + f, "w+") as dest_file:
|
||||
dest_file.writelines(output_lines)
|
||||
Reference in New Issue
Block a user