@@ -49,7 +49,7 @@ class gen_device:
|
||||
self.arg_member = []
|
||||
self.gen_class_name = gen_class_name
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
self.template_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"
|
||||
@@ -63,7 +63,7 @@ class gen_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)
|
||||
self.gen_kernel = gen_ker.gen_kernel(self.template_args, self.gen_class_name, self.b2b_num, output_dir, cutlass_deps_root, project_root)
|
||||
|
||||
|
||||
def __check_arg_type(self, temp_arg):
|
||||
@@ -126,7 +126,7 @@ class gen_device:
|
||||
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)
|
||||
gen_code = gen_ir.gen_template_class(self.gen_class_name, self.template_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:
|
||||
@@ -142,7 +142,7 @@ class gen_device:
|
||||
|
||||
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]])
|
||||
self.template_args.append([self.__check_arg_type(arg), arg, self.args[arg]])
|
||||
|
||||
def update_b2b_args(self):
|
||||
|
||||
|
||||
@@ -444,7 +444,7 @@ class gen_kernel:
|
||||
|
||||
self.gen_class_name = "B2bGemm"
|
||||
self.gen_kernel_name = gen_class_name + "Kernel"
|
||||
self.tempalte_args = []
|
||||
self.template_args = []
|
||||
|
||||
self.cutlass_deps_root = cutlass_deps_root
|
||||
self.project_root = project_root
|
||||
|
||||
@@ -957,13 +957,13 @@ public:\n\
|
||||
|
||||
def gen_code(self):
|
||||
|
||||
tempalte_arg = []
|
||||
template_arg = []
|
||||
for i in range(self.b2b_num):
|
||||
tempalte_arg.append(("typename", helper.var_idx("Shape", i)))
|
||||
template_arg.append(("typename", helper.var_idx("Shape", i)))
|
||||
for i in range(self.b2b_num):
|
||||
tempalte_arg.append(("typename", helper.var_idx("Policy", i)))
|
||||
template_arg.append(("typename", helper.var_idx("Policy", i)))
|
||||
for i in range(self.b2b_num):
|
||||
tempalte_arg.append((int, helper.var_idx("Stage", i)))
|
||||
template_arg.append((int, helper.var_idx("Stage", i)))
|
||||
|
||||
|
||||
|
||||
@@ -971,7 +971,7 @@ public:\n\
|
||||
code_body += self.gen_protected()
|
||||
code_body += self.gen_public_member()
|
||||
|
||||
class_code = gen_ir.gen_template_class("B2bMmaBase", tempalte_arg, code_body)
|
||||
class_code = gen_ir.gen_template_class("B2bMmaBase", template_arg, code_body)
|
||||
|
||||
code = self.gen_include_header() + gen_ir.gen_namespace("cutlass", gen_ir.gen_namespace("gemm", gen_ir.gen_namespace("threadblock", class_code)))
|
||||
|
||||
|
||||
@@ -68,7 +68,7 @@
|
||||
"source": [
|
||||
"## Define the epilogue visitor functor\n",
|
||||
"The epilogue functor can be defined as a simple Python function and a set of example tensors for inputs and outputs. The example below illustrates a complex epilogue under the directed acyclic graph structure (`F` is used twice). The epilogue takes source tensors in different ranks: `alpha`, `beta` are scalars, `bias` is a column vector to broadcast, and `C`, `aux` are matrices. It contains various math operations from basic arithmatic operations and built-in callable functions like `relu`. It also accomodates multiple outputs `D` and `F`. Note that there are some restrictions on syntax.\n",
|
||||
"* Each named variable must be assigned exactly once and defined before it it used.\n",
|
||||
"* Each named variable must be assigned exactly once and defined before it used.\n",
|
||||
"* Reserved names: `accum`, `C`, and `D` are reserved for accumulator, tensor_C, and tensor_D.\n",
|
||||
"* Return values must be a named variable.\n",
|
||||
"\n",
|
||||
|
||||
Reference in New Issue
Block a user