wehub-resource-sync
299 lines
12 KiB
Python
299 lines
12 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
head_template = """// auto-generated by generate_configs.py
|
|
|
|
#pragma once
|
|
|
|
#include "cutlass/gemm_coord.h"
|
|
|
|
namespace ap {
|
|
|
|
constexpr int kNumConfigsHalf = ${num_configs_fp16};
|
|
constexpr int kNumConfigsFloat = ${num_configs_fp32};
|
|
|
|
template <int SwizzleFactor, bool Batched> struct SwizzleWrapper {
|
|
using Type =
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
|
|
};
|
|
|
|
// template <int SwizzleFactor>
|
|
// struct SwizzleWrapper<SwizzleFactor, true> {
|
|
// using Type =
|
|
// cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
|
|
// };
|
|
"""
|
|
|
|
autotune_wrapper_template = """
|
|
#define AP_AUTOTUNE_${datatype}(func, stream, ...) { \\
|
|
using FuncType = decltype(func<0>); \\
|
|
static int selected_config_id = -1; \\
|
|
static std::vector<std::function<FuncType>> \\
|
|
matmul_functions = { \\
|
|
${repeat_functions} \\
|
|
}; \\
|
|
if (selected_config_id == -1) { \\
|
|
selected_config_id = ap::ProfileBestConfig(matmul_functions, stream, ##__VA_ARGS__); \\
|
|
} \\
|
|
matmul_functions[selected_config_id](__VA_ARGS__); \\
|
|
}
|
|
"""
|
|
|
|
fp16_config_template_0 = """
|
|
template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
|
|
struct GemmTuningConfigs {
|
|
using TShape = cutlass::gemm::GemmShape<${tshape}>;
|
|
using WShape = cutlass::gemm::GemmShape<${wshape}>;
|
|
using IShape = cutlass::gemm::GemmShape<${ishape}>;
|
|
static constexpr int kNumStages = ${stages};
|
|
|
|
using SwizzleThreadBlock =
|
|
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
|
|
static constexpr int kId = Id;
|
|
};
|
|
"""
|
|
|
|
fp16_config_template = """
|
|
template <typename ElementT, int SwizzleFactor, bool Batched>
|
|
struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ${config_id}> {
|
|
using TShape = cutlass::gemm::GemmShape<${tshape}>;
|
|
using WShape = cutlass::gemm::GemmShape<${wshape}>;
|
|
using IShape = cutlass::gemm::GemmShape<${ishape}>;
|
|
static constexpr int kNumStages = ${stages};
|
|
|
|
using SwizzleThreadBlock =
|
|
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
|
|
static constexpr int kId = ${config_id};
|
|
};
|
|
"""
|
|
|
|
fp32_config_template_0 = """
|
|
// Specialization for float
|
|
template <int SwizzleFactor, bool Batched, int Id>
|
|
struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
|
|
using TShape = cutlass::gemm::GemmShape<${tshape}>;
|
|
using WShape = cutlass::gemm::GemmShape<${wshape}>;
|
|
using IShape = cutlass::gemm::GemmShape<${ishape}>;
|
|
static constexpr int kNumStages = ${stages};
|
|
|
|
using SwizzleThreadBlock =
|
|
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
|
|
static constexpr int kId = Id;
|
|
};
|
|
"""
|
|
|
|
fp32_config_template = """
|
|
template <int SwizzleFactor, bool Batched>
|
|
struct GemmTuningConfigs<float, SwizzleFactor, Batched, ${config_id}> {
|
|
using TShape = cutlass::gemm::GemmShape<${tshape}>;
|
|
using WShape = cutlass::gemm::GemmShape<${wshape}>;
|
|
using IShape = cutlass::gemm::GemmShape<${ishape}>;
|
|
static constexpr int kNumStages = ${stages};
|
|
|
|
using SwizzleThreadBlock =
|
|
typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
|
|
static constexpr int kId = ${config_id};
|
|
};
|
|
"""
|
|
|
|
tail_code_str = """
|
|
} // namespace ap
|
|
"""
|
|
|
|
|
|
class GemmTuningConfig:
|
|
def __init__(self, tshape, wshape, ishape, stages, level):
|
|
self.tshape = tshape
|
|
self.wshape = wshape
|
|
self.ishape = ishape
|
|
self.stages = stages
|
|
self.level = level
|
|
|
|
def __eq__(self, other):
|
|
def check_shape(s1, s2):
|
|
assert len(s1) == len(s2)
|
|
res = True
|
|
for i in range(len(s1)):
|
|
if s1[i] != s2[i]:
|
|
res = False
|
|
break
|
|
return res
|
|
|
|
res = check_shape(self.tshape, other.tshape)
|
|
res = res and check_shape(self.wshape, other.wshape)
|
|
res = res and check_shape(self.ishape, other.ishape)
|
|
res = res and self.stages == other.stages
|
|
return res
|
|
|
|
|
|
def all_configs_sm80_fp16():
|
|
all_tuning_configs_fp16 = [
|
|
GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 2, 2),
|
|
GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 2, 1),
|
|
GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 3, 3),
|
|
GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 3, 2),
|
|
GemmTuningConfig([64, 128, 64], [32, 64, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 3, 3),
|
|
GemmTuningConfig([128, 64, 64], [64, 32, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 3, 2),
|
|
GemmTuningConfig([128, 256, 32], [64, 64, 32], [16, 8, 16], 3, 2),
|
|
GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 3, 2),
|
|
GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([256, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
|
|
GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 4, 3),
|
|
GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 4, 2),
|
|
GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 4, 2),
|
|
GemmTuningConfig([64, 256, 32], [64, 64, 32], [16, 8, 16], 4, 2),
|
|
GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 4, 3),
|
|
GemmTuningConfig([128, 32, 64], [32, 32, 64], [16, 8, 16], 4, 1),
|
|
GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 4, 1),
|
|
GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 4, 1),
|
|
GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 4, 2),
|
|
GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 4, 1),
|
|
GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 4, 1),
|
|
GemmTuningConfig([16, 64, 64], [16, 32, 64], [16, 8, 16], 5, 2),
|
|
GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 5, 3),
|
|
GemmTuningConfig([32, 64, 64], [16, 32, 64], [16, 8, 16], 5, 1),
|
|
GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 5, 3),
|
|
GemmTuningConfig([64, 64, 64], [32, 32, 64], [16, 8, 16], 5, 1),
|
|
GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 5, 3),
|
|
GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 5, 1),
|
|
GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 5, 1),
|
|
GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 5, 2),
|
|
GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 16], 6, 1),
|
|
GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 16], 6, 1),
|
|
GemmTuningConfig([128, 32, 32], [32, 32, 32], [16, 8, 16], 7, 1),
|
|
GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 16], 10, 1),
|
|
]
|
|
return all_tuning_configs_fp16
|
|
|
|
|
|
def all_configs_sm80_fp32():
|
|
all_tuning_configs_fp32 = [
|
|
GemmTuningConfig([64, 64, 16], [32, 32, 16], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([64, 256, 16], [32, 64, 16], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([64, 256, 32], [32, 64, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([128, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([256, 64, 16], [64, 32, 16], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([256, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
|
|
GemmTuningConfig([64, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
|
|
GemmTuningConfig([128, 64, 16], [64, 32, 16], [16, 8, 8], 4, 1),
|
|
GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
|
|
]
|
|
return all_tuning_configs_fp32
|
|
|
|
|
|
def generate_autotune_wrapper(datatype, num_configs):
|
|
repeat_func_strs = []
|
|
for i in range(num_configs):
|
|
repeat_func_strs.append(f"func<{i}>")
|
|
code_str = autotune_wrapper_template.replace(
|
|
"${datatype}", datatype
|
|
).replace("${repeat_functions}", ", \\\n ".join(repeat_func_strs))
|
|
return code_str
|
|
|
|
|
|
def get_configs(all_configs_list, level=3):
|
|
consigs_list = []
|
|
for i in range(len(all_configs_list)):
|
|
if all_configs_list[i].level <= level:
|
|
already_have = False
|
|
for config in consigs_list:
|
|
if config == all_configs_list[i]:
|
|
print(f"-- The {i}-th config is repeat.")
|
|
already_have = True
|
|
break
|
|
if not already_have:
|
|
consigs_list.append(all_configs_list[i])
|
|
return consigs_list
|
|
|
|
|
|
def generate_configs(configs_list, config_template_0, config_template):
|
|
code_str = ""
|
|
config_id = 0
|
|
for config in configs_list:
|
|
if config_id == 0:
|
|
config_code_str = (
|
|
config_template_0.replace(
|
|
"${tshape}", ", ".join(map(str, config.tshape))
|
|
)
|
|
.replace("${wshape}", ", ".join(map(str, config.wshape)))
|
|
.replace("${ishape}", ", ".join(map(str, config.ishape)))
|
|
.replace("${stages}", str(config.stages))
|
|
)
|
|
else:
|
|
config_code_str = (
|
|
config_template.replace(
|
|
"${tshape}", ", ".join(map(str, config.tshape))
|
|
)
|
|
.replace("${wshape}", ", ".join(map(str, config.wshape)))
|
|
.replace("${ishape}", ", ".join(map(str, config.ishape)))
|
|
.replace("${stages}", str(config.stages))
|
|
.replace("${config_id}", str(config_id))
|
|
)
|
|
code_str += config_code_str
|
|
config_id += 1
|
|
return len(configs_list), code_str
|
|
|
|
|
|
def main():
|
|
level = 1
|
|
num_fp16_configs, fp16_configs_code_str = generate_configs(
|
|
configs_list=get_configs(all_configs_sm80_fp16(), level=level),
|
|
config_template_0=fp16_config_template_0,
|
|
config_template=fp16_config_template,
|
|
)
|
|
num_fp32_configs, fp32_configs_code_str = generate_configs(
|
|
configs_list=get_configs(all_configs_sm80_fp32(), level=level),
|
|
config_template_0=fp32_config_template_0,
|
|
config_template=fp32_config_template,
|
|
)
|
|
print(
|
|
f"-- Total {num_fp16_configs} fp16 configs, {num_fp32_configs} fp32 configs"
|
|
)
|
|
head_code_str = head_template.replace(
|
|
"${num_configs_fp16}", str(num_fp16_configs)
|
|
).replace("${num_configs_fp32}", str(num_fp32_configs))
|
|
fp16_autotune_wrapper_code_str = generate_autotune_wrapper(
|
|
"half", num_fp16_configs
|
|
)
|
|
fp32_autotune_wrapper_code_str = generate_autotune_wrapper(
|
|
"float", num_fp32_configs
|
|
)
|
|
with open("all_tuning_configs.h", "w") as f:
|
|
f.write(head_code_str)
|
|
f.write(fp16_autotune_wrapper_code_str)
|
|
f.write(fp32_autotune_wrapper_code_str)
|
|
f.write(fp16_configs_code_str)
|
|
f.write(fp32_configs_code_str)
|
|
f.write(tail_code_str)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|