diff --git a/CMakeLists.txt b/CMakeLists.txt index e915f69d..855b5ae1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -48,4 +48,8 @@ else() message(STATUS "DLINFER_BUILD_DICP=OFF: skipping DICP vendor build (set env DLINFER_BUILD_DICP=ON to enable)") endif() +if(DEVICE STREQUAL "ascend") + add_subdirectory(dlinfer/vendor/ascend/csrc/grouped_matmul_direct) +endif() + install(CODE "message(STATUS \"Install completed for device: ${DEVICE}\")") diff --git a/benchmark/probe_grouped_matmul_direct.py b/benchmark/probe_grouped_matmul_direct.py new file mode 100644 index 00000000..2e5fea4b --- /dev/null +++ b/benchmark/probe_grouped_matmul_direct.py @@ -0,0 +1,71 @@ +"""Device correctness probe for DLInfer's bundled direct2560 extension.""" + +import argparse + +from dlinfer.vendor.ascend import grouped_matmul_direct +import torch +import torch_npu + + +def run(dtype: torch.dtype, iterations: int = 1) -> None: + if not grouped_matmul_direct.is_available(): + raise RuntimeError(grouped_matmul_direct.unavailable_reason()) + + experts, hidden_size, output_size, tokens = 2560, 16, 16, 8 + device = torch.device("npu:0") + x = torch.randn(tokens, hidden_size, device=device, dtype=dtype) + stored_weight = torch.randn( + experts, output_size, hidden_size, device=device, dtype=dtype + ) + weight = stored_weight.transpose(1, 2) + group_list = torch.zeros(experts, device=device, dtype=torch.int64) + group_list[:tokens] = 1 + + references = [] + row_start = 0 + for expert_start in range(0, experts, 1024): + expert_end = min(expert_start + 1024, experts) + chunk_groups = group_list[expert_start:expert_end] + row_end = row_start + int(chunk_groups.sum().cpu()) + if row_end > row_start: + references.append( + torch.ops.npu.npu_grouped_matmul( + [x[row_start:row_end]], + [weight[expert_start:expert_end]], + group_list=chunk_groups, + split_item=2, + group_type=0, + group_list_type=1, + )[0] + ) + row_start = row_end + + reference = torch.cat(references, dim=0) + for iteration in range(iterations): + print(f"ITERATION {iteration + 1}/{iterations} begin", flush=True) + direct = grouped_matmul_direct.grouped_matmul(x, stored_weight, group_list, 1) + print(f"ITERATION {iteration + 1}/{iterations} submitted", flush=True) + torch_npu.npu.synchronize() + print(f"ITERATION {iteration + 1}/{iterations} synchronized", flush=True) + torch.testing.assert_close(direct, reference, rtol=0, atol=0) + print( + f"PASS dtype={dtype} iterations={iterations} " + f"max_abs={(direct - reference).abs().max().item()}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dtype", choices=("fp16", "bf16", "both"), default="both") + parser.add_argument("--iterations", type=int, default=1) + args = parser.parse_args() + if args.iterations < 1: + parser.error("--iterations must be at least 1") + if args.dtype in ("fp16", "both"): + run(torch.float16, args.iterations) + if args.dtype in ("bf16", "both"): + run(torch.bfloat16, args.iterations) + + +if __name__ == "__main__": + main() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/CMakeLists.txt b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/CMakeLists.txt new file mode 100644 index 00000000..6f5ab20a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/CMakeLists.txt @@ -0,0 +1,48 @@ +include(ascend) + +# Build the direct-2560 operator from source as part of the Python wheel build. +set(ASCEND_CANN_PACKAGE_PATH "${ASCEND_TOOLKIT_HOME}") +set(ASCEND_COMPUTE_UNIT ascend910_93) +set(ASCEND_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") +add_subdirectory(opp) + +add_library(_grouped_matmul_direct MODULE grouped_matmul_direct.cpp) +add_dependencies(_grouped_matmul_direct cust_opapi) +find_library(TORCH_PYTHON_LIBRARY torch_python + PATHS "${TORCH_INSTALL_PREFIX}/lib" + REQUIRED + NO_DEFAULT_PATH +) +set_target_properties(_grouped_matmul_direct PROPERTIES + PREFIX "" + INSTALL_RPATH "\$ORIGIN/grouped_matmul_direct/op_api/lib" +) + +target_compile_definitions(_grouped_matmul_direct PRIVATE + TORCH_EXTENSION_NAME=_grouped_matmul_direct + _GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI} +) +target_compile_options(_grouped_matmul_direct PRIVATE -O2 -Wall -Wextra) +target_include_directories(_grouped_matmul_direct PRIVATE + ${DLINFER_GMM_OPP_BINARY_DIR}/autogen + ${TORCH_INCLUDE_DIRS} + ${CANN_INCLUDE_DIRS} + ${CANN_INCLUDE_DIRS}/aclnn +) +target_link_libraries(_grouped_matmul_direct PRIVATE + Python::Python + torch + ${TORCH_PYTHON_LIBRARY} + cust_opapi + ${CANN_LIBRARIES} +) + +install(TARGETS _grouped_matmul_direct + LIBRARY DESTINATION dlinfer/vendor/ascend +) +install(TARGETS cust_opapi + LIBRARY DESTINATION dlinfer/vendor/ascend/grouped_matmul_direct/op_api/lib +) +install(FILES ${DLINFER_GMM_OPP_BINARY_DIR}/autogen/aclnn_dlinfer_grouped_matmul_direct.h + DESTINATION dlinfer/vendor/ascend/grouped_matmul_direct/op_api/include +) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/grouped_matmul_direct.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/grouped_matmul_direct.cpp new file mode 100644 index 00000000..89ffd3df --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/grouped_matmul_direct.cpp @@ -0,0 +1,197 @@ +// Copied from DeepLink-org/DLBlas csrc/ascend/grouped_matmul_direct and +// adapted for DLInfer pybind packaging. +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int64_t kDirectGroups = 2560; + +at::Tensor GroupedMatmulDirect(const at::Tensor& x, const at::Tensor& weight, const at::Tensor& group_list, int64_t group_list_type, uint64_t stream_handle); + +aclDataType ToAclDataType(const at::ScalarType dtype) { + switch (dtype) { + case at::ScalarType::Half: + return ACL_FLOAT16; + case at::ScalarType::BFloat16: + return ACL_BF16; + case at::ScalarType::Long: + return ACL_INT64; + default: + TORCH_CHECK(false, "unsupported dtype for direct grouped matmul: ", dtype); + } +} + +struct AclTensorDeleter { + void operator()(aclTensor* tensor) const { + if (tensor != nullptr) { + aclDestroyTensor(tensor); + } + } +}; + +struct AclTensorListDeleter { + void operator()(aclTensorList* tensors) const { + if (tensors != nullptr) { + aclDestroyTensorList(tensors); + } + } +}; + +using AclTensorPtr = std::unique_ptr; +using AclTensorListPtr = std::unique_ptr; + +AclTensorPtr MakeAclTensor(const at::Tensor& tensor) { + std::vector shape(tensor.sizes().begin(), tensor.sizes().end()); + std::vector strides(tensor.strides().begin(), tensor.strides().end()); + auto* acl_tensor = aclCreateTensor(shape.data(), + shape.size(), + ToAclDataType(tensor.scalar_type()), + strides.data(), + tensor.storage_offset(), + ACL_FORMAT_ND, + shape.data(), + shape.size(), + tensor.data_ptr()); + TORCH_CHECK(acl_tensor != nullptr, "aclCreateTensor failed"); + return AclTensorPtr(acl_tensor); +} + +AclTensorPtr MakeTransposedWeightAclTensor(const at::Tensor& tensor) { + TORCH_CHECK(tensor.dim() == 3 && tensor.is_contiguous(), "stored weight must be contiguous [E, N, K]"); + const std::array view_shape{tensor.size(0), tensor.size(2), tensor.size(1)}; + // The kernel consumes the raw [E, N, K] storage with transposeWeight=true, + // while host tiling needs the logical [E, K, N] shape. Mark the ACL view + // contiguous so the generic individual-op executor does not insert a copy. + const std::array view_strides{ + view_shape[1] * view_shape[2], view_shape[2], 1}; + const std::array storage_shape{tensor.size(0), tensor.size(1), tensor.size(2)}; + auto* acl_tensor = aclCreateTensor(view_shape.data(), + view_shape.size(), + ToAclDataType(tensor.scalar_type()), + view_strides.data(), + tensor.storage_offset(), + ACL_FORMAT_ND, + storage_shape.data(), + storage_shape.size(), + tensor.data_ptr()); + TORCH_CHECK(acl_tensor != nullptr, "aclCreateTensor failed for weight"); + return AclTensorPtr(acl_tensor); +} + +AclTensorListPtr MakeAclTensorList(aclTensor* tensor) { + std::array tensors{tensor}; + auto* list = aclCreateTensorList(tensors.data(), tensors.size()); + TORCH_CHECK(list != nullptr, "aclCreateTensorList failed"); + return AclTensorListPtr(list); +} + +AclTensorPtr MakeEmptyAclTensor(const aclDataType dtype) { + constexpr std::array shape{0}; + constexpr std::array strides{1}; + auto* tensor = aclCreateTensor(shape.data(), + shape.size(), + dtype, + strides.data(), + 0, + ACL_FORMAT_ND, + shape.data(), + shape.size(), + nullptr); + TORCH_CHECK(tensor != nullptr, "aclCreateTensor failed for empty input"); + return AclTensorPtr(tensor); +} + +AclTensorListPtr MakeEmptyAclTensorList(const aclDataType dtype) { + auto tensor = MakeEmptyAclTensor(dtype); + auto* raw_tensor = tensor.get(); + auto* list = aclCreateTensorList(&raw_tensor, 1); + TORCH_CHECK(list != nullptr, "aclCreateTensorList failed for empty input"); + tensor.release(); + return AclTensorListPtr(list); +} + +at::Tensor GroupedMatmulDirect(const at::Tensor& x, const at::Tensor& weight, const at::Tensor& group_list, int64_t group_list_type, uint64_t stream_handle) { + TORCH_CHECK(x.device().type() == c10::DeviceType::PrivateUse1, "x must be an NPU tensor"); + TORCH_CHECK(weight.device() == x.device() && group_list.device() == x.device(), "x, weight and group_list must be on the same NPU device"); + TORCH_CHECK(x.dim() == 2, "x must have shape [M, K]"); + TORCH_CHECK(weight.dim() == 3 && weight.size(0) == kDirectGroups, "weight must have stored shape [2560, N, K]"); + TORCH_CHECK(group_list.dim() == 1 && group_list.numel() == kDirectGroups, "group_list must contain 2560 entries"); + TORCH_CHECK(group_list.scalar_type() == at::ScalarType::Long, "group_list must be int64"); + TORCH_CHECK(group_list_type == 1, "direct2560 only supports group_list_type=1"); + TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::BFloat16, "direct2560 only supports float16 and bfloat16"); + TORCH_CHECK(weight.scalar_type() == x.scalar_type(), "x and weight dtypes must match"); + TORCH_CHECK(x.size(1) == weight.size(2), "x K must match weight K"); + + auto out = at::empty({x.size(0), weight.size(1)}, x.options()); + auto workspace = at::Tensor(); + + auto x_acl = MakeAclTensor(x); + auto weight_acl = MakeTransposedWeightAclTensor(weight); + auto group_list_acl = MakeAclTensor(group_list); + auto out_acl = MakeAclTensor(out); + auto x_list = MakeAclTensorList(x_acl.get()); + auto weight_list = MakeAclTensorList(weight_acl.get()); + auto out_list = MakeAclTensorList(out_acl.get()); + auto empty_bias = MakeEmptyAclTensorList(ToAclDataType(x.scalar_type())); + auto empty_scale = MakeEmptyAclTensorList(ACL_UINT64); + auto empty_offset = MakeEmptyAclTensorList(ACL_FLOAT); + auto empty_antiquant_scale = MakeEmptyAclTensorList(ToAclDataType(x.scalar_type())); + auto empty_antiquant_offset = MakeEmptyAclTensorList(ToAclDataType(x.scalar_type())); + auto empty_per_token_scale = MakeEmptyAclTensor(ACL_FLOAT); + // aclTensorList owns the tensors passed to aclCreateTensorList. Transfer + // ownership to the lists so the individual guards do not destroy them a + // second time when this function returns. + x_acl.release(); + weight_acl.release(); + out_acl.release(); + + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + const auto status = aclnnDlinferGroupedMatmulDirectGetWorkspaceSize( + x_list.get(), + weight_list.get(), + empty_bias.get(), + empty_scale.get(), + empty_offset.get(), + empty_antiquant_scale.get(), + empty_antiquant_offset.get(), + group_list_acl.get(), + empty_per_token_scale.get(), + 2, + 0, + true, + false, + 0, + group_list_type, + 0, + nullptr, + out_list.get(), + &workspace_size, + &executor); + TORCH_CHECK(status == ACL_SUCCESS, "aclnnDlinferGroupedMatmulDirectGetWorkspaceSize failed, status=", status); + TORCH_CHECK(executor != nullptr, "GroupedMatmul returned a null executor"); + + void* workspace_addr = nullptr; + if (workspace_size > 0) { + workspace = at::empty({static_cast(workspace_size)}, x.options().dtype(at::ScalarType::Byte)); + workspace_addr = workspace.data_ptr(); + } + + TORCH_CHECK(stream_handle != 0, "current NPU stream handle must not be null"); + const auto stream = reinterpret_cast(stream_handle); + const auto execute_status = aclnnDlinferGroupedMatmulDirect(workspace_addr, workspace_size, executor, stream); + TORCH_CHECK(execute_status == ACL_SUCCESS, "aclnnDlinferGroupedMatmulDirect failed, status=", execute_status); + return out; +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { module.def("grouped_matmul", &GroupedMatmulDirect, "DLInfer bundled GroupedMatmul direct2560"); } diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/CMakeLists.txt b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/CMakeLists.txt new file mode 100644 index 00000000..2210718d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/CMakeLists.txt @@ -0,0 +1,43 @@ +cmake_minimum_required(VERSION 3.16.0) +project(opp) + +set(DLINFER_GMM_OPP_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") +set(DLINFER_GMM_OPP_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}" PARENT_SCOPE) +set(DLINFER_GMM_OPP_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}") + +include(cmake/config.cmake) +include(cmake/func.cmake) +include(cmake/intf.cmake) + +if(ENABLE_CROSS_COMPILE) + if(${CMAKE_SYSTEM_PROCESSOR} STREQUAL x86_64) + set(CROSS_COMPILE_PLATFORM aarch64) + else() + set(CROSS_COMPILE_PLATFORM x86_64) + endif() + set(PLATFORM ${CMAKE_SYSTEM_PROCESSOR}) + set(CMAKE_COMPILE_COMPILER_LIBRARY ${ASCEND_CANN_PACKAGE_PATH}/${PLATFORM}-linux/devlib/linux/${CROSS_COMPILE_PLATFORM}/) + set(CMAKE_COMPILE_RUNTIME_LIBRARY ${ASCEND_CANN_PACKAGE_PATH}/${PLATFORM}-linux/devlib/${CROSS_COMPILE_PLATFORM}/) + if(CMAKE_CROSS_LIBRARY_PATH) + set(CMAKE_COMPILE_COMPILER_LIBRARY ${CMAKE_CROSS_LIBRARY_PATH}) + set(CMAKE_COMPILE_RUNTIME_LIBRARY ${CMAKE_CROSS_LIBRARY_PATH}) + endif() + set(CMAKE_SYSTEM_PROCESSOR ${CROSS_COMPILE_PLATFORM}) + set(CMAKE_COMPILE ${CMAKE_CXX_COMPILER}) + set(CMAKE_CXX_COMPILER ${CMAKE_CROSS_PLATFORM_COMPILER}) +else() + set(CMAKE_COMPILE ${CMAKE_CXX_COMPILER}) +endif() + +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/framework) + add_subdirectory(framework) +endif() +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_host) + add_subdirectory(op_host) +endif() +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_kernel) + add_subdirectory(op_kernel) +endif() +if(ENABLE_TEST AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/testcases) + add_subdirectory(testcases) +endif() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/config.cmake b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/config.cmake new file mode 100755 index 00000000..76b02f51 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/config.cmake @@ -0,0 +1,52 @@ +set(CMAKE_CXX_FLAGS_DEBUG "") +set(CMAKE_CXX_FLAGS_RELEASE "") + +if (NOT DEFINED CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Release CACHE STRING "") +endif() +if (NOT DEFINED ASCEND_CANN_PACKAGE_PATH) + if (DEFINED ENV{ASCEND_TOOLKIT_HOME}) + set(ASCEND_CANN_PACKAGE_PATH "$ENV{ASCEND_TOOLKIT_HOME}" CACHE PATH "") + else() + set(ASCEND_CANN_PACKAGE_PATH /usr/local/Ascend/ascend-toolkit/latest CACHE PATH "") + endif() +endif() +if (NOT DEFINED ASCEND_PYTHON_EXECUTABLE) + set(ASCEND_PYTHON_EXECUTABLE python3 CACHE STRING "") +endif() +execute_process( + COMMAND ${ASCEND_PYTHON_EXECUTABLE} -S -c + "import sys; print(f'{sys.base_prefix}/lib/python{sys.version_info.major}.{sys.version_info.minor}/site-packages', end='')" + OUTPUT_VARIABLE ASCEND_SYSTEM_PYTHON_SITE +) +set(ASCEND_COMPILER_PYTHONPATH + "${ASCEND_CANN_PACKAGE_PATH}/python/site-packages:${ASCEND_CANN_PACKAGE_PATH}/opp/built-in/op_impl/ai_core/tbe:${ASCEND_SYSTEM_PYTHON_SITE}" +) +if (NOT DEFINED ASCEND_COMPUTE_UNIT) + set(ASCEND_COMPUTE_UNIT ascend910_93 CACHE STRING "") +endif() +if (NOT DEFINED ENABLE_TEST) + set(ENABLE_TEST FALSE CACHE BOOL "") +endif() +if (ASCEND_COMPUTE_UNIT STREQUAL "ascend610lite") + set(ENABLE_CROSS_COMPILE TRUE CACHE BOOL "") + set(CMAKE_CROSS_PLATFORM_COMPILER "clang-12" CACHE PATH "") +endif() +if (NOT DEFINED ENABLE_CROSS_COMPILE) + set(ENABLE_CROSS_COMPILE FALSE CACHE BOOL "") +endif() +if (NOT DEFINED CMAKE_CROSS_PLATFORM_COMPILER) + set(CMAKE_CROSS_PLATFORM_COMPILER "/usr/bin/aarch64-linux-gnu-g++" CACHE PATH "") +endif() +set(ASCEND_TENSOR_COMPILER_PATH ${ASCEND_CANN_PACKAGE_PATH}/compiler) +set(ASCEND_CCEC_COMPILER_PATH ${ASCEND_TENSOR_COMPILER_PATH}/ccec_compiler/bin) +set(ASCEND_AUTOGEN_PATH ${DLINFER_GMM_OPP_BINARY_DIR}/autogen) +set(ASCEND_AUTOGEN_GROUPPROTO_PATH ${DLINFER_GMM_OPP_BINARY_DIR}/autogen/group_proto) +file(MAKE_DIRECTORY ${ASCEND_AUTOGEN_PATH} ${ASCEND_AUTOGEN_GROUPPROTO_PATH}) +set(CUSTOM_COMPILE_OPTIONS "custom_compile_options.ini") +set(CUSTOM_OPC_OPTIONS "custom_opc_options.ini") +execute_process(COMMAND rm -rf ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} + COMMAND rm -rf ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + COMMAND touch ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} + COMMAND touch ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + ) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/func.cmake b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/func.cmake new file mode 100755 index 00000000..a1230cd8 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/func.cmake @@ -0,0 +1,165 @@ + +function(get_system_info SYSTEM_INFO) + if (UNIX) + execute_process(COMMAND grep -i ^id= /etc/os-release OUTPUT_VARIABLE TEMP) + string(REGEX REPLACE "\n|id=|ID=|\"" "" SYSTEM_NAME ${TEMP}) + set(${SYSTEM_INFO} ${SYSTEM_NAME}_${CMAKE_SYSTEM_PROCESSOR} PARENT_SCOPE) + elseif (WIN32) + message(STATUS "System is Windows. Only for pre-build.") + else () + message(FATAL_ERROR "${CMAKE_SYSTEM_NAME} not support.") + endif () +endfunction() + +function(opbuild) + message(STATUS "Opbuild generating sources") + cmake_parse_arguments(OPBUILD "" "OUT_DIR;PROJECT_NAME;ACCESS_PREFIX;ENABLE_SOURCE" "OPS_SRC" ${ARGN}) + execute_process(COMMAND ${CMAKE_COMPILE} -g -fPIC -shared -std=c++17 ${OPBUILD_OPS_SRC} -D_GLIBCXX_USE_CXX11_ABI=0 + -I ${ASCEND_CANN_PACKAGE_PATH}/include + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include/aclnn + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include/base + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc/op_common + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc/base + -I ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include/tiling + -I ${DLINFER_GMM_OPP_SOURCE_DIR}/common/include + -I ${DLINFER_GMM_OPP_SOURCE_DIR}/op_host + -I ${DLINFER_GMM_OPP_SOURCE_DIR}/op_host/op_tiling + -I ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel + -L ${ASCEND_CANN_PACKAGE_PATH}/lib64 -lexe_graph -lregister -ltiling_api + -o ${OPBUILD_OUT_DIR}/libascend_all_ops.so + RESULT_VARIABLE EXEC_RESULT + OUTPUT_VARIABLE EXEC_INFO + ERROR_VARIABLE EXEC_ERROR + ) + if (${EXEC_RESULT}) + message("build ops lib info: ${EXEC_INFO}") + message("build ops lib error: ${EXEC_ERROR}") + message(FATAL_ERROR "opbuild run failed!") + endif() + set(proj_env "") + set(prefix_env "") + if (NOT "${OPBUILD_PROJECT_NAME}x" STREQUAL "x") + set(proj_env "OPS_PROJECT_NAME=${OPBUILD_PROJECT_NAME}") + endif() + if (NOT "${OPBUILD_ACCESS_PREFIX}x" STREQUAL "x") + set(prefix_env "OPS_DIRECT_ACCESS_PREFIX=${OPBUILD_ACCESS_PREFIX}") + endif() + + set(ENV{OPS_PRODUCT_NAME} ${ASCEND_COMPUTE_UNIT}) + set(ENV{ENABLE_SOURCE_PACAKGE} ${OPBUILD_ENABLE_SOURCE}) + execute_process(COMMAND ${proj_env} ${prefix_env} ${ASCEND_CANN_PACKAGE_PATH}/toolkit/tools/opbuild/op_build + ${OPBUILD_OUT_DIR}/libascend_all_ops.so ${OPBUILD_OUT_DIR} + RESULT_VARIABLE EXEC_RESULT + OUTPUT_VARIABLE EXEC_INFO + ERROR_VARIABLE EXEC_ERROR + ) + unset(ENV{OPS_PRODUCT_NAME}) + unset(ENV{ENABLE_SOURCE_PACAKGE}) + if (${EXEC_RESULT}) + message("opbuild ops info: ${EXEC_INFO}") + message(FATAL_ERROR "opbuild ops error: ${EXEC_ERROR}") + endif() + message(STATUS "Opbuild generating sources - done") +endfunction() + +function(build_optiling_for_compile) + message(STATUS "building optiling so for compile") + cmake_parse_arguments(TILING_COMPILE "" "OUT_DIR" "OPS_SRC" ${ARGN}) + file(MAKE_DIRECTORY ${TILING_COMPILE_OUT_DIR}/op_impl/ai_core/tbe/op_tiling/) + execute_process(COMMAND ${CMAKE_COMPILE} -fPIC -shared -std=c++11 ${TILING_COMPILE_OPS_SRC} -D_GLIBCXX_USE_CXX11_ABI=0 + -I ${ASCEND_CANN_PACKAGE_PATH}/include -I ${CMAKE_CURRENT_SOURCE_DIR}/../op_host -L ${ASCEND_CANN_PACKAGE_PATH}/lib64 + -DOP_TILING_LIB -fvisibility=hidden -lexe_graph -lregister -Wl, --whole-archive -ltiling_api -lrt2_registry -Wl, --no-whole-archive + -o ${TILING_COMPILE_OUT_DIR}/op_impl/ai_core/tbe/op_tiling/liboptiling.so + RESULT_VARIABLE EXEC_RESULT + OUTPUT_VARIABLE EXEC_INFO + ERROR_VARIABLE EXEC_ERROR + ) + if (${EXEC_RESULT}) + message("build optiling lib for compile info: ${EXEC_INFO}") + message("build optiling lib for compile error: ${EXEC_ERROR}") + message(FATAL_ERROR "optiling lib for compile failed") + endif() + message(STATUS "building optiling so for compile - done") +endfunction() + +function(add_ops_compile_options OP_TYPE) + cmake_parse_arguments(OP_COMPILE "" "OP_TYPE" "COMPUTE_UNIT;OPTIONS" ${ARGN}) + execute_process(COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${DLINFER_GMM_OPP_SOURCE_DIR}/cmake/util/ascendc_gen_options.py + ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} ${OP_TYPE} ${OP_COMPILE_COMPUTE_UNIT} + ${OP_COMPILE_OPTIONS} + RESULT_VARIABLE EXEC_RESULT + OUTPUT_VARIABLE EXEC_INFO + ERROR_VARIABLE EXEC_ERROR) + if (${EXEC_RESULT}) + message("add ops compile options info: ${EXEC_INFO}") + message("add ops compile options error: ${EXEC_ERROR}") + message(FATAL_ERROR "add ops compile options failed!") + endif() +endfunction() + +function(add_kernel_compile op_type src) + cmake_parse_arguments(BINCMP "" "OPS_INFO;OUT_DIR;TILING_LIB" "COMPUTE_UNIT;OPTIONS;CONFIGS" ${ARGN}) + if (NOT DEFINED BINCMP_COMPUTE_UNIT) + set(BINCMP_COMPUTE_UNIT ${ASCEND_COMPUTE_UNIT}) + endif() + if (NOT DEFINED BINCMP_OUT_DIR) + set(BINCMP_OUT_DIR ${CMAKE_CURRENT_BINARY_DIR}) + endif() + set(BINCMP_OUT_DIR ${BINCMP_OUT_DIR}/kernel) + if (NOT DEFINED BINCMP_TILING_LIB) + set(BINCMP_TILING_LIB $) + endif() + if (NOT TARGET op_kernel_pack) + add_custom_target(op_kernel_pack + COMMAND ${CMAKE_COMMAND} -E env + "PYTHONDONTWRITEBYTECODE=1" + "PYTHONPATH=${ASCEND_COMPILER_PYTHONPATH}" + ${ASCEND_PYTHON_EXECUTABLE} -S + ${DLINFER_GMM_OPP_SOURCE_DIR}/cmake/util/ascendc_pack_kernel.py + --input-path=${BINCMP_OUT_DIR} + --output-path=${BINCMP_OUT_DIR}/library) + add_library(ascend_kernels INTERFACE) + target_link_libraries(ascend_kernels INTERFACE kernels) + target_link_directories(ascend_kernels INTERFACE ${BINCMP_OUT_DIR}/library) + target_include_directories(ascend_kernels INTERFACE ${BINCMP_OUT_DIR}/library) + add_dependencies(ascend_kernels op_kernel_pack) + endif() + + # add Environment Variable Configurations of ccache + set(_ASCENDC_ENV_VAR) + if(${CMAKE_CXX_COMPILER_LAUNCHER} MATCHES "ccache$") + list(APPEND _ASCENDC_ENV_VAR export ASCENDC_CCACHE_EXECUTABLE=${CMAKE_CXX_COMPILER_LAUNCHER} &&) + endif() + + foreach(compute_unit ${BINCMP_COMPUTE_UNIT}) + if (NOT DEFINED BINCMP_OPS_INFO) + set(BINCMP_OPS_INFO ${ASCEND_AUTOGEN_PATH}/aic-${compute_unit}-ops-info.ini) + endif() + add_custom_target(${op_type}_${compute_unit} + COMMAND ${_ASCENDC_ENV_VAR} ${CMAKE_COMMAND} -E env + "PYTHONDONTWRITEBYTECODE=1" + ${ASCEND_PYTHON_EXECUTABLE} ${DLINFER_GMM_OPP_SOURCE_DIR}/cmake/util/ascendc_compile_kernel.py + --op-name=${op_type} + --src-file=${src} + --compute-unit=${compute_unit} + --compile-options=\"${BINCMP_OPTIONS}\" + --debug-config=\"${BINCMP_CONFIGS}\" + --config-ini=${BINCMP_OPS_INFO} + --tiling-lib=${BINCMP_TILING_LIB} + --output-path=${BINCMP_OUT_DIR}) + add_dependencies(${op_type}_${compute_unit} cust_optiling) + add_dependencies(op_kernel_pack ${op_type}_${compute_unit}) + endforeach() +endfunction() + +function(add_cross_compile_target) + cmake_parse_arguments(CROSSMP "" "TARGET;OUT_DIR;INSTALL_DIR" "" ${ARGN}) + add_custom_target(${CROSSMP_TARGET} ALL + DEPENDS ${CROSSMP_OUT_DIR} + ) + install(DIRECTORY ${CROSSMP_OUT_DIR} + DESTINATION ${CROSSMP_INSTALL_DIR} + ) +endfunction() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/intf.cmake b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/intf.cmake new file mode 100755 index 00000000..2088fbd2 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/intf.cmake @@ -0,0 +1,38 @@ + +add_library(intf_pub INTERFACE) +target_compile_options(intf_pub INTERFACE + -fPIC + -fvisibility=hidden + -fvisibility-inlines-hidden + $<$:-O2> + $<$:-O0 -g> + $<$:-std=c++11> + $<$,$>:-ftrapv -fstack-check> + $<$:-pthread -Wfloat-equal -Wshadow -Wformat=2 -Wno-deprecated -Wextra> + $,-fstack-protector-strong,-fstack-protector-all> +) +target_compile_definitions(intf_pub INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:_FORTIFY_SOURCE=2> +) +target_include_directories(intf_pub INTERFACE ${ASCEND_CANN_PACKAGE_PATH}/include + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include/aclnn + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/include/base + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc/op_common + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/pkg_inc/base + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/asc/include/tiling + ${DLINFER_GMM_OPP_SOURCE_DIR}/common/include + ${DLINFER_GMM_OPP_SOURCE_DIR}/op_host + ${DLINFER_GMM_OPP_SOURCE_DIR}/op_host/op_tiling + ${CMAKE_CURRENT_SOURCE_DIR}/op_kernel +) +target_link_options(intf_pub INTERFACE + $<$,EXECUTABLE>:-pie> + $<$:-s> + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack +) +target_link_directories(intf_pub INTERFACE ${ASCEND_CANN_PACKAGE_PATH}/lib64) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_bin_param_build.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_bin_param_build.py new file mode 100755 index 00000000..3a8f26d4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_bin_param_build.py @@ -0,0 +1,548 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Created on Feb 28 20:56:45 2020 +Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved. +""" + +import argparse +import sys +import os +import json +import hashlib +import re +import copy +from collections import defaultdict +from typing import Dict, List, Set, Tuple, NamedTuple + +import const_var +import opdesc_parser + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + + +class ParamInfo(NamedTuple): + dtype_list: list + format_list: list + dtype_for_bin_list: dict + format_for_bin_list: dict + + +class BinParamBuilder(opdesc_parser.OpDesc): + def __init__(self: any, op_type: str): + super().__init__(op_type) + self.soc = '' + self.out_path = '' + self.tiling_keys = set() + self.op_debug_config = '' + self.op_super_config = [] + + def set_soc_version(self: any, soc: str): + self.soc = soc + + def set_out_path(self: any, out_path: str): + self.out_path = out_path + + def set_tiling_key(self: any, tiling_key_info: Set): + if tiling_key_info: + self.tiling_keys.update(tiling_key_info) + + def set_op_debug_config(self: any, op_debug_config: str): + if op_debug_config: + self.op_debug_config = op_debug_config + + def set_op_super_config(self: any, op_super_config: str): + if op_super_config: + self.op_super_config = op_super_config + + def get_full_list(self: any): + dtype_list = [] + for dtype_in in self.input_dtype: + dtype_list.append(dtype_in.split(',')) + for dtype_out in self.output_dtype: + dtype_list.append(dtype_out.split(',')) + + format_list = [] + for fmt_in in self.input_fmt: + format_list.append(fmt_in.split(',')) + for fmt_out in self.output_fmt: + format_list.append(fmt_out.split(',')) + + dtype_for_bin_list = [[] for _ in range(len(self.input_dtype) + len(self.output_dtype))] + format_for_bin_list = copy.deepcopy(dtype_for_bin_list) + + for key, value in self.input_dtype_for_bin.items(): + dtype_for_bin_list[key] = value.split(',') + for key, value in self.output_dtype_for_bin.items(): + dtype_for_bin_list[key + len(self.input_dtype)] = value.split(',') + for key, value in self.input_fmt_for_bin.items(): + format_for_bin_list[key] = value.split(',') + for key, value in self.output_fmt_for_bin.items(): + format_for_bin_list[key + len(self.input_dtype)] = value.split(',') + + return ParamInfo(dtype_list, format_list, dtype_for_bin_list, format_for_bin_list) + + + def gen_bin_cprs_list(self: any, param_info: ParamInfo): + combine_dict = {} + origin_combine_dict = {} + for cob_idx in range(0, len(self.input_dtype[0].split(','))): + origin_combine = "" + combine = "" + for param_idx in range(0, len(self.input_dtype) + len(self.output_dtype)): + if (param_info.dtype_for_bin_list[param_idx]): + combine += param_info.dtype_for_bin_list[param_idx][cob_idx] + else: + combine += param_info.dtype_list[param_idx][cob_idx] + origin_combine += param_info.dtype_list[param_idx][cob_idx] + if (param_info.format_for_bin_list[param_idx]): + combine += param_info.format_for_bin_list[param_idx][cob_idx] + else: + combine += param_info.format_list[param_idx][cob_idx] + origin_combine += param_info.format_list[param_idx][cob_idx] + if (combine not in combine_dict): + combine_dict[combine] = [] + combine_dict[combine].append(cob_idx) + origin_combine_dict[origin_combine] = cob_idx + for key, value in combine_dict.items(): + if (key not in origin_combine_dict): + print(f"WARNING: ForBinQuery {key} not in origin combine") + self.bin_save_list += value + continue + if len(value) == 1 and value[0] == origin_combine_dict[key]: + self.bin_save_list += value + continue + self.bin_cprs_head.append(origin_combine_dict[key]) + self.bin_cprs_list.append(value) + for index, sub_list in enumerate(self.bin_cprs_list): + if self.bin_cprs_head[index] not in self.bin_save_list: + continue + sub_list.append(self.bin_cprs_head[index]) + self.bin_save_list += self.bin_cprs_head + + + def gen_for_bin_list(self: any, param_info: ParamInfo): + combine_size = len(self.input_dtype[0].split(',')) + input_size = len(self.input_dtype) + output_size = len(self.output_dtype) + + self.input_dtype_for_bin_list = [[] for _ in range(input_size)] + self.output_dtype_for_bin_list = [[] for _ in range(output_size)] + for i in range(0, input_size): + self.input_dtype_for_bin_list[i] = [[] for _ in range(combine_size)] + for i in range(0, output_size): + self.output_dtype_for_bin_list[i] = [[] for _ in range(combine_size)] + self.input_fmt_for_bin_list = copy.deepcopy(self.input_dtype_for_bin_list) + self.output_fmt_for_bin_list = copy.deepcopy(self.output_dtype_for_bin_list) + + for index, sub_list in enumerate(self.bin_cprs_list): + head_idx = self.bin_cprs_head[index] + for cmb_idx in sub_list: + for i in range(0, input_size): + self.input_dtype_for_bin_list[i][head_idx].append(param_info.dtype_list[i][cmb_idx]) + self.input_fmt_for_bin_list[i][head_idx].append(param_info.format_list[i][cmb_idx]) + for i in range(0, output_size): + self.output_dtype_for_bin_list[i][head_idx].append(param_info.dtype_list[i + input_size][cmb_idx]) + self.output_fmt_for_bin_list[i][head_idx].append(param_info.format_list[i + input_size][cmb_idx]) + + + def rm_cprs_cmb(self: any, dtype_list, format_list, input_size, output_size): + for i in range(0, input_size): + self.input_dtype_for_bin_list[i] = [ + element for index, element in enumerate(self.input_dtype_for_bin_list[i]) + if index in self.bin_save_list + ] + self.input_fmt_for_bin_list[i] = [ + element for index, element in enumerate(self.input_fmt_for_bin_list[i]) + if index in self.bin_save_list + ] + new_dtype_list = [ + element for index, element in enumerate(dtype_list[i]) + if index in self.bin_save_list + ] + new_dtype_str = "" + for dtype in new_dtype_list: + new_dtype_str += f"{dtype}," + self.input_dtype[i] = new_dtype_str[:-1] + new_format_list = [ + element for index, element in enumerate(format_list[i]) + if index in self.bin_save_list + ] + new_format_str = "" + for fmt in new_format_list: + new_format_str += f"{fmt}," + self.input_fmt[i] = new_format_str[:-1] + for i in range(0, output_size): + self.output_dtype_for_bin_list[i] = [ + element for index, element in enumerate(self.output_dtype_for_bin_list[i]) + if index in self.bin_save_list + ] + self.output_fmt_for_bin_list[i] = [ + element for index, element in enumerate(self.output_fmt_for_bin_list[i]) + if index in self.bin_save_list + ] + new_dtype_list = [ + element for index, element in enumerate(dtype_list[i + input_size]) + if index in self.bin_save_list + ] + new_dtype_str = "" + for dtype in new_dtype_list: + new_dtype_str += f"{dtype}," + self.output_dtype[i] = new_dtype_str[:-1] + new_format_list = [ + element for index, element in enumerate(format_list[i + input_size]) + if index in self.bin_save_list + ] + new_format_str = "" + for fmt in new_format_list: + new_format_str += f"{fmt}," + self.output_fmt[i] = new_format_str[:-1] + + + def is_set_for_bin_query(self: any): + return any([ + self.input_dtype_for_bin, + self.output_dtype_for_bin, + self.input_fmt_for_bin, + self.output_fmt_for_bin, + ]) + + + def for_bin_list_match(self: any): + if not self.is_set_for_bin_query(): + return + input_size = len(self.input_dtype) + output_size = len(self.output_dtype) + param_info = self.get_full_list() + self.gen_bin_cprs_list(param_info) + self.gen_for_bin_list(param_info) + if len(self.bin_save_list) == len(self.input_dtype[0].split(',')): + print(f'WARNING: ForBinQuery can not compress number of bin file with this set, please check!!.') + return + self.rm_cprs_cmb(param_info.dtype_list, param_info.format_list, input_size, output_size) + + + def gen_input_json(self: any, auto_gen_path: str): + key_map = {} + self.for_bin_list_match() + if len(self.input_dtype) == 0: + count = len(self.output_dtype[0].split(',')) + else: + count = len(self.input_dtype[0].split(',')) + if count == 0: + raise RuntimeError(f'Op {self.op_type} must have at least one input or output') + required_parameters = set() + index_value = -1 + + for i in range(0, count): + inputs = [] + outputs = [] + attrs = [] + required_parameter = [] + op_node = {} + + for idx in range(0, len(self.input_name)): + idtypes = self.input_dtype[idx].split(',') + ifmts = self.input_fmt[idx].split(',') + itype = self.input_type[idx] + para = {} + para['name'] = self.input_name[idx][:-5] + para['index'] = idx + para['dtype'] = idtypes[i] + if self.is_set_for_bin_query() and self.input_dtype_for_bin_list[idx][i]: + para['dtypeForBinQuery'] = self.input_dtype_for_bin_list[idx][i] + para['format'] = ifmts[i] + if self.is_set_for_bin_query() and self.input_fmt_for_bin_list[idx][i]: + para['formatForBinQuery'] = self.input_fmt_for_bin_list[idx][i] + para['paramType'] = itype + para['shape'] = [-2] + para['format_match_mode'] = 'FormatAgnostic' + + input_parameter_key = (idtypes[i], ifmts[i]) + if itype == 'dynamic': + inputs.append([para]) + required_parameter.append(input_parameter_key) + elif itype == 'required': + inputs.append(para) + required_parameter.append(input_parameter_key) + else: + inputs.append(para) + + for idx in range(0, len(self.output_name)): + odtypes = self.output_dtype[idx].split(',') + ofmts = self.output_fmt[idx].split(',') + otype = self.output_type[idx] + para = {} + para['name'] = self.output_name[idx][:-5] + para['index'] = idx + para['dtype'] = odtypes[i] + if self.is_set_for_bin_query() and self.output_dtype_for_bin_list[idx][i]: + para['dtypeForBinQuery'] = self.output_dtype_for_bin_list[idx][i] + para['format'] = ofmts[i] + if self.is_set_for_bin_query() and self.output_fmt_for_bin_list[idx][i]: + para['formatForBinQuery'] = self.output_fmt_for_bin_list[idx][i] + para['paramType'] = otype + para['shape'] = [-2] + para['format_match_mode'] = 'FormatAgnostic' + output_parameter_key = (odtypes[i], ofmts[i]) + if otype == 'dynamic': + outputs.append([para]) + required_parameter.append(output_parameter_key) + elif otype == 'required': + outputs.append(para) + required_parameter.append(output_parameter_key) + else: + outputs.append(para) + + for attr in self.attr_list: + att = {} + att['name'] = attr + atype = self.attr_val.get(attr).get('type').lower() + att['dtype'] = atype + att['value'] = const_var.ATTR_DEF_VAL.get(atype) + attrs.append(att) + + required_parameter_tuple = tuple(required_parameter) + if required_parameter_tuple in required_parameters: + continue + else: + required_parameters.add(required_parameter_tuple) + index_value +=1 + + op_node['bin_filename'] = '' + op_node['inputs'] = inputs + op_node['outputs'] = outputs + if len(attrs) > 0: + op_node['attrs'] = attrs + + param = {} + param['op_type'] = self.op_type + param['op_list'] = [op_node] + objstr = json.dumps(param, indent=' ') + md5sum = hashlib.md5(objstr.encode('utf-8')).hexdigest() + while key_map.get(md5sum) is not None: + objstr += '1' + md5sum = hashlib.md5(objstr.encode('utf-8')).hexdigest() + key_map[md5sum] = md5sum + bin_file = self.op_type + '_' + md5sum + op_node['bin_filename'] = bin_file + param_file = os.path.join(self.out_path, bin_file + '_param.json') + param_file = os.path.realpath(param_file) + + self._write_build_json(param_file, param) + self._write_build_cmd(param_file, bin_file, index_value, auto_gen_path) + if self.op_super_config: + bin_file += "_relocatable" + op_node['bin_filename'] = bin_file + param_file = os.path.join(self.out_path, bin_file + '_param.json') + param_file = os.path.realpath(param_file) + self._write_build_json(param_file, param) + index_value += 1 + self._write_build_cmd(param_file, bin_file, index_value, auto_gen_path, True) + + def _write_build_json(self: any, param_file: str, param): + with os.fdopen(os.open(param_file, const_var.WFLAGS, const_var.WMODES), 'w') as fd: + json.dump(param, fd, indent=' ') + + def _generate_check_result(self: any, enable_tiling_keys: bool, bin_file: str): + check_result = "" + if enable_tiling_keys is False: + check_result += "echo \"${res}\"\n" + check_result += const_var.CHK_CMD.format(res_file=bin_file + '.json') + check_result += const_var.CHK_CMD.format(res_file=bin_file + '.o') + else: + check_result += "if [ $? -eq 1 ]; then\n" + check_result += " if echo \"${res}\" | \ +grep -q \"None of the given tiling keys are in the supported list\"; then\n" + check_result += " echo \"${res}\"\n" + check_result += " else\n" + check_result += " echo \"${res}\"\n" + check_result += " exit 1\n" + check_result += " fi\n" + check_result += "else\n" + check_result += "echo \"${res}\"\n" + check_result += const_var.CHK_CMD.format(res_file=bin_file + '.json') + check_result += const_var.CHK_CMD.format(res_file=bin_file + '.o') + check_result += "fi\n" + return check_result + + def _write_build_cmd(self: any, param_file: str, bin_file: str, index: int, auto_gen_path: str, super_mode=False): + hard_soc = const_var.conv_soc_ver(self.soc) + if not hard_soc: + hard_soc = self.soc.capitalize() + name_com = [self.op_type, self.op_file, str(index)] + compile_file = os.path.join(self.out_path, '-'.join(name_com) + '.sh') + compile_file = os.path.realpath(compile_file) + + bin_cmd_str = 'res=$(opc $1 --main_func={fun} --input_param={param} --soc_version={soc} \ + --output=$2 --impl_mode={impl} --simplified_key_mode=0 --op_mode=dynamic ' + + build_cmd_var = "#!/bin/bash\n" + build_cmd_var += f'echo "[{self.soc}] Generating {bin_file} ..."\n' + plog_level = os.environ.get("ASCEND_GLOBAL_LOG_LEVEL") + plog_stdout = os.environ.get("ASCEND_SLOG_PRINT_TO_STDOUT") + if plog_level is None: + build_cmd_var += const_var.SET_PLOG_LEVEL_ERROR + if plog_stdout is None: + build_cmd_var += const_var.SET_PLOG_STDOUT + build_cmd_var += const_var.SRC_ENV + if hard_soc == "Ascend610Lite": + build_cmd_var += f'export ASCEND_CUSTOM_OPP_PATH={auto_gen_path}:$ASCEND_CUSTOM_OPP_PATH \n' + build_cmd_var += bin_cmd_str.format(fun=self.op_intf, soc=hard_soc, param=param_file, + impl='high_performance,optional') + enable_tiling_keys = False + if self.tiling_keys: + tiling_keys_list = sorted(list(self.tiling_keys)) + tiling_key_str = ','.join([str(_key) for _key in tiling_keys_list]) + build_cmd_var += f' --tiling_key="{tiling_key_str}"' + enable_tiling_keys = True + + if self.op_debug_config: + op_debug_str = ','.join([str(_key) for _key in list(self.op_debug_config)]) + build_cmd_var += f' --op_debug_config={op_debug_str}' + + if super_mode and self.op_super_config: + op_super_config_str = ' '.join([str(_key) for _key in list(self.op_super_config)]) + build_cmd_var += f' {op_super_config_str}' + + build_cmd_var += ")\n" + build_cmd_var += "\n" + + check_result = self._generate_check_result(enable_tiling_keys, bin_file) + build_cmd_var += check_result + build_cmd_var += f'echo "[{self.soc}] Generating {bin_file} Done"\n' + + with os.fdopen(os.open(compile_file, const_var.WFLAGS, const_var.WMODES), 'w') as fd: + fd.write(build_cmd_var) + + +def get_tiling_keys(tiling_keys: str) -> Set: + all_tiling_keys = set() + if not tiling_keys: + return all_tiling_keys + + tiling_key_list = tiling_keys.split(';') + for tiling_key_value in tiling_key_list: + pattern = r"(? int(end): + continue + for i in range(int(start), int(end) + 1): + all_tiling_keys.add(i) + elif tiling_key_value.isdigit(): + all_tiling_keys.add(int(tiling_key_value)) + return all_tiling_keys + + +def trans_soc_verion(soc_ver: str): + low_soc_ver = soc_ver.lower() + if low_soc_ver not in opdesc_parser.SOC_TO_SHORT_SOC_MAP: + return low_soc_ver + return opdesc_parser.SOC_TO_SHORT_SOC_MAP[low_soc_ver] + + +def parse_op_debug_confg(opc_config_file: str, soc: str) -> Dict: + tiling_key_info = defaultdict(set) + op_debug_config = defaultdict(set) + if not opc_config_file: + return tiling_key_info, op_debug_config + + if not os.path.exists(opc_config_file): + return tiling_key_info, op_debug_config + + with open(opc_config_file, 'r') as file: + contents = file.readlines() + + for _content in contents: + content = _content.strip() + opc_configs = content.split('@') + if len(opc_configs) < 3: + continue + + op_type = opc_configs[0] + if not op_type: + continue + + compute_unit = opc_configs[1] + if compute_unit: + compute_unit_list = compute_unit.split(';') + soc_lists = [] + for soc_ver in compute_unit_list: + short_soc_ver = trans_soc_verion(soc_ver) + soc_lists.append(short_soc_ver) + if soc not in soc_lists: + continue + + for options in opc_configs[2:]: + if "--tiling_key" in options: + format_tiling_keys = get_tiling_keys(options.split('=')[1]) + if format_tiling_keys: + tiling_key_info[op_type].update(format_tiling_keys) + if "--op_debug_config" in options: + first_index = options.find('=') + if first_index != -1: + debug_config = options[first_index + 1:] + else: + debug_config = "" + + format_debug_config = set(debug_config.split(';')) + for _config in format_debug_config: + op_debug_config[op_type].add(_config) + return tiling_key_info, op_debug_config + + +def gen_bin_param_file(cfgfile: str, out_dir: str, soc: str, + opc_config_file: str = '', ops: list = None): + if not os.path.exists(cfgfile): + print(f'INFO: {cfgfile} does not exists in this project, skip generating compile commands.') + return + + debug_config = defaultdict(set) + super_config = defaultdict(set) + + op_descs = opdesc_parser.get_op_desc(cfgfile, [], [], BinParamBuilder, ops) + tiling_key_info, op_debug_config = parse_op_debug_confg(opc_config_file, soc) + for _op_type, _op_option in op_debug_config.items(): + for _option in _op_option: + if (_option.startswith("--op_relocatable_kernel_binary") + or _option.startswith("--op_super_kernel_options")): + super_config[_op_type].add(_option) + else: + debug_config[_op_type].add(_option) + + auto_gen_path_dir = os.path.dirname(cfgfile) + all_soc_key = "ALL" + for op_desc in op_descs: + op_desc.set_soc_version(soc) + op_desc.set_out_path(out_dir) + if op_desc.op_type in debug_config: + op_desc.set_op_debug_config(debug_config[op_desc.op_type]) + if all_soc_key in debug_config: + op_desc.set_op_debug_config(debug_config[all_soc_key]) + if op_desc.op_type in super_config: + op_desc.set_op_super_config(super_config[op_desc.op_type]) + if op_desc.op_type in tiling_key_info: + op_desc.set_tiling_key(tiling_key_info[op_desc.op_type]) + if all_soc_key in tiling_key_info: + op_desc.set_tiling_key(tiling_key_info[all_soc_key]) + op_desc.gen_input_json(auto_gen_path_dir) + + +def parse_args(argv): + """Command line parameter parsing""" + parser = argparse.ArgumentParser() + parser.add_argument('argv', nargs='+') + parser.add_argument('--opc-config-file', nargs='?', const='', default='') + return parser.parse_args(argv) + + +if __name__ == '__main__': + args = parse_args(sys.argv) + if len(args.argv) <= 3: + raise RuntimeError('arguments must greater than 3') + gen_bin_param_file(args.argv[1], + args.argv[2], + args.argv[3], + opc_config_file=args.opc_config_file) \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_compile_kernel.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_compile_kernel.py new file mode 100755 index 00000000..49a04f74 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_compile_kernel.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. +""" + +import sys +import os +import subprocess +import time +import glob +import shutil +import argparse +import const_var +import ascendc_impl_build +import ascendc_bin_param_build +import ascendc_op_info + + +class CompileKernel: + def __init__(self: any, args: any): + self.op_type = args.op_name + self.op_cpp_file = os.path.realpath(args.src_file) + self.op_soc_ver = args.compute_unit + self.compile_options = args.compile_options + self.op_debug_config = args.debug_config + self.op_cfg_ini = os.path.realpath(args.config_ini) + self.op_tiling = os.path.realpath(args.tiling_lib) + self.op_output = os.path.realpath(args.output_path) + self.op_impl_py = None + self.compile_sh = [] + self.working_dir = os.path.join( + os.getcwd(), + self.op_type + "_" + self.op_soc_ver, + ) + self.build_opp_path = os.path.join(self.working_dir, "customize") + os.makedirs(self.working_dir) + os.makedirs(self.op_output, exist_ok=True) + if args.dynamic_dir is not None and args.dynamic_dir != "": + self.dynamic_dir = os.path.realpath(args.dynamic_dir) + else: + self.dynamic_dir = None + if args.json_file is not None and args.json_file != "": + self.json_file = args.json_file + else: + self.json_file = None + + + def clean(self: any): + if 'dump_cce' not in self.op_debug_config: + shutil.rmtree(self.working_dir) + return + + def ascendc_gen_impl(self: any): + rep_cfg = {} + rep_cfg[const_var.REPLAY_BATCH] = "" + rep_cfg[const_var.REPLAY_ITERATE] = "" + cfg_dir = {} + cfg_dir[const_var.CFG_IMPL_DIR] = os.path.dirname(self.op_cpp_file) + cfg_dir[const_var.CFG_OUT_DIR] = os.path.join(self.working_dir, "dynamic") + os.makedirs(os.path.join(self.working_dir, "dynamic"), exist_ok=True) + cfg_dir[const_var.AUTO_GEN_DIR] = os.path.dirname(self.op_cfg_ini) + ascendc_impl_build.write_scripts( + self.op_cfg_ini, rep_cfg, cfg_dir, [self.op_type], self.compile_options + ) + py_files = glob.glob(os.path.join(self.working_dir, "dynamic", "*.py")) + if py_files is None or len(py_files) != 1: + self.clean() + raise RuntimeError("compile py file {} generated error!".format(py_files)) + self.op_impl_py = os.path.join( + self.working_dir, "dynamic", self.op_type + ".py" + ) + if self.dynamic_dir is not None: + shutil.copy(py_files[0], self.dynamic_dir) + os.rename(py_files[0], self.op_impl_py) + if not os.path.exists(self.op_impl_py): + self.clean() + raise RuntimeError( + "compile py file {} not generated!".format(self.op_impl_py) + ) + + def ascendc_gen_param(self: any): + bin_param_path = os.path.join(self.working_dir, "bin_param") + os.makedirs(bin_param_path) + base_dir = os.path.dirname(self.op_cfg_ini) + opc_config_file = os.path.join(base_dir, "custom_opc_options.ini") + ascendc_bin_param_build.gen_bin_param_file( + self.op_cfg_ini, bin_param_path, self.op_soc_ver, opc_config_file, [self.op_type] + ) + tiling_key_info, op_debug_config = ascendc_bin_param_build.parse_op_debug_confg(opc_config_file, self.op_type) + if self.op_type in op_debug_config: + self.op_debug_config = op_debug_config[self.op_type] + if "ALL" in op_debug_config: + self.op_debug_config = op_debug_config["ALL"] + bin_param_files = glob.glob(os.path.join(bin_param_path, "*.json")) + if bin_param_files is None or len(bin_param_files) <= 0: + self.clean() + raise RuntimeError("compile binary param json file not generated!") + self.compile_sh = glob.glob(os.path.join(bin_param_path, "*.sh")) + if self.compile_sh is None or len(self.compile_sh) != len(bin_param_files): + self.clean() + raise RuntimeError("compile binary shell file not generated!") + + def ascendc_put_tiling(self: any): + tiling_path = os.path.join( + self.build_opp_path, "op_impl", "ai_core", "tbe", "op_tiling" + ) + os.makedirs(tiling_path) + tiling_so = os.path.join(tiling_path, "liboptiling.so") + os.symlink(self.op_tiling, tiling_so) + if not os.path.exists(tiling_so): + self.clean() + raise RuntimeError("prepare tiling lib {} link failed!".format(tiling_so)) + + def ascendc_put_json(self: any): + if self.json_file is not None: + json_file_dir = os.path.join(self.build_opp_path, + "op_impl", + "ai_core", + "tbe", + "config", + self.op_soc_ver) + os.makedirs(json_file_dir) + shutil.copy(self.json_file, json_file_dir) + build_json_file = os.path.join(json_file_dir, "aic-{}-ops-info.json".format(self.op_soc_ver)) + if not os.path.exists(build_json_file): + self.clean() + raise RuntimeError("prepare json file aic-{}-ops-info.json failed!".format(self.op_soc_ver)) + + def ascendc_build(self: any): + op_info = ascendc_op_info.OpInfo(self.op_type, self.op_cfg_ini) + op_file = op_info.get_op_file() + op_bin_dir = os.path.join(self.op_output, self.op_soc_ver, op_file) + os.makedirs(op_bin_dir, exist_ok=True) + all_tar = [] + sub_cmd = [] + index = 0 + for sh in self.compile_sh: + tar = op_file + str(index) + build_path = os.path.join(self.working_dir, "kernel_" + str(index)) + os.makedirs(build_path) + all_tar.append(tar) + sub_cmd.append(tar + ":") + sub_cmd.append( + "\tcd {} && bash {} --kernel-src=$(CPP) $(PY) $(OUT) $(MAKE)".format( + build_path, sh + ) + ) + index += 1 + mkfile = os.path.join(self.working_dir, op_file + ".make") + with os.fdopen(os.open(mkfile, const_var.WFLAGS, const_var.WMODES), "w") as fd: + sub_cmd.insert(0, "all: " + " ".join(all_tar)) + fd.write("\n".join(sub_cmd)) + + python_version = "python{}.{}".format(sys.version_info.major, sys.version_info.minor) + system_site = os.path.join(sys.base_prefix, "lib", python_version, "site-packages") + compiler_python = "{} -S".format(sys.executable) + if os.getenv("TILINGKEY_PAR_COMPILE") is None: + cmd_str = ('export PYTHONPATH=$PYTHONPATH:{0} && export HI_PYTHON="{1}" && ' + 'export ASCEND_CUSTOM_OPP_PATH={{}} && export TILINGKEY_PAR_COMPILE=1' + '&& make -f {{}} PY={{}} OUT={{}} CPP={{}}').format(system_site, compiler_python) + else: + cmd_str = ('export PYTHONPATH=$PYTHONPATH:{0} && export HI_PYTHON="{1}" && ' + 'export ASCEND_CUSTOM_OPP_PATH={{}} && make -f {{}} PY={{}} OUT={{}} CPP={{}}').format( + system_site, compiler_python) + + if os.system(cmd_str.format(self.build_opp_path, mkfile, self.op_impl_py, op_bin_dir, self.op_cpp_file)) != 0: + raise RuntimeError('Kernel Compilation Error: OpType {} Kernel File {}!'.format( + self.op_type, self.op_cpp_file)) + + +def args_parse(): + parser = argparse.ArgumentParser() + parser.add_argument( + "-n", "--op-name", nargs="?", help="Op name(Camel string) to compile." + ) + parser.add_argument("-s", "--src-file", nargs="?", help="Op kernel source file.") + + parser.add_argument("-u", "--compute-unit", nargs="?", help="Compute unit.") + parser.add_argument( + "-c", "--compile-options", nargs="?", help="Compile options of compiler." + ) + parser.add_argument( + "-d", + "--debug-config", + nargs="?", + help="Debug config of op, ref opc op-debug-config.", + ) + parser.add_argument("-i", "--config-ini", nargs="?", help="Op config ini file.") + parser.add_argument( + "-t", "--tiling-lib", nargs="?", help="Tiling shared library file." + ) + + parser.add_argument( + "-o", "--output-path", nargs="?", help="Output path of compile result." + ) + parser.add_argument( + "-dy", "--dynamic-dir", nargs="?", default=None, help="dynamic path of source compile." + ) + parser.add_argument( + "-eb", "--enable-binary", nargs="?", default=None, help="whether binary compile is enabled." + ) + parser.add_argument( + "-j", "--json-file", nargs="?", default=None, help="aic--ops-info.json file path." + ) + # $(MAKE) is necessary for parallel compiling + parser.add_argument( + "-b", "--build-tool", nargs="?", default=None, help="build tool must be make." + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = args_parse() + kernel_builder = CompileKernel(args) + kernel_builder.clean() + if args.enable_binary == "False": + kernel_builder.ascendc_gen_impl() + kernel_builder.clean() + else: + kernel_builder.ascendc_gen_impl() + kernel_builder.ascendc_gen_param() + kernel_builder.ascendc_put_json() + kernel_builder.ascendc_put_tiling() + kernel_builder.ascendc_build() + kernel_builder.clean() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_gen_options.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_gen_options.py new file mode 100755 index 00000000..2684e9ad --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_gen_options.py @@ -0,0 +1,90 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +# Copyright (c) Huawei Technologies Co., Ltd. 2024. 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. +# ============================================================================ + +import sys +import stat +import os +import re +import json +import const_var + + +def write_options_to_file(file_name: str, options_str: str, \ + op_type: str, compute_unit: str, split_char: str): + flags = os.O_WRONLY | os.O_CREAT + modes = stat.S_IWUSR | stat.S_IRUSR + try: + with os.fdopen(os.open(file_name, flags, modes), 'a') as fd: + fd.write(op_type + split_char + compute_unit + split_char + options_str + '\n') + except Exception as err: + print("write compile options config file failed") + raise(err) + + +def gen_compile_options(compile_options_file: str, op_type: str, \ + compute_unit: str, compile_options: list): + base_dir = os.path.dirname(compile_options_file) + opc_config_file = os.path.join(base_dir, "custom_opc_options.ini") + compile_opt = [] + opc_debug_config = [] + opc_tiling_keys = "" + for opts in compile_options: + if "oom" in opts: + if opts == "--oom": + opc_debug_config.append("oom") + else: + raise RuntimeError(f"Unknown oom option format {opts}") + elif "--save-temp-files" in opts: + opc_debug_config.append("dump_cce") + elif opts.startswith("--op_relocatable_kernel_binary"): + opc_debug_config.append(opts) + elif opts.startswith("--op_super_kernel_options"): + opc_debug_config.append(opts) + elif "--tiling_key" in opts: + keys = opts.strip().split('=')[1].split(',') + keys_str = ";".join([key for key in keys]) + opc_tiling_keys = keys_str + else: + compile_opt.append(opts) + if len(compile_opt) > 0: + options_str = ';'.join([opt for opt in compile_opt]) + write_options_to_file(compile_options_file, options_str, op_type, compute_unit, ",") + opc_config_str = "" + if opc_debug_config: + opc_config_str = "--op_debug_config=" + ';'.join([opt for opt in opc_debug_config]) + if len(opc_tiling_keys) > 0: + if opc_config_str != "": + opc_config_str += "@" + opc_config_str += "--tiling_key=" + opc_tiling_keys + + if opc_config_str != "": + write_options_to_file(opc_config_file, opc_config_str, op_type, compute_unit, "@") + + +if __name__ == '__main__': + if len(sys.argv) < 4: + raise RuntimeError('arguments must greater than 4') + compute_soc = "" + comp_options = [] + for i in range(len(sys.argv) - 3): + if sys.argv[i + 3].upper().startswith("ASCEND"): + compute_soc += sys.argv[i + 3] + ";" + else: + comp_options.append(sys.argv[i + 3]) + if compute_soc != "": + compute_soc = compute_soc[0:-1] + gen_compile_options(sys.argv[1], sys.argv[2], compute_soc, comp_options) \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_impl_build.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_impl_build.py new file mode 100755 index 00000000..9c1cfaa4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_impl_build.py @@ -0,0 +1,698 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. +""" + +import argparse +import glob +import sys +import os +import re +import datetime +from typing import List +import json +import opdesc_parser +import const_var + + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + +IMPL_HEAD = '''#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Copyright (c) Huawei Technologies Co., Ltd. {}-{}. All rights reserved. +""" + +import re +import os, sys +import ctypes +import json +import shutil +from tbe.common.platform import get_soc_spec +from tbe.common.utils import para_check +from tbe.tikcpp import compile_op, replay_op, check_op_cap, generalize_op_params, get_code_channel, OpInfo +from tbe.tikcpp.compile_op import CommonUtility, AscendCLogLevel +from tbe.common.buildcfg import get_default_build_config +import tbe.common.register as tbe_register +from tbe.common.buildcfg import get_current_build_config +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + +DTYPE_MAP = {{"float32": ["DT_FLOAT", "float"], + "float16": ["DT_FLOAT16", "half"], + "int8": ["DT_INT8", "int8_t"], + "int16": ["DT_INT16", "int16_t"], + "int32": ["DT_INT32", "int32_t"], + "int64": ["DT_INT64", "int64_t"], + "uint1": ["DT_UINT1", "uint1b_t"], + "uint8": ["DT_UINT8", "uint8_t"], + "uint16": ["DT_UINT16", "uint16_t"], + "uint32": ["DT_UINT32", "uint32_t"], + "uint64": ["DT_UINT64", "uint64_t"], + "bool": ["DT_BOOL", "bool"], + "double": ["DT_DOUBLE", "double"], + "dual": ["DT_DUAL", "unknown"], + "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"], + "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"], + "string": ["DT_STRING", "unknown"], + "complex32": ["DT_COMPLEX32", "complex32"], + "complex64": ["DT_COMPLEX64", "complex64"], + "complex128": ["DT_COMPLEX128", "unknown"], + "qint8": ["DT_QINT8", "unknown"], + "qint16": ["DT_QINT16", "unknown"], + "qint32": ["DT_QINT32", "unknown"], + "quint8": ["DT_QUINT8", "unknown"], + "quint16": ["DT_QUINT16", "unknown"], + "resource": ["DT_RESOURCE", "unknown"], + "string_ref": ["DT_STRING_REF", "unknown"], + "int4": ["DT_INT4", "int4b_t"], + "bfloat16": ["DT_BF16", "bfloat16_t"], + "float8_e5m2": ["DT_FLOAT8_E5M2", "fp8_e5m2_t"], + "float8_e4m3fn": ["DT_FLOAT8_E4M3FN", "fp8_e4m3fn_t"], + "hifloat8":["DT_HIFLOAT8", "hifloat8_t"], + "float8_e8m0":["DT_FLOAT8_E8M0", "fp8_e8m0_t"], + "float4_e2m1":["DT_FLOAT4_E2M1", "fp4x2_e2m1_t"], + "float4_e1m2":["DT_FLOAT4_E1M2", "fp4x2_e1m2_t"], + "int2": ["DT_INT2", "int2b_t"]}} + +def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False): + options = [] + x_fmt = x.get("format") + x_dtype = x.get("dtype") + x_n_in_kernel = x_n + '_REF' if is_ref else x_n + options.append("-DDTYPE_{{n}}={{t}}".format(n=x_n_in_kernel, t=DTYPE_MAP.get(x_dtype)[1])) + options.append("-DORIG_DTYPE_{{n}}={{ot}}".format(n=x_n_in_kernel, ot=DTYPE_MAP.get(x_dtype)[0])) + options.append("-DFORMAT_{{n}}=FORMAT_{{f}}".format(n=x_n_in_kernel, f=x_fmt)) + return options + +def get_dtype_fmt_options(__inputs__, __outputs__): + options = [] + input_names = {} + output_names = {} + unique_param_name_set = set() + for idx, x in enumerate(__inputs__): + if x is None: + continue + x_n = input_names[idx].upper() + unique_param_name_set.add(x_n) + options += add_dtype_fmt_option_single(x, x_n) + + for idx, x in enumerate(__outputs__): + if x is None: + continue + x_n = output_names[idx].upper() + if x_n in unique_param_name_set: + options += add_dtype_fmt_option_single(x, x_n, True) + else: + options += add_dtype_fmt_option_single(x, x_n) + return options + +def load_dso(so_path): + try: + ctypes.CDLL(so_path) + except OSError as error : + CommonUtility.print_compile_log("", error, AscendCLogLevel.LOG_ERROR) + raise RuntimeError("cannot open %s" %(so_path)) + else: + msg = "load so succ " + so_path + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + +def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str): + compile_options = [] + if shortsoc in compile_option_list: + compile_options.extend(compile_option_list[shortsoc]) + if '__ALLSOC__' in compile_option_list: + compile_options.extend(compile_option_list['__ALLSOC__']) + return compile_options + +def get_kernel_source(src_file, dir_snake, dir_ex): + src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file) + if os.path.exists(src_ex): + return src_ex + src = os.environ.get('BUILD_KERNEL_SRC') + if src and os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file) + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, src_file) + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, dir_snake + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, dir_ex + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", os.path.splitext(src_file)[0], src_file) + if os.path.exists(src): + return src + return src_ex + +''' + +IMPL_API = ''' +@tbe_register.register_operator("{}", trans_bool_to_s8=False) +@para_check.check_op_params({}) +def {}({}, kernel_name="{}"{}): +{} + if get_current_build_config("enable_op_prebuild"): + return + __inputs__, __outputs__, __attrs__ = _build_args({}) + options = get_dtype_fmt_options(__inputs__, __outputs__) + options += ["-x", "cce"] + bisheng = os.environ.get('BISHENG_REAL_PATH') + if bisheng is None: + bisheng = shutil.which("bisheng") + if bisheng != None: + bisheng_path = os.path.dirname(bisheng) + tikcpp_path = os.path.realpath(os.path.join(bisheng_path, "..", "..", "tikcpp")) + else: + tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp") + options.append("-I" + tikcpp_path) + options.append("-I" + os.path.join(tikcpp_path, "..", "..", "include")) + options.append("-I" + os.path.join(tikcpp_path, "..", "..", "include", "ascendc")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "impl")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "interface")) + options.append("-I" + os.path.join(tikcpp_path, "..", "ascendc", "act")) + options.append("-I" + os.path.join(PYF_PATH, "..", "ascendc", "common")) + if "impl_mode" in locals(): + if impl_mode == "high_performance": + options.append("-DHIGH_PERFORMANCE=1") + elif impl_mode == "high_precision": + options.append("-DHIGH_PRECISION=1") + elif "high_precision" in impl_mode and "high_performance" in impl_mode: + options.append("-DHIGH_PRECISION=1 -DHIGH_PERFORMANCE=1") + if get_current_build_config("enable_deterministic_mode") == 1: + options.append("-DDETERMINISTIC_MODE=1") + else: + options.append("-DDETERMINISTIC_MODE=0") + ascendc_api_version_header_path = os.path.join(tikcpp_path, "tikcfw/lib/ascendc_api_version.h") + if os.path.exists(ascendc_api_version_header_path): + with open(ascendc_api_version_header_path, "r") as ascendc_api_version_file: + ascendc_api_version = re.findall(r"#define ASCENDC_API_VERSION (\d+)", ascendc_api_version_file.read()) + if ascendc_api_version: + options.append(f"-DASCENDC_API_VERSION={{ascendc_api_version[0]}}") + custom_compile_options = {}, + custom_all_compile_options = {}, + soc_version = get_soc_spec("SOC_VERSION") + soc_short = get_soc_spec("SHORT_SOC_VERSION").lower() + custom_compile_options_soc = get_shortsoc_compile_option(custom_compile_options[0], soc_short) + custom_all_compile_options_soc = get_shortsoc_compile_option(custom_all_compile_options[0], soc_short) + options += custom_all_compile_options_soc + options += custom_compile_options_soc + + origin_func_name = "{}" + ascendc_src_dir_ex = "{}" + ascendc_src_dir = "{}" + ascendc_src_file = "{}" + src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex) +''' + +REPLAY_OP_API = ''' + msg = "start replay Ascend C Operator {}, kernel name is {}" + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + tikreplay_codegen_path = tikcpp_path + "/tikreplaylib/lib" + tikreplay_stub_path = tikcpp_path + "/tikreplaylib/lib/" + soc_version + msg = "start load libtikreplaylib_codegen.so and libtikreplaylib_stub.so" + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + codegen_so_path = tikreplay_codegen_path + "/libtikreplaylib_codegen.so" + replaystub_so_path = tikreplay_stub_path + "/libtikreplaylib_stub.so" + if PYF_PATH.endswith("dynamic"): + op_replay_path = os.path.join(PYF_PATH, "..", "..", "op_replay") + else: + op_replay_path = os.path.join(PYF_PATH, "..", "op_replay") + replayapi_so_path = os.path.join(op_replay_path, "libreplay_{}_" + soc_short + ".so") + load_dso(codegen_so_path) + load_dso(replaystub_so_path) + load_dso(replayapi_so_path) + op_type = "{}" + entry_obj = os.path.join(op_replay_path, "{}_entry_" + soc_short + ".o") + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__, impl_mode = impl_mode, param_type_dynamic = {}) + res, msg = replay_op(op_info, entry_obj, code_channel, src, options) + if not res: + print("call replay op failed for %s and get into call compile op" %(msg)) + compile_op(src, origin_func_name, op_info, options, code_channel, '{}') +''' + +COMPILE_OP_API = ''' + msg = "start compile Ascend C Operator {}, kernel name is " + kernel_name + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + op_type = "{}" + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__ {}, origin_inputs=[{}], origin_outputs = [{}],\\ + param_type_dynamic = {}, mc2_ctx = {}, param_type_list = {}, init_value_list = {},\\ + output_shape_depend_on_compute = {}) + compile_op(src, origin_func_name, op_info, options, code_channel, '{}', {}) +''' +COMPILE_OP_API_BUILT_IN = ''' + msg = "start compile Ascend C Operator {}, kernel name is " + kernel_name + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + op_type = "{}" + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__ {}, origin_inputs=[{}], origin_outputs = [{}],\\ + param_type_dynamic = {}, mc2_ctx = {}, param_type_list = {}, init_value_list = {},\\ + output_shape_depend_on_compute = {}) + + op_compile_option = '{}' + opp_path = os.environ.get('ASCEND_OPP_PATH') + dat_path = os.path.realpath(os.path.join(opp_path, "built-in", "op_impl", "ai_core", "tbe", "ascendc_impl.dat")) + if opp_path and os.path.exists(dat_path): + # dat file exists: built in hidden src file online compiling process. append vfs compile option in compile_op + abs_rel_kernel_src_path = "{}" + extend_options = {} + extend_options['opp_kernel_hidden_dat_path'] = dat_path + compile_op(abs_rel_kernel_src_path, origin_func_name, op_info, options, code_channel, op_compile_option,\\ + extend_options) + else: + raise RuntimeError("built-in opp compile, ascendc_impl.dat file path does not exist: %s" %(dat_path)) +''' +SUP_API = ''' +def {}({}{}): + __inputs__, __outputs__, __attrs__ = _build_args({}) + ret_str = check_op_cap("{}", "{}", __inputs__, __outputs__, __attrs__) + ret_dict = json.loads(ret_str) + err_code = ret_dict.get("ret_code") + sup = "Unknown" + reason = "Unknown reason" + if err_code is not None: + if err_code == 0: + sup = "True" + reason = "" + elif err_code == 1: + sup = "False" + reason = ret_dict.get("reason") + else: + sup = "Unknown" + reason = ret_dict.get("reason") + return sup, reason +''' +CAP_API = ''' +def {}({}{}): + __inputs__, __outputs__, __attrs__ = _build_args({}) + result = check_op_cap("{}", "{}", __inputs__, __outputs__, __attrs__) + return result.decode("utf-8") +''' +GLZ_API = ''' +@tbe_register.register_param_generalization("{}") +def {}_generalization({}, generalize_config=None): + __inputs__, __outputs__, __attrs__ = _build_args({}) + ret_str = generalize_op_params("{}", __inputs__, __outputs__, __attrs__, generalize_config) + return [json.loads(ret_str)] +''' + +ATTR_DEFAULT = {'bool': 'False', 'int': '0', 'float': '0.0', 'list_int': '[]', + 'list_float': '[]', 'list_bool': '[]', 'list_list_int': '[[]]', 'str': ''} + + +def optype_snake(origin_str): + temp_str = origin_str[0].lower() + origin_str[1:] + new_str = re.sub(r'([A-Z])', r'_\1', temp_str).lower() + return new_str + + +def optype_snake_ex(s): + snake_case = "" + for i, c in enumerate(s): + if i == 0: + snake_case += c.lower() + elif c.isupper(): + if s[i - 1] != '_': + if not s[i - 1].isupper(): + snake_case += "_" + elif s[i - 1].isupper() and (i + 1) < len(s) and s[i + 1].islower(): + snake_case += "_" + snake_case += c.lower() + else: + snake_case += c + return snake_case + + +class AdpBuilder(opdesc_parser.OpDesc): + def __init__(self: any, op_type: str): + self.argsdefv = [] + self.op_compile_option:str = '{}' + super().__init__(op_type) + + + def write_adapt(self: any, impl_path, path: str, op_compile_option_all: list = None): + self._build_paradefault() + if os.environ.get('BUILD_BUILTIN_OPP') != '1' and impl_path != "": + src_file = os.path.join(impl_path, self.op_file + '.cpp') + if not os.path.exists(src_file): + print(f"[ERROR]: operator: {self.op_file} source file: {src_file} does not found, please check.") + return + out_path = os.path.abspath(path) + if self.dynamic_shape and not out_path.endswith('dynamic'): + out_path = os.path.join(path, 'dynamic') + os.makedirs(out_path, exist_ok=True) + adpfile = os.path.join(out_path, self.op_file + '.py') + self._gen_op_compile_option(op_compile_option_all) + with os.fdopen(os.open(adpfile, const_var.WFLAGS, const_var.WMODES), 'w') as fd: + self._write_head(fd) + self._write_argparse(fd) + self._get_impl_mode() + self._write_impl(fd, impl_path) + if self.op_chk_support: + self._write_cap('check_supported', fd) + self._write_cap('get_op_support_info', fd) + if self.op_fmt_sel: + self._write_cap('op_select_format', fd) + self._write_cap('get_op_specific_info', fd) + if self.op_range_limit == 'limited' or self.op_range_limit == 'dynamic': + self._write_glz(fd) + + + def _gen_op_compile_option(self: any, op_compile_option_all: list = None): + if op_compile_option_all is not None: + if self.op_type in op_compile_option_all: + self.op_compile_option = op_compile_option_all[self.op_type] + elif "__all__" in op_compile_option_all: + self.op_compile_option = op_compile_option_all["__all__"] + + + def _ip_argpack(self: any, default: bool = True) -> list: + args = [] + for i in range(len(self.input_name)): + arg = self.input_name[i] + if default and self.argsdefv[i] is not None: + arg += '=' + self.argsdefv[i] + args.append(arg) + return args + + def _op_argpack(self: any, default: bool = True) -> list: + args = [] + argidx = len(self.input_name) + for i in range(len(self.output_name)): + arg = self.output_name[i] + if default and self.argsdefv[i + argidx] is not None: + arg += '=' + self.argsdefv[i + argidx] + args.append(arg) + return args + + def _attr_argpack(self: any, default: bool = True) -> list: + args = [] + argidx = len(self.input_name) + len(self.output_name) + for i in range(len(self.attr_list)): + att = self.attr_list[i] + arg = att + if default and self.argsdefv[i + argidx] is not None: + if self.attr_val.get(att).get('type') == 'str': + arg += '="' + self.argsdefv[i + argidx] + '"' + elif self.attr_val.get(att).get('type') == 'bool': + arg += '=' + self.argsdefv[i + argidx].capitalize() + elif self.attr_val.get(att).get('type') == 'list_bool': + arg += '=' + "[" + ", ".join(word.strip().capitalize() \ + for word in self.argsdefv[i + argidx].strip('[]').split(',')) + "]" + else: + arg += '=' + self.argsdefv[i + argidx] + args.append(arg) + return args + + def _build_paralist(self: any, default: bool = True) -> str: + args = [] + args.extend(self._ip_argpack(default)) + args.extend(self._op_argpack(default)) + args.extend(self._attr_argpack(default)) + return ', '.join(args) + + def _io_parachk(self: any, types: list, type_name: str) -> list: + chk = [] + for iot in types: + if iot == 'optional': + ptype = 'OPTION' + else: + ptype = iot.upper() + chk.append('para_check.{}_{}'.format(ptype, type_name)) + return chk + + def _attr_parachk(self: any) -> list: + chk = [] + for att in self.attr_list: + att_type = self.attr_val.get(att).get('type').upper() + chk.append('para_check.{}_ATTR_{}'.format('OPTION', att_type)) + return chk + + def _build_parachk(self: any) -> str: + chk = [] + chk.extend(self._io_parachk(self.input_type, 'INPUT')) + chk.extend(self._io_parachk(self.output_type, 'OUTPUT')) + chk.extend(self._attr_parachk()) + chk.append('para_check.KERNEL_NAME') + return ', '.join(chk) + + def _build_virtual(self: any) -> str: + virt_exp = [] + for index in range(len(self.input_name)): + if self.input_virt.get(index) is None: + continue + val = [] + val.append('"param_name":"{}"'.format(self.input_name[index])) + val.append('"index":{}'.format(index)) + val.append('"dtype":"{}"'.format(self.input_dtype[index].split(',')[0])) + val.append('"format":"{}"'.format(self.input_fmt[index].split(',')[0])) + val.append('"ori_format":"{}"'.format(self.input_fmt[index].split(',')[0])) + val.append('"paramType":"optional"') + val.append('"shape":[1]') + val.append('"ori_shape":[1]') + virt_exp.append(' ' + self.input_name[index] + ' = {' + ','.join(val) + '}') + if len(virt_exp) > 0: + return '\n'.join(virt_exp) + else: + return ' # do ascendc build step' + + def _build_mc2_ctx(self: any): + if len(self.mc2_ctx) != 0: + return '["' + '", "'.join(self.mc2_ctx) + '"]' + return '[]' + + def _build_paradefault(self: any): + optional = False + argtypes = [] + argtypes.extend(self.input_type) + argtypes.extend(self.output_type) + in_idx = 0 + for atype in argtypes: + if atype == 'optional': + optional = True + if optional: + self.argsdefv.append('None') + else: + self.argsdefv.append(None) + in_idx += 1 + for attr in self.attr_list: + atype = self.attr_val.get(attr).get('paramType') + if atype == 'optional': + optional = True + attrval = self.attr_val.get(attr).get('defaultValue') + if attrval is not None: + optional = True + if type == "bool": + attrval = attrval.capitalize() + elif type == "str": + attrval = "\"" + attrval + "\"" + self.argsdefv.append(attrval) + continue + if optional: + self.argsdefv.append(ATTR_DEFAULT.get(self.attr_val.get(attr).get('type'))) + else: + self.argsdefv.append(None) + + def _write_head(self: any, fd: object): + now = datetime.datetime.now() + curr_year = now.year + former_year = curr_year - 1 + fd.write(IMPL_HEAD.format(former_year, curr_year, self.input_ori_name, self.output_ori_name)) + + def _write_argparse(self: any, fd: object): + args = self._build_paralist(False) + fd.write('def _build_args({}):\n'.format(args)) + fd.write(' __inputs__ = []\n') + fd.write(' for arg in [{}]:\n'.format(', '.join(self.input_name))) + fd.write(' if arg != None:\n') + fd.write(' if isinstance(arg, (list, tuple)):\n') + fd.write(' if len(arg) == 0:\n') + fd.write(' continue\n') + fd.write(' __inputs__.append(arg[0])\n') + fd.write(' else:\n') + fd.write(' __inputs__.append(arg)\n') + fd.write(' else:\n') + fd.write(' __inputs__.append(arg)\n') + fd.write(' __outputs__ = []\n') + fd.write(' for arg in [{}]:\n'.format(', '.join(self.output_name))) + fd.write(' if arg != None:\n') + fd.write(' if isinstance(arg, (list, tuple)):\n') + fd.write(' if len(arg) == 0:\n') + fd.write(' continue\n') + fd.write(' __outputs__.append(arg[0])\n') + fd.write(' else:\n') + fd.write(' __outputs__.append(arg)\n') + fd.write(' else:\n') + fd.write(' __outputs__.append(arg)\n') + fd.write(' __attrs__ = []\n') + for attr in self.attr_list: + fd.write(' if {} != None:\n'.format(attr)) + fd.write(' attr = {}\n') + fd.write(' attr["name"] = "{}"\n'.format(attr)) + fd.write(' attr["dtype"] = "{}"\n'.format(self.attr_val.get(attr).get('type'))) + fd.write(' attr["value"] = {}\n'.format(attr)) + fd.write(' __attrs__.append(attr)\n') + fd.write(' return __inputs__, __outputs__, __attrs__\n') + + def _get_kernel_source(self: any, kernel_src_dir, src_file, dir_snake, dir_ex): + src_ex = os.path.join(kernel_src_dir, dir_ex, src_file) + if os.path.exists(src_ex): + return src_ex + src = os.environ.get('BUILD_KERNEL_SRC') + if src and os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_snake, src_file) + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, src_file) + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_snake, dir_snake + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_ex, dir_ex + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, os.path.splitext(src_file)[0], src_file) + if os.path.exists(src): + return src + return src_ex + + def _get_impl_mode(self: any): + op_compile_options = json.loads(self.op_compile_option) + if "impl_mode" in op_compile_options: + if op_compile_options['impl_mode'] == "": + self.impl_mode = "" + self.impl_mode_op_info = "" + del op_compile_options['impl_mode'] + self.op_compile_option = json.dumps(op_compile_options) + else: + self.impl_mode = ", impl_mode ='" + op_compile_options["impl_mode"] + "'" + self.impl_mode_op_info = ", impl_mode ='" + op_compile_options["impl_mode"] + "'" + else: + self.impl_mode = ', impl_mode = ""' + self.impl_mode_op_info = ", impl_mode = impl_mode" + + def _write_impl(self: any, fd: object, impl_path: str = ""): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + pchk = self._build_parachk() + if len(self.kern_name) > 0: + kern_name = self.kern_name + else: + kern_name = self.op_intf + src = self.op_file + '.cpp' + virt_exprs = self._build_virtual() + fd.write(IMPL_API.format(self.op_type, pchk, \ + self.op_intf, argsdef, kern_name, self.impl_mode, virt_exprs, argsval,\ + self.custom_compile_options, self.custom_all_compile_options, self.op_intf,\ + optype_snake_ex(self.op_type), optype_snake(self.op_type), src)) + if self.op_replay_flag: + fd.write(REPLAY_OP_API.format(self.op_type, kern_name, self.op_file,\ + self.op_type, self.op_file, self.param_type_dynamic, self.op_compile_option)) + else: + value_depend_obj = {key: value for key, value in self.input_value_depend.items()} + extend_opt = {"valueDepend": value_depend_obj} + if os.environ.get('BUILD_BUILTIN_OPP') == '1': + relative_kernel_src_path = os.path.realpath(self._get_kernel_source(impl_path, src,\ + optype_snake(self.op_type), optype_snake_ex(self.op_type))) + # to match src path in .dat file system, turn relative path into absolute path + abs_rel_kernel_src_path = os.path.join("/", os.path.relpath(relative_kernel_src_path, impl_path)) + + # compiling hidden src file requires src path before packaging .dat file, + # hard code such src path to .py + fd.write(COMPILE_OP_API_BUILT_IN.format(self.op_type, self.op_type,\ + self.impl_mode_op_info, ', '.join(self.input_name), \ + ', '.join(self.output_name), self.param_type_dynamic,\ + self._build_mc2_ctx(), self.input_type + self.output_type, self.output_init_value,\ + self.output_shape_depend_on_compute, self.op_compile_option, abs_rel_kernel_src_path, repr(extend_opt))) + else: + fd.write(COMPILE_OP_API.format(self.op_type, + self.op_type, self.impl_mode_op_info, ', '.join(self.input_name), \ + ', '.join(self.output_name), self.param_type_dynamic, self._build_mc2_ctx(),\ + self.input_type + self.output_type, self.output_init_value, self.output_shape_depend_on_compute,\ + self.op_compile_option, repr(extend_opt))) + + def _write_cap(self: any, cap_name: str, fd: object): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + if cap_name == 'check_supported': + fd.write(SUP_API.format(cap_name, argsdef, self.impl_mode, argsval, cap_name, self.op_type)) + else: + fd.write(CAP_API.format(cap_name, argsdef, self.impl_mode, argsval, cap_name, self.op_type)) + + def _write_glz(self: any, fd: object): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + fd.write(GLZ_API.format(self.op_type, self.op_intf, argsdef, argsval, self.op_type)) + + +def write_scripts(cfgfile: str, cfgs: dict, dirs: dict, ops: list = None, op_compile_option:list = None): + batch_lists = cfgs.get(const_var.REPLAY_BATCH).split(';') + iterator_lists = cfgs.get(const_var.REPLAY_ITERATE).split(';') + file_map = {} + op_descs = opdesc_parser.get_op_desc(cfgfile, batch_lists, iterator_lists, AdpBuilder,\ + ops, dirs.get(const_var.AUTO_GEN_DIR)) + for op_desc in op_descs: + op_desc.write_adapt(dirs.get(const_var.CFG_IMPL_DIR), dirs.get(const_var.CFG_OUT_DIR), op_compile_option) + file_map[op_desc.op_type] = op_desc.op_file + return file_map + + +class OpFileNotExistsError(Exception): + """File does not exist error.""" + def __str__(self) -> str: + return f"File aic-*-ops-info.ini does not exist in directory {super().__str__()}" + + +def get_ops_info_files(opsinfo_dir: List[str]) -> List[str]: + """Get all ops info files.""" + ops_info_files = [] + for _dir in opsinfo_dir: + ops_info_files.extend(glob.glob(f'{_dir}/aic-*-ops-info.ini')) + return sorted(ops_info_files) + + +def parse_args(argv): + """Command line parameter parsing""" + parser = argparse.ArgumentParser() + parser.add_argument('argv', nargs='+') + parser.add_argument('--opsinfo-dir', nargs='*', default=None) + return parser.parse_args(argv) + + +if __name__ == '__main__': + args = parse_args(sys.argv) + + if len(args.argv) <= 6: + raise RuntimeError('arguments must greater equal than 6') + + rep_cfg = {} + rep_cfg[const_var.REPLAY_BATCH] = args.argv[2] + rep_cfg[const_var.REPLAY_ITERATE] = args.argv[3] + + cfg_dir = {} + cfg_dir[const_var.CFG_IMPL_DIR] = args.argv[4] + cfg_dir[const_var.CFG_OUT_DIR] = args.argv[5] + cfg_dir[const_var.AUTO_GEN_DIR] = args.argv[6] + + ops_infos = [] + if args.opsinfo_dir: + ops_infos.extend(get_ops_info_files(args.opsinfo_dir)) + if not ops_infos: + raise OpFileNotExistsError(args.opsinfo_dir) + else: + ops_infos.append(args.argv[1]) + + for ops_info in ops_infos: + write_scripts(cfgfile=ops_info, cfgs=rep_cfg, dirs=cfg_dir) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_op_info.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_op_info.py new file mode 100755 index 00000000..3a6057cc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_op_info.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. +""" + +import sys +import os +import opdesc_parser + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + + +class OpInfo: + def __init__(self: any, op_type: str, cfg_file: str): + op_descs = opdesc_parser.get_op_desc( + cfg_file, [], [], opdesc_parser.OpDesc, [op_type] + ) + if op_descs is None or len(op_descs) != 1: + raise RuntimeError("cannot get op info of {}".format(op_type)) + self.op_desc = op_descs[0] + + def get_op_file(self: any): + return self.op_desc.op_file + + def get_op_intf(self: any): + return self.op_desc.op_intf + + def get_inputs_name(self: any): + return self.op_desc.input_ori_name + + def get_outputs_name(self: any): + return self.op_desc.output_ori_name + + +if __name__ == "__main__": + if len(sys.argv) <= 2: + raise RuntimeError("arguments must greater than 2") + op_info = OpInfo(sys.argv[1], sys.argv[2]) + print(op_info.get_op_file()) + print(op_info.get_op_intf()) + print(op_info.get_inputs_name()) + print(op_info.get_outputs_name()) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_ops_config.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_ops_config.py new file mode 100755 index 00000000..c5fbf421 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_ops_config.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Created on Feb 28 20:56:45 2020 +Copyright (c) Huawei Technologies Co., Ltd. 2020-2024. All rights reserved. +""" + +import os +import glob +import json +import sys +import argparse +from typing import NamedTuple, Dict +import const_var + + +class OpConfig(NamedTuple): + op_type: str + support_info: Dict + core_type: str + task_ration: str + obj_file: str + + +def load_json(json_file: str): + with open(json_file, encoding='utf-8') as file: + json_content = json.load(file) + return json_content + + +def get_specified_suffix_file(root_dir, suffix): + specified_suffix = os.path.join(root_dir, '**/*{}'.format(suffix)) + all_suffix_files = glob.glob(specified_suffix, recursive=True) + return sorted(all_suffix_files) + + +def add_dict_key(dict_to_add, key, value): + if value is None: + return + dict_to_add[key] = value + + +def correct_format_mode(format_mode): + if format_mode == 'FormatDefault': + return 'nd_agnostic' + if format_mode == 'FormatAgnostic': + return 'static_nd_agnostic' + if format_mode == 'FormatFixed': + return 'normal' + return format_mode + + +def get_input_or_output_config(in_or_out): + param_dict = {} + name = in_or_out.get('name') + index = in_or_out.get('index') + param_type = in_or_out.get('paramType') + + format_match_mode = in_or_out.get('format_match_mode') + format_mode = correct_format_mode(format_match_mode) + + dtype_mode = in_or_out.get('dtype_match_mode') + if dtype_mode == 'DtypeByte': + dtype_mode = 'bit' + + add_dict_key(param_dict, 'name', name) + add_dict_key(param_dict, 'index', index) + add_dict_key(param_dict, 'paramType', param_type) + add_dict_key(param_dict, 'dtypeMode', dtype_mode) + add_dict_key(param_dict, 'formatMode', format_mode) + return param_dict + + +def get_inputs_or_outputs_config(inputs_or_outputs): + if inputs_or_outputs is None: + return None + inputs_or_outputs_list = [] + + for in_or_out in inputs_or_outputs: + if isinstance(in_or_out, dict): + dict_param_config = get_input_or_output_config(in_or_out) + inputs_or_outputs_list.append(dict_param_config) + elif isinstance(in_or_out, list): + param_info = in_or_out[0] + list_param_config = get_input_or_output_config(param_info) + tmp_list = [list_param_config] + inputs_or_outputs_list.append(tmp_list) + return inputs_or_outputs_list + + +def gen_attrs_config(attrs): + attrs_list = [] + for attr in attrs: + attrs_dict = {} + name = attr.get('name') + mode = attr.get('mode') + add_dict_key(attrs_dict, 'name', name) + add_dict_key(attrs_dict, 'mode', mode) + attrs_list.append(attrs_dict) + return attrs_list + + +def get_params_config(support_info): + params_dict = {} + + inputs = support_info.get('inputs') + inputs_list = get_inputs_or_outputs_config(inputs) + params_dict['inputs'] = inputs_list + + outputs = support_info.get('outputs') + outputs_list = get_inputs_or_outputs_config(outputs) + params_dict['outputs'] = outputs_list + + attrs = support_info.get('attrs') + if attrs is not None: + attrs_list = gen_attrs_config(attrs) + params_dict['attrs'] = attrs_list + + return params_dict + + +def add_simplified_config(op_info, binary_info_config, config): + simplified_key = op_info.support_info.get('simplifiedKey') + + json_path = op_info.obj_file.split('.')[0] + '.json' + + simple_cfg = config.get(binary_info_config) + op_cfg = simple_cfg.get(op_info.op_type) + if not op_cfg: + op_cfg = {'dynamicRankSupport': True} + + simplified_key_mode = op_info.support_info.get('simplifiedKeyMode') + add_dict_key(op_cfg, 'simplifiedKeyMode', simplified_key_mode) + + optional_input_mode = op_info.support_info.get('optionalInputMode') + optional_output_mode = op_info.support_info.get('optionalOutputMode') + add_dict_key(op_cfg, 'optionalInputMode', optional_input_mode) + if optional_output_mode is not None: + add_dict_key(op_cfg, 'optionalOutputMode', optional_output_mode) + + params_info = get_params_config(op_info.support_info) + op_cfg['params'] = params_info + op_cfg['binaryList'] = [] + simple_cfg[op_info.op_type] = op_cfg + + bin_list = op_cfg.get('binaryList') + if op_info.core_type == 0 and op_info.task_ration == "tilingKey": + bin_list.append({'coreType': op_info.core_type, 'simplifiedKey': simplified_key, + 'multiKernelType': 1, 'binPath': op_info.obj_file, 'jsonPath': json_path}) + else: + bin_list.append({'coreType': op_info.core_type, 'simplifiedKey': simplified_key, + 'binPath': op_info.obj_file, 'jsonPath': json_path}) + + +def add_op_config(op_file, bin_info, config): + op_cfg = config.get(op_file) + if not op_cfg: + op_cfg = {'binList': []} + config[op_file] = op_cfg + op_cfg.get('binList').append(bin_info) + + +def gen_ops_config(json_file, soc, binary_info_config, config): + core_type_map = {'MIX': 0, 'AiCore': 1, 'VectorCore': 2, 'MIX_AICORE': 3, 'MIX_VECTOR_CORE': 4, 'MIX_AIV': 4} + contents = load_json(json_file) + if ('binFileName' not in contents) or ('supportInfo' not in contents): + return + json_base_name = os.path.basename(json_file) + op_dir = os.path.basename(os.path.dirname(json_file)) + + support_info = contents.get('supportInfo') + bin_name = contents.get('binFileName') + bin_suffix = contents.get('binFileSuffix') + core_type = contents.get("coreType") + task_ration = contents.get("taskRation") + core_type = core_type_map.get(core_type, -1) + if core_type == -1 and soc != 'ascend310b': + raise Exception("[ERROR]: must set coreType in json when soc version is {soc}.") + + bin_file_name = bin_name + bin_suffix + op_type = bin_name.split('_')[0] + op_file = op_dir + '.json' + bin_info = {} + + add_dict_key(bin_info, 'implMode', support_info.get('implMode')) + add_dict_key(bin_info, 'int64Mode', support_info.get('int64Mode')) + add_dict_key(bin_info, 'simplifiedKeyMode', support_info.get('simplifiedKeyMode')) + + simplified_key = support_info.get('simplifiedKey') + if simplified_key is not None: + bin_info['simplifiedKey'] = simplified_key + obj_file = os.path.join(soc, op_dir, bin_file_name) + op_info = OpConfig( + op_type=op_type, + support_info=support_info, + core_type=core_type, + task_ration=task_ration, + obj_file=obj_file, + ) + add_simplified_config(op_info, binary_info_config, config) + + add_dict_key(bin_info, 'dynamicParamMode', support_info.get('dynamicParamMode')) + bin_info['staticKey'] = support_info.get('staticKey') + bin_info['inputs'] = support_info.get('inputs') + bin_info['outputs'] = support_info.get('outputs') + if support_info.get('attrs'): + bin_info['attrs'] = support_info.get('attrs') + + add_dict_key(bin_info, 'opMode', support_info.get('opMode')) + add_dict_key(bin_info, 'optionalInputMode', support_info.get('optionalInputMode')) + add_dict_key(bin_info, 'deterministic', support_info.get('deterministic')) + if support_info.get('optionalOutputMode') is not None: + add_dict_key(bin_info, 'optionalOutputMode', support_info.get('optionalOutputMode')) + + bin_info['binInfo'] = {'jsonFilePath': os.path.join(soc, op_dir, json_base_name)} + add_op_config(op_file, bin_info, config) + + +def check_single_op_is_void(root_dir): + for root, dirs, _ in os.walk(root_dir): + for sub_dir in dirs: + dir_path = os.path.join(root, sub_dir) + if len(os.listdir(dir_path)) == 0: + print(f"[ERROR] op {sub_dir}: not any obj compile success") + sys.exit(1) + + +def write_jsons(out_dir, file_list, config): + for json_name in file_list: + json_file = os.path.join(out_dir, json_name) + with os.fdopen(os.open(json_file, const_var.WFLAGS, const_var.WMODES), 'w') as fd: + json.dump(config.get(json_name), fd, indent=' ') + + +def generate_operator_cfg_file(json_files, + binary_info_config, + soc, + out_dir, + gen_json_status): + + if not json_files: + return + + if gen_json_status == "not_generated": + return + + json_files.sort() + config = {binary_info_config: {}} + for _json in json_files: + gen_ops_config(_json, soc, binary_info_config, config) + + if gen_json_status == "single_json": + file_list = [json_file for json_file in config.keys() if json_file != binary_info_config] + elif gen_json_status == "summary_json": + file_list = [binary_info_config] + else: + file_list = config.keys() + + write_jsons(out_dir, file_list, config) + + +def gen_all_config(root_dir, soc, out_dir, + skip_binary_info_config, op_range="all"): + if op_range != "relocatable": + check_single_op_is_void(root_dir) + all_json_files = get_specified_suffix_file(root_dir, '.json') + relocatable_json_files = get_specified_suffix_file(root_dir, '_relocatable.json') + normal_json_files = list(set(all_json_files) - set(relocatable_json_files)) + os.makedirs(out_dir, exist_ok=True) + + if op_range != "relocatable": + for _json in all_json_files: + file_path = soc + _json.split(soc, maxsplit=1)[1] + with open(_json, "r+") as f: + data = json.load(f) + data["filePath"] = file_path + f.seek(0) + json.dump(data, f, indent=" ") + f.truncate() + + if skip_binary_info_config: + gen_normale_json = "single_json" + gen_relocatable_json = "not_generated" + else: + gen_normale_json = "all_json" + gen_relocatable_json = "summary_json" + + # normal kernel + if op_range == "all" or op_range == "normal": + binary_info_config = "binary_info_config.json" + generate_operator_cfg_file(normal_json_files, binary_info_config, + soc, out_dir, gen_normale_json) + + # relocatable kernel + if op_range == "all" or op_range == "relocatable": + binary_info_config = "relocatable_kernel_info_config.json" + generate_operator_cfg_file(relocatable_json_files, binary_info_config, + soc, out_dir, gen_relocatable_json) + + +# Parse multiple soc_versions ops in single path. +def gen_all_soc_config(all_path): + soc_roots = glob.glob(os.path.join(all_path, "ascend*")) + + for soc_root in soc_roots: + soc = os.path.basename(soc_root) + gen_all_config(soc_root, soc, soc_root, True) + cfg_files = glob.glob(os.path.join(soc_root, "*.json")) + cfg_path = os.path.join(all_path, "config", soc) + os.makedirs(cfg_path, exist_ok=True) + for cfg_file in cfg_files: + new_file = os.path.join(cfg_path, os.path.basename(cfg_file)) + os.rename(cfg_file, new_file) + + +def args_prase(): + parser = argparse.ArgumentParser() + parser.add_argument('-p', + '--path', + nargs='?', + required=True, + help='Parse the path of the json file.') + + parser.add_argument('-s', + '--soc', + nargs='?', + required=True, + help='Parse the soc_version of ops.') + + parser.add_argument('-o', + '--out', + nargs='?', + help='Output directory.') + + parser.add_argument('--skip-binary-info-config', + action='store_true', + help='binary_info_config.json file is not parsed.') + + parser.add_argument('--op-range', + type=str, + choices=["all", "normal", "relocatable"], + default='all', + help='all operators/normal operators/relocatable operators.') + + return parser.parse_args() + + +def main(): + args = args_prase() + if args.out is None: + out_dir = args.path + else: + out_dir = args.out + + gen_all_config(args.path, args.soc, out_dir, + args.skip_binary_info_config, args.op_range) + + +if __name__ == '__main__': + main() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_pack_kernel.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_pack_kernel.py new file mode 100755 index 00000000..f50f1a0e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/ascendc_pack_kernel.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +# Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + +import os +import sys +import subprocess +import json +import glob +import argparse +import math +import const_var +import ascendc_ops_config +from tbe.tikcpp.log_utils import LogUtil, AscendCLogLevel + + +class PackKernel: + def __init__(self: any, args: any): + self.in_path = os.path.realpath(args.input_path) + self.out_path = os.path.realpath(args.output_path) + self.is_lib = args.enable_library + self.platform = args.platform + self.op_info = {} + self.file_info = {} + try: + os.makedirs(self.out_path, exist_ok=True) + except Exception as e: + LogUtil.print_compile_log("", f"make {self.out_path} error: {e}!", + AscendCLogLevel.LOG_ERROR, LogUtil.Option.NON_SOC) + + def load_json(self: any, json_file: str): + with open(json_file, encoding="utf-8") as file: + json_content = json.load(file) + return json_content + + def get_symbol(self: any, name: str): + name = name.replace("/", "_") + return name.replace(".", "_") + + def ascendc_gen_object(self: any, in_file: str, soc: str): + sym = self.get_symbol("_binary_" + in_file) + out_file = os.path.join(self.out_path, sym + ".o") + #ascend610lite only supoort aarch64 + if soc == 'ascend610lite': + try: + subprocess.run(['llvm-objcopy', '--input-target', 'binary', '--output-target', 'elf64-littleaarch64', + '--binary-architecture', 'aarch64', in_file, out_file]) + except Exception as e: + LogUtil.print_compile_log("", " ascend610lite execute objcopy fail!", + AscendCLogLevel.LOG_ERROR, LogUtil.Option.NON_SOC) + return None + return [sym + "_start", sym + "_end"] + uname = os.popen("uname -m").read().strip() + if self.platform is not None: + target_platform = self.platform + else: + target_platform = uname + try: + if target_platform == "x86_64": + subprocess.run(['llvm-objcopy', '--input-target', 'binary', '--output-target', 'elf64-x86-64', + '--binary-architecture', 'i386', in_file, out_file]) + elif target_platform == "aarch64": + subprocess.run(['llvm-objcopy', '--input-target', 'binary', '--output-target', 'elf64-littleaarch64', + '--binary-architecture', 'aarch64', in_file, out_file]) + else: + subprocess.run(['echo', 'unsported environment!']) + except Exception as e: + LogUtil.print_compile_log("", f"{target_platform} execute objcopy error: {e}!", + AscendCLogLevel.LOG_ERROR, LogUtil.Option.NON_SOC) + return None + return [sym + "_start", sym + "_end"] + + def ascendc_get_config(self: any): + os.chdir(self.in_path) + soc_vers = os.listdir("config") + for soc in soc_vers: + bin_infos = glob.glob(os.path.join("config", soc, "*.json")) + cfgs = {} + for bin_info in bin_infos: + if bin_info.find("binary_info_config.json") > 0: + continue + jobj = self.load_json(bin_info) + for bin_cfg in jobj.get("binList"): + js_cfg = bin_cfg.get("binInfo").get("jsonFilePath") + op_type = os.path.basename(js_cfg).split("_")[0] + if cfgs.get(op_type) is None: + op_obj = {} + op_obj["obj"] = [] + op_obj["cfg"] = bin_info + cfgs[op_type] = op_obj + op_obj = cfgs.get(op_type) + op_obj.get("obj").append(js_cfg[:-5]) + self.file_info[soc] = cfgs + + def ascendc_pack_kernel(self: any): + for soc in self.file_info.keys(): + os.chdir(self.in_path) + op_cfgs = self.file_info.get(soc) + for op_type in op_cfgs.keys(): + op_obj = op_cfgs.get(op_type) + if self.op_info.get(op_type) is None: + op_info = {} + op_info["op_fun"] = ["nullptr", "nullptr"] + op_info["op_bin"] = {} + op_info["op_rkb"] = [] + self.op_info[op_type] = op_info + op_info = self.op_info.get(op_type) + op_bin = op_info.get("op_bin") + if op_bin.get(soc) is None: + op_bin[soc] = [] + op_bin[soc].append(self.ascendc_gen_object(op_obj["cfg"], soc)) + op_soc = op_bin.get(soc) + for objs in op_obj["obj"]: + op_soc.append(self.ascendc_gen_object(objs + ".json", soc)) + op_soc.append(self.ascendc_gen_object(objs + ".o", soc)) + + def ascendc_gen_header(self: any): + for op_type in self.op_info.keys(): + op_obj = self.op_info.get(op_type) + macro_op = "#define {}_OP_RESOURCES std::make_tuple, \\\n" \ + " std::map>>, \\\n" \ + " std::vector>>({{{}}}, \\\n".format( + op_type, ", ".join(op_obj.get("op_fun")) + ) + op_bin = op_obj.get("op_bin") + socs_res = [] + op_syms = [] + for soc in op_bin.keys(): + soc_res = '{{ "{}", {{'.format(soc) + soc_syms = op_bin.get(soc) + soc_pairs = [] + for pair_addr in soc_syms: + pair_addr1 = ["&" + s for s in pair_addr] + op_syms += pair_addr + soc_pairs.append( + " {{ {} }} ".format(", \\\n ".join(pair_addr1)) + ) + soc_res += ", \\\n ".join(soc_pairs) + soc_res += " } }" + socs_res.append(soc_res) + macro_op += " {{ {} }}, \\\n".format(", \\\n ".join(socs_res)) + macro_op += " {{ {} }})\n\n".format(", ".join(op_obj.get("op_rkb"))) + macro_str = '#define {}_RESOURCES {{{{"{}", {}}}}}'.format( + op_type, op_type, "{}_OP_RESOURCES".format(op_type) + ) + var_str = ("extern gert::OpImplRegisterV2 op_impl_register_optiling_{};\n".format(op_type)) + if len(op_syms) > 0: + var_str += ('extern uint8_t ' + ";\nextern uint8_t ".join(op_syms) + ";\n") + head_file = os.path.join(self.out_path, "{}_op_resource.h".format(op_type)) + try: + with os.fdopen( + os.open(head_file, const_var.WFLAGS, const_var.WMODES), "w" + ) as fd: + fd.write("#include \n") + fd.write("#include \n") + fd.write("#include \n") + fd.write("#include \n") + fd.write('#include "graph/ascend_string.h"\n') + fd.write('#include "register/op_impl_registry.h"\n\n') + fd.write(var_str) + fd.write('\n') + fd.write(macro_op) + fd.write(macro_str) + except Exception as e: + LogUtil.print_compile_log("", f"{op_type}_op_resource.h create error: {e}!", + AscendCLogLevel.LOG_ERROR, LogUtil.Option.NON_SOC) + + def ascendc_gen_lib(self: any): + out_lib = os.path.join(self.out_path, "libkernels.a") + if os.path.exists(out_lib): + os.remove(out_lib) + objs = glob.glob(os.path.join(self.out_path, "*.o")) + start = 0 + batch_size = 100 + for _ in range(math.ceil(len(objs) / batch_size)): + sub_objs = objs[start : start + batch_size] + start += batch_size + try: + subprocess.run(['ar', 'qc', out_lib] + sub_objs) + subprocess.run(['ranlib', out_lib]) + except Exception as e: + LogUtil.print_compile_log("", f"execute ar/ranlib command error: {e}!", + AscendCLogLevel.LOG_ERROR, LogUtil.Option.NON_SOC) + + def ascendc_gen_opsinfo(self: any): + ascendc_ops_config.gen_all_soc_config(self.in_path) + + +def args_parse(): + parser = argparse.ArgumentParser() + parser.add_argument( + "-i", "--input-path", nargs="?", help="Input path of compile result." + ) + parser.add_argument( + "-o", "--output-path", nargs="?", help="Output path of compile result." + ) + parser.add_argument( + "-l", "--enable-library", nargs="?", default=None, help="Whether library is enabled." + ) + parser.add_argument( + "-p", "--platform", nargs="?", default=None, help="target platform is x86_64 or aarch64." + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = args_parse() + kernel_packer = PackKernel(args) + if kernel_packer.is_lib is None: + kernel_packer.ascendc_gen_opsinfo() + kernel_packer.ascendc_get_config() + kernel_packer.ascendc_pack_kernel() + kernel_packer.ascendc_gen_header() + kernel_packer.ascendc_gen_lib() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/const_var.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/const_var.py new file mode 100755 index 00000000..e7195b12 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/const_var.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Function: +The replay funtion entry +Copyright Information: +Huawei Technologies Co., Ltd. All Rights Reserved © 2020 +""" + +import os +import stat + + +REPLAY_BATCH = 'batch' +REPLAY_ITERATE = 'iterate' +CFG_IMPL_DIR = 'impl_dir' +CFG_OUT_DIR = 'out_dir' +AUTO_GEN_DIR = 'auto_gen_dir' +WFLAGS = os.O_WRONLY | os.O_CREAT | os.O_TRUNC +WMODES = stat.S_IWUSR | stat.S_IRUSR +SOC_MAP_EXT = {'ascend310p': 'Ascend310P3', 'ascend310b': 'Ascend310B1', + 'ascend910': 'Ascend910A', 'ascend910b': 'Ascend910B1', + 'ascend910_93': 'Ascend910_9391', 'ascend610lite': 'Ascend610Lite', + 'ascend910_95': 'Ascend910_9599'} +BIN_CMD = 'opc $1 --main_func={fun} --input_param={param} --soc_version={soc} \ +--output=$2 --impl_mode={impl} --simplified_key_mode=0 --op_mode=dynamic\n' +SET_PLOG_LEVEL_ERROR = "export ASCEND_GLOBAL_LOG_LEVEL=3\n" +SET_PLOG_STDOUT = "export ASCEND_SLOG_PRINT_TO_STDOUT=1\n" +SRC_ENV = ''' +while true; do + case "$1" in + --kernel-src=*) + export BUILD_KERNEL_SRC=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + -*) + shift + ;; + *) + break + ;; + esac +done +''' +CHK_CMD = ''' +if ! test -f $2/{res_file} ; then + echo "$2/{res_file} not generated!" + exit 1 +fi +''' +ATTR_DEF_VAL = {'str' : '', 'int': 0, 'float': 0.0, 'bool': False, 'list_bool': [], + 'list_int': [], 'list_float': [], 'list_list_int': [[]]} + + +def conv_soc_ver(ver: str): + return SOC_MAP_EXT.get(ver) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/COPYING b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/COPYING new file mode 100644 index 00000000..d159169d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/COPYING @@ -0,0 +1,339 @@ + GNU GENERAL PUBLIC LICENSE + Version 2, June 1991 + + Copyright (C) 1989, 1991 Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The licenses for most software are designed to take away your +freedom to share and change it. By contrast, the GNU General Public +License is intended to guarantee your freedom to share and change free +software--to make sure the software is free for all its users. This +General Public License applies to most of the Free Software +Foundation's software and to any other program whose authors commit to +using it. (Some other Free Software Foundation software is covered by +the GNU Lesser General Public License instead.) You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +this service if you wish), that you receive source code or can get it +if you want it, that you can change the software or use pieces of it +in new free programs; and that you know you can do these things. + + To protect your rights, we need to make restrictions that forbid +anyone to deny you these rights or to ask you to surrender the rights. +These restrictions translate to certain responsibilities for you if you +distribute copies of the software, or if you modify it. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must give the recipients all the rights that +you have. You must make sure that they, too, receive or can get the +source code. And you must show them these terms so they know their +rights. + + We protect your rights with two steps: (1) copyright the software, and +(2) offer you this license which gives you legal permission to copy, +distribute and/or modify the software. + + Also, for each author's protection and ours, we want to make certain +that everyone understands that there is no warranty for this free +software. If the software is modified by someone else and passed on, we +want its recipients to know that what they have is not the original, so +that any problems introduced by others will not reflect on the original +authors' reputations. + + Finally, any free program is threatened constantly by software +patents. We wish to avoid the danger that redistributors of a free +program will individually obtain patent licenses, in effect making the +program proprietary. To prevent this, we have made it clear that any +patent must be licensed for everyone's free use or not licensed at all. + + The precise terms and conditions for copying, distribution and +modification follow. + + GNU GENERAL PUBLIC LICENSE + TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + + 0. This License applies to any program or other work which contains +a notice placed by the copyright holder saying it may be distributed +under the terms of this General Public License. The "Program", below, +refers to any such program or work, and a "work based on the Program" +means either the Program or any derivative work under copyright law: +that is to say, a work containing the Program or a portion of it, +either verbatim or with modifications and/or translated into another +language. (Hereinafter, translation is included without limitation in +the term "modification".) Each licensee is addressed as "you". + +Activities other than copying, distribution and modification are not +covered by this License; they are outside its scope. The act of +running the Program is not restricted, and the output from the Program +is covered only if its contents constitute a work based on the +Program (independent of having been made by running the Program). +Whether that is true depends on what the Program does. + + 1. You may copy and distribute verbatim copies of the Program's +source code as you receive it, in any medium, provided that you +conspicuously and appropriately publish on each copy an appropriate +copyright notice and disclaimer of warranty; keep intact all the +notices that refer to this License and to the absence of any warranty; +and give any other recipients of the Program a copy of this License +along with the Program. + +You may charge a fee for the physical act of transferring a copy, and +you may at your option offer warranty protection in exchange for a fee. + + 2. You may modify your copy or copies of the Program or any portion +of it, thus forming a work based on the Program, and copy and +distribute such modifications or work under the terms of Section 1 +above, provided that you also meet all of these conditions: + + a) You must cause the modified files to carry prominent notices + stating that you changed the files and the date of any change. + + b) You must cause any work that you distribute or publish, that in + whole or in part contains or is derived from the Program or any + part thereof, to be licensed as a whole at no charge to all third + parties under the terms of this License. + + c) If the modified program normally reads commands interactively + when run, you must cause it, when started running for such + interactive use in the most ordinary way, to print or display an + announcement including an appropriate copyright notice and a + notice that there is no warranty (or else, saying that you provide + a warranty) and that users may redistribute the program under + these conditions, and telling the user how to view a copy of this + License. (Exception: if the Program itself is interactive but + does not normally print such an announcement, your work based on + the Program is not required to print an announcement.) + +These requirements apply to the modified work as a whole. If +identifiable sections of that work are not derived from the Program, +and can be reasonably considered independent and separate works in +themselves, then this License, and its terms, do not apply to those +sections when you distribute them as separate works. But when you +distribute the same sections as part of a whole which is a work based +on the Program, the distribution of the whole must be on the terms of +this License, whose permissions for other licensees extend to the +entire whole, and thus to each and every part regardless of who wrote it. + +Thus, it is not the intent of this section to claim rights or contest +your rights to work written entirely by you; rather, the intent is to +exercise the right to control the distribution of derivative or +collective works based on the Program. + +In addition, mere aggregation of another work not based on the Program +with the Program (or with a work based on the Program) on a volume of +a storage or distribution medium does not bring the other work under +the scope of this License. + + 3. You may copy and distribute the Program (or a work based on it, +under Section 2) in object code or executable form under the terms of +Sections 1 and 2 above provided that you also do one of the following: + + a) Accompany it with the complete corresponding machine-readable + source code, which must be distributed under the terms of Sections + 1 and 2 above on a medium customarily used for software interchange; or, + + b) Accompany it with a written offer, valid for at least three + years, to give any third party, for a charge no more than your + cost of physically performing source distribution, a complete + machine-readable copy of the corresponding source code, to be + distributed under the terms of Sections 1 and 2 above on a medium + customarily used for software interchange; or, + + c) Accompany it with the information you received as to the offer + to distribute corresponding source code. (This alternative is + allowed only for noncommercial distribution and only if you + received the program in object code or executable form with such + an offer, in accord with Subsection b above.) + +The source code for a work means the preferred form of the work for +making modifications to it. For an executable work, complete source +code means all the source code for all modules it contains, plus any +associated interface definition files, plus the scripts used to +control compilation and installation of the executable. However, as a +special exception, the source code distributed need not include +anything that is normally distributed (in either source or binary +form) with the major components (compiler, kernel, and so on) of the +operating system on which the executable runs, unless that component +itself accompanies the executable. + +If distribution of executable or object code is made by offering +access to copy from a designated place, then offering equivalent +access to copy the source code from the same place counts as +distribution of the source code, even though third parties are not +compelled to copy the source along with the object code. + + 4. You may not copy, modify, sublicense, or distribute the Program +except as expressly provided under this License. Any attempt +otherwise to copy, modify, sublicense or distribute the Program is +void, and will automatically terminate your rights under this License. +However, parties who have received copies, or rights, from you under +this License will not have their licenses terminated so long as such +parties remain in full compliance. + + 5. You are not required to accept this License, since you have not +signed it. However, nothing else grants you permission to modify or +distribute the Program or its derivative works. These actions are +prohibited by law if you do not accept this License. Therefore, by +modifying or distributing the Program (or any work based on the +Program), you indicate your acceptance of this License to do so, and +all its terms and conditions for copying, distributing or modifying +the Program or works based on it. + + 6. Each time you redistribute the Program (or any work based on the +Program), the recipient automatically receives a license from the +original licensor to copy, distribute or modify the Program subject to +these terms and conditions. You may not impose any further +restrictions on the recipients' exercise of the rights granted herein. +You are not responsible for enforcing compliance by third parties to +this License. + + 7. If, as a consequence of a court judgment or allegation of patent +infringement or for any other reason (not limited to patent issues), +conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot +distribute so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you +may not distribute the Program at all. For example, if a patent +license would not permit royalty-free redistribution of the Program by +all those who receive copies directly or indirectly through you, then +the only way you could satisfy both it and this License would be to +refrain entirely from distribution of the Program. + +If any portion of this section is held invalid or unenforceable under +any particular circumstance, the balance of the section is intended to +apply and the section as a whole is intended to apply in other +circumstances. + +It is not the purpose of this section to induce you to infringe any +patents or other property right claims or to contest validity of any +such claims; this section has the sole purpose of protecting the +integrity of the free software distribution system, which is +implemented by public license practices. Many people have made +generous contributions to the wide range of software distributed +through that system in reliance on consistent application of that +system; it is up to the author/donor to decide if he or she is willing +to distribute software through any other system and a licensee cannot +impose that choice. + +This section is intended to make thoroughly clear what is believed to +be a consequence of the rest of this License. + + 8. If the distribution and/or use of the Program is restricted in +certain countries either by patents or by copyrighted interfaces, the +original copyright holder who places the Program under this License +may add an explicit geographical distribution limitation excluding +those countries, so that distribution is permitted only in or among +countries not thus excluded. In such case, this License incorporates +the limitation as if written in the body of this License. + + 9. The Free Software Foundation may publish revised and/or new versions +of the General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + +Each version is given a distinguishing version number. If the Program +specifies a version number of this License which applies to it and "any +later version", you have the option of following the terms and conditions +either of that version or of any later version published by the Free +Software Foundation. If the Program does not specify a version number of +this License, you may choose any version ever published by the Free Software +Foundation. + + 10. If you wish to incorporate parts of the Program into other free +programs whose distribution conditions are different, write to the author +to ask for permission. For software which is copyrighted by the Free +Software Foundation, write to the Free Software Foundation; we sometimes +make exceptions for this. Our decision will be guided by the two goals +of preserving the free status of all derivatives of our free software and +of promoting the sharing and reuse of software generally. + + NO WARRANTY + + 11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY +FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN +OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES +PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED +OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS +TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE +PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, +REPAIR OR CORRECTION. + + 12. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY AND/OR +REDISTRIBUTE THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, +INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING +OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED +TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY +YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER +PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE +POSSIBILITY OF SUCH DAMAGES. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +convey the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software; you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation; either version 2 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License along + with this program; if not, write to the Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + +Also add information on how to contact you by electronic and paper mail. + +If the program is interactive, make it output a short notice like this +when it starts in an interactive mode: + + Gnomovision version 69, Copyright (C) year name of author + Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, the commands you use may +be called something other than `show w' and `show c'; they could even be +mouse-clicks or menu items--whatever suits your program. + +You should also get your employer (if you work as a programmer) or your +school, if any, to sign a "copyright disclaimer" for the program, if +necessary. Here is a sample; alter the names: + + Yoyodyne, Inc., hereby disclaims all copyright interest in the program + `Gnomovision' (which makes passes at compilers) written by James Hacker. + + , 1 April 1989 + Ty Coon, President of Vice + +This General Public License does not permit incorporating your program into +proprietary programs. If your program is a subroutine library, you may +consider it more useful to permit linking proprietary applications with the +library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/Makefile b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/Makefile new file mode 100644 index 00000000..46382a78 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/Makefile @@ -0,0 +1,18 @@ +.PHONY: all clean test help + +VERSION := $(shell cat VERSION) +OUTPUT := makeself-$(VERSION).run + +all: $(OUTPUT) + +$(OUTPUT): makeself.sh makeself-header.sh VERSION + ./make-release.sh + +clean: + $(RM) makeself-*.run + +test: + ./test/run-tests.sh + +help: + $(info Targets: all $(OUTPUT) clean test help) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/README.md b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/README.md new file mode 100644 index 00000000..d112f2c1 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/README.md @@ -0,0 +1,253 @@ +[![License: GPL v2](https://img.shields.io/badge/License-GPL%20v2-blue.svg)](https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html) +![Build Status](https://github.com/megastep/makeself/workflows/CI/badge.svg) + +# makeself - Make self-extractable archives on Unix + +[makeself.sh][1] is a small shell script that generates a self-extractable +compressed tar archive from a directory. The resulting file appears as a shell script +(many of those have a **.run** suffix), and can be launched as is. The archive +will then uncompress itself to a temporary directory and an optional arbitrary +command will be executed (for example an installation script). This is pretty +similar to archives generated with WinZip Self-Extractor in the Windows world. +Makeself archives also include checksums for integrity self-validation (CRC +and/or MD5/SHA256 checksums). + +The makeself.sh script itself is used only to create the archives from a +directory of files. The resultant archive is actually a compressed (using +gzip, bzip2, or compress) TAR archive, with a small shell script stub at the +beginning. This small stub performs all the steps of extracting the files, +running the embedded command, and removing the temporary files when done. +All the user has to do to install the software contained in such an +archive is to "run" the archive, i.e **sh nice-software.run**. I recommend +using the ".run" (which was introduced by some Makeself archives released by +Loki Software) or ".sh" suffix for such archives not to confuse the users, +so that they will know they are actually shell scripts (with quite a lot of binary data +attached to them though!). + +I am trying to keep the code of this script as portable as possible, i.e it is +not relying on any bash-specific features and only calls commands that are +installed on any functioning UNIX-compatible system. This script as well as +the archives it generates should run on any Unix flavor, with any compatible +Bourne shell, provided of course that the compression programs are available. + +Makeself has been rewritten and tested on the following platforms : + +* Linux (all distributions) +* Sun Solaris (8 and above) +* HP-UX (tested on 11.0 and 11i on HPPA RISC) +* SCO OpenUnix and OpenServer +* IBM AIX +* macOS (Darwin) +* SGI IRIX 6.5 +* FreeBSD +* OpenBSD +* NetBSD +* UnicOS / Cray +* Windows (Cygwin, WSL) + +If you successfully run Makeself and/or archives created with it on another +system, then please [let me know][2]! + +Examples of publicly available archives made using makeself are : + +* Game patches and installers for [Id Software][3] games like Quake 3 for Linux or Return To Castle Wolfenstein ; +* All game patches released by [Loki Software][4] for the Linux version of popular games ; +* The [nVidia drivers][5] for Linux +* The installer for the Linux version of [Google Earth][6] +* The [VirtualBox][7] installers for Linux +* The [Makeself][1] distribution itself ;-) +* and countless others... + +**Important note for Apache users:** By default, most Web servers will think that Makeself archives are regular text files and thus they may show up as text in a Web browser. The correct way to prevent this is to add a MIME type for this file format, like so (in httpd.conf) : + +`AddType application/x-makeself .run` + +**Important note for certain GNU/Linux distributions:** Archives created with Makeself prior to v2.1.2 were using an old syntax for the _head_ and _tail_ Unix commands that is being progressively obsoleted in their GNU forms. Therefore you may have problems uncompressing some of these archives. A workaround for this is to set the environment variable $_POSIX2_VERSION to enable the old syntax, i.e. : + +`export _POSIX2_VERSION=199209` + +## Usage + +The syntax of makeself is the following: + +```sh +makeself.sh [args] archive_dir file_name label startup_script [script_args] +``` + + * _args_ are optional options for Makeself. The available ones are : + + * **`--version`** : Prints the version number on stdout, then exits immediately + * **`--gzip`** : Use gzip for compression (the default on platforms on which gzip is commonly available, like Linux) + * **`--bzip2`** : Use bzip2 instead of gzip for better compression. The bzip2 command must be available in the command path. It is recommended that the archive extension be set to something like '.bz2.run', so that potential users know that they'll need bzip2 to extract it. + * **`--bzip3`** : Use bzip3 instead of gzip for better compression. + * **`--pbzip2`** : Use pbzip2 instead of gzip for better and faster compression on machines having multiple CPUs. The pbzip2 command must be available in the command path. It is recommended that the archive extension be set to something like '.bz2.run', so that potential users know that they'll need bzip2 to extract it. + * **`--xz`** : Use xz instead of gzip for better compression. The xz command must be available in the command path. It is recommended that the archive extension be set to something like '.xz.run' for the archive, so that potential users know that they'll need xz to extract it. + * **`--lzo`** : Use lzop instead of gzip for better compression. The lzop command must be available in the command path. It is recommended that the archive extension be set to something like `.lzo.run` for the archive, so that potential users know that they'll need lzop to extract it. + * **`--lz4`** : Use lz4 instead of gzip for better compression. The lz4 command must be available in the command path. It is recommended that the archive extension be set to something like '.lz4.run' for the archive, so that potential users know that they'll need lz4 to extract it. + * **`--zstd`** : Use zstd instead of gzip for better compression. The zstd command must be available in the command path. It is recommended that the archive extension be set to something like '.zstd.run' for the archive, so that potential users know that they'll need zstd to extract it. + * **`--pigz`** : Use pigz for compression. + * **`--base64`** : Encode the archive to ASCII in Base64 format instead of compressing (base64 command required). + * **`--gpg-encrypt`** : Encrypt the archive using `gpg -ac -z $COMPRESS_LEVEL`. This will prompt for a password to encrypt with. Assumes that potential users have `gpg` installed. + * **`--ssl-encrypt`** : Encrypt the archive using `openssl aes-256-cbc -a -salt`. This will prompt for a password to encrypt with. Assumes that the potential users have the OpenSSL tools installed. + * **`--compress`** : Use the UNIX `compress` command to compress the data. This should be the default on all platforms that don't have gzip available. + * **`--nocomp`** : Do not use any compression for the archive, which will then be an uncompressed TAR. + * **`--complevel`** : Specify the compression level for gzip, bzip2, pbzip2, zstd, xz, lzo or lz4. (defaults to 9) + * **`--threads`** : Specify the number of threads to be used by compressors that support parallelization. Omit to use compressor's default. Most useful (and required) for opting into xz's threading, usually with `--threads=0` for all available cores. pbzip2 and pigz are parallel by default, and setting this value allows limiting the number of threads they use. + * **`--notemp`** : The generated archive will not extract the files to a temporary directory, but in a new directory created in the current directory. This is better to distribute software packages that may extract and compile by themselves (i.e. launch the compilation through the embedded script). + * **`--current`** : Files will be extracted to the current directory, instead of in a subdirectory. This option implies `--notemp` above. + * **`--follow`** : Follow the symbolic links inside of the archive directory, i.e. store the files that are being pointed to instead of the links themselves. + * **`--append`** _(new in 2.1.x)_: Append data to an existing archive, instead of creating a new one. In this mode, the settings from the original archive are reused (compression type, label, embedded script), and thus don't need to be specified again on the command line. + * **`--header`** : Makeself uses a separate file to store the header stub, called `makeself-header.sh`. By default, it is assumed that it is stored in the same location as makeself.sh. This option can be used to specify its actual location if it is stored someplace else. + * **`--cleanup`** : Specify a script that is run when execution is interrupted or finishes successfully. The script is executed with the same environment and initial `script_args` as `startup_script`. + * **`--copy`** : Upon extraction, the archive will first extract itself to a temporary directory. The main application of this is to allow self-contained installers stored in a Makeself archive on a CD, when the installer program will later need to unmount the CD and allow a new one to be inserted. This prevents "Filesystem busy" errors for installers that span multiple CDs. + * **`--nox11`** : Disable the automatic spawning of a new terminal in X11. + * **`--nowait`** : When executed from a new X11 terminal, disable the user prompt at the end of the script execution. + * **`--nomd5`** and **`--nocrc`** : Disable the creation of a MD5 / CRC checksum for the archive. This speeds up the extraction process if integrity checking is not necessary. + * **`--sha256`** : Adds a SHA256 checksum for the archive. This is in addition to the MD5 / CRC checksums unless `--nomd5` is also used. + * **`--lsm` _file_** : Provide a Linux Software Map (LSM) file to makeself, that will be embedded in the generated archive. LSM files are describing a software package in a way that is easily parseable. The LSM entry can then be later retrieved using the `--lsm` argument to the archive. An example of a LSM file is provided with Makeself. + * **`--tar-format opt`** : Specify the tar archive format (default is ustar); you may use any value accepted by your tar command (such as posix, v7, etc). + * **`--tar-extra opt`** : Append more options to the tar command line. + + For instance, in order to exclude the `.git` directory from the packaged archive directory using the GNU `tar`, one can use `makeself.sh --tar-extra "--exclude=.git" ...` + + * **`--keep-umask`** : Keep the umask set to shell default, rather than overriding when executing the self-extracting archive. + * **`--packaging-date date`** : Use provided string as the packaging date instead of the current date. + * **`--license` _file_** : Append a license file. + * **`--nooverwrite`** : Do not extract the archive if the specified target directory already exists. + * **`--help-header` _file_** : Add a header to the archive's `--help` output. + * `archive_dir` is the name of the directory that contains the files to be archived + * `file_name` is the name of the archive to be created + * `label` is an arbitrary text string describing the package. It will be displayed while extracting the files. + * `startup_script` is the command to be executed _from within_ the directory of extracted files. Thus, if you wish to execute a program contained in this directory, you must prefix your command with `./`. For example, `./program` will be fine. The `script_args` are additional arguments for this command. + Note that `startup_script` and its arguments are not strictly required for archives that don't extract in a temporary directory (i.e. when using `--notemp`). + +Here is an example, assuming the user has a package image stored in a **/home/joe/mysoft**, and he wants to generate a self-extracting package named +**mysoft.sh**, which will launch the "setup" script initially stored in /home/joe/mysoft : + +```sh +makeself.sh /home/joe/mysoft mysoft.sh "Joe's Nice Software Package" ./setup +``` + +Here is also how I created the [makeself.run][9] archive which contains the Makeself distribution : + +`makeself.sh --notemp makeself makeself.run "Makeself by Stephane Peter" echo "Makeself has extracted itself"` + +Archives generated with Makeself can be passed the following arguments: + +* **`--keep`** : Prevent the files to be extracted in a temporary directory that will be removed after the embedded script's execution. The files will then be extracted in the current working directory and will stay here until you remove them. +* **`--verbose`** : Will prompt the user before executing the embedded command +* **`--target dir`** : Allows to extract the archive in an arbitrary place. +* **`--nox11`** : Do not spawn a X11 terminal. +* **`--confirm`** : Prompt the user for confirmation before running the embedded command. +* **`--info`** : Print out general information about the archive (does not extract). +* **`--lsm`** : Print out the LSM entry, if it is present. +* **`--list`** : List the files in the archive. +* **`--check`** : Check the archive for integrity using the embedded checksums. Does not extract the archive. +* **`--nochown`** : By default, a `chown -R` command is run on the target directory after extraction, so that all files belong to the current user. This is mostly needed if you are running as root, as tar will then try to recreate the initial user ownerships. You may disable this behavior with this flag. +* **`--tar`** : Run the tar command on the contents of the archive, using the following arguments as parameter for the command. +* **`--noexec`** : Do not run the embedded script after extraction. +* **`--noexec-cleanup`** : Do not run the embedded cleanup script. +* **`--nodiskspace`** : Do not check for available disk space before attempting to extract. +* **`--cleanup-args`** : Specify arguments to be passed to the cleanup script. Wrap value in quotes to specify multiple arguments. + +Any subsequent arguments to the archive will be passed as additional arguments to the embedded command. You must explicitly use the `--` special command-line construct before any such options to make sure that Makeself will not try to interpret them. + +## Startup Script + +The startup script must be a regular Shell script. + +Within the startup script, you can use the `$USER_PWD` variable to get the path of the folder from which the self-extracting script is executed. This is especially useful to access files that are located in the same folder as the script, as shown in the example below. + +```sh +my-self-extracting-script.sh --fooBarFileParameter foo.bar +``` + +## Building and Testing + +Clone the git repo and execute `git submodule update --init --recursive` to obtain all submodules. + +* To make a release: `make` +* To run all tests: `make test` + +## Maven Usage + +Makeself is now supported by the following maven plugin [makeself-maven-plugin](https://github.com/hazendaz/makeself-maven-plugin). Please refer to project for usage and report any bugs in regards to maven plugin on that project. + +## License + +Makeself itself is covered by the [GNU General Public License][8] (GPL) version 2 and above. Archives generated by Makeself don't have to be placed under this license (although I encourage it ;-)), since the archive itself is merely data for Makeself. + +## Contributing + +I will gladly consider merging your pull requests on the [GitHub][10] repository. However, please keep the following in mind: + +* One of the main purposes of Makeself is portability. Do not submit patches that will break supported platforms. The more platform-agnostic, the better. +* Please explain clearly what the purpose of the patch is, and how you achieved it. + +## Download + +Get the latest official distribution [here][9] (version 2.5.0). + +The latest development version can be grabbed from [GitHub][10]. Feel free to submit any patches there through the fork and pull request process. + +## Version history + +* **v1.0:** Initial public release +* **v1.1:** The archive can be passed parameters that will be passed on to the embedded script, thanks to John C. Quillan +* **v1.2:** Cosmetic updates, support for bzip2 compression and non-temporary archives. Many ideas thanks to Francois Petitjean. +* **v1.3:** More patches from Bjarni R. Einarsson and Francois Petitjean: Support for no compression (`--nocomp`), script is no longer mandatory, automatic launch in an xterm, optional verbose output, and -target archive option to indicate where to extract the files. +* **v1.4:** Many patches from Francois Petitjean: improved UNIX compatibility, automatic integrity checking, support of LSM files to get info on the package at run time.. +* **v1.5.x:** A lot of bugfixes, and many other patches, including automatic verification through the usage of checksums. Version 1.5.5 was the stable release for a long time, even though the Web page didn't get updated ;-). Makeself was also officially made a part of the [Loki Setup installer][11], and its source is being maintained as part of this package. +* **v2.0:** Complete internal rewrite of Makeself. The command-line parsing was vastly improved, the overall maintenance of the package was greatly improved by separating the stub from makeself.sh. Also Makeself was ported and tested to a variety of Unix platforms. +* **v2.0.1:** First public release of the new 2.0 branch. Prior versions are officially obsoleted. This release introduced the `--copy` argument that was introduced in response to a need for the [UT2K3][12] Linux installer. +* **v2.1.0:** Big change : Makeself can now support multiple embedded tarballs, each stored separately with their own checksums. An existing archive can be updated with the `--append` flag. Checksums are also better managed, and the `--nochown` option for archives appeared. +* **v2.1.1:** Fixes related to the Unix compression (compress command). Some Linux distributions made the insane choice to make it unavailable, even though gzip is capable of uncompressing these files, plus some more bugfixes in the extraction and checksum code. +* **v2.1.2:** Some bug fixes. Use head -n to avoid problems with POSIX conformance. +* **v2.1.3:** Bug fixes with the command line when spawning terminals. Added `--tar`, `--noexec` for archives. Added `--nomd5` and `--nocrc` to avoid creating checksums in archives. The embedded script is now run through "eval". The `--info` output now includes the command used to create the archive. A man page was contributed by Bartosz Fenski. +* **v2.1.4:** Fixed `--info` output. Generate random directory name when extracting files to . to avoid problems. Better handling of errors with wrong permissions for the directory containing the files. Avoid some race conditions, Unset the $CDPATH variable to avoid problems if it is set. Better handling of dot files in the archive directory. +* **v2.1.5:** Made the md5sum detection consistent with the header code. Check for the presence of the archive directory. Added `--encrypt` for symmetric encryption through gpg (Eric Windisch). Added support for the digest command on Solaris 10 for MD5 checksums. Check for available disk space before extracting to the target directory (Andreas Schweitzer). Allow extraction to run asynchronously (patch by Peter Hatch). Use file descriptors internally to avoid error messages (patch by Kay Tiong Khoo). +* **v2.1.6:** Replaced one dot per file progress with a realtime progress percentage and a spinning cursor. Added `--noprogress` to prevent showing the progress during the decompression. Added `--target` dir to allow extracting directly to a target directory. (Guy Baconniere) +* **v2.2.0:** First major new release in years! Includes many bugfixes and user contributions. Please look at the [project page on Github][10] for all the details. +* **v2.3.0:** Support for archive encryption via GPG or OpenSSL. Added LZO and LZ4 compression support. Options to set the packaging date and stop the umask from being overriden. Optionally ignore check for available disk space when extracting. New option to check for root permissions before extracting. +* **v2.3.1:** Various compatibility updates. Added unit tests for Travis CI in the GitHub repo. New `--tar-extra`, `--untar-extra`, `--gpg-extra`, `--gpg-asymmetric-encrypt-sign` options. +* **v2.4.0:** Added optional support for SHA256 archive integrity checksums. +* **v2.4.2:** New --cleanup and --cleanup-args arguments for cleanup scripts. Added threading support for supported compressors. Now supports zstd compression. +* **v2.4.3:** Make explicit POSIX tar archives for increased compatibility. +* **v2.4.4:** Fixed various compatibility issues (no longer use POSIX tar archives), Github Actions to check on Solaris and FreeBSD. +* **v2.4.5:** Added `--tar-format` option to set the tar archive format (default is ustar) +* **v2.5.0:** Expended support to NetBSD, OpenBSD, Busybox and other minimal distributions such as Alpine Linux. Added bzip3 compression support and expanded GPG arguments. + +## Links + +* Check out the ["Loki Setup"][11] installer, used to install many Linux games and other applications, and of which I am the co-author. Since the demise of Loki, I am now the official maintainer of the project, and it is now being hosted here on GitHub. +* Bjarni R. Einarsson also wrote the **setup.sh** installer script, inspired by Makeself. [Check it out !][14] + +## Contact + +This script was written by [Stéphane Peter][15] (megastep at megastep.org). Any enhancements and suggestions are welcome. + +Contributions were included from John C. Quillan, Bjarni R. Einarsson, +Francois Petitjean, Ryan C. Gordon, and many contributors on GitHub. If you think I forgot +your name, don't hesitate to contact me. + +This project is now hosted on GitHub. Feel free to submit patches and bug reports on the [project page][10]. + +* * * + +[Stephane Peter][2] + + [1]: http://makeself.io/ + [2]: mailto:megastep@megastep.org + [3]: http://www.idsoftware.com/ + [4]: http://www.lokigames.com/products/myth2/updates.php3 + [5]: http://www.nvidia.com/ + [6]: http://earth.google.com/ + [7]: http://www.virtualbox.org/ + [8]: http://www.gnu.org/copyleft/gpl.html + [9]: https://github.com/megastep/makeself/releases/download/release-2.5.0/makeself-2.5.0.run + [10]: https://github.com/megastep/makeself + [11]: https://github.com/megastep/loki_setup/ + [12]: http://www.unrealtournament2003.com/ + [13]: http://www.icculus.org/ + [14]: http://bre.klaki.net/programs/setup.sh/ + [15]: https://stephanepeter.com/ diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/VERSION b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/VERSION new file mode 100644 index 00000000..437459cd --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/VERSION @@ -0,0 +1 @@ +2.5.0 diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/make-release.sh b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/make-release.sh new file mode 100755 index 00000000..df8dc4b9 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/make-release.sh @@ -0,0 +1,8 @@ +#!/bin/sh +# +# Create a distributable archive of the current version of Makeself + +VER=`cat VERSION` +mkdir -p /tmp/makeself-$VER release +cp -pPR makeself* README.md COPYING VERSION /tmp/makeself-$VER/ +./makeself.sh --notemp /tmp/makeself-$VER release/makeself-$VER.run "Makeself v$VER" echo "Makeself has extracted itself" diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself-header.sh b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself-header.sh new file mode 100755 index 00000000..48256c01 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself-header.sh @@ -0,0 +1,681 @@ +cat << EOF > "$archname" +#!/bin/bash +# This script was generated using Makeself $MS_VERSION +# The license covering this archive and its contents, if any, is wholly independent of the Makeself license (GPL) + +ORIG_UMASK=\`umask\` + +SHA="$SHAsum" +SHA_ARG="" +SIGNATURE="$Signature" +TMPROOT=\${TMPDIR:="/home/tmpdir"} + +if ! test -d "\$TMPROOT"; then + TMPROOT="\$HOME" +fi +if ! test -d "\$TMPROOT"; then + TMPROOT="\$PWD" +fi +export TMPDIR="\$TMPROOT" +USER_PWD="\$PWD" +if ! test -d "\$USER_PWD"; then + exit 1 +fi +export USER_PWD +ARCHIVE_DIR=\`dirname "\$0"\` +export ARCHIVE_DIR + +name_of_file="\$0 " +package_name=\`echo \$name_of_file | cut -d "/" -f2 | sed "s/.run//g" \` +pwd_of_file="\$PWD" +label="$LABEL" +script="$SCRIPT" +scriptargs="$SCRIPTARGS" +cleanup_script="${CLEANUP_SCRIPT}" +licensetxt="$LICENSE" +helpheader='$HELPHEADER' +targetdir="$archdirname" +filesizes="$filesizes" +totalsize="$totalsize" +keep="$KEEP" +nooverwrite="$NOOVERWRITE" +quiet="n" +accept="n" +nodiskspace="n" +export_conf="$EXPORT_CONF" +decrypt_cmd="$DECRYPT_CMD" +skip="$SKIP" +PACKAGE_LOG_NAME=makeself +readonly user_n=\$(whoami) +info_record_path="\$HOME/log/makeself" +info_record_file="makeself.log" +info_record_file_bak="makeself.log.bak" +log_file=\$info_record_path/\$info_record_file +LOG_SIZE_THRESHOLD=1024000 + +print_cmd_arg="" +if type printf > /dev/null; then + print_cmd="printf" +elif test -x /usr/ucb/echo; then + print_cmd="/usr/ucb/echo" +else + print_cmd="echo" +fi + +if test -d /usr/xpg4/bin; then + PATH=/usr/xpg4/bin:\$PATH + export PATH +fi + +if test -d /usr/sfw/bin; then + PATH=\$PATH:/usr/sfw/bin + export PATH +fi + +unset CDPATH + +function rotate_log() { + check_path "\$log_file" + mv -f "\$log_file" "\$info_record_path/\$info_record_file_bak" + touch "\$log_file" 2>/dev/null + check_path "\$info_record_path/\$info_record_file_bak" + chmod 440 "\$info_record_path/\$info_record_file_bak" + check_path "\$log_file" + chmod 640 "\$log_file" +} +function check_path() { + if [ "\$1" != \$(readlink -f "\$1") ]; then + echo >&2 + echo "Log file is not support symlink, exiting!" >&2 + exit 1 + fi +} + +function log_check() { + local log_size + log_size=\$(find \$log_file -exec ls -l {} \; | awk '{ print \$5 }') + if [[ "\${log_size}" -ge "\${LOG_SIZE_THRESHOLD}" ]];then + rotate_log + fi +} + +# usage log "INFO" "this is message" +function log() { + if [ ! -d "\$info_record_path" ];then + mkdir -p "\$info_record_path" + chmod 750 "\$info_record_path" + fi + if [[ ! -f "\$log_file" ]];then + touch "\$log_file" + chmod 640 "\$log_file" + fi + # print log to log file + if [ -f "\$log_file" ]; then + log_check "\$log_file" + if ! echo -e "[\${PACKAGE_LOG_NAME}] [\${package_name}][\$(date +%Y%m%d-%H:%M:%S)] [\$user_n] [\$1] \$2" >>"\$log_file" + then + echo "can not write log, exiting!" >&2 + exit 1 + fi + else + echo "log file not exist, exiting!" >&2 + exit 1 + fi +} + +MS_Printf() +{ + \$print_cmd \$print_cmd_arg "\$1" +} + +MS_PrintLicense() +{ + PAGER=\${PAGER:=more} + if test x"\$licensetxt" != x; then + PAGER_PATH=\`exec <&- 2>&-; which \$PAGER || command -v \$PAGER || type \$PAGER\` + if test -x "\$PAGER_PATH"; then + echo "\$licensetxt" | \$PAGER + else + echo "\$licensetxt" + fi + if test x"\$accept" != xy; then + while true + do + MS_Printf "Please type y to accept, n otherwise: " + read yn + if test x"\$yn" = xn; then + keep=n + eval \$finish; exit 1 + break; + elif test x"\$yn" = xy; then + break; + fi + done + fi + fi +} + +MS_diskspace() +{ + ( + df -kP "\$1" | tail -1 | awk '{ if (\$4 ~ /%/) {print \$3} else {print \$4} }' + ) +} + +MS_dd() +{ + blocks=\`expr \$3 / 1024\` + bytes=\`expr \$3 % 1024\` + # Test for ibs, obs and conv feature + if dd if=/dev/zero of=/dev/null count=1 ibs=512 obs=512 conv=sync 2> /dev/null; then + dd if="\$1" ibs=\$2 skip=1 obs=1024 conv=sync 2> /dev/null | \\ + { test \$blocks -gt 0 && dd ibs=1024 obs=1024 count=\$blocks ; \\ + test \$bytes -gt 0 && dd ibs=1 obs=1024 count=\$bytes ; } 2> /dev/null + else + dd if="\$1" bs=\$2 skip=1 2> /dev/null + fi +} + +MS_dd_Progress() +{ + if test x"\$noprogress" = xy; then + MS_dd "\$@" + return \$? + fi + file="\$1" + offset=\$2 + length=\$3 + pos=0 + bsize=4194304 + while test \$bsize -gt \$length; do + bsize=\`expr \$bsize / 4\` + done + blocks=\`expr \$length / \$bsize\` + bytes=\`expr \$length % \$bsize\` + ( + dd ibs=\$offset skip=1 2>/dev/null + pos=\`expr \$pos \+ \$bsize\` + MS_Printf " 0%% " 1>&2 + if test \$blocks -gt 0; then + while test \$pos -le \$length; do + dd bs=\$bsize count=1 2>/dev/null + pcent=\`expr \$length / 100\` + pcent=\`expr \$pos / \$pcent\` + if test \$pcent -lt 100; then + MS_Printf "\b\b\b\b\b\b\b" 1>&2 + if test \$pcent -lt 10; then + MS_Printf " \$pcent%% " 1>&2 + else + MS_Printf " \$pcent%% " 1>&2 + fi + fi + pos=\`expr \$pos \+ \$bsize\` + done + fi + if test \$bytes -gt 0; then + dd bs=\$bytes count=1 2>/dev/null + fi + MS_Printf "\b\b\b\b\b\b\b" 1>&2 + MS_Printf " 100%% " 1>&2 + ) < "\$file" +} + +MS_Help() +{ + cat << EOH >&2 +Usage: \$0 [options] +Options: + --help | -h Print this message + --info Print embedded info : title, default target directory, embedded script ... + --list Print the list of files in the archive + --check Checks integrity and version dependency of the archive + --quiet | -q Quiet install mode, skip human-computer interactions + If the package requires an EULA, quiet installations are useful for scripting the installation + Using this option means accepting the EULA + --nox11 Do not spawn an xterm + --noexec Do not run embedded script + --extract= Extract directly to a target directory (absolute or relative) + Usually used with --noexec to just extract files without running + --tar arg1 [arg2 ...] Access the contents of the archive through the tar command +\${helpheader} +EOH +} + +MS_Verify_Sig() +{ + GPG_PATH=\`exec <&- 2>&-; which gpg || command -v gpg || type gpg\` + MKTEMP_PATH=\`exec <&- 2>&-; which mktemp || command -v mktemp || type mktemp\` + test -x "\$GPG_PATH" || GPG_PATH=\`exec <&- 2>&-; which gpg || command -v gpg || type gpg\` + test -x "\$MKTEMP_PATH" || MKTEMP_PATH=\`exec <&- 2>&-; which mktemp || command -v mktemp || type mktemp\` + offset=\`head -n "\$skip" "\$1" | wc -c | tr -d " "\` + temp_sig=\`mktemp -t XXXXX\` + echo \$SIGNATURE | base64 --decode > "\$temp_sig" + gpg_output=\`MS_dd "\$1" \$offset \$totalsize | LC_ALL=C "\$GPG_PATH" --verify "\$temp_sig" - 2>&1\` + gpg_res=\$? + rm -f "\$temp_sig" + if test \$gpg_res -eq 0 && test \`echo \$gpg_output | grep -c Good\` -eq 1; then + if test \`echo \$gpg_output | grep -c \$sig_key\` -eq 1; then + test x"\$quiet" = xn && echo "GPG signature is good" >&2 + else + echo "GPG Signature key does not match" >&2 + exit 2 + fi + else + test x"\$quiet" = xn && echo "GPG signature failed to verify" >&2 + exit 2 + fi +} + +MS_Check() +{ + SHA_PATH=\`exec <&- 2>&-; which shasum || command -v shasum || type shasum\` + test -x "\$SHA_PATH" || SHA_PATH=\`exec <&- 2>&-; which sha256sum || command -v sha256sum || type sha256sum\` + + if ! test -x "\$SHA_PATH"; then + echo "Command sha256sum not found, please install it first." + log "ERROR" "Command sha256sum not found, please install it first." + exit 2 + fi + + if test x"\$quiet" = xn; then + MS_Printf "Verifying archive integrity..." + fi + offset=\`head -n "\$skip" "\$1" | wc -c | tr -d " "\` + fsize=\`cat "\$1" | wc -c | tr -d " "\` + if test \$totalsize -ne \`expr \$fsize - \$offset\`; then + echo " Unexpected archive size." >&2 + exit 2 + fi + verb=\$2 + i=1 + for s in \$filesizes + do + if test -x "\$SHA_PATH"; then + if test x"\`basename \$SHA_PATH\`" = xshasum; then + SHA_ARG="-a 256" + fi + sha=\`echo \$SHA | cut -d" " -f\$i\` + if test x"\$sha" = x0000000000000000000000000000000000000000000000000000000000000000; then + test x"\$verb" = xy && echo " \$1 does not contain an embedded SHA256 checksum." >&2 + else + shasum=\`MS_dd_Progress "\$1" \$offset \$s | eval "\$SHA_PATH \$SHA_ARG" | cut -b-64\`; + if test x"\$shasum" != x"\$sha"; then + echo "Error in SHA256 checksums: \$shasum is different from \$sha" >&2 + log "ERROR" "Error in SHA256 checksums: \$shasum is different from \$sha" + exit 2 + elif test x"\$quiet" = xn; then + MS_Printf " SHA256 checksums are OK." >&2 + log "INFO" "SHA256 checksums are OK." + fi + fi + fi + i=\`expr \$i + 1\` + offset=\`expr \$offset + \$s\` + done + if test x"\$quiet" = xn; then + echo " All good." + fi +} + +MS_Decompress() +{ + if test x"\$decrypt_cmd" != x""; then + { eval "\$decrypt_cmd" || echo " ... Decryption failed." >&2; } | eval "$GUNZIP_CMD" + else + eval "$GUNZIP_CMD" + fi + + if test \$? -ne 0; then + echo " ... Decompression failed." >&2 + log "ERROR" "Decompression failed." + fi +} + +UnTAR() +{ + if test x"\$quiet" = xn; then + tar \$1vf - $UNTAR_EXTRA 2>&1 || { echo " ... Extraction failed." >&2; kill -15 \$$; } + else + tar \$1f - $UNTAR_EXTRA 2>&1 || { echo Extraction failed. >&2; kill -15 \$$; } + fi +} + +MS_exec_cleanup() { + if test x"\$cleanup" = xy && test x"\$cleanup_script" != x""; then + cleanup=n + cd "\$tmpdir" + eval "\"\$cleanup_script\" \$scriptargs \$cleanupargs" + fi +} + +MS_cleanup() +{ + echo 'Signal caught, cleaning up' >&2 + MS_exec_cleanup + cd "\$TMPROOT" + rm -rf "\$tmpdir" + eval \$finish; exit 15 +} + +MS_check_user() +{ + userid=\`id -u\` + tmpdir_uid=\`stat -c %u \$tmpdir\` + user_name=\`stat -c %U \$tmpdir\` + if test x"\$userid" != x"\$tmpdir_uid"; then + echo "Run package was modified by user \$user_name, please check security." + exit 1 + fi +} + +Script_Args_Check() +{ + script_supported_args=\$(echo \${helpheader} | grep -o -E "[-][-][^ ]+" | awk -F"=" {'print \$1'}) + arg_to_test=\$(echo \$1|awk -F"=" {'print \$1'}) + + for arg in \${script_supported_args}; + do + if test x"\$arg_to_test" = x"\$arg" ;then + return + fi + done + + MS_Help + exit 1 +} + +finish=true +xterm_loop= +noprogress=$NOPROGRESS +nox11=$NOX11 +copy=$COPY +ownership=$OWNERSHIP +verbose=n +cleanup=y +cleanupargs= +sig_key= + +initargs="\$@" + +while [ -n "\$*" ] +do + case "\$1" in + -h | --help) + MS_Help + exit 0 + ;; + -q | --quiet) + quiet=y + noprogress=y + shift + ;; + --info) + echo Identification: "\$label" + echo Target directory: "\$targetdir" + echo Uncompressed size: $USIZE KB + echo Compression: $COMPRESS + if test x"$ENCRYPT" != x""; then + echo Encryption: $ENCRYPT + fi + echo Date of packaging: $DATE + echo Built with Makeself version $MS_VERSION + echo Build command was: "$MS_COMMAND" + if test x"\$script" != x; then + echo Script run after extraction: + echo " " \$script \$scriptargs + fi + if test x"$copy" = xcopy; then + echo "Archive will copy itself to a temporary location" + fi + if test x"$NEED_ROOT" = xy; then + echo "Root permissions required for extraction" + fi + if test x"$KEEP" = xy; then + echo "directory \$targetdir is permanent" + else + echo "\$targetdir will be removed after extraction" + fi + exit 0 + ;; + --list) + echo Target directory: \$targetdir + offset=\`head -n "\$skip" "\$0" | wc -c | tr -d " "\` + for s in \$filesizes + do + MS_dd "\$0" \$offset \$s | MS_Decompress | UnTAR t + offset=\`expr \$offset + \$s\` + done + exit 0 + ;; + --tar) + offset=\`head -n "\$skip" "\$0" | wc -c | tr -d " "\` + arg1="\$2" + shift 2 || { MS_Help; exit 1; } + log "INFO" "Start --tar process." + echo "Makeself logfile: \$log_file" + for s in \$filesizes + do + MS_dd "\$0" \$offset \$s | MS_Decompress | tar "\$arg1" - "\$@" + offset=\`expr \$offset + \$s\` + done + exit 0 + ;; + --check) + echo "Makeself logfile: \$log_file" + MS_Check "\$0" y + scriptargs="\$scriptargs \$1" + shift + ;; + --noexec) + script="" + cleanup_script="" + shift + ;; + --extract=*) + keep=y + targetdir=\`echo \$1 | cut -d"=" -f2 \` + if ! shift; then MS_Help; exit 1; fi + log "INFO" "Extract files to targetdir." + echo "Makeself logfile: \$log_file" + ;; + --nox11) + nox11=y + shift + ;; + --xwin) + if test "$NOWAIT" = n; then + finish="echo Press Return to close this window...; read junk" + fi + xterm_loop=1 + shift + ;; + --phase2) + copy=phase2 + shift + ;; + *) + Script_Args_Check \$1 + scriptargs="\$scriptargs '\$1'" + shift + ;; + esac +done + +quiet_para="" +if test x"\$quiet" = xy; then + quiet_para="--quiet " +fi +scriptargs="--\$name_of_file""--\"\$pwd_of_file\""" \$quiet_para""\$scriptargs" + +if test x"\$quiet" = xy -a x"\$verbose" = xy; then + echo Cannot be verbose and quiet at the same time. >&2 + exit 1 +fi + +if test x"$NEED_ROOT" = xy -a \`id -u\` -ne 0; then + echo "Administrative privileges required for this archive (use su or sudo)" >&2 + exit 1 +fi + +if test x"\$copy" \!= xphase2; then + MS_PrintLicense +fi + +case "\$copy" in +copy) + tmpdir="\$TMPROOT"/makeself.\$RANDOM.\`date +"%y%m%d%H%M%S"\`.\$\$ + mkdir "\$tmpdir" || { + echo "Could not create temporary directory \$tmpdir" >&2 + exit 1 + } + SCRIPT_COPY="\$tmpdir/makeself" + echo "Copying to a temporary location..." >&2 + cp "\$0" "\$SCRIPT_COPY" + chmod +x "\$SCRIPT_COPY" + cd "\$TMPROOT" + exec "\$SCRIPT_COPY" --phase2 -- \$initargs + ;; +phase2) + finish="\$finish ; rm -rf \`dirname \$0\`" + ;; +esac + +if test x"\$targetdir" = x.; then + tmpdir="." +else + if test x"\$keep" = xy; then + if test x"\$nooverwrite" = xy && test -d "\$targetdir"; then + echo "Target directory \$targetdir already exists, aborting." >&2 + exit 1 + fi + if test x"\$quiet" = xn; then + echo "Creating directory \$targetdir" >&2 + fi + tmpdir="\$targetdir" + dashp="-p" + else + tmpdir="\$TMPROOT/selfgz\$\$\$RANDOM" + dashp="" + fi + if [ -L "\$tmpdir" ]; then + tmpdir=\`readlink -f \$tmpdir\` + fi + if [ ! -d "\$tmpdir" ]; then + mkdir \$dashp "\$tmpdir" || { + echo 'Cannot create target directory' \$tmpdir >&2 + echo 'You should try option --extract=' >&2 + eval \$finish + exit 1 + } + fi +fi +tmpdir=\`readlink -f \$tmpdir\` + +location="\`pwd\`" +if test x"\$SETUP_NOCHECK" != x1; then + MS_Check "\$0" +fi +offset=\`head -n "\$skip" "\$0" | wc -c | tr -d " "\` + +if test x"\$verbose" = xy; then + MS_Printf "About to extract $USIZE KB in \$tmpdir ... Proceed ? [Y/n] " + read yn + if test x"\$yn" = xn; then + eval \$finish; exit 1 + fi +fi + +if test x"\$quiet" = xn; then + # Decrypting with openssl will ask for password, + # the prompt needs to start on new line + if test x"$ENCRYPT" = x"openssl"; then + echo "Decrypting and uncompressing \$label..." + else + MS_Printf "Uncompressing \$label" + fi +fi +res=3 +if test x"\$keep" = xn; then + trap MS_cleanup 1 2 3 15 +fi + +if test x"\$nodiskspace" = xn; then + leftspace=\`MS_diskspace "\$tmpdir"\` + if test -n "\$leftspace"; then + if test "\$leftspace" -lt $USIZE; then + echo + echo "Not enough space left in "\`dirname \$tmpdir\`" (\$leftspace KB) to decompress \$0 ($USIZE KB)" >&2 + if test x"\$keep" = xn; then + echo "Use the (export TMPDIR=) command to set a decompressed directory with more free space." + fi + eval \$finish; exit 1 + fi + fi +fi + +for s in \$filesizes +do + if MS_dd_Progress "\$0" \$offset \$s | MS_Decompress | ( cd "\$tmpdir"; umask \$ORIG_UMASK ; UnTAR xp ) 1>/dev/null; then + if test x"\$ownership" = xy; then + (cd "\$tmpdir"; chown -R \`id -u\` .; chgrp -R \`id -g\` .) + fi + else + echo >&2 + echo "Unable to decompress \$0" >&2 + eval \$finish; exit 1 + fi + offset=\`expr \$offset + \$s\` +done +if test x"\$quiet" = xn; then + echo +fi + +cd "\$tmpdir" +res=0 +if test x"\$script" != x; then + if test x"\$export_conf" = x"y"; then + MS_BUNDLE="\$0" + MS_LABEL="\$label" + MS_SCRIPT="\$script" + MS_SCRIPTARGS="\$scriptargs" + MS_ARCHDIRNAME="\$archdirname" + MS_KEEP="\$KEEP" + MS_NOOVERWRITE="\$NOOVERWRITE" + MS_COMPRESS="\$COMPRESS" + MS_CLEANUP="\$cleanup" + export MS_BUNDLE MS_LABEL MS_SCRIPT MS_SCRIPTARGS + export MS_ARCHDIRNAME MS_KEEP MS_NOOVERWRITE MS_COMPRESS + fi + + if test x"\$verbose" = x"y"; then + yn="x" + while test x"\$yn" != x -a x"\$yn" != xy -a x"\$yn" != xY -a x"\$yn" != xn -a x"\$yn" != xN + do + MS_Printf "OK to execute: \$script \$scriptargs \$* ? [Y/n] " + read yn + if test x"\$yn" = x -o x"\$yn" = xy -o x"\$yn" = xY; then + MS_check_user + eval "\"\$script\" \$scriptargs \"\\\$@\""; res=\$?; + elif test x"\$yn" = xn -o x"\$yn" = xN; then + echo "Unable to decompress \$script ,because of aborting! ";res=\$? + else + echo "Input value is unacceptable,please try again." + fi + done + else + MS_check_user + eval "\"\$script\" \$scriptargs \"\\\$@\""; res=\$? + fi + if test "\$res" -ne 0; then + test x"\$verbose" = xy && echo "The program '\$script' returned an error code (\$res)" >&2 + fi +fi + +MS_exec_cleanup + +if test x"\$keep" = xn; then + cd "\$TMPROOT" + rm -rf "\$tmpdir" +fi +eval \$finish; exit \$res +EOF diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.1 b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.1 new file mode 100755 index 00000000..39efd9b7 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.1 @@ -0,0 +1,158 @@ +.TH "MAKESELF" "1" "2.5.0" +.SH "NAME" +makeself \- An utility to generate self-extractable archives. +.SH "SYNTAX" +.B makeself [\fIoptions\fP] archive_dir file_name label +.B [\fIstartup_script\fP] [\fIargs\fP] +.SH "DESCRIPTION" +This program is a free (GPL) shell utility designed to create self-extractable +compressed archives from a directory. The resulting file appears as a shell script, and can be launched as is. The archive +will then uncompress itself to a temporary directory and an optional arbitrary +command will be executed (for example an installation script). +.TP +Makeself archives also include checksums for integrity self-validation (CRC and/or MD5/SHA256 checksums). +.SH "OPTIONS" +The following options are supported: +.TP 15 +.B -v, --version +Prints out the makeself version number and exits. +.TP +.B -h, --help +Print out help information. +.TP +.B --tar-quietly +Suppress verbose output from the tar command +.TP +.B --quiet +Do not print any messages other than errors +.TP +.B --gzip +Compress using gzip (default if detected). +.TP +.B --bzip2 +Compress using bzip2. +.TP +.B --bzip3 +Compress using bzip3. +.TP +.B --pbzip2 +Compress using pbzip2. +.TP +.B --xz +Compress using xz. +.TP +.B --lzo +Compress using lzop. +.TP +.B --lz4 +Compress using lz4. +.TP +.B --pigz +Compress using pigz. +.TP +.B --zstd +Compress using zstd. +.TP +.B --base64 +Encode the archive to ASCII in Base64 format instead of compressing (base64 command required). +.TP +.B --gpg-encrypt +Encrypt the archive using GPG. This will prompt for a password to encrypt with. +.TP +.B --ssl-encrypt +Encrypt the archive using OpenSSL. This will prompt for a password to encrypt with. +.TP +.B --keep-umask +Keep the umask set to shell default, rather than overriding when executing the self-extracting archive. +.TP +.B --compress +Compress using the UNIX 'compress' command. +.TP +.B --nocomp +Do not compress the data. +.TP +.B --complevel lvl +Specify the compression level for gzip, bzip2, pbzip2, xz, zstd, lzo or lz4. Defaults to 9. +.TP +.B --threads num +Specify the number of threads to be used by compressors that support parallelization. +.TP +.B --tar-format opt + Specify the tar archive format (default is ustar); you may use any value accepted by your tar command (such as posix, v7, etc). +.TP +.B --tar-extra opt +Append more options to the tar command line. +.TP +.B --notemp +The archive will create archive_dir in the current directory and +uncompress in ./archive_dir. +.TP +.B --copy +Upon extraction, the archive will first copy itself to a temporary directory. +.TP +.B --append +Append more files to an existing makeself archive. The label and startup scripts will then be ignored. +.TP +.B --current +Files will be extracted to the current directory. Both --current and --target dir imply --notemp. +.TP +.B --target dir +Extract directly to a target directory. Directory path can be either absolute or relative. +.TP +.B --header file +Specify location of the header script. +.TP +.B --help-header file +Add a header to the archive's help output. +.TP +.B --cleanup file +Specify a cleanup script that executes on interrupt and when finished successfully. +.TP +.B --follow +Follow the symlinks in the archive. +.TP +.B --noprogress +Do not show the progress during the decompression. +.TP +.B --nooverwrite +Do not extract the archive if the target directory already exists. +.TP +.B --nox11 +Disable automatic spawn of an xterm if running in X11. +.TP +.B --nowait +Do not wait for user input after executing embedded program from an xterm. +.TP +.B --nomd5 +Do not create a MD5 checksum for the archive. +.TP +.B --sha256 +Adds a SHA256 checksum for the archive. +.TP +.B --nocrc +Do not create a CRC32 checksum for the archive. +.TP +.B --lsm file +LSM file describing the package. +.TP +.B --license file +Append a license file. +.TP +.B --packaging-date date +Use provided string as the packaging date instead of the current date. +.TP +.SH "EXAMPLES" +Here is an example, assuming the user has a package image stored in a /home/joe/mysoft, +and he wants to generate a self-extracting package named mysoft.sh, which will launch +the "setup" script initially stored in /home/joe/mysoft: +.TP +makeself.sh /home/joe/mysoft mysoft.sh "Joe's Nice Software Package" ./setup +.TP +Here is also how I created the makeself.run archive which contains the Makeself distribution: +.TP +makeself.sh --notemp makeself makeself.run "Makeself by Stephane Peter" echo "Makeself has extracted itself" +.SH "AUTHORS" +Makeself has been written by Stephane Peter . +.BR +This man page was originally written by Bartosz Fenski for the +Debian GNU/Linux distribution (but it may be used by others). diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.sh b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.sh new file mode 100755 index 00000000..588d3c45 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/makeself/makeself.sh @@ -0,0 +1,776 @@ +#!/bin/sh +# +# Makeself version 2.5.x +# by Stephane Peter +# +# Utility to create self-extracting tar.gz archives. +# The resulting archive is a file holding the tar.gz archive with +# a small Shell script stub that uncompresses the archive to a temporary +# directory and then executes a given script from withing that directory. +# +# Makeself home page: https://makeself.io/ - Version history available on GitHub +# +# (C) 1998-2023 by Stephane Peter +# +# This software is released under the terms of the GNU GPL version 2 and above +# Please read the license at http://www.gnu.org/copyleft/gpl.html +# Self-extracting archives created with this script are explictly NOT released under the term of the GPL +# + +MS_VERSION=2.5.0 +MS_COMMAND="$0" +unset CDPATH + +for f in ${1+"$@"}; do + MS_COMMAND="$MS_COMMAND \\\\ + \\\"$f\\\"" +done + +# For Solaris systems +if test -d /usr/xpg4/bin; then + PATH=/usr/xpg4/bin:$PATH + export PATH +fi + +# Procedures + +MS_Usage() +{ + echo "Usage: $0 [args] archive_dir file_name label startup_script [script_args]" + echo "args can be one or more of the following :" + echo " --version | -v : Print out Makeself version number and exit" + echo " --help | -h : Print out this help message" + echo " --tar-quietly : Suppress verbose output from the tar command" + echo " --quiet | -q : Do not print any messages other than errors." + echo " --gzip : Compress using gzip (default if detected)" + echo " --pigz : Compress with pigz" + echo " --zstd : Compress with zstd" + echo " --bzip2 : Compress using bzip2 instead of gzip" + echo " --pbzip2 : Compress using pbzip2 instead of gzip" + echo " --bzip3 : Compress using bzip3 instead of gzip" + echo " --xz : Compress using xz instead of gzip" + echo " --lzo : Compress using lzop instead of gzip" + echo " --lz4 : Compress using lz4 instead of gzip" + echo " --compress : Compress using the UNIX 'compress' command" + echo " --complevel lvl : Compression level for gzip pigz zstd xz lzo lz4 bzip2 pbzip2 and bzip3 (default 9)" + echo " --threads thds : Number of threads to be used by compressors that support parallelization." + echo " Omit to use compressor's default. Most useful (and required) for opting" + echo " into xz's threading, usually with '--threads=0' for all available cores." + echo " pbzip2 and pigz are parallel by default, and setting this value allows" + echo " limiting the number of threads they use." + echo " --base64 : Instead of compressing, encode the data using base64" + echo " --gpg-encrypt : Instead of compressing, encrypt the data using GPG" + echo " --gpg-asymmetric-encrypt-sign" + echo " : Instead of compressing, asymmetrically encrypt and sign the data using GPG" + echo " --gpg-extra opt : Append more options to the gpg command line" + echo " --ssl-encrypt : Instead of compressing, encrypt the data using OpenSSL" + echo " --ssl-passwd pass : Use the given password to encrypt the data using OpenSSL" + echo " --ssl-pass-src src : Use the given src as the source of password to encrypt the data" + echo " using OpenSSL. See \"PASS PHRASE ARGUMENTS\" in man openssl." + echo " If this option is not supplied, the user will be asked to enter" + echo " encryption password on the current terminal." + echo " --ssl-no-md : Do not use \"-md\" option not supported by older OpenSSL." + echo " --nochown : Do not give the target folder to the current user (default)" + echo " --chown : Give the target folder to the current user recursively" + echo " --nocomp : Do not compress the data" + echo " --notemp : The archive will create archive_dir in the current directory" + echo " and uncompress in ./archive_dir" + echo " Note: persistent archives do not strictly require a startup_script" + echo " --needroot : Check that the root user is extracting the archive before proceeding" + echo " --copy : Upon extraction, the archive will first copy itself to" + echo " a temporary directory" + echo " --append : Append more files to an existing Makeself archive" + echo " The label and startup scripts will then be ignored" + echo " --target dir : Extract directly to a target directory" + echo " directory path can be either absolute or relative" + echo " --current : Files will be extracted to the current directory" + echo " Both --current and --target imply --notemp, and do not require a startup_script" + echo " --nooverwrite : Do not extract the archive if the specified target directory exists" + echo " --tar-format opt : Specify a tar archive format (default is ustar)" + echo " --tar-extra opt : Append more options to the tar command line" + echo " --untar-extra opt : Append more options to the during the extraction of the tar archive" + echo " --nomd5 : Don't calculate an MD5 for archive" + echo " --nocrc : Don't calculate a CRC for archive" + echo " --sha256 : Compute a SHA256 checksum for the archive" + echo " --header file : Specify location of the header script" + echo " --cleanup file : Specify a cleanup script that executes on interrupt and when finished successfully." + echo " --follow : Follow the symlinks in the archive" + echo " --noprogress : Do not show the progress during the decompression" + echo " --nox11 : Disable automatic spawn of a xterm" + echo " --nowait : Do not wait for user input after executing embedded" + echo " program from an xterm" + echo " --sign passphrase : Signature private key to sign the package with" + echo " --lsm file : LSM file describing the package" + echo " --license file : Append a license file" + echo " --help-header file : Add a header to the archive's --help output" + echo " --packaging-date date" + echo " : Use provided string as the packaging date" + echo " instead of the current date." + echo + echo " --keep-umask : Keep the umask set to shell default, rather than overriding when executing self-extracting archive." + echo " --export-conf : Export configuration variables to startup_script" + echo + echo "Do not forget to give a fully qualified startup script name" + echo "(i.e. with a ./ prefix if inside the archive)." + exit 1 +} + +# Default settings +if type gzip >/dev/null 2>&1; then + COMPRESS=gzip +elif type compress >/dev/null 2>&1; then + COMPRESS=compress +else + echo "ERROR: missing commands: gzip, compress" >&2 + MS_Usage +fi +ENCRYPT=n +PASSWD="" +PASSWD_SRC="" +OPENSSL_NO_MD=n +COMPRESS_LEVEL=9 +DEFAULT_THREADS=123456 # Sentinel value +THREADS=$DEFAULT_THREADS +KEEP=n +CURRENT=n +NOX11=n +NOWAIT=n +APPEND=n +TAR_QUIETLY=n +KEEP_UMASK=n +QUIET=n +NOPROGRESS=n +COPY=none +NEED_ROOT=n +TAR_ARGS=rvf +TAR_FORMAT=ustar +TAR_EXTRA="" +GPG_EXTRA="" +DU_ARGS=-ks +HEADER=`dirname "$0"`/makeself-header.sh +SIGNATURE="" +TARGETDIR="" +NOOVERWRITE=n +DATE=`LC_ALL=C date` +EXPORT_CONF=n +SHA256=n +OWNERSHIP=n +SIGN=n +GPG_PASSPHRASE="" + +# LSM file stuff +LSM_CMD="echo No LSM. >> \"\$archname\"" + +while true +do + case "$1" in + --version | -v) + echo Makeself version $MS_VERSION + exit 0 + ;; + --pbzip2) + COMPRESS=pbzip2 + shift + ;; + --bzip3) + COMPRESS=bzip3 + shift + ;; + --bzip2) + COMPRESS=bzip2 + shift + ;; + --gzip) + COMPRESS=gzip + shift + ;; + --pigz) + COMPRESS=pigz + shift + ;; + --zstd) + COMPRESS=zstd + shift + ;; + --xz) + COMPRESS=xz + shift + ;; + --lzo) + COMPRESS=lzo + shift + ;; + --lz4) + COMPRESS=lz4 + shift + ;; + --compress) + COMPRESS=compress + shift + ;; + --base64) + COMPRESS=base64 + shift + ;; + --gpg-encrypt) + COMPRESS=gpg + shift + ;; + --gpg-asymmetric-encrypt-sign) + COMPRESS=gpg-asymmetric + shift + ;; + --gpg-extra) + GPG_EXTRA="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --ssl-encrypt) + ENCRYPT=openssl + shift + ;; + --ssl-passwd) + PASSWD=$2 + shift 2 || { MS_Usage; exit 1; } + ;; + --ssl-pass-src) + PASSWD_SRC=$2 + shift 2 || { MS_Usage; exit 1; } + ;; + --ssl-no-md) + OPENSSL_NO_MD=y + shift + ;; + --nocomp) + COMPRESS=none + shift + ;; + --complevel) + COMPRESS_LEVEL="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --threads) + THREADS="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --nochown) + OWNERSHIP=n + shift + ;; + --chown) + OWNERSHIP=y + shift + ;; + --notemp) + KEEP=y + shift + ;; + --copy) + COPY=copy + shift + ;; + --current) + CURRENT=y + KEEP=y + shift + ;; + --tar-format) + TAR_FORMAT="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --tar-extra) + TAR_EXTRA="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --untar-extra) + UNTAR_EXTRA="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --target) + TARGETDIR="$2" + KEEP=y + shift 2 || { MS_Usage; exit 1; } + ;; + --sign) + SIGN=y + GPG_PASSPHRASE="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --nooverwrite) + NOOVERWRITE=y + shift + ;; + --needroot) + NEED_ROOT=y + shift + ;; + --header) + HEADER="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --cleanup) + CLEANUP_SCRIPT="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --license) + # We need to escape all characters having a special meaning in double quotes + LICENSE=$(sed 's/\\/\\\\/g; s/"/\\\"/g; s/`/\\\`/g; s/\$/\\\$/g' "$2") + shift 2 || { MS_Usage; exit 1; } + ;; + --follow) + TAR_ARGS=rvhf + DU_ARGS=-ksL + shift + ;; + --noprogress) + NOPROGRESS=y + shift + ;; + --nox11) + NOX11=y + shift + ;; + --nowait) + NOWAIT=y + shift + ;; + --nomd5) + NOMD5=y + shift + ;; + --sha256) + SHA256=y + shift + ;; + --nocrc) + NOCRC=y + shift + ;; + --append) + APPEND=y + shift + ;; + --lsm) + LSM_CMD="awk 1 \"$2\" >> \"\$archname\"" + shift 2 || { MS_Usage; exit 1; } + ;; + --packaging-date) + DATE="$2" + shift 2 || { MS_Usage; exit 1; } + ;; + --help-header) + HELPHEADER=`sed -e "s/'/'\\\\\''/g" $2` + shift 2 || { MS_Usage; exit 1; } + [ -n "$HELPHEADER" ] && HELPHEADER="$HELPHEADER +" + ;; + --tar-quietly) + TAR_QUIETLY=y + shift + ;; + --keep-umask) + KEEP_UMASK=y + shift + ;; + --export-conf) + EXPORT_CONF=y + shift + ;; + -q | --quiet) + QUIET=y + shift + ;; + -h | --help) + MS_Usage + ;; + -*) + echo Unrecognized flag : "$1" + MS_Usage + ;; + *) + break + ;; + esac +done + +if test $# -lt 1; then + MS_Usage +else + if test -d "$1"; then + archdir="$1" + else + echo "Directory $1 does not exist." >&2 + exit 1 + fi +fi +archname="$2" + +if test "$QUIET" = "y" || test "$TAR_QUIETLY" = "y"; then + if test "$TAR_ARGS" = "rvf"; then + TAR_ARGS="rf" + elif test "$TAR_ARGS" = "rvhf"; then + TAR_ARGS="rhf" + fi +fi + +if test "$APPEND" = y; then + if test $# -lt 2; then + MS_Usage + fi + + # Gather the info from the original archive + OLDENV=`sh "$archname" --dumpconf` + if test $? -ne 0; then + echo "Unable to update archive: $archname" >&2 + exit 1 + else + eval "$OLDENV" + OLDSKIP=`expr $SKIP + 1` + fi +else + if test "$KEEP" = n -a $# = 3; then + echo "ERROR: Making a temporary archive with no embedded command does not make sense!" >&2 + echo >&2 + MS_Usage + fi + # We don't want to create an absolute directory unless a target directory is defined + if test "$CURRENT" = y; then + archdirname="." + elif test x"$TARGETDIR" != x; then + archdirname="$TARGETDIR" + else + archdirname=`basename "$1"` + fi + + if test $# -lt 3; then + MS_Usage + fi + + LABEL="$3" + SCRIPT="$4" + test "x$SCRIPT" = x || shift 1 + shift 3 + SCRIPTARGS="$*" +fi + +if test "$KEEP" = n -a "$CURRENT" = y; then + echo "ERROR: It is A VERY DANGEROUS IDEA to try to combine --notemp and --current." >&2 + exit 1 +fi + +case $COMPRESS in +gzip) + GZIP_CMD="gzip -c$COMPRESS_LEVEL" + GUNZIP_CMD="gzip -cd" + ;; +pigz) + GZIP_CMD="pigz -$COMPRESS_LEVEL" + if test $THREADS -ne $DEFAULT_THREADS; then # Leave as the default if threads not indicated + GZIP_CMD="$GZIP_CMD --processes $THREADS" + fi + GUNZIP_CMD="gzip -cd" + ;; +zstd) + GZIP_CMD="zstd -$COMPRESS_LEVEL" + if test $THREADS -ne $DEFAULT_THREADS; then # Leave as the default if threads not indicated + GZIP_CMD="$GZIP_CMD --threads=$THREADS" + fi + GUNZIP_CMD="zstd -cd" + ;; +pbzip2) + GZIP_CMD="pbzip2 -c$COMPRESS_LEVEL" + if test $THREADS -ne $DEFAULT_THREADS; then # Leave as the default if threads not indicated + GZIP_CMD="$GZIP_CMD -p$THREADS" + fi + GUNZIP_CMD="bzip2 -d" + ;; +bzip3) + # Map the compression level to a block size in MiB as 2^(level-1). + BZ3_COMPRESS_LEVEL=`echo "2^($COMPRESS_LEVEL-1)" | bc` + GZIP_CMD="bzip3 -b$BZ3_COMPRESS_LEVEL" + if test $THREADS -ne $DEFAULT_THREADS; then # Leave as the default if threads not indicated + GZIP_CMD="$GZIP_CMD -j$THREADS" + fi + JOBS=`echo "10-$COMPRESS_LEVEL" | bc` + GUNZIP_CMD="bzip3 -dj$JOBS" + ;; +bzip2) + GZIP_CMD="bzip2 -$COMPRESS_LEVEL" + GUNZIP_CMD="bzip2 -d" + ;; +xz) + GZIP_CMD="xz -c$COMPRESS_LEVEL" + # Must opt-in by specifying a value since not all versions of xz support threads + if test $THREADS -ne $DEFAULT_THREADS; then + GZIP_CMD="$GZIP_CMD --threads=$THREADS" + fi + GUNZIP_CMD="xz -d" + ;; +lzo) + GZIP_CMD="lzop -c$COMPRESS_LEVEL" + GUNZIP_CMD="lzop -d" + ;; +lz4) + GZIP_CMD="lz4 -c$COMPRESS_LEVEL" + GUNZIP_CMD="lz4 -d" + ;; +base64) + GZIP_CMD="base64" + GUNZIP_CMD="base64 --decode -i -" + ;; +gpg) + GZIP_CMD="gpg $GPG_EXTRA -ac -z$COMPRESS_LEVEL" + GUNZIP_CMD="gpg -d" + ENCRYPT="gpg" + ;; +gpg-asymmetric) + GZIP_CMD="gpg $GPG_EXTRA -z$COMPRESS_LEVEL -es" + GUNZIP_CMD="gpg --yes -d" + ENCRYPT="gpg" + ;; +compress) + GZIP_CMD="compress -fc" + GUNZIP_CMD="(type compress >/dev/null 2>&1 && compress -fcd || gzip -cd)" + ;; +none) + GZIP_CMD="cat" + GUNZIP_CMD="cat" + ;; +esac + +if test x"$ENCRYPT" = x"openssl"; then + if test x"$APPEND" = x"y"; then + echo "Appending to existing archive is not compatible with OpenSSL encryption." >&2 + fi + + ENCRYPT_CMD="openssl enc -aes-256-cbc -salt" + DECRYPT_CMD="openssl enc -aes-256-cbc -d" + + if test x"$OPENSSL_NO_MD" != x"y"; then + ENCRYPT_CMD="$ENCRYPT_CMD -md sha256" + DECRYPT_CMD="$DECRYPT_CMD -md sha256" + fi + + if test -n "$PASSWD_SRC"; then + ENCRYPT_CMD="$ENCRYPT_CMD -pass $PASSWD_SRC" + elif test -n "$PASSWD"; then + ENCRYPT_CMD="$ENCRYPT_CMD -pass pass:$PASSWD" + fi +fi + +tmpfile="${TMPDIR:-/tmp}/mkself$$" + +if test -f "$HEADER"; then + oldarchname="$archname" + archname="$tmpfile" + # Generate a fake header to count its lines + SKIP=0 + . "$HEADER" + SKIP=`cat "$tmpfile" |wc -l` + # Get rid of any spaces + SKIP=`expr $SKIP` + rm -f "$tmpfile" + if test "$QUIET" = "n"; then + echo "Header is $SKIP lines long" >&2 + fi + archname="$oldarchname" +else + echo "Unable to open header file: $HEADER" >&2 + exit 1 +fi + +if test "$QUIET" = "n"; then + echo +fi + +if test "$APPEND" = n; then + if test -f "$archname"; then + echo "WARNING: Overwriting existing file: $archname" >&2 + fi +fi + +USIZE=`du $DU_ARGS "$archdir" | awk '{print $1}'` + +if test "." = "$archdirname"; then + if test "$KEEP" = n; then + archdirname="makeself-$$-`date +%Y%m%d%H%M%S`" + fi +fi + +test -d "$archdir" || { echo "Error: $archdir does not exist."; rm -f "$tmpfile"; exit 1; } +if test "$QUIET" = "n"; then + echo "About to compress $USIZE KB of data..." + echo "Adding files to archive named \"$archname\"..." +fi + +# See if we have GNU tar +TAR=`exec <&- 2>&-; which gtar || command -v gtar || type gtar` +test -x "$TAR" || TAR=`exec <&- 2>&-; which bsdtar || command -v bsdtar || type bsdtar` +test -x "$TAR" || TAR=tar + +tmparch="${TMPDIR:-/tmp}/mkself$$.tar" +( + if test "$APPEND" = "y"; then + tail -n "+$OLDSKIP" "$archname" | eval "$GUNZIP_CMD" > "$tmparch" + fi + cd "$archdir" + # "Determining if a directory is empty" + # https://www.etalabs.net/sh_tricks.html + find . \ + \( \ + ! -type d \ + -o \ + \( -links 2 -exec sh -c ' + is_empty () ( + cd "$1" + set -- .[!.]* ; test -f "$1" && return 1 + set -- ..?* ; test -f "$1" && return 1 + set -- * ; test -f "$1" && return 1 + return 0 + ) + is_empty "$0"' {} \; \ + \) \ + \) -print \ + | LC_ALL=C sort \ + | sed 's/./\\&/g' \ + | xargs $TAR $TAR_EXTRA --format $TAR_FORMAT -$TAR_ARGS "$tmparch" +) || { + echo "ERROR: failed to create temporary archive: $tmparch" + rm -f "$tmparch" "$tmpfile" + exit 1 +} + +USIZE=`du $DU_ARGS "$tmparch" | awk '{print $1}'` + +eval "$GZIP_CMD" <"$tmparch" >"$tmpfile" || { + echo "ERROR: failed to create temporary file: $tmpfile" + rm -f "$tmparch" "$tmpfile" + exit 1 +} +rm -f "$tmparch" + +if test x"$ENCRYPT" = x"openssl"; then + echo "About to encrypt archive \"$archname\"..." + { eval "$ENCRYPT_CMD -in $tmpfile -out ${tmpfile}.enc" && mv -f ${tmpfile}.enc $tmpfile; } || \ + { echo Aborting: could not encrypt temporary file: "$tmpfile".; rm -f "$tmpfile"; exit 1; } +fi + +fsize=`cat "$tmpfile" | wc -c | tr -d " "` + +# Compute the checksums + +shasum=0000000000000000000000000000000000000000000000000000000000000000 +md5sum=00000000000000000000000000000000 +crcsum=0000000000 + +if test "$NOCRC" = y; then + if test "$QUIET" = "n"; then + echo "skipping crc at user request" + fi +else + crcsum=`CMD_ENV=xpg4 cksum < "$tmpfile" | sed -e 's/ /Z/' -e 's/ /Z/' | cut -dZ -f1` + if test "$QUIET" = "n"; then + echo "CRC: $crcsum" + fi +fi + +if test "$SHA256" = y; then + SHA_PATH=`exec <&- 2>&-; which shasum || command -v shasum || type shasum` + if test -x "$SHA_PATH"; then + shasum=`eval "$SHA_PATH -a 256" < "$tmpfile" | cut -b-64` + else + SHA_PATH=`exec <&- 2>&-; which sha256sum || command -v sha256sum || type sha256sum` + shasum=`eval "$SHA_PATH" < "$tmpfile" | cut -b-64` + fi + if test "$QUIET" = "n"; then + if test -x "$SHA_PATH"; then + echo "SHA256: $shasum" + else + echo "SHA256: none, SHA command not found" + fi + fi +fi +if test "$NOMD5" = y; then + if test "$QUIET" = "n"; then + echo "Skipping md5sum at user request" + fi +else + # Try to locate a MD5 binary + OLD_PATH=$PATH + PATH=${GUESS_MD5_PATH:-"$OLD_PATH:/bin:/usr/bin:/sbin:/usr/local/ssl/bin:/usr/local/bin:/opt/openssl/bin"} + MD5_ARG="" + MD5_PATH=`exec <&- 2>&-; which md5sum || command -v md5sum || type md5sum` + test -x "$MD5_PATH" || MD5_PATH=`exec <&- 2>&-; which md5 || command -v md5 || type md5` + test -x "$MD5_PATH" || MD5_PATH=`exec <&- 2>&-; which digest || command -v digest || type digest` + PATH=$OLD_PATH + if test -x "$MD5_PATH"; then + if test `basename ${MD5_PATH}`x = digestx; then + MD5_ARG="-a md5" + fi + md5sum=`eval "$MD5_PATH $MD5_ARG" < "$tmpfile" | cut -b-32` + if test "$QUIET" = "n"; then + echo "MD5: $md5sum" + fi + else + if test "$QUIET" = "n"; then + echo "MD5: none, MD5 command not found" + fi + fi +fi +if test "$SIGN" = y; then + GPG_PATH=`exec <&- 2>&-; which gpg || command -v gpg || type gpg` + if test -x "$GPG_PATH"; then + SIGNATURE=`$GPG_PATH --pinentry-mode=loopback --batch --yes $GPG_EXTRA --passphrase "$GPG_PASSPHRASE" --output - --detach-sig $tmpfile | base64 | tr -d \\\\n` + if test "$QUIET" = "n"; then + echo "Signature: $SIGNATURE" + fi + else + echo "Missing gpg command" >&2 + fi +fi + +totalsize=0 +for size in $fsize; +do + totalsize=`expr $totalsize + $size` +done + +if test "$APPEND" = y; then + mv "$archname" "$archname".bak || exit + + # Prepare entry for new archive + filesizes="$fsize" + CRCsum="$crcsum" + MD5sum="$md5sum" + SHAsum="$shasum" + Signature="$SIGNATURE" + # Generate the header + . "$HEADER" + # Append the new data + cat "$tmpfile" >> "$archname" + + chmod +x "$archname" + rm -f "$archname".bak + if test "$QUIET" = "n"; then + echo "Self-extractable archive \"$archname\" successfully updated." + fi +else + filesizes="$fsize" + CRCsum="$crcsum" + MD5sum="$md5sum" + SHAsum="$shasum" + Signature="$SIGNATURE" + + # Generate the header + . "$HEADER" + + # Append the compressed tar data after the stub + if test "$QUIET" = "n"; then + echo + fi + cat "$tmpfile" >> "$archname" + chmod +x "$archname" + if test "$QUIET" = "n"; then + echo Self-extractable archive \"$archname\" successfully created. + fi +fi +rm -f "$tmpfile" diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/opdesc_parser.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/opdesc_parser.py new file mode 100755 index 00000000..c187ed02 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/opdesc_parser.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Created on Feb 28 20:56:45 2020 +Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved. +""" + +import sys +import os + + +OP_ALL = '__ALLOP__' +SOC_ALL = '__ALLSOC__' +SOC_TO_SHORT_SOC_MAP = { + "ascend910a": "ascend910", + "ascend910proa": "ascend910", + "ascend910b": "ascend910", + "ascend910prob": "ascend910", + "ascend910premiuma": "ascend910", + "ascend910b1": "ascend910b", + "ascend910b2": "ascend910b", + "ascend910b2c": "ascend910b", + "ascend910b3": "ascend910b", + "ascend910b4": "ascend910b", + "ascend910b4-1": "ascend910b", + "ascend910_9391": "ascend910_93", + "ascend910_9381": "ascend910_93", + "ascend910_9372": "ascend910_93", + "ascend910_9392": "ascend910_93", + "ascend910_9382": "ascend910_93", + "ascend910_9362": "ascend910_93", + "ascend310p1": "ascend310p", + "ascend310p3": "ascend310p", + "ascend310p5": "ascend310p", + "ascend310p7": "ascend310p", + "ascend310p3vir01": "ascend310p", + "ascend310p3vir02": "ascend310p", + "ascend310p3vir04": "ascend310p", + "ascend310p3vir08": "ascend310p", + "ascend310b1": "ascend310b", + "bs9sx1aa": "bs9sx1a", + "ascend610lite": "ascend610lite", + "ascend910_9599": "ascend910_95" +} +CONFLICT_KEYWORDS = { + "and", "as", "assert", "break", "class", "continue", "def", "del", "elif", "else", + "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda", + "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "False", + "None", "True", "nonlocal", "arg", "__inputs__", "__outputs__", "options", "bisheng", + "bisheng_path", "tikcpp_path", "impl_mode", "custom_compile_options", + "custom_all_compile_options", "soc_version", "soc_short", "custom_compile_options_soc", + "custom_all_compile_options_soc", "origin_func_name", "ascendc_src_dir_ex", + "ascendc_src_dir", "ascendc_src_file", "src", "op_type", "code_channel", "op_info", + "compile_op", "get_code_channel", "result", "__attrs__", "isinstance", "attr", + "get_current_build_config", "_build_args", "get_dtype_fmt_options", "shutil", "os", + "get_kernel_source" +} + + +class OpDesc: + def __init__(self: any, op_type: str): + self.op_type = op_type + self.attr_list = [] + self.attr_val = {} + self.input_name = [] + self.input_ori_name = [] + self.input_type = [] + self.input_dtype = [] + self.input_dtype_for_bin_list = [] + self.input_dtype_for_bin = {} + self.input_fmt = [] + self.input_fmt_for_bin_list = [] + self.input_fmt_for_bin = {} + self.input_virt = {} + self.input_value_depend = {} + self.output_name = [] + self.output_ori_name = [] + self.output_type = [] + self.output_dtype = [] + self.output_dtype_for_bin_list = [] + self.output_dtype_for_bin = {} + self.output_fmt = [] + self.output_fmt_for_bin_list = [] + self.output_fmt_for_bin = {} + self.output_init_value = [] + self.output_shape_depend_on_compute = [] + self.op_fmt_sel = False + self.op_chk_support = False + self.op_intf = '' + self.kern_name = '' + self.op_file = '' + self.op_replay_flag = False + self.op_replay_batch = False + self.input_idx = -1 + self.output_idx = -1 + self.max_block_dim = 32 + self.max_shape_size = 268435456 + self.dynamic_shape = False + self.op_range_limit = '' + self.custom_compile_options = {} + self.custom_all_compile_options = {} + self.param_type_dynamic = False + self.mc2_ctx = [] + self.bin_cprs_list = [] + self.bin_cprs_head = [] + self.bin_save_list = [] + + @staticmethod + def _parse_digit(conf: str) -> int: + return int(conf.split('=')[1]) + + @staticmethod + def _parse_flag(conf: str) -> bool: + if 'true' == conf.split('=')[1]: + return True + return False + + @staticmethod + def _parse_str(conf: str) -> str: + return conf.split('=')[1] + + @staticmethod + def _parse_list(conf: str) -> list: + return conf.split('=')[1].split(',') + + def parse_input(self: any, conf: str): + if conf.startswith('input{}.name'.format(int(self.input_idx) + 1)): + self.input_idx += 1 + self.input_ori_name.append(self._parse_str(conf)) + self.input_name.append(self.input_ori_name[-1] + '_in__') + elif conf.startswith('input{}.paramType'.format(int(self.input_idx))): + param_type = self._parse_str(conf) + self.input_type.append(param_type) + if param_type == "dynamic": + self.param_type_dynamic = True + elif conf.startswith('input{}.dtype'.format(int(self.input_idx))): + self.input_dtype.append(self._parse_str(conf)) + elif conf.startswith('input{}.for_bin_dtype'.format(int(self.input_idx))): + self.input_dtype_for_bin.update({self.input_idx : self._parse_str(conf)}) + elif conf.startswith('input{}.format'.format(int(self.input_idx))): + self.input_fmt.append(self._parse_str(conf)) + elif conf.startswith('input{}.for_bin_format'.format(int(self.input_idx))): + self.input_fmt_for_bin.update({self.input_idx : self._parse_str(conf)}) + elif conf.startswith('input{}.virtual'.format(int(self.input_idx))): + self.input_virt[self.input_idx] = self._parse_str(conf) + elif conf.startswith('input{}.valueDepend'.format(int(self.input_idx))): + self.input_value_depend[self.input_idx] = self._parse_str(conf) + elif conf.startswith('input{}.initValue'.format(int(self.input_idx))): + raise Exception(f'[ERROR]: Op: {{\'{self.op_type}\'}} input {self.input_ori_name[int(self.input_idx)]}\ + has InitValue, which is not support!') + else: + return + + def parse_output(self: any, conf: str): + if conf.startswith('output{}.name'.format(int(self.output_idx) + 1)): + self.output_idx += 1 + self.output_ori_name.append(self._parse_str(conf)) + self.output_name.append(self.output_ori_name[-1] + '_out_') + self.output_init_value.append(None) + elif conf.startswith('output{}.paramType'.format(int(self.output_idx))): + param_type = self._parse_str(conf) + self.output_type.append(param_type) + if param_type == "dynamic": + self.param_type_dynamic = True + elif conf.startswith('output{}.dtype'.format(int(self.output_idx))): + self.output_dtype.append(self._parse_str(conf)) + elif conf.startswith('output{}.for_bin_dtype'.format(int(self.output_idx))): + self.output_dtype_for_bin.update({self.output_idx : self._parse_str(conf)}) + elif conf.startswith('output{}.format'.format(int(self.output_idx))): + self.output_fmt.append(self._parse_str(conf)) + elif conf.startswith('output{}.for_bin_format'.format(int(self.output_idx))): + self.output_fmt_for_bin.update({self.output_idx : self._parse_str(conf)}) + elif conf.startswith('output{}.initValue'.format(int(self.output_idx))): + self.output_init_value[int(self.output_idx)] = self._parse_str(conf) + elif conf.startswith('output{}.outputShapeDependOnCompute=true'.format(int(self.output_idx))): + self.output_shape_depend_on_compute.append(int(self.output_idx)) + else: + return + + def parse_op_format(self: any, conf: str): + self.op_fmt_sel = self._parse_flag(conf) + + def parse_check_support(self: any, conf: str): + self.op_chk_support = self._parse_flag(conf) + + def parse_range_limit(self: any, conf: str): + self.op_range_limit = self._parse_str(conf) + + def parse_kern_name(self: any, conf: str): + self.kern_name = self._parse_str(conf) + + def parse_op_intf(self: any, conf: str): + self.op_intf = self._parse_str(conf) + + def parse_op_file(self: any, conf: str): + self.op_file = self._parse_str(conf) + + def parse_dynamic_shape(self: any, conf: str): + self.dynamic_shape = self._parse_flag(conf) + + def parse_attr_list(self: any, conf: str): + self.attr_list = self._parse_list(conf) + intersection_element = set(self.attr_list) & CONFLICT_KEYWORDS + if intersection_element: + raise Exception(f'[ERROR]: The attribute name: {intersection_element} in op: {{\'{self.op_type}\'}} \ +conflicts with the built-in variable name. Use a complex name or prefix the operator name.') + + def parse_mc2_ctx(self: any, conf: str): + self.mc2_ctx = self._parse_list(conf) + + @staticmethod + def _camel_to_snake(camel_case_str: str): + snake_case_str = '' + for i, c in enumerate(camel_case_str): + if i == 0: + snake_case_str += c.lower() + elif c.isupper(): + snake_case_str += '_' + c.lower() + else: + snake_case_str += c + return snake_case_str + + def parse_attr_val(self: any, conf: str): + for attr in self.attr_list: + if self.attr_val.get(attr) is None: + self.attr_val[attr] = {} + if conf.startswith('attr_{}.type'.format(attr)): + self.attr_val.get(attr)['type'] = self._camel_to_snake(self._parse_str(conf)) + elif conf.startswith('attr_{}.paramType'.format(attr)): + self.attr_val.get(attr)['paramType'] = self._parse_str(conf) + elif conf.startswith('attr_{}.defaultValue'.format(attr)): + self.attr_val.get(attr)['defaultValue'] = self._parse_str(conf) + + def parse_replay_val(self: any, batch_list: list, iterator_list: list): + if self.op_type in batch_list: + self.op_replay_flag = True + self.op_replay_batch = True + elif self.op_type in iterator_list: + self.op_replay_flag = True + self.op_replay_batch = False + + +def _is_op_type_in_opdesc(op_descs: list, op_type: str): + for op in op_descs: + if op_type == op.op_type: + return True + return False + + +def _set_all_options_to_opdescs(op_descs, soc_ver_compile_options): + for op in op_descs: + op.custom_all_compile_options = soc_ver_compile_options + + +def _set_options_to_opdesc(op_descs, op_type, soc_ver_compile_options): + for op in op_descs: + if op.op_type != op_type: + continue + op.custom_compile_options.update(soc_ver_compile_options) + + +def _trans_soc_ver_to_short(soc_ver: str): + low_soc_ver = soc_ver.lower() + if low_soc_ver not in SOC_TO_SHORT_SOC_MAP: + print(f'WARNING: caution: {soc_ver} will trans into ascend910, if not your intention,' + f'use ascend910b1~4 instead') + return SOC_TO_SHORT_SOC_MAP[low_soc_ver] + + +def _get_op_custom_options(op_descs: list, auto_gen_dir: str): + if auto_gen_dir is None: + return {} + file = os.path.join(auto_gen_dir, "custom_compile_options.ini") + if not os.path.exists(file): + print(f'WARNING: cannot find {auto_gen_dir}/custom_compile_options.ini') + return {} + with open (file, 'r') as fd: + lines = fd.readlines() + for line in lines: + param_list = str.split(line.rstrip('\n'), ',') + if len(param_list) != 3: + raise Exception(f'ERROR: custom compile option {param_list} len is not 3') + op_type = param_list[0] + if op_type.upper() == 'ALL': + op_type = OP_ALL + if op_type != OP_ALL and _is_op_type_in_opdesc(op_descs, op_type) == False: + continue + soc_ver_compile_options = {} + soc_ver = param_list[1] + options_str = param_list[2] + options = str.split(options_str, ';') + if soc_ver == '': + soc_ver_compile_options[SOC_ALL] = options + else: + soc_ver_list = str.split(soc_ver, ';') + for ver in soc_ver_list: + short_ver = _trans_soc_ver_to_short(ver) + soc_ver_compile_options[short_ver] = options + if op_type == OP_ALL: + _set_all_options_to_opdescs(op_descs, soc_ver_compile_options) + else: + _set_options_to_opdesc(op_descs, op_type, soc_ver_compile_options) + + +def get_op_desc(file: str, batch_list: list, iterator_list: list, builder: any, + op_type: list, auto_gen_dir: str = None) -> list: + op_descs = [] + op_match = False + with open (file, 'r') as fd: + lines = fd.readlines() + for line in lines: + line = line.strip() + if line.startswith('['): + name = line[1:-1] + if op_type is None or name in op_type: + op_match = True + op_desc = builder(name) + op_desc.parse_replay_val(batch_list, iterator_list) + op_descs.append(op_desc) + else: + op_match = False + if op_type is not None and len(op_descs) == len(op_type): + break + continue + if not op_match: + continue + if line.startswith('input'): + op_desc.parse_input(line) + elif line.startswith('output'): + op_desc.parse_output(line) + elif line.startswith('dynamicFormat.flag'): + op_desc.parse_op_format(line) + elif line.startswith('needCheckSupport.flag'): + op_desc.parse_check_support(line) + elif line.startswith('rangeLimit.value'): + op_desc.parse_range_limit(line) + elif line.startswith('opInterface.value'): + op_desc.parse_op_intf(line) + elif line.startswith('kernel.name'): + op_desc.parse_kern_name(line) + elif line.startswith('opFile.value'): + op_desc.parse_op_file(line) + elif line.startswith('dynamicShapeSupport.flag'): + op_desc.parse_dynamic_shape(line) + elif line.startswith('mc2.ctx'): + op_desc.parse_mc2_ctx(line) + elif line.startswith('attr.list'): + op_desc.parse_attr_list(line) + elif line.startswith('attr_'): + op_desc.parse_attr_val(line) + _get_op_custom_options(op_descs, auto_gen_dir) + return op_descs \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/preset_parse.py b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/preset_parse.py new file mode 100755 index 00000000..0d25f09a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/cmake/util/preset_parse.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +# Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. + +import json +import sys +import os + + +def read_json(file): + with open(file, 'r') as fd: + config = json.load(fd) + return config + + +def get_config_opts(file): + config = read_json(file) + + src_dir = os.path.abspath(os.path.dirname(file)) + opts = '' + + for conf in config: + if conf == 'configurePresets': + for node in config[conf]: + macros = node.get('cacheVariables') + if macros is not None: + for key in macros: + opts += '-D{}={} '.format(key, macros[key]['value']) + + opts = opts.replace('${sourceDir}', src_dir) + print(opts) + + +if __name__ == "__main__": + get_config_opts(sys.argv[1]) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/op_api_def.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/op_api_def.h new file mode 100644 index 00000000..ca83c7bf --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/op_api_def.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_api_def.h + * \brief + */ + +#ifndef Transformer_COMMON_OP_API_DEF_H +#define Transformer_COMMON_OP_API_DEF_H + +namespace op { + constexpr size_t MAX_SUPPORT_DIMS_NUMS = 8; + constexpr size_t BN_MIN_SUPPORT_DIMS_NUMS = 2; + constexpr int8_t FP16FP32_KEEP_DTYPE = -1; + constexpr int8_t KEEP_DTYPE = 0; + constexpr int8_t ALLOW_FP32_DOWN_PRECISION = 1; + constexpr int8_t USE_FP16 = 2; + constexpr int8_t USE_HF32 = 3; + constexpr size_t MAX_MASK_LEN64 = 64; +} // namespace op +#endif // Transformer_COMMON_OP_API_DEF_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.cpp new file mode 100644 index 00000000..56bab699 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.cpp @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "tensor_util.h" +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/transpose.h" +#include "aclnn_kernels/reshape.h" +#include "aclnn_kernels/cast.h" +#include "aclnn_kernels/contiguous.h" +#include "level0/unsqueeze.h" +#include "level0/squeeze.h" +#include "level0/fill.h" +#include "aclnn/aclnn_base.h" + +namespace op { +const aclIntArray* getAllDims(const aclTensor* self, aclOpExecutor* executor) { + auto input_shape = self->GetViewShape(); + const size_t input_dim_num = input_shape.GetDimNum(); + std::vector dims(input_dim_num); + for (size_t idx = 0; idx < input_dim_num; idx++) { + dims[idx] = idx; + } + return executor->AllocIntArray(dims.data(), input_dim_num); +} + +constexpr size_t MAX_DIM_CNT = 5; +const aclTensor* ResizeFrom1D(const aclTensor* cdim, const aclTensor* input, bool isSupportNcdhw, aclOpExecutor* executor) { + auto cdimContiguous = l0op::Contiguous(cdim, executor); + if (cdimContiguous == nullptr) { + return cdimContiguous; + } + + auto cdimCast = l0op::Cast(cdimContiguous, DataType::DT_FLOAT, executor); + if (cdimCast == nullptr) { + return cdimCast; + } + + size_t inputDim = input->GetViewShape().GetDimNum(); + + const int64_t appendDim[] = {0, 2, 3}; + aclIntArray* newShape = executor->AllocIntArray(appendDim, sizeof(appendDim) / sizeof(int64_t)); + if (inputDim == MAX_DIM_CNT) { + const int64_t value[] = {0, 2, 3, 4}; + newShape = executor->AllocIntArray(value, sizeof(value) / sizeof(int64_t)); + } + auto cdimUnsqueeze = l0op::UnsqueezeNd(cdimCast, newShape, executor); + if (cdimUnsqueeze == nullptr) { + return cdimUnsqueeze; + } + + op::Format format = inputDim == MAX_DIM_CNT ? Format::FORMAT_NCDHW : Format::FORMAT_NCHW; + auto cdimFormat = l0op::ReFormat(cdimUnsqueeze, format); + if (cdimFormat == nullptr) { + return cdimFormat; + } + + if ((inputDim == MAX_DIM_CNT) && !isSupportNcdhw) { + return l0op::TransDataSpecial(cdimFormat, Format::FORMAT_NDC1HWC0, 0, executor); + } + + return cdimFormat; +} + +const aclTensor* ResizeTo1D(const aclTensor* result, const aclTensor* output, bool isSupportNcdhw, aclOpExecutor* executor) { + auto resultTransdata = result; + size_t resultDim = result->GetViewShape().GetDimNum(); + if (resultDim >= MAX_DIM_CNT && !isSupportNcdhw) { + resultTransdata = l0op::TransDataSpecial(result, Format::FORMAT_NCDHW, 0, executor); + if (resultTransdata == nullptr) { + return resultTransdata; + } + } + + const int64_t appendDim[] = {0, 2, 3}; + aclIntArray* newShape = executor->AllocIntArray(appendDim, sizeof(appendDim) / sizeof(int64_t)); + if (resultTransdata->GetViewShape().GetDimNum() == MAX_DIM_CNT) { + const int64_t value[] = {0, 2, 3, 4}; + newShape = executor->AllocIntArray(value, sizeof(value) / sizeof(int64_t)); + } + auto resultNchw = l0op::SqueezeNd(resultTransdata, newShape, executor); + if (resultNchw == nullptr) { + return resultNchw; + } + + auto resultNd = l0op::ReFormat(resultNchw, Format::FORMAT_ND); + if (resultNd == nullptr) { + return resultNd; + } + + auto resultCast = l0op::Cast(resultNd, output->GetDataType(), executor); + if (resultCast == nullptr) { + return resultCast; + } + + return l0op::ViewCopy(resultCast, output, executor); +} + +const aclTensor* ResizeFromND(const aclTensor* input, aclOpExecutor* executor) { + const int nchw_dims = 4; + auto inputShape = input->GetViewShape(); + int64_t nchwShape[nchw_dims]; + for (size_t i = 0; i < nchw_dims; i++) { + nchwShape[i] = i < inputShape.GetDimNum() ? inputShape[i] : 1; + } + aclIntArray* nchwArray = executor->AllocIntArray(nchwShape, nchw_dims); + + auto inputReshape = l0op::Reshape(input, nchwArray, executor); + if (inputReshape == nullptr) { + return inputReshape; + } + + return l0op::ReFormat(inputReshape, Format::FORMAT_NCHW); +} + +const aclTensor* ResizeToND(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + size_t dimNum = inputShape.GetDimNum(); + + int64_t ndShape[dimNum]; + for (size_t i = 0; i < inputShape.GetDimNum(); i++) { + ndShape[i] = inputShape[i]; + } + aclIntArray* ndArray = executor->AllocIntArray(ndShape, dimNum); + + auto outputReshape = l0op::Reshape(output, ndArray, executor); + if (outputReshape == nullptr) { + return outputReshape; + } + + return l0op::ReFormat(outputReshape, input->GetViewFormat()); +} + +const aclTensor* ResizeFrom5D(const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + // NCDHW -> NDCHW + const int64_t value[] = {0, 2, 1, 3, 4}; + aclIntArray* ndchwShape = executor->AllocIntArray(value, MAX_DIM_CNT); + auto inputTranspose = l0op::Transpose(input, ndchwShape, executor); + if (inputTranspose == nullptr) { + return inputTranspose; + } + + // NDCHW -> NCHW + const int64_t nchwShape[] = {inputShape[0] * inputShape[2], inputShape[1], inputShape[3], inputShape[4]}; + aclIntArray* nchwArray = executor->AllocIntArray(nchwShape, sizeof(nchwShape) / sizeof(int64_t)); + auto inputReshape = l0op::Reshape(inputTranspose, nchwArray, executor); + if (inputReshape == nullptr) { + return inputReshape; + } + + return l0op::ReFormat(inputReshape, Format::FORMAT_NCHW); +} + +const aclTensor* ResizeTo5D(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + // nchw -> ndchw + const int64_t ndchwShape[] = {inputShape[0], inputShape[2], inputShape[1], inputShape[3], inputShape[4]}; + aclIntArray* ndchwArray = executor->AllocIntArray(ndchwShape, MAX_DIM_CNT); + auto outputReshape = l0op::Reshape(output, ndchwArray, executor); + if (outputReshape == nullptr) { + return outputReshape; + } + + auto outputFormat = l0op::ReFormat(outputReshape, Format::FORMAT_NCDHW); + if (outputFormat == nullptr) { + return outputFormat; + } + // ndchw -> ncdhw + const int64_t ncdhwShape[] = {0, 2, 1, 3, 4}; + aclIntArray* ncdhwArray = executor->AllocIntArray(ncdhwShape, MAX_DIM_CNT); + return l0op::Transpose(outputFormat, ncdhwArray, executor); +} + +aclTensor* FillScalar(int64_t dim, int value, aclOpExecutor* executor) { + const aclScalar* dimScalar = executor->AllocScalar(dim); + const aclTensor* dimTensor = executor->ConvertToTensor(dimScalar, op::DataType::DT_INT32); + aclIntArray* outShape = executor->AllocIntArray(&dim, 1); + + const aclScalar* valueScalar = executor->AllocScalar(value); + const aclTensor* valueTensor = executor->ConvertToTensor(valueScalar, op::DataType::DT_FLOAT); + + auto fillTensor = l0op::Fill(dimTensor, valueTensor, outShape, executor); + if (fillTensor == nullptr) { + return nullptr; + } + + return const_cast(fillTensor); +} + +aclTensor* FillVector(const op::Shape dstShape, const aclTensor* src, float value, aclOpExecutor* executor) { + op::FVector fillDims = op::ToShapeVector(dstShape); + auto shapes = executor->AllocIntArray(fillDims.data(), src->GetViewShape().GetDimNum()); + const aclTensor* dimTensor = executor->ConvertToTensor(shapes, op::DataType::DT_INT32); + const aclScalar* valueScalar = executor->AllocScalar(value); + const aclTensor* valueTensor = executor->ConvertToTensor(valueScalar, src->GetDataType()); + auto fillTensor = l0op::Fill(dimTensor, valueTensor, shapes, executor); + if (fillTensor == nullptr) { + return nullptr; + } + fillTensor = l0op::ReFormat(fillTensor, op::Format::FORMAT_ND); + return const_cast(fillTensor); +} + +aclnnStatus ProcessEmptyTensorWithValue(aclTensor* src, float initValue, aclOpExecutor* executor) { + auto srcShape = src->GetViewShape(); + auto dst = FillVector(srcShape, src, initValue, executor); + auto dstCopyResult = l0op::ViewCopy(dst, src, executor); + CHECK_RET(dstCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + return ACLNN_SUCCESS; +} + +op::DataType CombineCategories(op::DataType higher, op::DataType lower) { + if (IsFloatingType(higher)) { + return higher; + } + + if (IsFloatingType(lower) || higher == op::DataType::DT_BOOL) { + return op::PromoteType(higher, lower); + } + + return (higher != op::DataType::DT_UNDEFINED) ? higher : lower; +} +} // namespace op \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.h new file mode 100644 index 00000000..8120e73a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/common/tensor_util.h @@ -0,0 +1,47 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn/aclnn_base.h" +#include "opdev/common_types.h" + +namespace op { +const aclIntArray* getAllDims(const aclTensor* self, aclOpExecutor* executor); + +const aclTensor* ResizeFrom1D(const aclTensor* cdim, const aclTensor* input, bool isSupportNcdhw, + aclOpExecutor* executor); + +const aclTensor* ResizeTo1D(const aclTensor* result, const aclTensor* output, bool isSupportNcdhw, + aclOpExecutor* executor); + +const aclTensor* ResizeFromND(const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeToND(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeFrom5D(const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeTo5D(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor); + +aclTensor* FillScalar(int64_t dim, int value, aclOpExecutor* executor); + +aclnnStatus ProcessEmptyTensorWithValue(aclTensor* src, float initValue, aclOpExecutor* executor); + +op::DataType CombineCategories(op::DataType higher, op::DataType lower); +} // namespace op + +#ifdef __cplusplus +extern "C" { +#endif + +aclnnStatus BatchNorm(const aclTensor* input, const aclTensor* weight, const aclTensor* bias, aclTensor* runningMean, + aclTensor* runningVar, bool training, float momentum, float eps, aclTensor** output, + aclTensor* saveMean, aclTensor* saveInvstd, aclOpExecutor* executor); + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/err/ops_err.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/err/ops_err.h new file mode 100644 index 00000000..e83b8984 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/err/ops_err.h @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file ops_err.h + * \brief + */ + +#ifndef Transformer_COMMON_OPS_ERR_H +#define Transformer_COMMON_OPS_ERR_H + +#include "log/log.h" + +#define OPS_INNER_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...) \ + do { \ + OpLogSub(OP, DLOG_ERROR, OPS_DESC, FMT, ##__VA_ARGS__); \ + REPORT_INNER_ERR_MSG(ERR_CODE_STR, FMT, ##__VA_ARGS__); \ + } while (0) + + +/* 基础报错 */ +#define OPS_REPORT_VECTOR_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E89999", OPS_DESC, __VA_ARGS__) +#define OPS_REPORT_CUBE_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E69999", OPS_DESC, __VA_ARGS__) + +#endif // Transformer_COMMON_OPS_ERR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/cast.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/cast.h new file mode 100644 index 00000000..25f5c91f --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/cast.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Cast(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); + +// 专攻卷积反向定制 +const aclTensor* CastOnlyForConvBackward(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/common/op_error_check.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/common/op_error_check.h new file mode 100644 index 00000000..48f41354 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/common/op_error_check.h @@ -0,0 +1,251 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_ERROR_CHECK_H__ +#define OP_ERROR_CHECK_H__ + +#include "opdev/op_log.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/shape_utils.h" + +const int32_t NCHW_N_DIM = 0; +const int32_t NCHW_C_DIM = 1; +const int32_t NHWC_N_DIM = 0; +const int32_t NHWC_C_DIM = 3; + +static inline bool IsNullptr(const aclTensor *tensor, const char *name) { + if (tensor == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper Tensor but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclTensorList *tensorList, const char *name) { + if (tensorList == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper TensorList but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclScalar *scalar, const char *name) { + if (scalar == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type number for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclIntArray *intArr, const char *name) { + if (intArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[int] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclBoolArray *boolArr, const char *name) { + if (boolArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[bool] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclFloatArray *floatArr, const char *name) { + if (floatArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[float] for argument %s but instead found type \ + null.", name); + return true; + } + return false; +} + +static inline bool CheckDims(const aclTensor *tensor) { + const auto& xShape = tensor->GetViewShape(); + for(size_t i = 0; i < xShape.GetDimNum(); i++) { + if (xShape.GetDim(i) > INT32_MAX) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The tensor's shape cannot be larger than %d.", INT32_MAX); + return false; + } + } + return true; +} + +static inline bool CheckReduceOutShape(const aclTensor *inferOut, const aclTensor *out) +{ + auto const &xShape = inferOut->GetViewShape(); + auto const &yShape = out->GetViewShape(); + if (xShape != yShape) { + if (!(xShape.GetShapeSize() == 1 && yShape.GetShapeSize() == 1)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The out tensor's shape[%s] is not equal with inferOut shape[%s].", + op::ToString(out->GetViewShape()).GetString(), op::ToString(inferOut->GetViewShape()).GetString()); + return false; + } + } + return true; +} + +static inline bool CheckNCDimValid(const aclTensor *self, const aclTensor *out) { + auto format = self->GetStorageFormat(); + int64_t selfDimN = 0; + int64_t selfDimC = 0; + int64_t outDimN = 0; + int64_t outDimC = 0; + if (format == op::Format::FORMAT_NCHW) { + selfDimN = self->GetViewShape().GetDim(NCHW_N_DIM); + selfDimC = self->GetViewShape().GetDim(NCHW_C_DIM); + outDimN = out->GetViewShape().GetDim(NCHW_N_DIM); + outDimC = out->GetViewShape().GetDim(NCHW_C_DIM); + } else if (format == op::Format::FORMAT_NHWC) { + selfDimN = self->GetViewShape().GetDim(NHWC_N_DIM); + selfDimC = self->GetViewShape().GetDim(NHWC_C_DIM); + outDimN = out->GetViewShape().GetDim(NHWC_N_DIM); + outDimC = out->GetViewShape().GetDim(NHWC_C_DIM); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Input and output format only support [NCHW, NHWC] format ."); + return false; + } + if ((selfDimN != outDimN) || (selfDimC != outDimC)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The selfDimN[%ld]/outDimN[%ld] or selfDimC[%ld]/outDimC[%ld] not equal .", + selfDimN, outDimN, selfDimC, outDimC); + return false; + } + return true; +} + + +#define OP_CHECK_NULL(param, retExpr) \ + if (IsNullptr(param, #param)) { \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SUPPORT(tensor, supportList, retExpr) \ + if (!CheckType(tensor->GetDataType(), supportList)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s not implemented for %s, should be in dtype support list %s.", \ + #tensor, op::ToString(tensor->GetDataType()).GetString(), op::ToString(supportList).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_MATCH(tensor, expectedDtype, retExpr) \ + if (tensor->GetDataType() != expectedDtype) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s expected dtype is %s but found %s.", \ + #tensor, op::ToString(expectedDtype).GetString(), op::ToString(tensor->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SAME(tensor1, tensor2, retExpr) \ + if (tensor1->GetDataType() != tensor2->GetDataType()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected both tensors to have same dtype, but found %s %s and %s %s.", \ + #tensor1, op::ToString(tensor1->GetDataType()).GetString(), \ + #tensor2, op::ToString(tensor2->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_RESULT_DTYPE_CAST_FAILED(dtype, desiredDtype, retExpr); \ + if (!CanCast(dtype, desiredDtype)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Result type %s can't be cast to the desired output type %s.", \ + op::ToString(dtype).GetString(), op::ToString(desiredDtype).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST(tensor1, tensor2, retExpr) \ + if (!CheckBroadcastShape(tensor1->GetViewShape(), tensor2->GetViewShape())) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_WITH_SHAPE(tensor, shape, retExpr) \ + if (!CheckBroadcastShape(tensor->GetViewShape(), shape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size %s.", \ + #tensor, op::ToString(tensor->GetViewShape()).GetString(), op::ToString(shape).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_AND_INFER_SHAPE(tensor1, tensor2, retShape, retExpr) \ + if (!BroadcastInferShape(tensor1->GetViewShape(), tensor2->GetViewShape(), retShape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL(tensor1, tensor2, retExpr) \ + if (tensor1->GetViewShape() != tensor2->GetViewShape()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as tensor for %s, but %s does not " \ + "equal %s.", #tensor1, #tensor2, op::ToString(tensor1->GetViewShape()).GetString(), \ + op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(tensor, shape, retExpr) \ + if (tensor->GetViewShape() != shape) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as %s, but got %s.", \ + #tensor, op::ToString(shape).GetString(), op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_WRONG_DIMENSION(tensor, expectedDimNum, retExpr) \ + if (tensor->GetViewShape().GetDimNum() != expectedDimNum) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected %zu dimension input, but got %s with sizes %s.", \ + static_cast(expectedDimNum), #tensor, op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_MAX_DIM(tensor, maxDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() > static_cast(maxDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor cannot be larger than %zu dimensions.", \ + #tensor, static_cast(maxDim)); \ + retExpr; \ + } + +#define OP_CHECK_MIN_DIM(tensor, minDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() < static_cast(minDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor must have at least %zu dimensions.", \ + #tensor, static_cast(minDim)); \ + retExpr; \ + } + +#define OP_CHECK_COMM_INPUT(workspaceSize, executor) \ + if (workspaceSize == nullptr || executor == nullptr) { \ + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "The workspaceSize or executor is nullptr."); \ + return ACLNN_ERR_PARAM_NULLPTR; \ + } + +#define OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_STATIC_WORKSPACE_INVALID, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_INFERSHAPE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_INFERSHAPE_ERROR, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_TENSORLIST_SIZE_EQUAL(tensorlist1, tensorlist2, retExpr) \ + if ((tensorlist1)->Size() != (tensorlist2)->Size()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, \ + "The %s tensorlist and %s tensorlist must have the same number of tensors, but got %ld and %ld.", \ + #tensorlist1, #tensorlist2, (tensorlist1)->Size(), (tensorlist2)->Size()); \ + retExpr; \ + } + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/contiguous.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/contiguous.h new file mode 100644 index 00000000..0c49bfdc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/contiguous.h @@ -0,0 +1,89 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H + +#include "opdev/op_def.h" +#include "opdev/common_types.h" + +namespace l0op { + +typedef struct { + // 每个op::Shape 18ns + int64_t viewOffset; + + // Transpose + op::Shape transposeSrcShape; + op::Shape transposeDstShape; + op::FVector perm; + + // broadcast to + op::Shape broadcastSrcShape; + op::Shape broadcastDstShape; + op::FVector shape; + + // slice + op::Shape sliceSrcShape; + op::Shape sliceDstShape; + op::FVector offset; + op::FVector size; + + // strided slice + op::Shape stridedsliceSrcShape; + op::Shape stridedsliceDstShape; + op::FVector begin; + op::FVector end; + op::FVector strides; + + // optimizer + bool mayBroadcast; + bool mayTranspose; + bool maySlice; + bool mayStridedslice; +} ContiguousParam; + +/** + * @brief 将非连续Tensor转换为连续Tensor + * @param x + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* Contiguous(const aclTensor* x, aclOpExecutor* executor); + +/** + * @brief 将连续tensor拷贝到非连续的tensor上 + * @param x + * @param y + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* ViewCopy(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +/** + * @brief 对Tensor创建一个View，要求Tensor满足PickView的条件 + * @param x 输入Tensor，可以是一整块的非连续Tensor + * @param executor + * @return 输出Shape是一个连续Tensor + */ +const aclTensor* PickViewAsContiguous(const aclTensor* x, aclOpExecutor* executor); + +const aclTensor* ReViewToOut(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +// ============内部接口============= +bool CanOptimizeContiguous( + const op::Shape& viewShape, const op::Strides& strides, int64_t offset, int64_t storageSize, + ContiguousParam& param); + +bool CanOptimizeView(const op::Shape& viewShape, const op::Strides& strides, int64_t offset, ContiguousParam& param); +// ============内部接口============= +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/pad.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/pad.h new file mode 100644 index 00000000..ab740a43 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/pad.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Pad(const aclTensor* self, const aclTensor* paddings, aclOpExecutor* executor); +} +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/reshape.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/reshape.h new file mode 100644 index 00000000..6a9a62d2 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/reshape.h @@ -0,0 +1,37 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H + +#include "opdev/shape_utils.h" +#include "opdev/op_def.h" + +namespace l0op { +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor.ldd + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const op::Shape& shape, aclOpExecutor* executor); + +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor. + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const aclIntArray* shape, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/slice.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/slice.h new file mode 100644 index 00000000..3594f478 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/slice.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Slice( + const aclTensor* x, const aclTensor* y, const aclTensor* offset, const aclTensor* size, aclOpExecutor* executor); + +const aclTensor* Slice( + const aclTensor* x, const aclIntArray* offsets, const aclIntArray* size, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transdata.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transdata.h new file mode 100644 index 00000000..33e7bd7d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transdata.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor* ReFormat(const aclTensor* x, const op::Format& format, aclOpExecutor* executor = nullptr); + +/** + * TransData + * Formal Transdata. Set the c0 size strictly based on the data type and chip block size. + * support data type as follows: fp16,fp32,int32,uint32,int8,uint8 + * fp16: block_size/2 + * fp32/int32/uint32: block_size/4 (this is different from `TransDataSpecial`) + * int8/uint8: block_size/1 + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransData(const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); +/** + * Special Transdata. Set the c0 size strictly based on the data type and chip block size. + * this transdata c0 size rule: + * fp16: block_size/2 + * fp32/int32/uint32: block_size/2 + * int8/uint8: block_size/1 + * bool not supported, should do: + * (NCHW, bool)-> cast -> (NCHW, fp16) -> TransDataSpecial -> (5HD, fp16) -> cast -> (5HD, bool) + * (5HD, bool)-> cast -> (5HD, fp16) -> TransDataSpecial -> (NCHW, fp16) -> cast -> (NCHW, bool) + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransDataSpecial( + const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); + +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transpose.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transpose.h new file mode 100644 index 00000000..b67fa20a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_kernels/transpose.h @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Transpose(const aclTensor* x, const aclTensor* y, const aclTensor* perm, aclOpExecutor* executor); +const aclTensor* Transpose(const aclTensor* x, const aclIntArray* perm, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_util.h new file mode 100644 index 00000000..1b46d2df --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/external/aclnn_util.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_util.h + * \brief + */ +#ifndef Transformer_COMMON_ACLNN_UTIL_H +#define Transformer_COMMON_ACLNN_UTIL_H + +#define ACLNN_API __attribute__((visibility("default"))) + +#endif // Transformer_COMMON_ACLNN_UTIL_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback.h new file mode 100644 index 00000000..4fb9fc69 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback.h @@ -0,0 +1,499 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback.h + * \brief + */ + +#ifndef ACLNNFALLBACK_OPAPI_H_ +#define ACLNNFALLBACK_OPAPI_H_ + +#include + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "fallback/fallback_comm.h" +#include "mc2_log.h" +#include "runtime/base.h" +#include "log/log.h" + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; +using namespace std; + +namespace std_utils { + template + struct index_sequence {}; + + template + struct make_index_sequence_helper : make_index_sequence_helper {}; + + template + struct make_index_sequence_helper<0, Is...> { + using type = index_sequence; + }; + + template + using make_index_sequence = typename make_index_sequence_helper::type; +} + +using aclOpExecutor = struct aclOpExecutor; +using aclTensor = struct aclTensor; +using aclScalar = struct aclScalar; +using aclIntArray = struct aclIntArray; +using aclFloatArray = struct aclFloatArray; +using aclBoolArray = struct aclBoolArray; +using aclTensorList = struct aclTensorList; + +using _aclCreateTensor = aclTensor* (*)(const int64_t* view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t* stride, int64_t offset, aclFormat format, + const int64_t* storage_dims, uint64_t storage_dims_num, void* tensor_data); + +using _aclCreateScalar = aclScalar* (*)(void* value, aclDataType data_type); +using _aclCreateIntArray = aclIntArray* (*)(const int64_t* value, uint64_t size); +using _aclCreateFloatArray = aclFloatArray* (*)(const float* value, uint64_t size); +using _aclCreateBoolArray = aclBoolArray* (*)(const bool* value, uint64_t size); +using _aclCreateTensorList = aclTensorList* (*)(const aclTensor* const *value, uint64_t size); + +using _aclDestroyTensor = int (*)(const aclTensor* tensor); +using _aclDestroyScalar = int (*)(const aclScalar* scalar); +using _aclDestroyIntArray = int (*)(const aclIntArray* array); +using _aclDestroyFloatArray = int (*)(const aclFloatArray* array); +using _aclDestroyBoolArray = int (*)(const aclBoolArray* array); +using _aclDestroyTensorList = int (*)(const aclTensorList* array); + +#define GET_OP_API_FUNC(apiName) reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +inline const char* GetOpApiLibName(void) { + return "libopapi.so"; +} + +inline const char* GetCustOpApiLibName(void) { + return "libcust_opapi.so"; +} + +inline void* GetOpApiFuncAddrInLib(void* handler, const char* libName, const char* apiName) { + auto funcAddr = dlsym(handler, apiName); + if (funcAddr == nullptr) { + OP_LOGW("aclnnfallback", "dlsym %s from %s failed, error:%s.", apiName, libName, dlerror()); + } + return funcAddr; +} + +inline void* GetOpApiLibHandler(const char* libName) { + auto handler = dlopen(libName, RTLD_LAZY); + if (handler == nullptr) { + OP_LOGW("aclnnfallback", "dlopen %s failed, error:%s.", libName, dlerror()); + } + return handler; +} + +inline void* GetAclnnArrdByApiName(const char *apiName) { + vector libs = {"libaclnn_ops_infer.so", "libaclnn_ops_train.so", "libaclnn_math.so", + "libaclnn_rand.so", "libaclnn_sparse.so", "libaclnn_fft.so"}; + for (const auto &libName : libs) { + static auto libHandler = GetOpApiLibHandler(libName.c_str()); + if (libHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(libHandler, libName.c_str(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + } + OP_LOGE("aclnnfallback", "api %s can't find in any aclnn lib.", apiName); + return nullptr; +} + +inline void* GetOpApiFuncAddr(const char* apiName) { + static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (custOpApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + OP_LOGD("aclnnfallback", "opapi lib is not exist,will use aclnn lib."); + return GetAclnnArrdByApiName(apiName); +} + +inline aclTensor* ConvertType(aclTensor* ge_tensor) { + return ge_tensor; +} + +inline aclIntArray* ConvertType(const std::vector &arr) { + if (arr.empty()) { + return nullptr; + } + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + auto array = aclCreateIntArray(arr.data(), arr.size()); + return array; +} + +inline aclDataType GetConvertType(const gert::Tensor* ge_tensor) { + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = aclDataType::ACL_FLOAT16; + if (dataType_ge == DT_FLOAT) { + dataType = aclDataType::ACL_FLOAT; + } else if (dataType_ge == DT_BF16) { + dataType = aclDataType::ACL_BF16; + } else if (dataType_ge == DT_BOOL) { + dataType = aclDataType::ACL_BOOL; + } else if (dataType_ge == DT_INT64) { + dataType = aclDataType::ACL_INT64; + } else if (dataType_ge == DT_INT32) { + dataType = aclDataType::ACL_INT32; + } else if (dataType_ge == DT_UINT64) { + dataType = aclDataType::ACL_UINT64; + } else if (dataType_ge == DT_UINT32) { + dataType = aclDataType::ACL_UINT32; + } else if (dataType_ge == DT_INT8) { + dataType = aclDataType::ACL_INT8; + } else if (dataType_ge == DT_UINT8) { + dataType = aclDataType::ACL_UINT8; + } else if (dataType_ge == DT_INT4) { + dataType = aclDataType::ACL_INT4; + } else if (dataType_ge == DT_FLOAT8_E4M3FN) { + dataType = aclDataType::ACL_FLOAT8_E4M3FN; + } else { + dataType = aclDataType::ACL_FLOAT16; + } + + return dataType; +} + +inline aclTensor* ConvertType(const gert::Tensor* ge_tensor) { + if (ge_tensor == nullptr) { + return nullptr; + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OP_CHECK_IF(aclCreateTensor == nullptr, OP_LOGE("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = nullptr; + device_addr = const_cast(ge_tensor->GetAddr()); + + auto dataType = GetConvertType(ge_tensor); + + OP_LOGD("aclnnfallback", "aclCreateTensor: tensor type is %d", dataType); + + // convert shape + auto gert_shape = ge_tensor->GetStorageShape(); + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + aclTensor* out = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), + 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), device_addr); + + OP_CHECK_IF(out == nullptr, + OP_LOGE("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline aclTensorList* ConvertType(std::vector& ge_tenserList) { + OP_CHECK_IF(ge_tenserList.size() == 0, + OP_LOGE("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + OP_CHECK_IF(aclCreateTensorList == nullptr, + OP_LOGE("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + std::vector tmp; + for (size_t i = 0; i < ge_tenserList.size(); i++) { + auto t_acl = ConvertType(ge_tenserList[i]); + tmp.push_back(t_acl); + } + + aclTensorList* tensorList = aclCreateTensorList(tmp.data(), tmp.size()); + return tensorList; +} + +template +inline aclScalar* ConvertScalarType(T value) { + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + OP_CHECK_IF(aclCreateScalar == nullptr, + OP_LOGE("aclnnfallback", "aclCreateScalar nullptr"), return nullptr); + if (typeid(value) == typeid(float)) { + return aclCreateScalar(&value, aclDataType::ACL_FLOAT); + } + return nullptr; +} + +template +T ConvertType(T value) { + return value; +} + +inline aclTensor* ConvertMmType(const gert::Tensor* ge_tensor, bool transpose, bool enable_NZ=false) { + if (ge_tensor == nullptr) { + return nullptr; + } + auto gert_shape = ge_tensor->GetStorageShape(); + if (gert_shape.GetDimNum() <= 1) { + return ConvertType(ge_tensor); + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OP_CHECK_IF(aclCreateTensor == nullptr, OP_LOGE("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = const_cast(ge_tensor->GetAddr()); + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = ToAclDataType(dataType_ge); + // convert shape + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + auto viewShape = shape; + // 对于transpose后的tensor对后两维度进行strides, viewShape转换 + if (transpose) { + // dimM 为倒数第二维， dimN 为倒数第一维度 + auto dimM = shape.size() - 2; + auto dimN = shape.size() - 1; + auto swap = strides[dimN]; + strides[dimN] = strides[dimM]; + strides[dimM] = swap; + // 修改viewShape + viewShape[dimN] = shape[dimM]; + viewShape[dimM] = shape[dimN]; + } + auto acl_format = aclFormat::ACL_FORMAT_ND; + if (enable_NZ && GetPrimaryFormat(ge_tensor->GetStorageFormat()) == ge::Format::FORMAT_FRACTAL_NZ) { + acl_format = aclFormat::ACL_FORMAT_FRACTAL_NZ; + } + aclTensor* out = aclCreateTensor(viewShape.data(), shape.size(), dataType, strides.data(), + 0, acl_format, shape.data(), shape.size(), device_addr); + OP_CHECK_IF(out == nullptr, OP_LOGE("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline void Release(aclTensor* p) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + OP_CHECK_IF(aclDestroyTensor == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensor is null"), return); + aclDestroyTensor(p); +} + +inline void Release(aclScalar* p) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + OP_CHECK_IF(aclDestroyScalar == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyScalar is null"), return); + aclDestroyScalar(p); +} + +inline void Release(aclIntArray* p) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + OP_CHECK_IF(aclDestroyIntArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyIntArray is null"), return); + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray* p) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + OP_CHECK_IF(aclDestroyBoolArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyBoolArray is null"), return); + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList* p) { + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + OP_CHECK_IF(aclDestroyTensorList == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensorList is null"), return); + aclDestroyTensorList(p); +} + +template +void Release(T value) { + (void)value; +} + +template +void CallRelease(Tuple t, std_utils::index_sequence) { + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple& t) { + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std_utils::make_index_sequence{}); +} + +template +auto ConvertTypes(Ts&... args) -> decltype(std::make_tuple(ConvertType(args)...)) { + auto tp = std::make_tuple(ConvertType(args)...); + return tp; +} + +template +auto call(Function f, Tuple t, std_utils::index_sequence) -> int { + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) -> int { + static constexpr auto size = std::tuple_size::value; + return call(f, t, std_utils::make_index_sequence{}); +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr, std_utils::index_sequence) + -> int (*)(typename std::decay(params))>::type...) { + using LocalOpApiFunc = int (*)(typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr) + -> typename std::enable_if::value != 0, + decltype(ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence::value>{}))>::type { + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence{}); +} + +template +class ConvertedParams { + public: + ConvertedParams(Tuple&& convertedParams) : convertedParams_(std::move(convertedParams)){}; + ConvertedParams(ConvertedParams&& other) : convertedParams_(std::move(other.convertedParams_)) { + other.validParams_ = false; + }; + ConvertedParams& operator=(ConvertedParams&& other) { + if (this == &other) { + return *this; + } + + convertedParams_ = std::move(other.convertedParams_); + validParams_ = true; + other.validParams_ = false; + return *this; + } + + ConvertedParams() = delete; + ConvertedParams(const ConvertedParams& other) = delete; + ConvertedParams& operator=(const ConvertedParams& other) = delete; + + ~ConvertedParams() { + if (validParams_) { + ReleaseConvertTypes(convertedParams_); + } + } + + const Tuple& GetConvertedParams() const { + return convertedParams_; + } + + private: + Tuple convertedParams_; + bool validParams_{true}; +}; + +using InitHugeMemThreadLocal = int (*)(void*, bool); +using UnInitHugeMemThreadLocal = void (*)(void*, bool); +using ReleaseHugeMem = void (*)(void*, bool); +using PTAGetExecCache = aclOpExecutor* (*)(uint64_t, uint64_t*); +using InitPTACacheThreadLocal = void (*)(); +using SetPTAHashKey = void (*)(uint64_t); +using CanUsePTACache = bool (*)(const char*); + +using ResetCacheThreadLocal = void (*)(); + +#define EXEC_OPAPI_CMD(aclnn_api, ...) \ + ({ \ + static auto ret = GRAPH_SUCCESS; \ + do { \ + static const auto ResetCacheThreadLocalAddr = GetOpApiFuncAddr("ResetCacheThreadLocal"); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + if (getWorkspaceSizeFuncAddr == nullptr || opApiFuncAddr == nullptr || ResetCacheThreadLocalAddr == nullptr) { \ + OP_LOGE("aclnnfallback", "%s or %s not in %s or %s or ResetCacheThreadLocal not found.", \ + #aclnn_api "GetWorkspaceSize", #aclnn_api, GetOpApiLibName(), GetOpApiLibName()); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + auto ResetCacheThreadLocalFunc = reinterpret_cast(ResetCacheThreadLocalAddr); \ + ResetCacheThreadLocalFunc(); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + aclOpExecutor* executor = nullptr; \ + aclOpExecutor** executor_addr = &executor; \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + if (workspace_status != 0) { \ + OP_LOGE("aclnnfallback", "call %s failed:", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + void* workspace_addr = nullptr; \ + if (workspace_size > 0) { \ + workspace_addr = host_api_ctx->MallocWorkspace(workspace_size); \ + if (workspace_addr == nullptr) { \ + OP_LOGE("aclnnfallback", "call %s allocate workspace failed", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + } \ + auto acl_stream = host_api_ctx->GetStream(); \ + auto acl_call = [converted_params, workspace_addr, workspace_size, host_api_ctx, acl_stream, \ + executor]() -> int { \ + using OpApiFunc = int (*)(void*, uint64_t, aclOpExecutor*, const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret_inner = opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + ReleaseConvertTypes(converted_params); \ + host_api_ctx->FreeWorkspace(); \ + if (api_ret_inner != 0) { \ + OP_LOGE("aclnnfallback", "call %s allocate workspace failed api_ret_inner: %d", #aclnn_api, api_ret_inner); \ + return GRAPH_FAILED; \ + } \ + return api_ret_inner; \ + }; \ + \ + ret = acl_call(); \ + } while (false); \ + (ret); \ + }) + +} // namespace fallback + +#endif // ACLNNFALLBACK_OPAPI_H_ diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_2stages.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_2stages.h new file mode 100644 index 00000000..8564403a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_2stages.h @@ -0,0 +1,126 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ +#define ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ + +#include + +#include +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "fallback.h" +#include "fallback_comm.h" +#include "fallback_comm_2stages.h" +#include "log/log.h" +#include "mc2_log.h" + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; + +inline void Collect(aclTensor *p, std::vector ¶ms) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + OPS_ERR_IF(aclDestroyTensor == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensor is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyTensor(static_cast(param));}}); +} + +inline void Collect(aclScalar *p, std::vector ¶ms) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + OPS_ERR_IF(aclDestroyScalar == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyScalar is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyScalar(static_cast(param));}}); +} + +inline void Collect(aclIntArray *p, std::vector ¶ms) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + OPS_ERR_IF(aclDestroyIntArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyIntArray is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyIntArray(static_cast(param));}}); +} + +inline void Collect(aclBoolArray *p, std::vector ¶ms) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + OPS_ERR_IF(aclDestroyBoolArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyBoolArray is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyBoolArray(static_cast(param));}}); +} + +inline void Collect(aclTensorList *p, std::vector ¶ms) { + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + OPS_ERR_IF(aclDestroyTensorList == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensorList is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyTensorList(static_cast(param));}}); +} + +template +void Collect(T value, std::vector ¶ms) { + (void)value; + params.emplace_back(OpApiAnyValue{nullptr, nullptr}); +} + +template +void CallCollect(Tuple t, std_utils::index_sequence, std::vector ¶ms) { + (void)std::initializer_list{(Collect(std::get(t), params), 0)...}; +} + +template +void CollectConvertedTypes(Tuple &t, std::vector ¶ms) { + static constexpr auto size = std::tuple_size::value; + CallCollect(t, std_utils::make_index_sequence{}, params); +} + +#define EXEC_OPAPI_PREPARE_CMD(aclnn_api, ...) \ + ({ \ + static auto ret = GRAPH_SUCCESS; \ + do { \ + static const auto ResetCacheThreadLocalAddr = GetOpApiFuncAddr("ResetCacheThreadLocal"); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + if (getWorkspaceSizeFuncAddr == nullptr || opApiFuncAddr == nullptr || ResetCacheThreadLocalAddr == nullptr) { \ + OP_LOGE("aclnnfallback", "%s or %s not in %s or %s or ResetCacheThreadLocal not found.", \ + #aclnn_api "GetWorkspaceSize", #aclnn_api, GetOpApiLibName(), GetOpApiLibName()); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + auto *op_api_params = new (std::nothrow) OpApiParams(); \ + auto ResetCacheThreadLocalFunc = reinterpret_cast(ResetCacheThreadLocalAddr); \ + ResetCacheThreadLocalFunc(); \ + op_api_params->op_api_func = reinterpret_cast(opApiFuncAddr); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + aclOpExecutor** executor_addr = &op_api_params->executor; \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + using TupleT = decltype(converted_params); \ + constexpr size_t tuple_size = std::tuple_size::value; \ + op_api_params->converted_params.reserve(tuple_size); \ + CollectConvertedTypes(converted_params, op_api_params->converted_params); \ + host_api_ctx->SetOpApiParamsWithDefaultDeleter(op_api_params); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + if (workspace_status != 0) { \ + OP_LOGE("aclnnfallback", "call %s failed:", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + ret = host_api_ctx->SetWorkspaceSizes({workspace_size}); \ + } while (false); \ + (ret); \ + }) + +} // namespace fallback + +#endif // ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm.h new file mode 100644 index 00000000..778c568c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm.h @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm.h + * \brief + */ + +#ifndef INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ +#define INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ + +#include "aclnn/aclnn_base.h" +#include "exe_graph/runtime/op_execute_context.h" +#include "exe_graph/runtime/tensor.h" +#include "register/op_impl_registry.h" +#if __has_include("runtime/base.h") +#include "runtime/base.h" +#else +#include "runtime/rt_external_base.h" +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { + +aclDataType ToAclDataType(ge::DataType dtype); +} // namespace fallback + +#ifdef __cplusplus +} +#endif + +#endif // INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm_2stages.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm_2stages.h new file mode 100644 index 00000000..dbd741ec --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/fallback/fallback_comm_2stages.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ +#define INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ + +#include "aclnn/aclnn_base.h" +#include "aclnn/acl_meta.h" +#include "exe_graph/runtime/op_execute_context.h" +#include "exe_graph/runtime/op_execute_prepare_context.h" +#include "exe_graph/runtime/op_execute_launch_context.h" +#include "exe_graph/runtime/tensor.h" +#include "register/op_impl_kernel_registry.h" +#include "register/op_impl_registry.h" +#if __has_include("runtime/base.h") +#include "runtime/base.h" +#else +#include "runtime/rt_external_base.h" +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { + +using OpApiAnyValueDeleter = void (*)(void *); +typedef struct { + void *pointer; + OpApiAnyValueDeleter deleter; +} OpApiAnyValue; + +// aclnn算子params结构体，用于传递算子一阶段到二阶段的参数，定义在算子仓，由算子感知，GE框架不感知 +using OpApiFunc = int (*)(void *, uint64_t, aclOpExecutor *, const aclrtStream); +struct OpApiParams { + std::vector converted_params; // 算子下发依赖的参数 + aclOpExecutor *executor = nullptr; // aclOpExecutor指针 + OpApiFunc op_api_func = nullptr; // aclnnxx函数指针，实现算子launch下发 +}; + +// aclnn算子注册的二阶段launch func，函数实现可以与算子类型无关，所有算子使用同一个二阶段注册接口 +ge::graphStatus ExecuteOpLaunch(gert::OpExecuteLaunchContext *context); +} // namespace fallback + +#ifdef __cplusplus +} +#endif + +#endif // INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/framework/onnx_common.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/framework/onnx_common.h new file mode 100644 index 00000000..f88455e8 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/framework/onnx_common.h @@ -0,0 +1,85 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file onnx_common.h + * \brief + */ + +#ifndef MATH_COMMON_ONNX_COMMON_H +#define MATH_COMMON_ONNX_COMMON_H + +#include +#include +#include + +#include "stub_ops.h" +#include "register/register.h" +#include "graph/operator.h" +#include "graph/graph.h" +#include "base/err_msg.h" +#include "log/log.h" +#include "onnx/proto/ge_onnx.pb.h" + +namespace domi { +template +inline std::string GetOpName(const T& op) +{ + ge::AscendString op_ascend_name; + ge::graphStatus ret = op.GetName(op_ascend_name); + if (ret != ge::GRAPH_SUCCESS) { + std::string op_name = "None"; + return op_name; + } + return op_ascend_name.GetString(); +} + +template +inline ge::Tensor Vec2Tensor(vector& vals, const vector& dims, ge::DataType dtype, ge::Format format = ge::FORMAT_ND) { + ge::Shape shape(dims); + ge::TensorDesc desc(shape, format, dtype); + ge::Tensor tensor(desc, reinterpret_cast(vals.data()), vals.size() * sizeof(T)); + return tensor; +} + +template +inline ge::Tensor CreateScalar(T val, ge::DataType dtype, ge::Format format = ge::FORMAT_ND) { + vector dims_scalar = {}; + ge::Shape shape(dims_scalar); + ge::TensorDesc desc(shape, format, dtype); + ge::Tensor tensor(desc, reinterpret_cast(&val), sizeof(T)); + return tensor; +} + +inline Status ChangeFormatFromOnnx(ge::Operator& op, const int idx, ge::Format format, bool is_input) { + if (is_input) { + ge::TensorDesc org_tensor = op.GetInputDesc(idx); + org_tensor.SetOriginFormat(format); + org_tensor.SetFormat(format); + auto ret = op.UpdateInputDesc(idx, org_tensor); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "change input format failed."); + return FAILED; + } + } else { + ge::TensorDesc org_tensor_y = op.GetOutputDesc(idx); + org_tensor_y.SetOriginFormat(format); + org_tensor_y.SetFormat(format); + auto ret_y = op.UpdateOutputDesc(idx, org_tensor_y); + if (ret_y != ge::GRAPH_SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "change output format failed."); + return FAILED; + } + } + return SUCCESS; +} +} // namespace domi + +#endif // MATH_COMMON_ONNX_COMMON_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common.h new file mode 100644 index 00000000..a0d6be23 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ + +#ifndef INCLUDE_COMMON_H +#define INCLUDE_COMMON_H + +#define CONST_2 2 + +#define SET_FLAG(trigger, waiter, e) AscendC::SetFlag((e)) +#define WAIT_FLAG(trigger, waiter, e) AscendC::WaitFlag((e)) +#define PIPE_BARRIER(pipe) AscendC::PipeBarrier() + +#ifndef FORCE_INLINE +#define FORCE_INLINE inline __attribute__((always_inline)) +#endif + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common_func.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common_func.h new file mode 100644 index 00000000..983a5229 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/common_func.h @@ -0,0 +1,117 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common_func.h + * \brief + */ + + #ifndef INCLUDE_COMMON_FUNC_H + #define INCLUDE_COMMON_FUNC_H + + #include + #include + + #ifdef __CCE_KT_TEST__ + #include "stub_def.h" + #include "stub_fun.h" + #else + #include "kernel_operator.h" + #endif + + template + inline __aicore__ T RoundUp(const T val) + { + static_assert(ALIGN != 0, "align must not be zero"); + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + T align = ALIGN; + if (val + align - 1 < val) { + return val; + } + return (val + align - 1) / align * align; + } + + template + inline __aicore__ T RoundUp(const T val, const T align) + { + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + if (align == 0 || val + align - 1 < val) { + return val; + } + return (val + align - 1) / align * align; + } + + template + inline __aicore__ T CeilDiv(const T dividend) + { + static_assert(DIVISOR != 0, "align must not be zero"); + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + T divisor = DIVISOR; + if (dividend + divisor - 1 < dividend) { + return dividend; + } + return (dividend + divisor - 1) / divisor; + } + + template + constexpr T T_MAX = std::numeric_limits::max(); + + template + inline __aicore__ T CeilDiv(const T dividend, const T divisor) + { + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + if (divisor == 0 || dividend + divisor - 1 < dividend) { + return T_MAX; + } + return (dividend + divisor - 1) / divisor; + } + + template + __aicore__ inline T Min(const T lhs, const T rhs) + { + return lhs < rhs ? lhs : rhs; + } + + template __aicore__ __attribute__((always_inline)) inline uint32_t BlockSize() + { + return 32 / sizeof(Dtype); + } + + template __aicore__ __attribute__((always_inline)) inline uint32_t MatrixSize() + { + return 512 / sizeof(Dtype); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t BlockSizeRoundUp(uint64_t num) + { + return (num + BlockSize() - 1) / BlockSize() * BlockSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t NumBlocksRoundUp(uint64_t num) + { + return (num + BlockSize() - 1) / BlockSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t MatrixSizeRoundUp(uint64_t num) + { + return (num + MatrixSize() - 1) / MatrixSize() * MatrixSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t NumMatrixsRoundUp(uint64_t num) + { + return (num + MatrixSize() - 1) / MatrixSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t L0HalfSize() + { + return 32 * 1024 / sizeof(Dtype); + } + + #endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/dropmask.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/dropmask.h new file mode 100644 index 00000000..b57b0126 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/dropmask.h @@ -0,0 +1,121 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dropmask.h + * \brief + */ + +#ifndef DROPMASK_H +#define DROPMASK_H + +#include "util.h" + +using AscendC::DROPOUT_MODE_BIT_MISALIGN; +using AscendC::DropOutShapeInfo; +using AscendC::DropOut; + +struct DropMaskInfo { + // for compute dropout mask offset + // 参数按B N G S1 S2全部切分设置进行偏移计算，没有切分的轴对应的参数设置为合适的0或者原始值 + int64_t n2G; // n2 * g + int64_t gSize; // g + int64_t s1Size; // s1 + int64_t s2Size; // s2 + int64_t gOutIdx; // g out index + int64_t bSSOffset; // boidx * s1 * s2 ===bSSOffset + int64_t n2OutIdx; // n out index + int64_t s1OutIdx; // s1 out index ===s1oIdx + int64_t s1InnerIdx; // s1 inner index, 配比 ===loopIdx + int64_t s1BaseSize; // S1基本块大小 + int64_t splitS1BaseSize; // s1 split size ===vec1S1BaseSize + int64_t s2StartIdx; // s2 start index + int64_t s2Idx; // s2 index =====s2LoopCount + int64_t s2BaseNratioSize; // s2的配比长度: s2BaseSize(S2基本块大小) * nRatio + + // for copy in dropout mask + uint32_t s1CopySize; + uint32_t s2CopySize; + int64_t s2TotalSize; + + // for compute dropout mask + uint32_t firstAxis; + uint32_t lstAxis; + uint32_t maskLstAxis; + int64_t vecCoreOffset = 0; + float keepProb; + + bool boolMode; +}; + +template +__aicore__ inline int64_t ComputeDropOffset(DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + // boidx * n2 * g* s1 * s2 + int64_t bOffset = dropMaskInfo.bSSOffset * dropMaskInfo.n2G; + // n2oIdx * g * s1 *s2 + int64_t n2Offset = dropMaskInfo.n2OutIdx * dropMaskInfo.gSize * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // goIdx * s1 * s2 + int64_t gOffset = dropMaskInfo.gOutIdx * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // s1oIdx * s1BaseSize * s2Size + s1innerindex * vec1S1BaseSize * s2Size + int64_t s1Offset = (dropMaskInfo.s1OutIdx * dropMaskInfo.s1BaseSize + dropMaskInfo.vecCoreOffset + + dropMaskInfo.s1InnerIdx * dropMaskInfo.splitS1BaseSize) * dropMaskInfo.s2Size; + // s2StartIdx + s2index * s2BaseNratioSize + int64_t s2Offset = dropMaskInfo.s2StartIdx + dropMaskInfo.s2Idx * dropMaskInfo.s2BaseNratioSize; + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template +__aicore__ inline void CopyInDropMask(LocalTensor&dstTensor, GlobalTensor& srcBoolTensor, + GlobalTensor& srcByteTensor, DropMaskInfo &dropMaskInfo, int64_t alignedSize = blockBytes) +{ + if constexpr (hasDrop == true) { + int64_t dropMaskOffset = ComputeDropOffset(dropMaskInfo); + if (unlikely(dropMaskInfo.boolMode)) { + BoolCopyIn(dstTensor, srcBoolTensor, dropMaskOffset, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } else { + Bit2Int8CopyIn(dstTensor, srcByteTensor, dropMaskOffset, 1, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } + return; + } +} + +template +__aicore__ inline void ComputeDropMask(LocalTensor& dstTensor, LocalTensor& srcTensor, + LocalTensor& dropoutBuffer, LocalTensor& tmpDropBuffer, DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + DropOutShapeInfo dropOutShapeInfo; + dropOutShapeInfo.firstAxis = dropMaskInfo.firstAxis; + dropOutShapeInfo.srcLastAxis = dropMaskInfo.lstAxis; + + if (unlikely(dropMaskInfo.boolMode)) { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis, blockBytes) * blockBytes; + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis / byteBitRatio, blockBytes) * blockBytes; + if (likely(dropMaskInfo.lstAxis / byteBitRatio % blockBytes == 0)) { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, + dropMaskInfo.keepProb, dropOutShapeInfo); + } + } + return; + } +} + +#endif // DROPMASK_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_l1_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_l1_iterator.h new file mode 100644 index 00000000..401a7a93 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_l1_iterator.h @@ -0,0 +1,169 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_to_l1_iterator.h + * \brief + */ + +#ifndef GM_TO_L1_ITERATOR_H +#define GM_TO_L1_ITERATOR_H + +#include "iterator.h" + +constexpr uint32_t STRIDE_LIMIT_H = 65536; + +// Partial specialization for V220, ND_in, ND_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::DataCopyParams(1, // nBurst + CeilDiv(nTileActual * dTileActual), // lenBurst + 0, // srcGap + 0)); // dstGap + }; +}; + +// Partial specialization for NZ_in, NZ_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + uint64_t srcStride = nTileCeil - nTileActual; + if (srcStride < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, gmTensor, + AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst + nTileActual, // lenBurst + nTileCeil - nTileActual, // srcGap + 0)); // dstGap + } else { + for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) { + uint64_t dstOffset = i * nTileActual * BLOCK_SIZE; + uint64_t srcOffset = i * nTileCeil * BLOCK_SIZE; + AscendC::DataCopy(l1Tensor[dstOffset], gmTensor[srcOffset], + AscendC::DataCopyParams(1, // nBurst + nTileActual, // lenBurst + 0, // srcGap + 0)); // dstGap + } + } + }; +}; + +// Partial specialization for V220, ND_in, ND_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + if (dVal < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + nTileActual, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + dVal, // srcDValue + nTileCeil, // dstNzC0Stride + 1, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } else { + for (uint32_t i = 0; i < nTileActual; i++) { + AscendC::DataCopy(l1Tensor[i * BLOCK_SIZE], + gmTensor[i * dVal], + AscendC::Nd2NzParams(1, // ndNum + 1, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + 0, // srcDValue + nTileCeil, // dstNzC0Stride + 0, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } + } + }; +}; + +// Partial specialization for V220, ND_in, NZ_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + if (dVal < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + nTileActual, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + dVal, // srcDValue + nTileCeil, // dstNzC0Stride + 1, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } else { + for (uint32_t i = 0; i < nTileActual; ++i) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + 1, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + 0, // srcDValue + nTileCeil, // dstNzC0Stride + 0, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } + } + }; +}; + +#endif // GM_TO_L1_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_ub_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_ub_iterator.h new file mode 100644 index 00000000..a906bc00 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/gm_to_ub_iterator.h @@ -0,0 +1,97 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_to_ub_iterator.h + * \brief + */ + +#ifndef GM_TO_UB_ITERATOR_H +#define GM_TO_UB_ITERATOR_H + +#include "iterator.h" + +constexpr uint32_t STRIDE_LIMIT_I = 65536; + +template struct gm_to_ub { + __aicore__ inline gm_to_ub(AscendC::LocalTensor dstTensor, AscendC::GlobalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template struct gm_to_ub_align { + __aicore__ inline gm_to_ub_align(AscendC::LocalTensor dstTensor, AscendC::GlobalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum, + uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap) + { + AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0), + AscendC::DataCopyPadExtParams(false, leftPaddingNum, rightPaddingNum, 0)); + }; +}; + +template struct ub_to_ub { + __aicore__ inline ub_to_ub(AscendC::LocalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template +struct ub_to_gm { + __aicore__ inline ub_to_gm(AscendC::GlobalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template struct ub_to_gm { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ ub_to_gm(AscendC::GlobalTensor gmTensor, AscendC::LocalTensor l1Tensor, + uint32_t nTileActual, uint32_t nTileCeil, uint32_t nVal, uint32_t dTileActual, + uint32_t dTileCeil, uint32_t dVal) + { + uint64_t dstStride = nTileCeil - nTileActual; + if (dstStride < STRIDE_LIMIT_I) { + AscendC::DataCopy(gmTensor, l1Tensor, + AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst + nTileActual, // lenBurst + 0, // srcGap + dstStride)); // dstGap + } else { + for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) { + uint64_t srcOffset = i * nTileActual * BLOCK_SIZE; + uint64_t dstOffset = i * nTileCeil * BLOCK_SIZE; + AscendC::DataCopy(gmTensor[dstOffset], l1Tensor[srcOffset], + AscendC::DataCopyParams(1, // nBurst + nTileActual, // lenBurst + 0, // srcGap + 0)); // dstGap + } + } + }; +}; + +template struct ub_to_gm_align { + __aicore__ inline ub_to_gm_align(AscendC::GlobalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum, + uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap) + { + AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0)); + }; +}; + +#endif // GM_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/hardware.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/hardware.h new file mode 100644 index 00000000..c5fa1156 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/hardware.h @@ -0,0 +1,40 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hardware.h + * \brief + */ + +#ifndef INCLUDE_HARDWARE_H +#define INCLUDE_HARDWARE_H + +enum class ArchType { ASCEND_V220, ASCEND_V200, ASCEND_M200 }; + +template +struct HardwareInfo { + static uint32_t const l2BW = 5; + static uint32_t const hbmBW = 1; + static uint32_t const supportMix = 0; + static uint32_t const l1Size = 512 * 1024; + static uint32_t const l0ASize = 64 * 1024; + static uint32_t const l0BSize = 64 * 1024; + static uint32_t const l0CSize = 128 * 1024; + static uint32_t const l2Size = 192 * 1024 * 1024; + static uint32_t const biasSize = 1024; + static uint32_t const fixBufSize = 7 * 1024; + static uint32_t const ubSize = 192 * 1024; + static uint32_t const fractalSize = 512; + static uint32_t const l1l0BlockSize = 32; + static uint32_t const btBlockSize = 64; + static uint32_t const fbBlockSize = 128; +}; + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/iterator.h new file mode 100644 index 00000000..95180180 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/iterator.h @@ -0,0 +1,123 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file iterator.h + * \brief + */ + +#ifndef INCLUDE_ITERTOR_H +#define INCLUDE_ITERTOR_H + +#include "common_func.h" +#include "hardware.h" +#include "kernel_operator.h" +#include "layout.h" +#include "mem.h" + +///////////////////////////////////////////////////// +// gm_to_l1 +///////////////////////////////////////////////////// +template +struct gm_to_l1 { + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_a +///////////////////////////////////////////////////// +template +struct l1_to_l0_a { + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// +template +struct l1_to_l0_b { + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) {}; +}; + +// l1_to_l0_a +///////////////////////////////////////////////////// +template +struct l1_to_l0_a_v1 { + __aicore__ l1_to_l0_a_v1(AscendC::LocalTensor l0_tensor, + AscendC::LocalTensor l1_tensor, + uint32_t m_tile_ceil, + uint32_t k_tile_ceil, + uint32_t k_part, + uint32_t k_part_ceil, + uint32_t k_part_idx) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// +template +struct l1_to_l0_b_v1 { + __aicore__ l1_to_l0_b_v1(AscendC::LocalTensor l0_tensor, + AscendC::LocalTensor l1_tensor, + int32_t n_tile_ceil, + int32_t k_tile_ceil, + int32_t k_part_ceil, + int32_t k_part_idx) {}; +}; + +///////////////////////////////////////////////////// +// l0c_to_gm +///////////////////////////////////////////////////// +template +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) {}; +}; + +///////////////////////////////////////////////////// +// l0c_to_l1 +///////////////////////////////////////////////////// +template +struct l0c_to_l1 { + __aicore__ l0c_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor deqTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) {}; +}; + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_gm_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_gm_iterator.h new file mode 100644 index 00000000..03c332a7 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_gm_iterator.h @@ -0,0 +1,213 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_gm_iterator.h + * \brief + */ + +#ifndef L0C_TO_GM_ITERATOR_H +#define L0C_TO_GM_ITERATOR_H + +#ifdef __CCE_KT_TEST__ +#define __bf16 bfloat16_t +#endif + +#include "iterator.h" +constexpr uint32_t BLOCK_NUM = 16; +constexpr uint32_t BLOCK_SIZE_INT8 = 32; + +template <> +struct l0c_to_gm { + /** + * @brief Copy data from L0C buffer to global memory, partial specialized for + * + * @param gmTensor the destination tensor on global memory, which is stored in ND format. + * @param l0cTensor the source tensor on L0C buffer, which is stored in FRACTAL_NZ format. + * @param mTileActual the m-direction size of the matrix in L0C buffer. + * @param nTileActual the n-direction size of the matrix in L0C buffer. + * @param srcStride the source stride between the adjacent fractal matrics along n-direction in unit of C0_SIZE. + * @param dstStride the leading dimension of the destination matrix in unit of element. + */ + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322F16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::F322F16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::VDEQF16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::VDEQF16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor<__bf16> gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322BF16; + AscendC::Fixpipe<__bf16, float, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::F322BF16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +// Partial specialization ND, float +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::NoQuant; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::NoQuant}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322F16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride - (nTileActual * sizeof(half) / sizeof(float))); + intriParams.quantParams = {QuantMode_t::F322F16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride){ +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::NoQuant; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif +}; +}; + +#endif // L0C_TO_GM_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_l1_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_l1_iterator.h new file mode 100644 index 00000000..9f3fc5c3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_l1_iterator.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_l1_iterator.h + * \brief + */ + +#ifndef L0C_TO_L1_ITERATOR_H +#define L0C_TO_L1_ITERATOR_H + +#include "iterator.h" +///////////////////////////////////////////////////// +// l0c_to_l1 +///////////////////////////////////////////////////// + +// Partial specialization ZN, half, int32_t +template +struct l0c_to_l1 { + using ElementOut = half; + using ElementIn = int32_t; + __aicore__ l0c_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor deqTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) + { + constexpr uint32_t BLOCK_NUM = 16; + constexpr uint32_t BLOCK_SIZE = 32; + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE), + 0, + mTileCeil - static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE) * + sizeof(ElementOut) / sizeof(ElementIn)); + intriParams.nz2ndParams = {false, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::VDEQF16}; + AscendC::Fixpipe(l1Tensor, l0cTensor, deqTensor, intriParams); + }; +}; + +#endif // L0C_TO_L1_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_ub_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_ub_iterator.h new file mode 100644 index 00000000..5d2ea8e7 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l0c_to_ub_iterator.h @@ -0,0 +1,72 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_ub_iterator.h + * \brief + */ +#ifndef L0C_TO_UB_ITERATOR_H +#define L0C_TO_UB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l0c_to_ub +///////////////////////////////////////////////////// + +// Partial specialization ZN, half, int32_t +template struct l0c_to_ub { + __aicore__ l0c_to_ub(AscendC::LocalTensor ubTensor, AscendC::LocalTensor l0cTensor, + uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + constexpr auto mode = + MatrixMode ? AscendC::BlockMode::BLOCK_MODE_MATRIX : AscendC::BlockMode::BLOCK_MODE_VECTOR; + AscendC::DataCopy(ubTensor, l0cTensor, + AscendC::DataCopyParams(nBurst, // count + lenBurst, // len + srcStride, // srcStrideIn + dstStride), // dstStrideIn + AscendC::DataCopyEnhancedParams(mode, // blockModeIn + AscendC::DeqScale::DEQ_NONE, // deqScaleIn + 0, // deqValueIn + 0, // sidStoreModeIn + false, // isReluIn + pad_t::PAD_NONE, // padModeIn + 0) // padValueIn + ); + }; +}; + +template +struct l0c_to_ub { + __aicore__ l0c_to_ub(AscendC::LocalTensor ubTensor, + AscendC::LocalTensor l0cTensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(ubTensor, l0cTensor, + AscendC::DataCopyParams(nBurst, // count + lenBurst, // len + srcStride, // srcStrideIn + dstStride), // dstStrideIn + AscendC::DataCopyEnhancedParams(AscendC::BlockMode::BLOCK_MODE_MATRIX, // blockModeIn + AscendC::DeqScale::VDEQ16, // deqScaleIn + 0, // deqValueIn + 0, // sidStoreModeIn + false, // isReluIn + pad_t::PAD_NONE, // padModeIn + 0) // padValueIn + ); + }; +}; + +#endif // L0C_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_bt_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_bt_iterator.h new file mode 100644 index 00000000..a2db67a3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_bt_iterator.h @@ -0,0 +1,39 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_bt_iterator.h + * \brief + */ +#ifndef L1_TO_BT_ITERATOR_H +#define L1_TO_BT_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_bt +///////////////////////////////////////////////////// + +// Partial specialization for V220 +template +struct l1_to_bt { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::btBlockSize / sizeof(DataType); + + __aicore__ l1_to_bt(AscendC::LocalTensor biasTableTensor, + AscendC::LocalTensor biasL1Tensor, + uint32_t ntileActual) + { + AscendC::DataCopy( + biasTableTensor, biasL1Tensor, {1, static_cast(CeilDiv(ntileActual)), 0, 0}); + }; +}; + +#endif // L1_TO_BT_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_fb_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_fb_iterator.h new file mode 100644 index 00000000..8f75074a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_fb_iterator.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_fb_iterator.h + * \brief + */ + +#ifndef L1_TO_FB_ITERATOR_H +#define L1_TO_FB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_fb +///////////////////////////////////////////////////// + +// Partial specialization for V220 +template +struct l1_to_fb { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::fbBlockSize / sizeof(DataType); + + __aicore__ + l1_to_fb(AscendC::LocalTensor fbTensor, AscendC::LocalTensor l1Tensor, uint32_t ntileActual) + { + copy_cbuf_to_fbuf((__fbuf__ DataType *)fbTensor.GetPhyAddr(), + (__cbuf__ DataType *)l1Tensor.GetPhyAddr(), + 1, + CeilDiv(ntileActual), + 0, + 0); + }; +}; + +#endif // L1_TO_FB_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_l0_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_l0_iterator.h new file mode 100644 index 00000000..72e58ddc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_l0_iterator.h @@ -0,0 +1,259 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_l0_iterator.h + * \brief + */ + +#ifndef L1_TO_L0_ITERATOR_H +#define L1_TO_L0_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_l0_a +///////////////////////////////////////////////////// + +// Partial specialization for vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + AscendC::LoadData(l0Tensor, + l1Tensor, + AscendC::LoadData2dParams(0, // baseIdx + kPartCeil, // repeat + kSrcStride, // srcStride + 0, // sid + kDstStride, // dstStride + IsTransPose, // transpose + 0)); // addrCalMode + }; +}; + +// Partial specialization for no transpose, not vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / BLOCK_NUM_PER_FRACTAL; i++) { + AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE], + l1Tensor[i * mSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, // baseIdx + static_cast(kPartCeil / BLOCK_SIZE), // repeat + kSrcStride, // srcStride + 0, // sid + kDstStride - 1, // dstStride + false, // transpose + 0)); // addrCalMode + } + }; +}; + +// Partial specialization for transpose, not vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / BLOCK_SIZE; i++) { + AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE], + l1Tensor[i * mSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, + static_cast(kPartCeil / BLOCK_NUM_PER_FRACTAL), + kSrcStride, + 0, + kDstStride - 1, + true, + 0)); + } + }; +}; + +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + // 16 * 32 + static constexpr uint32_t ROW_BLOCK_SIZE = 16; + static constexpr uint32_t COL_BLOCK_SIZE = 32 / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / ROW_BLOCK_SIZE; i++) { + AscendC::LoadData(l0Tensor[i * ROW_BLOCK_SIZE * kPartCeil], + l1Tensor[i * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, + static_cast(kPartCeil / COL_BLOCK_SIZE), + mTileCeil / ROW_BLOCK_SIZE, + 0, + 0, + false, + 0)); + } + }; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// + +// Partial specialization for vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + AscendC::LoadData( + l0Tensor, l1Tensor, AscendC::LoadData2dParams(0, kPartCeil, kSrcStride, 0, kDstStride, IsTransPose, 0)); + }; +}; + +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + using DataType = int8_t; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < nTileCeil / BLOCK_SIZE; i++) { + AscendC::LoadDataWithTranspose(l0Tensor[i * kPartCeil * BLOCK_SIZE], + l1Tensor[i * BLOCK_SIZE * BLOCK_SIZE], + AscendC::LoadData2dTransposeParams(0, // startIndexIn + kPartCeil / BLOCK_SIZE, // repeatTimesIn + nTileCeil / BLOCK_SIZE, // srcStrideIn + 1, // dstGapIn + 0, // dstfracGapIn + 0) // addrModeIn + ); + } + }; +}; + +// Partial specialization for no transpose, not vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < kPartCeil / BLOCK_NUM_PER_FRACTAL; i++) { + AscendC::LoadData(l0Tensor[i * kDstStride * FRACTAL_SIZE], + l1Tensor[i * kSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, // baseIdx + static_cast(nTileCeil / BLOCK_SIZE), // repeat + nSrcStride, // srcStride + 0, // sid + nDstStride - 1, // dstStride + true, // transpose + 0)); // addrCalMode + } + }; +}; + +// Partial specialization for transpose, not vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + AscendC::LoadData( + l0Tensor, + l1Tensor, + AscendC::LoadData2dParams(0, // baseIdx + static_cast(kPartCeil * nTileCeil / FRACTAL_SIZE), // repeat + 1, // srcStride + 0, // sid + 0, // dstStride + false, // transpose + 0)); // addr_cal_mode_t + }; +}; + +#endif // L1_TO_L0_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_ub_iterator.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_ub_iterator.h new file mode 100644 index 00000000..2446d110 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/l1_to_ub_iterator.h @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_ub_iterator.h + * \brief + */ + +#ifndef L1_TO_UB_ITERATOR_H +#define L1_TO_UB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_ub +///////////////////////////////////////////////////// +template +struct l1_to_ub { + __aicore__ l1_to_ub(AscendC::LocalTensor ubTensor, + AscendC::LocalTensor l1Tensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(ubTensor, l1Tensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +///////////////////////////////////////////////////// +// ub_to_l1 +///////////////////////////////////////////////////// +template +struct ub_to_l1 { + __aicore__ ub_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor ubTensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(l1Tensor, ubTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; +#endif // L1_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/layout.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/layout.h new file mode 100644 index 00000000..bdd3f024 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/layout.h @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file layout.h + * \brief + */ + +#ifndef INCLUDE_LAYOUT_H +#define INCLUDE_LAYOUT_H + +enum class DataFormatT { + ND = 0, + NZ, + ZN, + ZZ, + NN, + VECTOR +}; + +#endif // INCLUDE_LAYOUT_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mem.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mem.h new file mode 100644 index 00000000..467d1454 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mem.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file mem.h + * \brief + */ + +#ifndef INCLUDE_MEM_H +#define INCLUDE_MEM_H + +#include "hardware.h" +#include "kernel_operator.h" +#include "kernel_tensor.h" + +enum class BufferType { ASCEND_UB, ASCEND_CB, ASCEND_L0A, ASCEND_L0B, ASCEND_L0C, ASCEND_MAX }; + +template +__aicore__ constexpr AscendC::TPosition GetPosition() +{ + if constexpr (BufferType_ == BufferType::ASCEND_UB) { + return AscendC::TPosition::VECIN; + } else if constexpr (BufferType_ == BufferType::ASCEND_CB) { + return AscendC::TPosition::A1; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0A) { + return AscendC::TPosition::A2; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0B) { + return AscendC::TPosition::B2; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0C) { + return AscendC::TPosition::CO1; + } + return AscendC::TPosition::GM; +} + +template +struct AsdopsBuffer { +public: + __aicore__ AsdopsBuffer() + { + constexpr uint32_t bufferSize[(uint32_t)BufferType::ASCEND_MAX] = {HardwareInfo::ubSize, + HardwareInfo::l1Size, + HardwareInfo::l0ASize, + HardwareInfo::l0BSize, + HardwareInfo::l0CSize}; +#ifdef __DAV_C220_VEC__ + tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]); +#elif __DAV_C220_CUBE__ + tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]); + tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]); + tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]); + tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]); +#else +#ifndef __clang__ + tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]); + tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]); + tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]); + tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]); + tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]); +#endif +#endif + }; + + template + __aicore__ AscendC::LocalTensor GetBuffer(const uint32_t offset) const + { + return tensor[(uint32_t)BufferType_][offset].template ReinterpretCast(); + } + +public: + AscendC::LocalTensor tensor[(uint32_t)BufferType::ASCEND_MAX]; +}; +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mma.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mma.h new file mode 100644 index 00000000..72ab9fd3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/mma.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file mma.h + * \brief + */ + +#ifndef INCLUDE_MMA_H +#define INCLUDE_MMA_H + +#include "hardware.h" +#include "kernel_tensor.h" + +template +struct mmad { + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) {}; + + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint64_t biasBt, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) {}; +}; + +// Partial specialization for V220, int8_t, not_vector_A, not TransposeA +template +struct mmad { + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) + { + AscendC::Mmad(l0cTensor, + l0aTensor, + l0bTensor, + AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC)); + }; + + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint64_t biasBt, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) + { + AscendC::LocalTensor biasTensor; + biasTensor.InitBuffer(biasBt, mTileActual); + biasTensor.address_.logicPos = static_cast(AscendC::TPosition::C2); + AscendC::Mmad(l0cTensor, + l0aTensor, + l0bTensor, + biasTensor, + AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC)); + }; +}; + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/pse.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/pse.h new file mode 100644 index 00000000..185a7ea3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/pse.h @@ -0,0 +1,483 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file pse.h + * \brief + */ + +#ifndef FLASH_ATTENTION_SCORE_PSE_H +#define FLASH_ATTENTION_SCORE_PSE_H + +#include "kernel_operator.h" +#include "util.h" + +constexpr static int64_t pseS1S2 = 0; +constexpr static int64_t pse1S2 = 1; +constexpr static int64_t pseSlopeBn = 2; +constexpr static int64_t pseSlopeN = 3; + +constexpr static uint8_t pseEncodeALibiS2Full = 0x11; + +enum class PseTypeEnum { + PSE_OUTER_MUL_ADD_TYPE = 0, // default + PSE_OUTER_ADD_MUL_TYPE, + PSE_INNER_MUL_ADD_TYPE, + PSE_INNER_MUL_ADD_SQRT_TYPE, + PSE_INVALID_TYPE +}; + +struct PseInfo { + int64_t blockCount; + int64_t bSSOffset; // boidx * s1 * s2 + int64_t boIdx; + int64_t gSize; + int64_t goIdx; + int64_t loopIdx; + int64_t n2G; + int64_t n2oIdx; + int64_t pseBSize; + int64_t pseS1Size; // for alibi + int64_t pseS2ComputeSize; // for alibi, do not need assignment + int64_t pseS2Size; // for alibi + uint32_t pseShapeType; + int64_t readS2Size; // for alibi, do not need assignment + int64_t s1BaseSize; + int64_t s1Size; + int64_t s1oIdx; + int64_t s2AlignedSize; + int64_t s2BaseNratioSize; + int64_t s2LoopCount; + int64_t s2RealSize; + int64_t s2Size; + int64_t s2SizeAcc; // accumulated sum of s2 size + int64_t s2StartIdx; + int64_t vec1S1BaseSize; + int64_t vec1S1RealSize; + uint32_t pseEncodeType; // for distinguish alibi + uint32_t pseType; // 0: outer, mul-add 1:outer, add-mul 2:inner, mul-add 3:inner, mul-add-sqrt + int64_t pseAlibiBaseS1; + int64_t pseAlibiBaseS2; + int64_t qStartIdx; + int64_t kvStartIdx; + int64_t vecCoreOffset = 0; + bool needCast; + bool align8 = false; + bool pseEndogenous = false; +}; + +template +__aicore__ inline void DataCopyInCommon(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int32_t dtypeSize, + int32_t alignedS2Size) +{ + if constexpr (hasPse == true) { + uint32_t shapeArray[] = {static_cast(s1Size), static_cast(alignedS2Size)}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.blockLen = CeilDiv(s2Size * dtypeSize, blockBytes); // 单位32B + dataCopyParams.dstStride = alignedS2Size * dtypeSize / blockBytes - dataCopyParams.blockLen; // gap + if (actualS2Len * dtypeSize % blockBytes == 0) { + dataCopyParams.srcStride = + (actualS2Len * dtypeSize - dataCopyParams.blockLen * blockBytes) / blockBytes; // srcGap + DataCopy(dstTensor, srcTensor[offset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size * dtypeSize; // 单位Byte + dataCopyParams.srcStride = (actualS2Len * dtypeSize - dataCopyParams.blockLen); + dataCopyParams.dstStride = (alignedS2Size - s2Size) * dtypeSize / blockBytes; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = false; + DataCopyPad(dstTensor, srcTensor[offset], dataCopyParams, dataCopyPadParams); + } + } +} + +template +__aicore__ inline void DataCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + int32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +template +__aicore__ inline void DataCopyInAlign8(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + if (dtypeSize == 0){ + return; + } + int32_t alignedS2Size = CeilDiv(s2Size, 32 / dtypeSize) * (32 / dtypeSize); + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +/* +dst = BroadcastAdd(src0, src1) +src0 shape: (s1, s2) +src1 shape: (1, s2) +dst shape: (s1, s2) +*/ +template +__aicore__ inline void BroadcastAdd(const LocalTensor &src0Tensor, const LocalTensor &src1Tensor, + int64_t src0Offset, int32_t src1Size, int32_t repeatTimes) +{ + if constexpr (hasPse == true) { + /* Total data number of single step should be smaller than 256bytes. + * If larger, we need to do add multiple times. */ + int32_t innerLoop = src1Size / repeatMaxSize; // s2轴整块计算次数 + int32_t innerRemain = src1Size % repeatMaxSize; // s2轴尾块计算量 + BinaryRepeatParams binaryRepeatParams; + binaryRepeatParams.src0BlkStride = 1; + binaryRepeatParams.src0RepStride = src1Size / blockSize; + binaryRepeatParams.src1BlkStride = 1; + binaryRepeatParams.src1RepStride = 0; + binaryRepeatParams.dstRepStride = binaryRepeatParams.src0RepStride; + binaryRepeatParams.blockNumber = binaryRepeatParams.src0RepStride; + + for (int32_t j = 0; j < innerLoop; j++) { + auto innerOffset = j * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], repeatMaxSize, repeatTimes, + binaryRepeatParams); + } + if (innerRemain > 0) { + auto innerOffset = innerLoop * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], innerRemain, repeatTimes, + binaryRepeatParams); + } + } +} + +template +__aicore__ inline void PseBroadcastAdd(int32_t s1Size, int32_t s2Size, int32_t computeSize, const LocalTensor &pseUb, + const LocalTensor &dstTensor, uint32_t pseShapeType) +{ + if constexpr (hasPse == true) { + if (pseShapeType == pseS1S2 || pseShapeType == pseSlopeBn || pseShapeType == pseSlopeN) { + Add(dstTensor, dstTensor, pseUb, computeSize); + } else { + /* Total repeated times should be <= repeatMaxTimes. If larger, + * we need to do multiple inner loops. */ + int32_t s1OuterLoop = s1Size / repeatMaxTimes; + int32_t s1OuterRemain = s1Size % repeatMaxTimes; + for (int32_t s1OuterIdx = 0; s1OuterIdx < s1OuterLoop; s1OuterIdx++) { + int32_t s1OuterOffset = s1OuterIdx * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, repeatMaxTimes); + } + if (s1OuterRemain > 0) { + int32_t s1OuterOffset = s1OuterLoop * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, s1OuterRemain); + } + } + } +} +template __aicore__ inline int64_t PseComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = 0; + int64_t s1Offset = 0; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t gOffset = 0; + if (pseInfo.pseShapeType == pseS1S2) { + // b, n2, g, s1, s2 + bOffset = pseInfo.bSSOffset * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s1Size * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s1Size * pseInfo.s2Size; + s1Offset = (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize) * pseInfo.s2Size; + } else if (pseInfo.pseShapeType == pse1S2) { + // b, n2, g, 1, s2 + bOffset = pseInfo.s2SizeAcc * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s2Size; + } + if (pseInfo.pseBSize == 1) { + bOffset = 0; + } + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline int64_t PseAlibiComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = (pseInfo.boIdx % pseInfo.pseBSize) * pseInfo.n2G * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t gOffset = pseInfo.goIdx * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t row = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t column = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t m = 0; + int64_t k = 0; + if constexpr (layOutType != LayOutTypeEnum::LAYOUT_TND) { + int64_t threshold = pseInfo.s1Size - pseInfo.pseS1Size; + if (row >= threshold) { + m = row - threshold; + k = column; + } else { + m = row % pseInfo.pseS1Size; + k = pseInfo.pseS2Size - (row - column) - (pseInfo.pseS1Size - m); + } + } else { + int64_t threshold = pseInfo.pseS2Size - pseInfo.pseS1Size; + int64_t posVal = row - column - threshold; + if (threshold >= 0) { + if (posVal >= 0) { + m = posVal; + k = 0; + } else { + m = 0; + k = -posVal; + } + } else { + m = posVal; + k = 0; + } + } + int64_t s1Offset = m * pseInfo.pseS2Size; + int64_t s2Offset = k; + pseInfo.readS2Size = Min(pseInfo.s2AlignedSize, pseInfo.pseS2Size - k); + pseInfo.pseS2ComputeSize = Align(pseInfo.readS2Size); + + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline bool NeedPseAlibiCompute(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + // Alibi编码只计算下三角 + if (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + (pseInfo.loopIdx + 1) * pseInfo.vec1S1BaseSize <= + pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize) { + return false; + } + return true; + } else { + return false; + } +} + +template +__aicore__ inline void PseAlibiCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + int64_t offset = PseAlibiComputeOffset(pseInfo); + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, + pseInfo.readS2Size, pseInfo.pseS2Size); + } + return; + } + + DataCopyIn(tmpTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + } + return; + } +} + +template +__aicore__ inline void PseSlopeCopyIn(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, GlobalTensor &alibiGm, PseInfo &pseInfo, + int64_t alignedSize = 16) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + + DataCopyIn(helpTensor, alibiGm, 0, pseInfo.vec1S1RealSize, + pseInfo.s2RealSize, pseInfo.pseAlibiBaseS2, alignedSize); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + + if (pseInfo.needCast) { + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } + } +} + +template +__aicore__ inline void PseSlopeCast(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } +} + +template +__aicore__ inline void PseCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCopyIn(dstTensor, tmpTensor, srcTensor, pseInfo, alignedSize); + } + int64_t offset = PseComputeOffset(pseInfo); + int64_t s1Size = pseInfo.pseShapeType == pse1S2 ? (pseInfo.blockCount == 0 ? 1 : pseInfo.blockCount) : + pseInfo.vec1S1RealSize; + + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, + pseInfo.s2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size); + } + return; + } + DataCopyIn(tmpTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size, + alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, s1Size * pseInfo.s2AlignedSize); + } + return; + } +} + +template +__aicore__ inline void PseAlibiCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + Add(dstTensor, dstTensor, pseTensor, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + return; + } +} + +template +__aicore__ inline void PseCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCompute(dstTensor, pseTensor, pseInfo); + } + int64_t computeSize = (pseInfo.pseShapeType == pseS1S2 || pseInfo.pseShapeType == pseSlopeBn || + pseInfo.pseShapeType == pseSlopeN) + ? pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize + : pseInfo.s2AlignedSize; + PseBroadcastAdd(pseInfo.vec1S1RealSize, pseInfo.s2AlignedSize, computeSize, pseTensor, + dstTensor, pseInfo.pseShapeType); + return; + } +} + +template +__aicore__ inline void PseInnerAlibiCreate(GlobalTensor &dstTensor, LocalTensor &helpTensor, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + if (pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_TYPE && pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + return; + } + event_t eventIdMte3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + event_t eventIdMte3ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + float tmpValue = -1.0; + + for (int64_t i = 0; i < pseInfo.pseAlibiBaseS1; i++) { + CreateVecIndex(helpTensor, (half)(i * tmpValue), pseInfo.pseAlibiBaseS2); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopy(dstTensor[i * pseInfo.pseAlibiBaseS2], helpTensor, pseInfo.pseAlibiBaseS2); + SetFlag(eventIdMte3ToV); + WaitFlag(eventIdMte3ToV); + SetFlag(eventIdMte3ToS); + WaitFlag(eventIdMte3ToS); + } + } +} +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/simd.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/simd.h new file mode 100644 index 00000000..de67a52a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/simd.h @@ -0,0 +1,433 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file simd.h + * \brief + */ + +#ifndef INCLUDE_SIMD_H +#define INCLUDE_SIMD_H + +#ifdef __CCE_KT_TEST__ +#define __bf16 bfloat16_t +#endif + +#include "hardware.h" +#include "kernel_operator.h" + +///////////////////////////////////////////////////// +// vadd +///////////////////////////////////////////////////// +template +__aicore__ inline void add_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Add( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vadds +///////////////////////////////////////////////////// +template +__aicore__ inline void adds_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + DType scalarValue, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t srcBlockStride, + uint8_t dstRepeatStride, + uint8_t srcRepeatStride) +{ + AscendC::Adds( + dst, + src, + scalarValue, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vcadd +///////////////////////////////////////////////////// +template +__aicore__ inline void cadd_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstRepeatStride, + uint16_t srcBlockStride, + uint16_t srcRepeatStride) +{ + AscendC::RepeatReduceSum(dst, src, repeat, 0, 0, srcBlockStride, dstRepeatStride, srcRepeatStride); +} +///////////////////////////////////////////////////// +// vbrcb +///////////////////////////////////////////////////// +template +__aicore__ inline void brcb_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint16_t dstBlockStride, + uint16_t dstRepeatStride, + uint8_t repeat) +{ + AscendC::Brcb(dst, src, repeat, AscendC::BrcbRepeatParams(dstBlockStride, dstRepeatStride)); +} + +///////////////////////////////////////////////////// +// vcmax +///////////////////////////////////////////////////// +template +__aicore__ inline void cmax_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstRepeatStride, + uint16_t srcBlockStride, + uint16_t srcRepeatStride) +{ +#if defined(__DAV_C220_VEC__) + AscendC::WholeReduceMax( + dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride, OrderType); +#else + AscendC::WholeReduceMax( + dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride); +#endif +} + +///////////////////////////////////////////////////// +// vconv +///////////////////////////////////////////////////// +template +__aicore__ inline void conv_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + if constexpr (std::is_same::value && std::is_same::value) { + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); + } else { + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); + } +} + +///////////////////////////////////////////////////// +// vconv_f322bf16r +///////////////////////////////////////////////////// +template +__aicore__ inline void convr_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vdiv +///////////////////////////////////////////////////// +template +__aicore__ inline void div_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Div( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vexp +///////////////////////////////////////////////////// +template +__aicore__ inline void exp_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Exp( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vmax +///////////////////////////////////////////////////// +template +__aicore__ inline void max_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Max( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmul +///////////////////////////////////////////////////// +template +__aicore__ inline void mul_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Mul( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmuls +///////////////////////////////////////////////////// +template +__aicore__ inline void muls_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Muls( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vsub +///////////////////////////////////////////////////// +template +__aicore__ inline void sub_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Sub( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmaxs +///////////////////////////////////////////////////// +template +__aicore__ inline void maxs_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Maxs( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vmins +///////////////////////////////////////////////////// +template +__aicore__ inline void mins_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Mins( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vsqrt +///////////////////////////////////////////////////// +template +__aicore__ inline void sqrt_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Sqrt( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vln +///////////////////////////////////////////////////// +template +__aicore__ inline void ln_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Ln( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vtranspose +///////////////////////////////////////////////////// +template +__aicore__ inline void tranpose_v(AscendC::LocalTensor dst, AscendC::LocalTensor src) +{ + AscendC::Transpose(dst, src); +} + +///////////////////////////////////////////////////// +// vcgmax +///////////////////////////////////////////////////// +template +__aicore__ inline void cgmax_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + const int32_t repeat, + const int32_t dstRepStride, + const int32_t srcBlkStride, + const int32_t srcRepStride) +{ + AscendC::BlockReduceMax(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride); +} + +///////////////////////////////////////////////////// +// vcgadd +///////////////////////////////////////////////////// +template +__aicore__ inline void cgadd_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + const int32_t repeat, + const int32_t dstRepStride, + const int32_t srcBlkStride, + const int32_t srcRepStride) +{ + AscendC::BlockReduceSum(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride); +} +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/util.h new file mode 100644 index 00000000..71ac733b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/kernel/util.h @@ -0,0 +1,159 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file util.h + * \brief + */ + +#ifndef FLASH_ATTENTION_UTIL_H +#define FLASH_ATTENTION_UTIL_H + +constexpr int32_t blockBytes = 32; +constexpr int32_t byteBitRatio = 8; +constexpr int64_t prefixAttenMaskDownHeight = 1024; +constexpr static int32_t blockSize = blockBytes / 4; // 4 means sizeof(T) +constexpr static int32_t repeatMaxBytes = 256; +constexpr static int32_t repeatMaxTimes = 255; +constexpr static int32_t repeatMaxSize = repeatMaxBytes / 4; // 4 means sizeof(T) + +using AscendC::LocalTensor; +using AscendC::GlobalTensor; +using AscendC::DataFormat; +using AscendC::ShapeInfo; +using AscendC::DataCopyParams; +using AscendC::DataCopyExtParams; +using AscendC::DataCopyPadParams; +using AscendC::DataCopyPadExtParams; +using AscendC::BinaryRepeatParams; +using AscendC::IsSameType; +using AscendC::HardEvent; +using AscendC::SetFlag; +using AscendC::WaitFlag; + +enum class LayOutTypeEnum { None = 0, LAYOUT_BSH = 1, LAYOUT_SBH = 2, LAYOUT_BNSD = 3, LAYOUT_TND = 4, LAYOUT_NTD_TND = 5}; + +namespace math { +template __aicore__ inline T Ceil(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template __aicore__ inline T Align(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} +} + +template +__aicore__ inline T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +__aicore__ inline void BoolCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t s1Size, uint32_t s2Size, int64_t totalS2Size, int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + uint32_t shapeArray[] = {s1Size, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.dstStride = 0; + if (totalS2Size == blockBytes && alignedSize == 64) { // totalS2Size < 64 && totalS2Size % blockBytes == 0 + dataCopyParams.dstStride = 1; + alignedSize = blockBytes; + alignedS2Size = CeilDiv(s2Size, blockBytes) * blockBytes; + } + if (likely(totalS2Size - s2Size <= UINT16_MAX)) { + if (totalS2Size % alignedSize == 0) { + dataCopyParams.blockLen = alignedS2Size / blockBytes; + dataCopyParams.srcStride = (totalS2Size - alignedS2Size) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size; + dataCopyParams.srcStride = totalS2Size - s2Size; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes); + dataCopyPadParams.paddingValue = 1; + DataCopyPad(dstTensor, srcTensor[srcOffset], dataCopyParams, dataCopyPadParams); + } + } else { + DataCopyExtParams extParams; + extParams.blockCount = s1Size; + extParams.dstStride = 0; + extParams.blockLen = s2Size; + extParams.srcStride = totalS2Size - s2Size; + DataCopyPadExtParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes); + dataCopyPadParams.paddingValue = 1; + DataCopyPad(dstTensor, srcTensor[srcOffset], extParams, dataCopyPadParams); + } +} + +__aicore__ inline void Bit2Int8CopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t batchSize, uint32_t s1BaseSize, uint32_t s2BaseSize, int64_t s2TotalSize, + int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2BaseSize / byteBitRatio, alignedSize) * alignedSize; + uint32_t shapeArray[] = {batchSize * s1BaseSize, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(batchSize * s1BaseSize * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = batchSize * s1BaseSize; + dataCopyParams.blockLen = CeilDiv(s2BaseSize / byteBitRatio, blockBytes); + dataCopyParams.dstStride = 0; + if (s2TotalSize / byteBitRatio % alignedSize == 0 && s2BaseSize / byteBitRatio % alignedSize == 0) { + dataCopyParams.srcStride = + (s2TotalSize / byteBitRatio - dataCopyParams.blockLen * blockBytes) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams); + } else { + dataCopyParams.blockLen = CeilDiv(s2BaseSize , byteBitRatio); + dataCopyParams.srcStride = (s2TotalSize - s2BaseSize) / byteBitRatio; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = 0; + dataCopyPadParams.paddingValue = 0; + DataCopyPad(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams, dataCopyPadParams); + } +} + +__aicore__ inline int32_t Align(int32_t shape) +{ + int32_t alignFactor = 16; + int32_t alignedSize = CeilDiv(shape, alignFactor) * alignFactor; + return alignedSize; +} + +#endif // FLASH_ATTENTION_UTIL_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/op_graph/op_transformer_proto_extend.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/op_graph/op_transformer_proto_extend.h new file mode 100644 index 00000000..b779f14a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/op_graph/op_transformer_proto_extend.h @@ -0,0 +1,100 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_transformer_proto_extend.h + * \brief + */ +#ifndef OPS_OP_MATH_PROTO_EXTEND_H_ +#define OPS_OP_MATH_PROTO_EXTEND_H_ + +#include "graph/operator_reg.h" + +namespace ge { +/** +* @brief swin_transformer model specific structure.Operator only supports swin_transformer. + +* @par Inputs: +* Three inputs, including: +* @li x: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be (B*W, N, S1, S2) or (B, W, N, S1, S2). +* @li atten_mask: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be (W, S1, S2) or (W, 1, S1, S2) or (1, W, 1, S1, S2) +* @li relative_pos_bias: An ND Tensor. Must be one of the following types: float16, float, bfloat16. + the shape sholud be (N, S1, S2) or (1, N, S1, S2) or (1, 1, N, S1, S2) + +* @par Attributes: +* @li scale_value: A optional attribute, the type is float. Defaults to 1.0. +* @li inner_precision_mode: A optional attribute, the type is int. Defaults to 0, reserved field. + +* @par Outputs: +* One output, including: +* @li y: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be same with x. +*/ +REG_OP(MaskedSoftmaxWithRelPosBias) + .INPUT(x, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .OPTIONAL_INPUT(atten_mask, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .INPUT(relative_pos_bias, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .OUTPUT(y, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .ATTR(scale_value, Float, 1.0) + .ATTR(inner_precision_mode, Int, 0) + .OP_END_FACTORY_REG(MaskedSoftmaxWithRelPosBias) + +/** +* @brief AttentionScore's forward calculation. + +* @par Inputs: +* six inputs, including: +* @li query: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li key: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li value: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li padding_mask: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li scale: A scalar. The type only support float16. Enter a 4D Tensor. +* @li drop_mask: A matrix Tensor. An optional input parameter. The type only support uint8. Enter a 4D Tensor. + +* @par Attributes: +* @li keep_prob: A float. The keep probability of dropout. Default: 1.0. +* @li query_transpose: A bool. If True, changes the shape of "query" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li key_transpose: A bool. If True, changes the shape of "key" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li bmm_score_transpose_a: A bool. If True, changes the shape of "mid_data" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li bmm_score_transpose_b: A bool. If True, changes the shape of "value" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li softmax_axes: A list of int. The dimension softmax would be performed on. Defaults to "[-1]". + +* @par Outputs: +* attention_score: The result matrix Tensor. The type only support float16. The output shape is the same as query. +* softmax_output: The result matrix Tensor. The type only support float16. The output shape is the same as query. + +* @par Restrictions: +* Warning: THIS FUNCTION IS EXPERIMENTAL. Please do not use. +*/ +REG_OP(AttentionScore) + .INPUT(query, TensorType({DT_FLOAT16})) + .INPUT(key, TensorType({DT_FLOAT16})) + .INPUT(value, TensorType({DT_FLOAT16})) + .INPUT(padding_mask, TensorType({DT_FLOAT16})) + .INPUT(scale, TensorType({DT_FLOAT16})) + .OPTIONAL_INPUT(drop_mask, TensorType({DT_INT8})) + .OUTPUT(attention_score, TensorType({DT_FLOAT16})) + .OUTPUT(softmax_output, TensorType({DT_FLOAT16})) + .ATTR(keep_prob, Float, 1.0) + .ATTR(query_transpose, Bool, false) + .ATTR(key_transpose, Bool, false) + .ATTR(bmm_score_transpose_a, Bool, false) + .ATTR(bmm_score_transpose_b, Bool, false) + .ATTR(softmax_axes, ListInt, {-1}) + .OP_END_FACTORY_REG(AttentionScore) +} + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/op_resource.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/op_resource.h new file mode 100644 index 00000000..2ad2ed1a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/op_resource.h @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_resource.h + * \brief + */ +#ifndef COMMON_NN_OP_RESOURCE_H +#define COMMON_NN_OP_RESOURCE_H + +#define EXTERN_OP_RESOURCE(kernelName) \ +namespace l0op { \ + extern void * kernelName##TilingRegisterResource(); \ + extern void * kernelName##InferShapeRegisterResource(); \ + extern void * kernelName##TuningRegisterResource(); \ + extern const OP_BINARY_RES& kernelName##KernelResource(); \ + extern const OP_RUNTIME_KB_RES& kernelName##TuningResource(); \ + [[maybe_unused]] uint32_t kernelName##_kernelName_Be_Defined_Multi_Times___; \ +} + +#define AUTO_GEN_OP_RESOURCE(kernelName) {{ #kernelName, \ + {{l0op::kernelName##TilingRegisterResource(), l0op::kernelName##InferShapeRegisterResource(), l0op::kernelName##TuningRegisterResource()}, \ + l0op::kernelName##KernelResource(), l0op::kernelName##TuningResource()}}} \ + +#endif // COMMON_NN_OP_RESOURCE_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/static_space.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/static_space.h new file mode 100644 index 00000000..3491adc5 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/static/static_space.h @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file static_space.h + * \brief + */ +#ifndef CANN_OPS_STATIC_SPACE_H_ +#define CANN_OPS_STATIC_SPACE_H_ +#include "base/registry/op_impl_space_registry_v2.h" + +class StaticSpaceInitializer { +public: + static StaticSpaceInitializer& GetInstance() { + static StaticSpaceInitializer instance; + return instance; + } +private: + StaticSpaceInitializer () { + auto space_registry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry(); + if (space_registry == nullptr) { + space_registry = std::make_shared(); + gert::DefaultOpImplSpaceRegistryV2::GetInstance().SetSpaceRegistry(space_registry); + } + } +}; +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling.h new file mode 100644 index 00000000..2f2c3e4a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include "data_copy_transpose_tiling_def.h" + +namespace optiling { + +inline void GetDataCopyTransposeTiling(const ge::Shape &dstShape, const ge::Shape &srcShape, const uint32_t typeSize, + optiling::CopyTransposeTiling &tiling) +{ + constexpr int64_t B_INDEX = 0; + constexpr int64_t N_INDEX = 1; + constexpr int64_t S_INDEX = 2; + constexpr int64_t H_INDEX = 3; + std::vector dstShapeInfo = dstShape.GetDims(); + std::vector srcShapeInfo = srcShape.GetDims(); + + tiling.set_dstShapeB(dstShapeInfo[B_INDEX]); + tiling.set_dstShapeN(dstShapeInfo[N_INDEX]); + tiling.set_dstShapeS(dstShapeInfo[S_INDEX]); + tiling.set_dstShapeH(dstShapeInfo[H_INDEX]); + tiling.set_dstShapeHN(tiling.get_dstShapeH() / tiling.get_dstShapeN()); + + tiling.set_srcShapeB(srcShapeInfo[B_INDEX]); + tiling.set_srcShapeN(srcShapeInfo[N_INDEX]); + tiling.set_srcShapeS(srcShapeInfo[S_INDEX]); + tiling.set_srcShapeHN(srcShapeInfo[H_INDEX]); + tiling.set_originalShapeNLen(tiling.get_srcShapeHN() * typeSize); + tiling.set_shapeSHValue(tiling.get_dstShapeS() * tiling.get_dstShapeH()); + tiling.set_shapeNsValue(tiling.get_dstShapeN() * tiling.get_dstShapeS()); + tiling.set_shapeNsnValue(tiling.get_dstShapeN() * tiling.get_srcShapeS() * tiling.get_srcShapeN()); + tiling.set_shapeBHValue(tiling.get_dstShapeB() * tiling.get_dstShapeH()); +} + +} // namespace optiling diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling_def.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling_def.h new file mode 100644 index 00000000..891a566a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/data_copy_transpose_tiling_def.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling_def.h + * \brief + */ + +#pragma once + +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyTransposeTiling) +TILING_DATA_FIELD_DEF(uint32_t, dstShapeB); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeS); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeH); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeB); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeN); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeS); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, originalShapeNLen); +TILING_DATA_FIELD_DEF(uint32_t, shapeSHValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsnValue); +TILING_DATA_FIELD_DEF(uint32_t, invalidParamCopyTransposeTiling); +TILING_DATA_FIELD_DEF(uint32_t, shapeBHValue); +TILING_DATA_FIELD_DEF(uint32_t, paramsAlign); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(CopyTransposeTilingOp, CopyTransposeTiling) + +} // namespace optiling diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_base.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_base.h new file mode 100644 index 00000000..9c64f217 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_base.h @@ -0,0 +1,256 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "tiling/platform/platform_ascendc.h" +#include "log/log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +struct AiCoreParams { + uint64_t ubSize = 0; + uint64_t blockDim = 0; + uint64_t aicNum = 0; + uint64_t l1Size = 0; + uint64_t l0aSize = 0; + uint64_t l0bSize = 0; + uint64_t l0cSize = 0; +}; + +struct CompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +struct FlashAttentionScoreGradCompileInfo { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + platform_ascendc::SocVersion socVersion; +}; + +struct FACompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +class TilingBaseClass { +public: + explicit TilingBaseClass(gert::TilingContext* context) : context_(context) + {} + + virtual ~TilingBaseClass() = default; + + // Tiling执行框架 + // 1、GRAPH_SUCCESS: 成功，并且不需要继续执行后续Tiling类的实现 + // 2、GRAPH_FAILED: 失败，中止整个Tiling流程 + // 3、GRAPH_PARAM_INVALID: 本类不支持，需要继续往下执行其他Tiling类的实现 + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // 更新 context + virtual void Reset(gert::TilingContext* context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2、获取INPUT/OUTPUT/ATTR信息 + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3、计算数据切分TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4、计算高阶API的TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5、计算TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6、计算Workspace 大小 + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7、保存Tiling数据 + virtual ge::graphStatus PostTiling() = 0; + // 8、Dump Tiling数据 + virtual void DumpTilingInfo() + { + int32_t enable = CheckLogLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t*)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context_->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template + [[nodiscard]] std::string GetShapeDebugStr(const T& shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr( + const gert::StorageShape* shape, const gert::CompileTimeTensorDesc* tensor) + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_key.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_key.h new file mode 100644 index 00000000..3c256093 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_key.h @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_key.h + * \brief + */ + +#pragma once + +#include + +namespace Ops { +namespace Transformer { +namespace OpTiling { +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr uint64_t kBase = 10; // 10进制进位基数 +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + kBase * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputed types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputed types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace Optiling +} // namespace Transformer +} // namespace Ops diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_templates_registry.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_templates_registry.h new file mode 100644 index 00000000..c62a09ab --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_templates_registry.h @@ -0,0 +1,350 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "exe_graph/runtime/tiling_context.h" +#include "tiling_base/tiling_base.h" +#include "log/log.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +template +std::unique_ptr TILING_CLASS(gert::TilingContext* context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext*); + +class TilingCases { +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + void AddTiling(int32_t priority) + { + OP_CHECK_IF( + cases_.find(priority) != cases_.end(), OP_LOGE(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OP_CHECK_IF( + cases_[priority] == nullptr, + OP_LOGE(op_type_, "Register op tiling func failed, please check the class name."), return); + } + + const std::map& GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +// --------------------------------Interfacce with soc version -------------------------------- +class TilingRegistryNew { +public: + TilingRegistryNew() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryNew& GetInstance(); +#else + static TilingRegistryNew& GetInstance() + { + static TilingRegistryNew registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + if (soc_iter == registry_map_.end()) { + std::map> op_type_map; + op_type_map[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + registry_map_[soc_version] = op_type_map; + } else { + if (soc_iter->second.find(op_type) == soc_iter->second.end()) { + soc_iter->second[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + } + + OP_CHECK_IF( + registry_map_[soc_version][op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[soc_version][op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t soc_version = (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION; + const char* op_type = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = static_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + if (soc_version == (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION) { + OP_LOGE(op_type, "Do op tiling failed, cannot find soc version."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + int32_t soc_version; + const char* op_type = context->GetNodeType(); + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = reinterpret_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + } + + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto priority_id : priorities) { + auto tilingCaseIter = tilingTemplateRegistryMap.find(priority_id); + if (tilingCaseIter != tilingTemplateRegistryMap.end()) { + auto templateFunc = tilingCaseIter->second(context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priority_id); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priority_id); + } + } + } + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + OP_CHECK_IF( + soc_iter == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the soc version %d", soc_version), + return empty_tiling_case_); + auto op_iter = soc_iter->second.find(op_type); + OP_CHECK_IF( + op_iter == soc_iter->second.end(), OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return op_iter->second->GetTilingCases(); + } + +private: + std::map>> registry_map_; // key is socversion + const std::map empty_tiling_case_{}; +}; + +class RegisterNew { +public: + explicit RegisterNew(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + RegisterNew& tiling(int32_t priority, int32_t soc_version) + { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterNew& tiling(int32_t priority, const std::vector& soc_versions) + { + for (int32_t soc_version : soc_versions) { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string op_type_; +}; + +// --------------------------------Interfacce without soc version -------------------------------- +class TilingRegistry { +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry& GetInstance(); +#else + static TilingRegistry& GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OP_CHECK_IF( + registry_map_[op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do op tiling failed"); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priorityId); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type) + { + OP_CHECK_IF( + registry_map_.find(op_type) == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_; +}; + +class Register { +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + Register& tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops + +// op_type: 算子名称， class_name: 注册的 tiling 类, soc_version：芯片版本号 +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(op_type, class_name, soc_versions, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_versions) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + static Ops::Transformer::OpTiling::Register VAR_UNUSED##op_type_##class_name##priority_register = \ + Ops::Transformer::OpTiling::Register(op_type).tiling(priority) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// soc_version: soc版本，用于区分不同的soc +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_NEW(op_type, class_name, soc_version, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_version) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +// 取代 REGISTER_TILING_TEMPLATE , 传入的op_type如果是字符串常量，需要去掉引号 +#define REGISTER_OPS_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register \ + __attribute__((unused)) tiling_##op_type##_##class_name##_##priority##_register = \ + Ops::Transformer::OpTiling::Register(#op_type).tiling(priority) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_type.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_type.h new file mode 100644 index 00000000..be3e4fcc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_type.h @@ -0,0 +1,139 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_type.h + * \brief + */ + +#pragma once + +#include + +namespace optiling { + +enum class AxisEnum { + B = 0, + N2 = 1, + G = 2, + S1 = 3, + S2 = 4, + D = 5, + NONE = 9, +}; + +enum class DtypeEnum { + FLOAT16 = 0, + FLOAT32 = 1, + BFLOAT16 = 2, + FLOAT16_PRECISION = 3, +}; + +enum class PerformanceOrientedEnum { + BIG_BUFFER = 1, + BIG_DOUBLE_BUFFER = 2, +}; + +enum class MatmulConfig { + NULL_CONFIG = 0, + NORMAL_CONFIG = 1, + MDL_CONFIG = 2 +}; + +enum class PseConfig { + NO_PSE = 0, + EXIST_PSE = 1 +}; + +enum class AttenMaskConfig { + NO_ATTEN_MASK = 0, + EXIST_ATTEN_MASK = 1 +}; + +enum class DropOutConfig { + NO_DROP_OUT = 0, + EXIST_DROP_OUT = 1 +}; + +enum class CubeFormatEnum { + ND = 0, + NZ = 1 +}; +enum class LayoutEnum { + BSND = 0, + SBND = 1, + BNSD = 2, + TND = 3, + NTD_TND = 4 +}; + +enum class CubeInputSourceEnum { + GM = 0, + L1 = 1 +}; + +enum class OptionEnum { + DISABLE = 0, + ENABLE = 1 +}; + +enum class SparseEnum { + ALL = 0, + NONE = 1, + ANY = 2, + CAUSAL = 3, + BAND = 4, + PREFIX = 5, + BAND_COMPRESS = 6, + RIGHT_DOWN_CAUSAL = 7, + RIGHT_DOWN_CAUSAL_BAND = 8, + BAND_LEFT_UP_CAUSAL = 9 +}; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr int64_t base10Multiplier = 10; + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + base10Multiplier * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputed types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputed types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace optiling diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_util.h new file mode 100644 index 00000000..602c34ce --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_base/tiling_util.h @@ -0,0 +1,30 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.h + * \brief + */ + +#pragma once + +#include "register/op_impl_registry.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { +bool IsRegbaseSocVersion(const gert::TilingParseContext* context); + +bool IsRegbaseSocVersion(const gert::TilingContext* context); + +const gert::Shape& EnsureNotScalar(const gert::Shape& inShape); +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/device_op_impl_registry_impl.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/device_op_impl_registry_impl.h new file mode 100644 index 00000000..994d4a81 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/device_op_impl_registry_impl.h @@ -0,0 +1,49 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file device_op_impl_registry_impl.h + * \brief + */ + +#ifndef OP_TILING_DEVICE_OP_IMPL_REGISTRY_IMPL_H +#define OP_TILING_DEVICE_OP_IMPL_REGISTRY_IMPL_H + +#include +#include +#include "register/device_op_impl_registry.h" + +namespace optiling { +class DeviceOpImplRegistry { + public: + static DeviceOpImplRegistry& GetSingleton(); + void RegisterSinkTiling(std::string &opType, SinkTilingFunc& func); + SinkTilingFunc GetSinkTilingFunc(std::string &opType); + + private: + DeviceOpImplRegistry() = default; + ~DeviceOpImplRegistry() = default; + + private: + std::map sinkTilingFuncsMap_; +}; + +class DeviceOpImplRegisterImpl { + public: + DeviceOpImplRegisterImpl() = default; + ~DeviceOpImplRegisterImpl(); + std::string& GetOpType(); + + private: + std::string opType_ = ""; +}; +} // namespace optiling + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/tiling_aicpu_task.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/tiling_aicpu_task.h new file mode 100644 index 00000000..24e01332 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/include/tiling_sink/tiling_aicpu_task.h @@ -0,0 +1,30 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_aicpu_task.h + * \brief + */ + +#ifndef TILING_SINK_TILING_AICPU_TASK_H_ +#define TILING_SINK_TILING_AICPU_TASK_H_ +#include "exe_graph/runtime/tiling_context.h" + +namespace tilingsink { +struct TilingAicpuTask { + gert::TilingContext *tilingContext; + const char *opType; + uint64_t notifyAddr; + uint64_t workspaceAddr; + uint64_t workspaceSize; +}; +} // namespace optiling + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/tiling_util.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/tiling_util.cpp new file mode 100644 index 00000000..d6757874 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/common/tiling_util.cpp @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.cpp + * \brief + */ + +#include "tiling_base/tiling_util.h" +#include "platform/platform_ascendc.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { +static const gert::Shape g_vec_1_shape = {1}; + +static bool IsRegbaseSocVersion(platform_ascendc::SocVersion version) +{ + const static std::set regbaseSocVersions = { + platform_ascendc::SocVersion::ASCEND910_95}; + + return regbaseSocVersions.find(version) != regbaseSocVersions.end(); +} + +bool IsRegbaseSocVersion(const gert::TilingParseContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto socVersion = ascendcPlatform.GetSocVersion(); + return IsRegbaseSocVersion(socVersion); +} + +bool IsRegbaseSocVersion(const gert::TilingContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto socVersion = ascendcPlatform.GetSocVersion(); + return IsRegbaseSocVersion(socVersion); +} + +const gert::Shape &EnsureNotScalar(const gert::Shape &inShape) { + if (inShape.IsScalar()) { + return g_vec_1_shape; + } + return inShape; +} +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/CMakeLists.txt b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/CMakeLists.txt new file mode 100644 index 00000000..4b743003 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/CMakeLists.txt @@ -0,0 +1,59 @@ + +file(GLOB_RECURSE ops_srcs CONFIGURE_DEPENDS + ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) +list(FILTER ops_srcs EXCLUDE REGEX "generated.*disabled$") +list(FILTER ops_srcs EXCLUDE REGEX "/op_api/") +list(FILTER ops_srcs EXCLUDE REGEX "dlinfer_grouped_matmul_direct_def.cpp$") +list(APPEND ops_srcs ${DLINFER_GMM_OPP_SOURCE_DIR}/common/tiling_util.cpp) + +opbuild(OPS_SRC ${ops_srcs} + OUT_DIR ${ASCEND_AUTOGEN_PATH} + ENABLE_SOURCE ${ENABLE_SOURCE_PACAKGE} +) + +if(ENABLE_CROSS_COMPILE) + build_optiling_for_compile( + OPS_SRC ${ops_srcs} + OUT_DIR ${ASCEND_AUTOGEN_PATH} + ) +endif() + +add_library(cust_optiling SHARED ${ops_srcs}) +target_compile_definitions(cust_optiling PRIVATE OP_TILING_LIB) +target_compile_options(cust_optiling PRIVATE + -fvisibility=hidden +) +if(ENABLE_CROSS_COMPILE) + target_link_directories(cust_optiling PRIVATE + ${CMAKE_COMPILE_COMPILER_LIBRARY} + ${CMAKE_COMPILE_RUNTIME_LIBRARY} + ) +endif() +target_link_libraries(cust_optiling PRIVATE + intf_pub + exe_graph + register + -Wl,--whole-archive + tiling_api + rt2_registry + -Wl,--no-whole-archive +) +set_target_properties(cust_optiling PROPERTIES OUTPUT_NAME + cust_opmaster_rt2.0 +) + +file(GLOB aclnn_src ${ASCEND_AUTOGEN_PATH}/aclnn_*.cpp) +file(GLOB aclnn_inc ${ASCEND_AUTOGEN_PATH}/aclnn_*.h) +add_library(cust_opapi SHARED + ${ops_srcs} + ${aclnn_src} +) +target_compile_definitions(cust_opapi PRIVATE ACLNN_WITH_BINARY OP_TILING_LIB) +include_directories(cust_opapi ${DLINFER_GMM_OPP_BINARY_DIR}/op_kernel/kernel/library) +if(ENABLE_CROSS_COMPILE) + target_link_directories(cust_opapi PRIVATE + ${CMAKE_COMPILE_COMPILER_LIBRARY} + ${CMAKE_COMPILE_RUNTIME_LIBRARY} + ) +endif() +target_link_libraries(cust_opapi PRIVATE intf_pub register ascend_kernels nnopbase ascendcl exe_graph tiling_api) diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/OWNERS b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/OWNERS new file mode 100644 index 00000000..81119c9e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/OWNERS @@ -0,0 +1,67 @@ +approvers: +- andylhy +- xiaolong_han +- huxiaobang +- fanqirui +- luanma_bl +- chenbinbin199309 +- chen-kang30 +- chq317 +- chaotang233 +- fang-guocan +- xu-binglin +- yangyang016 +- wuyi_huawei +- zhao-yingchao +- shi-ray +- crystalhu +- yu-xinjie62 +- zcc9707 +- liuzhuheng +- juyangokok + +reviewers: +- li-xulong +- Allan_Yu +- mabing726 +- miao-fangzheng +- chengsheng304064 +- miao-fangzheng +- wang-fei6 +- shawn-hu +- li-shengxian +- xig514 +- chen-vvjob +- renkyk +- yang-binrong +- song-jionghui +- wang-zhe123456789 +- huangli70 +- huangwei791 +- shunqi +- wangjun-qt +- jiang-lirui +- realmadrid1016 +- monologue815 +- zhangtj0209 +- GodantShen +- Liexss +- fzzach +- zhanglei_hw +- songkai-huawei +- yangxu19921206 +- crystalhu +- yu-xinjie62 +- zhu-yijun-Julius + +files: + "aclnn_grouped_matmul.h": + approvers: + - Allan_Yu + - wang-yongguang + - juyangokok + "grouped_matmul_def.cpp": + approvers: + - Allan_Yu + - wang-yongguang + - juyangokok \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_minimal_def.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_minimal_def.cpp new file mode 100644 index 00000000..5d171024 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_minimal_def.cpp @@ -0,0 +1,52 @@ +#include "register/op_def_registry.h" + +namespace ops { +class DlinferGroupedMatmulDirect : public OpDef { +public: + explicit DlinferGroupedMatmulDirect(const char *name) : OpDef(name) + { + const std::initializer_list floatTypes = { + ge::DT_FLOAT16, ge::DT_BF16}; + const std::initializer_list ndFormats = { + ge::FORMAT_ND, ge::FORMAT_ND}; + + this->Input("x").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + this->Input("weight").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + this->Input("bias").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + this->Input("scale").ParamType(DYNAMIC) + .DataType({ge::DT_UINT64, ge::DT_UINT64}).Format(ndFormats); + this->Input("offset").ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}).Format(ndFormats); + this->Input("antiquant_scale").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + this->Input("antiquant_offset").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + this->Input("group_list").ParamType(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64}).Format(ndFormats); + this->Input("per_token_scale").ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}).Format(ndFormats); + this->Output("y").ParamType(DYNAMIC).DataType(floatTypes).Format(ndFormats); + + this->Attr("split_item").AttrType(OPTIONAL).Int(0); + this->Attr("dtype").AttrType(OPTIONAL).Int(0); + this->Attr("transpose_weight").AttrType(OPTIONAL).Bool(false); + this->Attr("transpose_x").AttrType(OPTIONAL).Bool(false); + this->Attr("group_type").AttrType(OPTIONAL).Int(-1); + this->Attr("group_list_type").AttrType(OPTIONAL).Int(0); + this->Attr("act_type").AttrType(OPTIONAL).Int(0); + this->Attr("tuning_config").AttrType(OPTIONAL).ListInt({0}); + + OpAICoreConfig config; + config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("prebuildPattern.value", "Opaque") + .ExtendCfgInfo("coreType.value", "AiCore") + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); + this->AICore().AddConfig("ascend910_93", config); + } +}; + +OP_ADD(DlinferGroupedMatmulDirect); +} // namespace ops diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_tiling.h new file mode 100644 index 00000000..d5e3db24 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/dlinfer_grouped_matmul_direct_tiling.h @@ -0,0 +1,10 @@ + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(DlinferGroupedMatmulDirectTilingData) + TILING_DATA_FIELD_DEF(uint32_t, size); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(DlinferGroupedMatmulDirect, DlinferGroupedMatmulDirectTilingData) +} diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_host_util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_host_util.h new file mode 100644 index 00000000..0387627c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_host_util.h @@ -0,0 +1,183 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_host_util.h + * \brief + */ + +#ifndef GROUPED_MATMUL_HOST_UTIL_H +#define GROUPED_MATMUL_HOST_UTIL_H + +#include + +namespace DlinferGroupedMatmulDirect { +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t BIAS_INDEX = 2; +constexpr uint32_t SCALE_INDEX = 3; +constexpr uint32_t OFFSET_INDEX = 4; +constexpr uint32_t ANTIQUANT_SCALE_INDEX = 5; +constexpr uint32_t GROUPLIST_INDEX = 7; +constexpr uint32_t PER_TOKEN_SCALE_INDEX = 8; +constexpr uint32_t Y_INDEX = 0; +constexpr uint64_t BEST_L1_PARTA = 256UL * 1024UL; +constexpr uint64_t BEST_L1_PARTB = 128UL * 1024UL; +constexpr uint64_t L1_PARTA_SIZE = 256UL * 1024UL; +constexpr int32_t BEST_BASEN = 256; +constexpr int32_t BEST_BASEN_A4W4 = 512; +constexpr int32_t BEST_BASEN_QUANT_ONE_GROUP = 128; +constexpr int32_t BEST_BASEM_QUANT_ONE_GROUP = 256; +constexpr int32_t BEST_BASEK_QUANT_ONE_GROUP = 128; +constexpr int32_t BEST_BASEN_MSD = 512; +constexpr int32_t BEST_UB_BASEK = 256; +constexpr int32_t BEST_UB_BASEN = 512; +constexpr int32_t MAX_BASEM = 256; +constexpr uint32_t MIN_UB_BASEN = 128; +constexpr uint32_t BASIC_BLOCK_SIZE_128 = 128; +constexpr uint32_t BASIC_BLOCK_SIZE_256 = 256; +constexpr uint32_t BASIC_BLOCK_SIZE_512 = 512; +constexpr uint32_t A16W8_MSD_STEP = 2; +constexpr uint32_t A16W8_MSD_KN_BASE_BLOCK = 128; +constexpr uint32_t A16W8_MSD_AVERAGE_TOKEN_NUM = 64; +constexpr uint32_t A16W8_MSD_MAX_K = 12U * 1024U; +constexpr uint32_t A16W8_MSD_MIN_N = 1024; +constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data +constexpr uint32_t UB_ANTIQUANT_PER_BLOCK_ALIGN = 4U * 1024U; +constexpr uint32_t UB_A16W8_BLOCK_NUM_FP16 = 6; // 2 * sizeof(int8) + 2 * sizeof(half) +constexpr uint32_t UB_A16W8_IO_USED_BLOCK_FP16 = 6; +constexpr uint32_t UB_A16W8_BLOCK_NUM_BF16 = 8; // tmpUb used 2 blks +constexpr uint32_t UB_A16W8_IO_USED_BLOCK_BF16 = 6; +constexpr uint32_t UB_A16W4_BLOCK_NUM_FP16 = 5; // 2 * sizeof(int4) + 2 * sizeof(half) +constexpr uint32_t UB_A16W4_IO_USED_BLOCK_FP16 = 5; +constexpr uint32_t UB_A16W4_BLOCK_NUM_BF16 = 7; // tmpUb used 2 blks +constexpr uint32_t UB_A16W4_IO_USED_BLOCK_BF16 = 5; +constexpr uint32_t UB_A4W4_BLOCK_NUM = 16; +constexpr uint32_t UB_A4W4_IO_USED_BLOCK_HALF = 8; // 2 * sizeof(fp16) + 2 * sizeof(fp16/bf16) +constexpr uint32_t UB_A4W4_PER_BLOCK_ALIGN = 8U * 1024U; +constexpr uint32_t UB_DYNAMIC_QUANT_BLOCK_NUM = 28; +constexpr uint32_t UB_DUNAMIC_QUANT_IO_USED_BLOCK = 12; +constexpr uint32_t UB_QUANT_BLOCK_ALIGN = 2U * 1024U; +constexpr uint32_t UB_A16W8_MSD_BLOCK_NUM = 30; +constexpr uint32_t UB_A16W8_MSD_IO_USED_BLOCK = 6; +constexpr uint32_t UB_A16W8_MSD_BLOCK_ALIGN = 512; +constexpr uint32_t UB_STATIC_QUANT_BLOCK_NUM_BF16 = 20; +constexpr uint32_t UB_STATIC_QUANT_BLOCK_NUM_FP16 = 24; +constexpr uint32_t UB_STATIC_QUANT_IO_USED_BLOCK = 12; +constexpr uint32_t QUEUE_DOUBLE_BUFFER = 2; +constexpr uint32_t FP32_DATATYPE_SIZE = 4; +constexpr uint64_t TILING_KEY = 0; +constexpr uint64_t TILING_KEY_TRANS_X = 1; +constexpr uint64_t TILING_KEY_TRANS_W = 2; +constexpr uint64_t TILING_KEY_ANTIQUANT_PERFORMANCE = 3; +constexpr uint64_t TILING_KEY_QUANT_2VECTOR = 4; +constexpr uint64_t TILING_KEY_QUANT_2VECTOR_TRANS_W = 5; +constexpr uint64_t TILING_KEY_A16W8_MSD = 6; +constexpr uint64_t TILING_KEY_A16W8_MSD_TRANS_W = 7; +constexpr uint64_t TILING_KEY_A8W4_MSD = 8; +constexpr uint64_t TILING_KEY_A8W4_MSD_NEW = 12; +constexpr uint64_t TILING_KEY_A8W4 = 18; // per group +constexpr uint64_t TILING_KEY_A8W4_FAKE_A8W8 = 17; //per channel +constexpr uint64_t TILING_KEY_A8W4_AUTOTILING_A8W4 = 21; +constexpr uint64_t TILING_KEY_A8W8_SPARSE_M = 9; +constexpr uint64_t TILING_KEY_A8W8_SPARSE_M_TRANS_W = 10; +constexpr uint64_t TILING_KEY_STATIC_TILING_OFFSET = 13; +constexpr uint64_t ATTR_INDEX_SPLIT_ITEM = 0; +constexpr uint64_t ATTR_INDEX_TRANS_W = 2; +constexpr uint64_t ATTR_INDEX_TRANS_X = 3; +constexpr uint64_t ATTR_INDEX_GROUPTYPE = 4; +constexpr uint32_t ATTR_INDEX_GROUP_LIST_TYPE = 5; +constexpr uint64_t ATTR_INDEX_ACT_TYPE = 6; +constexpr uint64_t ATTR_INDEX_TUNING_CONFIG = 7; +constexpr uint64_t DOUBLE_BUFFER_L0A_L0B = 2; +constexpr uint64_t DOUBLE_BUFFER_STEPKA_STEPKB = 2; +constexpr uint32_t SYS_WORKSPACE_SIZE = 16U * 1024U * 1024U; +constexpr uint32_t GROUPLIST_TYPE_SPARSE_M = 2; +constexpr int32_t NO_SPLIT = -1; +constexpr int32_t SPLIT_M = 0; +constexpr int32_t SPLIT_K = 2; +// used for whether going into performance branch in antiquant case, by experiment +constexpr int64_t ANTIQUANT_PERFORMANCE_THRESHOLD = 5L * 1024L * 1024L; +constexpr int64_t ACT_TYPE_GELU = 2; +constexpr uint16_t MAX_TENSOR_CONT = 128; +constexpr int64_t FULL_K_SINGLE_N = 1280; // used for fullload k case, by experiment +constexpr int64_t FULL_K_N_THRESHOLD = 2560; // used for fullload k case, by experiment +constexpr int64_t FULL_K_M_THRESHOLD = 2048; // used for fullload k case, by experiment +constexpr int64_t FULL_K_M_E_THRESHOLD = 256; // used for fullload k case, by experiment +constexpr int64_t FULL_K_MAX_K_THRESHOLD = 384; // used for fullload k case, by experiment +constexpr int64_t FULL_K_MIN_K_THRESHOLD = 320; // used for fullload k case, by experiment +constexpr uint32_t MIN_NZ_DIM = 4; +constexpr uint32_t MIN_ND_DIM = 2; +constexpr uint32_t A3_AIC_NUM = 24; +constexpr uint32_t OFFSET_DIM_A8W4 = 3; +constexpr float INT4_DATA_TYPE_SIZE = 0.5; +// used for static tiling template +constexpr int32_t STATIC_TILING_DEPTH_A1_B1 = 8; +constexpr int32_t STATIC_TILING_STEP_KA_KB = 4; +constexpr int32_t STATIC_TILING_MAX_K = 8192; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +template +constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + 2U * RecursiveSum(templateIds...); +} + +const std::map, std::array> A8W8_PRETILING_WHITE_LIST = { // used for A8W8 preTiling, by experiment + {{576, 7168, 4096, 0}, {128, 512}}, + {{576, 2048, 7168, 1}, {96, 1792}} +}; + +const std::map, int64_t> A8W4_PRETILING_WHITE_LIST = { // used for A8W4 preTiling, by experiment + {{1, 16, 256, 512, 1}, 1}, + {{256, 1024, 512, 32768, 1}, 1} +}; + +template +auto CeilDiv(T1 a, T2 b) -> T1 +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +auto CeilDiv(T a, T b) -> T +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +template +auto CeilAlign(T a, T b) -> T +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} + +/** + * if align is 0, return 0 + */ +template +auto FloorAlign(T x, T align) -> typename std::enable_if::value, T>::type { + return align == 0 ? 0 : x / align * align; +} +} // namespace DlinferGroupedMatmulDirect + +#endif // GROUPED_MATMUL_HOST_UTIL_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape.cpp new file mode 100644 index 00000000..6d9226c5 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape.cpp @@ -0,0 +1,1754 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + +/*! + * \file grouped_matmul_infershape.cpp + * \brief + */ +#include + +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "platform/platform_info.h" +#include "grouped_matmul_infershape_weight_quant_checker.h" +#include "grouped_matmul_infershape_quant_checker.h" +#include "grouped_matmul_infershape_common_util.h" + +using namespace ge; +namespace ops { + +static std::set GmmDavidSupportSoc = {"Ascend910_95"}; + +enum class PlatformID : std::uint8_t { + UNKNOWN, + ASCEND310P, + ASCEND910B, + ASCEND910_95 +}; + +struct GMMParamsInfo { + size_t numX; + size_t numWeight; + size_t numY; + int64_t lenGroupList; + size_t groupNum; + size_t numScale; + size_t numOffset; + size_t numAntiquantScale; + size_t numAntiquantOffset; + PlatformID platform; +}; + +struct GMMSetOutputParams { + bool isSingleX; + bool isSingleY; + size_t xDimM; + size_t weightDimN; + int64_t lenGroupList; + size_t numWeight; + size_t numX; +}; + +static inline std::string ToString(const std::int64_t value) { + return std::to_string(value); +} + +static ge::graphStatus CheckSplitItem(int64_t splitItem) { + if (splitItem == GMM_X_Y_SEPARATED || splitItem == GMM_NO_SEPARATED || + splitItem == GMM_X_SEPARATED || splitItem == GMM_Y_SEPARATED) { + return GRAPH_SUCCESS; + } else { + return GRAPH_FAILED; + } +} + +static bool IsTensorListNullOrEmpty(const gert::InferShapeContext* context, size_t index) { + auto shape = context->GetDynamicInputShape(index, 0); + if (shape == nullptr) { + return true; + } + if (shape->GetDimNum() == 0 || (shape->GetDimNum() == 1 && shape->GetDim(0) == 0)) { + if (context->GetDynamicInputShape(index, 1) == nullptr) { + return true; + } + } + return false; +} + +static ge::graphStatus CheckGroupType(const gert::InferShapeContext* context, int64_t groupType) { + if (groupType == GMM_NO_SPLIT || groupType == GMM_SPLIT_M || groupType == GMM_SPLIT_K) { + return GRAPH_SUCCESS; + } else if (groupType == GMM_SPLIT_N) { + OP_LOGE(context->GetNodeName(), "Splitting tensor along the N-axis is not supported yet."); + return GRAPH_FAILED; + } else { + OP_LOGE(context->GetNodeName(), "GroupType can only be -1/0/2 now, but actually %ld is given.", groupType); + return GRAPH_FAILED; + } +} + +static ge::graphStatus UpdateShapeYMultiDim(gert::InferShapeContext* context, size_t idxY, const gert::Shape* xShape, + const gert::Shape* weightShape) { + gert::Shape* yShape = context->GetOutputShape(idxY); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + *yShape = *xShape; + size_t dimY = yShape->GetDimNum(); + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const bool* transposeWPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_W); + const bool* transposeXPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_X); + + OP_CHECK_NULL_WITH_CONTEXT(context, weightShape); + if (transposeWPtr != nullptr && *transposeWPtr) { + yShape->SetDim(dimY - 1, weightShape->GetDim(weightShape->GetDimNum() - 2)); // -2: transpose weight + } else { + yShape->SetDim(dimY - 1, weightShape->GetDim(weightShape->GetDimNum() - 1)); + } + if (transposeXPtr != nullptr && *transposeXPtr) { + yShape->SetDim(dimY - 2, xShape->GetDim(xShape->GetDimNum() - 1)); // -2: last two dim of Y + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus UpdateShapeY(gert::InferShapeContext* context, size_t idxY, std::vector yDims) { + gert::Shape* yShape = context->GetOutputShape(idxY); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + yShape->SetDimNum(yDims.size()); + for (size_t dim = 0; dim < yDims.size(); ++dim) { + yShape->SetDim(dim, yDims[dim]); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus UpdateMultipleShapeY(gert::InferShapeContext* context, const gert::Tensor* groupListTensor, + size_t weightDimN, bool isXTransposed, size_t xDimM) { + auto groupListData = groupListTensor->GetData(); + OP_CHECK_IF(groupListData == nullptr, + OP_LOGE(context->GetNodeName(), "Failed to obtain necessary data from groupListTensor."), + return GRAPH_FAILED); + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* groupListTypePtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_LIST_TYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListTypePtr); + const gert::Shape* x0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, x0Shape); + const gert::Shape* weight0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weight0Shape); + int64_t preOffset = 0; + for (int idx = 0; idx < groupListTensor->GetShapeSize(); ++idx) { + const gert::Shape* weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idx); + if (weightShape == nullptr) { + weightShape = weight0Shape; + } + if (isXTransposed) { + const gert::Shape* xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, idx); + if (xShape == nullptr) { + xShape = x0Shape; + } + std::vector yDims = {xShape->GetDim(xDimM), weightShape->GetDim(weightDimN)}; + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y + idx, yDims) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Failed to update shape of y."), return GRAPH_FAILED); + } else { + std::vector yDims; + if (*groupListTypePtr == 0) { + yDims = {groupListData[idx] - preOffset, weightShape->GetDim(weightDimN)}; + } else if (*groupListTypePtr == 1) { + yDims = {groupListData[idx], weightShape->GetDim(weightDimN)}; + } else { + OP_LOGE(context->GetNodeName(), "Invalid groupListType = %ld", *groupListTypePtr); + return GRAPH_FAILED; + } + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y + idx, yDims) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Failed to update shape of y."), return GRAPH_FAILED); + preOffset = groupListData[idx]; + } + } + + return GRAPH_SUCCESS; +} + +static ge::graphStatus MultiInMultiOutWithoutGroupList(gert::InferShapeContext* context) { + size_t idx = 0; + size_t idw = 0; + const gert::Shape* w0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Shape); + while (true) { + const gert::Shape* xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, idx); + if (xShape == nullptr) { + break; + } + ++idx; + const gert::Shape* wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idw); + if (wShape) { + ++idw; + } else { + wShape = w0Shape; + } + OP_CHECK_IF(UpdateShapeYMultiDim(context, GMM_INDEX_OUT_Y + idx - 1, xShape, wShape) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* groupTypePtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_TYPE); + bool success = true; + if (w0Shape->GetDimNum() == 2) { // 2 two-dim weight tensor + if (groupTypePtr != nullptr && *groupTypePtr == 2) { + success = true; + } else { + success = idx == idw; + } + } else { + success = static_cast(idx) == w0Shape->GetDim(0); + } + OP_CHECK_IF(!success, + OP_LOGE(context->GetNodeName(), + "x tensorList's length[%zu] != weight tensor's first dim[%ld] and length[%zu]", + idx, w0Shape->GetDim(0), idw), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus MultiWeightMultiOutWithoutGroupList(gert::InferShapeContext* context) { + size_t idx = 0; + const gert::Shape* x0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, x0Shape); + while (true) { + const gert::Shape* wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idx); + if (!wShape) { + break; + } + ++idx; + OP_CHECK_IF(UpdateShapeYMultiDim(context, GMM_INDEX_OUT_Y + idx - 1, x0Shape, wShape) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } + + return GRAPH_SUCCESS; +} + +template +static ge::graphStatus GetAttrsValue(T context, GMMAttrs &gmmAttrs) +{ + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + + const int64_t *splitItemPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_SPLIT_ITEM); + OP_CHECK_NULL_WITH_CONTEXT(context, splitItemPtr); + gmmAttrs.splitItem = *splitItemPtr; + OP_LOGI(context->GetNodeName(), "Attr splitItem = %ld", gmmAttrs.splitItem); + + const int64_t *dtypePtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_OUTPUT_DTYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, dtypePtr); + gmmAttrs.outputDtype = *dtypePtr; + OP_LOGI(context->GetNodeName(), "Attr dtype = %ld", gmmAttrs.outputDtype); + + const auto tuningConfigPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TUNING_CONFIG); + gmmAttrs.tuningConfig = (tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > 0) ? + (reinterpret_cast(tuningConfigPtr->GetData()))[0] : 0; + OP_LOGI(context->GetNodeName(), "Attr tuningConfig = %ld", gmmAttrs.tuningConfig); + + const int64_t *groupTypePtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_TYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, groupTypePtr); + gmmAttrs.groupType = *groupTypePtr; + OP_LOGI(context->GetNodeName(), "Attr groupType = %ld", gmmAttrs.groupType); + + const bool *transposeWPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_W); + OP_CHECK_NULL_WITH_CONTEXT(context, transposeWPtr); + gmmAttrs.transposeWeight = *transposeWPtr; + OP_LOGI(context->GetNodeName(), "Attr isWeightTransposed = %d", gmmAttrs.transposeWeight); + + const bool *transposeXPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_X); + OP_CHECK_NULL_WITH_CONTEXT(context, transposeXPtr); + gmmAttrs.transposeX = *transposeXPtr; + OP_LOGI(context->GetNodeName(), "Attr isXTransposed = %d", gmmAttrs.transposeX); + + const int64_t *activeType = attrs->GetInt(GMM_INDEX_ATTR_ACT_TYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, activeType); + gmmAttrs.activeType = *activeType; + OP_LOGI(context->GetNodeName(), "Attr activeType = %ld", gmmAttrs.activeType); + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckAttrs(gert::InferShapeContext* context, GMMAttrs& gmmAttrs) { + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + OP_CHECK_IF(CheckSplitItem(gmmAttrs.splitItem) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Invalid splitItem, which can only be one of 0/1/2/3."), return GRAPH_FAILED); + OP_CHECK_IF(CheckGroupType(context, gmmAttrs.groupType) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid groupType."), return GRAPH_FAILED); + const int64_t* activeType = attrs->GetInt(GMM_INDEX_ATTR_ACT_TYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, activeType); + OP_CHECK_IF(*activeType < 0 || *activeType >= static_cast(GMMActType::END_ACT_TYPE_ENUM), + OP_LOGE(context->GetNodeName(), "activeType must be no less than 0 and smaller than 6"), + return GRAPH_FAILED); + OP_CHECK_IF(*activeType == static_cast(GMMActType::GMM_ACT_TYPE_GELU_ERR_FUNC), + OP_LOGE(context->GetNodeName(), "Activation function not support GELU_ERR_FUNC now."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus GetNumOfInputs(const gert::InferShapeContext* context, size_t& numX, + size_t& numWeight, int64_t& lenGroupList) { + ge::graphStatus res = GRAPH_SUCCESS; + const gert::Shape* shape = nullptr; + while (true) { + shape = context->GetDynamicInputShape(GMM_INDEX_IN_X, numX); + if (shape == nullptr) { // last shape + break; + } + for (size_t i = 0; i < shape->GetDimNum(); ++i) { + if (shape->GetDim(i) < 0) { // shape dim cannot be smaller than 0 + res = GRAPH_FAILED; + break; + } + } + ++numX; + } + OP_LOGI(context->GetNodeName(), "numX = %lu", numX); + + while (true) { + shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, numWeight); + if (shape == nullptr) { // last shape + break; + } + for (size_t i = 0; i < shape->GetDimNum(); ++i) { + if (shape->GetDim(i) < 0) { // shape dim cannot be smaller than 0 + res = GRAPH_FAILED; + break; + } + } + ++numWeight; + } + OP_LOGI(context->GetNodeName(), "numWeight = %lu", numWeight); + + const gert::Tensor* groupListTensor = context->GetOptionalInputTensor(GMM_INDEX_IN_GROUP_LIST); + if (groupListTensor != nullptr) { + lenGroupList = groupListTensor->GetStorageShape().GetDim(0); // groupListType 2 shape is [e, 2] + if (lenGroupList < 0) { // lenGroupList cannot be smaller than 0 + res = GRAPH_FAILED; + } + } + OP_LOGI(context->GetNodeName(), "lenGroupList = %ld", lenGroupList); + + return res; +} + +static int64_t GetDim0(const gert::InferShapeContext* context, bool isXTransposed, size_t numX, size_t xDimM) { + int64_t dim0 = 0; + if (isXTransposed) { + const gert::Shape* x0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + dim0 = (x0Shape == nullptr ? 0 : x0Shape->GetDim(xDimM)); + } else { + for (size_t idx = 0; idx < numX; ++idx) { + const gert::Shape* xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, idx); + int64_t tmpDim0 = (xShape == nullptr ? 0 : xShape->GetDim(0)); + if(tmpDim0 >= 0) { + dim0 += tmpDim0; + } else { + return tmpDim0; + } + } + } + + return dim0; +} + +static bool inline IsNonEmpty(const gert::Shape* shape) { + return (shape != nullptr && !(shape->GetDimNum() == 1 && shape->GetDim(0) == 0)); +} + +static ge::graphStatus IsGmmAntiQuantEmpty(gert::InferShapeContext* context) { + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_ANTIQUANT_SCALE), + OP_LOGE(context->GetNodeName(), "antiquantScale is not null or empty!"), + return GRAPH_FAILED); + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_ANTIQUANT_OFFSET), + OP_LOGE(context->GetNodeName(), "antiquantOffset is not null or empty!"), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus IsGmmQuantEmpty(gert::InferShapeContext* context) { + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_SCALE), + OP_LOGE(context->GetNodeName(), "scale is not null or empty!"), + return GRAPH_FAILED); + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_OFFSET), + OP_LOGE(context->GetNodeName(), "offset is not null or empty!"), + return GRAPH_FAILED); + const gert::Shape* pertokenQuantScale0Shape = context->GetOptionalInputShape(GMM_INDEX_IN_PERTOKEN_SCALE); + OP_CHECK_IF(IsNonEmpty(pertokenQuantScale0Shape), + OP_LOGE(context->GetNodeName(), "pertokenQuant scale is not null or empty!"), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckNonQuant(gert::InferShapeContext* context) { + OP_CHECK_IF(IsGmmQuantEmpty(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Detected nonquant, but quant inputs is not empty!"), + return GRAPH_FAILED); + OP_CHECK_IF(IsGmmAntiQuantEmpty(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Detected nonquant, but antiquant inputs is not empty!"), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus GetGroupSize(const gert::InferShapeContext* context, GMMParamsInfo& paramsInfo) { + size_t groupNum = 1; + size_t maxGroupNum = GMM_MAX_GROUP_LIST_SIZE_ARRAY; // init max value + if (paramsInfo.numX > 1UL) { + groupNum = paramsInfo.numX; + } else if (paramsInfo.numWeight > 1UL) { + groupNum = paramsInfo.numWeight; + } else if (paramsInfo.numY > 1UL) { + groupNum = paramsInfo.numY; + } else if (paramsInfo.lenGroupList > 0) { + groupNum = static_cast(paramsInfo.lenGroupList); + maxGroupNum = static_cast(GMM_MAX_GROUP_LIST_SIZE_TENSOR); // only this case allows GMM_MAX_GROUP_LIST_SIZE_TENSOR size + } + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* groupListType = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_LIST_TYPE); + if (groupListType != nullptr && *groupListType == 2L && paramsInfo.numWeight == 1UL) { + const gert::Shape* weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weightShape); + groupNum = static_cast(weightShape->GetDim(0)); + } + OP_CHECK_IF(groupNum > maxGroupNum, + OP_LOGE(context->GetNodeName(), "groupNum[%zu] is larger than %zu.", + groupNum, maxGroupNum), + return GRAPH_FAILED); + paramsInfo.groupNum = groupNum; + return GRAPH_SUCCESS; +} + +static graphStatus CheckDimNumAndPerGroupNum(const gert::InferShapeContext* context, bool isAntiquantInt4, + const std::tuple& dimData, const gert::Shape* tensorShape, const std::string& tensorType) { + size_t tensorDimNum = std::get<0>(dimData); + size_t expectedDimNum = std::get<1>(dimData); // 1: the sceond element + int64_t weightKDimValue = std::get<2>(dimData); // 2: the third element + if (isAntiquantInt4) { + if (tensorDimNum == expectedDimNum) { + int64_t perGroupNum = tensorShape->GetDim(tensorDimNum - 2); // 2: the last 2-th index + OP_CHECK_IF(!(perGroupNum > 0 && weightKDimValue % perGroupNum == 0), + OP_LOGE(context->GetNodeName(), "perGroupNum must be larger than 0, and can evenly divided " + "by K[%ld] in A16W4-pergroup case, but now perGroupNum is %ld.", weightKDimValue, perGroupNum), + return GRAPH_FAILED); + } else { + OP_CHECK_IF(tensorDimNum != expectedDimNum - 1, + OP_LOGE(context->GetNodeName(), "%s Dim must be %zu for in perchannel case or " + "%zu for pergroup case in A16W4, but now is %zu.", + tensorType.c_str(), expectedDimNum - 1, expectedDimNum, tensorDimNum), + return GRAPH_FAILED); + } + } else { + OP_CHECK_IF(tensorDimNum != expectedDimNum - 1, + OP_LOGE(context->GetNodeName(), "%s Dim must be %zu, but now is %zu.", + tensorType.c_str(), expectedDimNum - 1, tensorDimNum), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckOptionalTensorList(gert::InferShapeContext* context, const std::string tensorType, + const GMMParamsInfo& paramsInfo, const GMMAttrs& gmmAttrs, size_t nodeIdx) { + // check bias，scale, antiquant scale or antiquant offset's size，tensor dimension and shape. + const size_t& groupNum = paramsInfo.groupNum; + size_t tensorSize = 0; + while (context->GetDynamicInputShape(nodeIdx, tensorSize) != nullptr) { + ++tensorSize; + } + uint64_t weightGroupedSize = static_cast(paramsInfo.numWeight); + const int64_t& groupType = gmmAttrs.groupType; + auto shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, shape); + uint64_t weightNDimIdx = shape->GetDimNum() - (gmmAttrs.transposeWeight ? 2 : 1); + auto tensor0Shape = context->GetDynamicInputShape(nodeIdx, 0); + // tensorList size should equals with weight's size + OP_CHECK_IF(tensorSize != weightGroupedSize, OP_LOGE(context->GetNodeName(), + "%s size[%lu] must be equal with weight size[%lu].", tensorType.c_str(), tensorSize, weightGroupedSize), return GRAPH_FAILED); + bool isSingleWeight = (weightGroupedSize == 1 && groupType != GMM_NO_SPLIT); + auto w0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + bool isAntiquantInt4 = (w0Desc->GetDataType() == DT_INT4 && tensorType.find("antiquant") != std::string::npos); + if (isSingleWeight) { // In this case, nodeIdx should have only single tensor, its dim should be 2. + OP_CHECK_IF(IsTensorListNullOrEmpty(context, nodeIdx), OP_LOGE(context->GetNodeName(), + "%s must not be nullptr or empty, but now is nullptr or empty.", tensorType.c_str()), return GRAPH_FAILED); + size_t tensorDimNum = tensor0Shape->GetDimNum(); + int64_t k = shape->GetDim(shape->GetDimNum() - (gmmAttrs.transposeWeight ? 1 : 2)); // 2: axis index + // 3: shape is (E,G,N),G is the perGroupNum + OP_CHECK_IF(CheckDimNumAndPerGroupNum(context, isAntiquantInt4, {tensorDimNum, 3, k}, tensor0Shape, tensorType) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckDimNumAndPerGroupNum failed."), return GRAPH_FAILED); + OP_CHECK_IF(static_cast(tensor0Shape->GetDim(0)) != groupNum, OP_LOGE(context->GetNodeName(), "%s batch size[%ld] should be " + "euqal with groupList length[%lu].", tensorType.c_str(), tensor0Shape->GetDim(0), groupNum), return GRAPH_FAILED); + // tensor's N axis size should equal with weight's N axis. + int64_t weightNDimValue = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0)->GetDim(weightNDimIdx); + int64_t tensorNDimValue = tensor0Shape->GetDim(tensorDimNum - 1); + OP_CHECK_IF(tensorNDimValue != weightNDimValue, OP_LOGE(context->GetNodeName(), + "NDim[%ld] of %s should be equal with NDim[%ld] of weight.", tensorNDimValue, tensorType.c_str(), weightNDimValue), + return GRAPH_FAILED); + } else { + for (uint64_t i = 0; i < groupNum; i++) { + auto tensorShape = context->GetDynamicInputShape(nodeIdx, i); + OP_CHECK_IF(tensorShape == nullptr, OP_LOGE(context->GetNodeName(), + "%s[%lu] must not be nullptr, but now is nullptr.", tensorType.c_str(), i), return GRAPH_FAILED); + // check each of tensor's dim to be 1 + size_t tensorDimNum = tensorShape->GetDimNum(); + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + int64_t k = wShape->GetDim(wShape->GetDimNum() - (gmmAttrs.transposeWeight ? 1 : 2)); // 2: axis index + // 2: shape is (G,N), G is the perGroupNum + OP_CHECK_IF(CheckDimNumAndPerGroupNum(context, isAntiquantInt4, {tensorDimNum, 2, k}, tensorShape, tensorType) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckDimNumAndPerGroupNum failed."), return GRAPH_FAILED); + int64_t weightNDimValue = wShape->GetDim(weightNDimIdx); + int64_t tensorNDimValue = tensorShape->GetDim(tensorDimNum - 1); + OP_CHECK_IF(tensorNDimValue != weightNDimValue, OP_LOGE(context->GetNodeName(), "NDim[%ld] of %s[%lu] should be equal with " + "NDim[%ld] of weight[%lu].", tensorNDimValue, tensorType.c_str(), i, weightNDimValue, i), return GRAPH_FAILED); + } + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckPerTokenScale(const gert::InferShapeContext* context, const GMMParamsInfo& paramsInfo) { + // check pertoken scale's size, tensor dimension and shape + const size_t& xGroupedSize = paramsInfo.numX; + const size_t& weightGroupedSize = paramsInfo.numWeight; + const size_t& yGroupedSize = paramsInfo.numY; + uint64_t xMDimIdx = 0; + // check pertoken scale's size to be equal with x's + if ((xGroupedSize == 1UL) && (yGroupedSize == 1UL)) { + auto perTokenScale0Shape = context->GetOptionalInputShape(GMM_INDEX_IN_PERTOKEN_SCALE); + OP_CHECK_IF(perTokenScale0Shape == nullptr, + OP_LOGE(context->GetNodeName(), "perTokenScaleOptional must not be nullptr, but now is nullptr."), + return GRAPH_FAILED); + // tensor dimension of pertoken_scale should be 1. + size_t tensorDimNum = perTokenScale0Shape->GetDimNum(); + OP_CHECK_IF(tensorDimNum != 1, + OP_LOGE(context->GetNodeName(), + "perTokenScaleOptional dim num must be 1 when x is single tensor, but now is %zu.", tensorDimNum), + return GRAPH_FAILED); + // check pertoken_scale's tensor shape size to be equal with M axis size of x. + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + int64_t xMDimValue = xShape->GetDim(xMDimIdx); + int64_t tensorMDimValue = perTokenScale0Shape->GetDim(tensorDimNum - 1); + OP_CHECK_IF(tensorMDimValue != xMDimValue, + OP_LOGE(context->GetNodeName(), + "MDim[%ld] of perTokenScaleOptional should be equal with MDim[%ld] of x.", + tensorMDimValue, xMDimValue), + return GRAPH_FAILED); + } else { + OP_LOGE(context->GetNodeName(), "per-token quant case is only supported " + "when x, weight and y are all single tensor, but now x size is %zu, weight size is %zu, y size is %zu", + xGroupedSize, weightGroupedSize, yGroupedSize); + return GRAPH_FAILED; + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckDlinferGroupedMatmulDirectQuant(gert::InferShapeContext* context, const GMMAttrs& gmmAttrs, + const GMMParamsInfo& paramsInfo) { + OP_CHECK_IF(paramsInfo.platform == PlatformID::ASCEND310P, + OP_LOGE(context->GetNodeName(), "quant cases do not support on Ascend310P."), + return GRAPH_FAILED); + OP_CHECK_IF(gmmAttrs.groupType == GMM_SPLIT_K, + OP_LOGE(context->GetNodeName(), "quant cases do not support splited axis is K."), + return GRAPH_FAILED); + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_OFFSET), + OP_LOGE(context->GetNodeName(), "offset must be nullptr in quant, but now is not nullptr."), + return GRAPH_FAILED); + if (gmmAttrs.outputDtype != GMM_OUT_DTYPE_INT32) { // output dtype is int32, this scene does not need scale + OP_CHECK_IF(IsTensorListNullOrEmpty(context, GMM_INDEX_IN_SCALE), + OP_LOGE(context->GetNodeName(), "scale must not be nullptr in quant, but now is nullptr."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckOptionalTensorList(context, "scale", paramsInfo, gmmAttrs, GMM_INDEX_IN_SCALE) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid scale."), + return GRAPH_FAILED); + } + bool isPerTokenQuant = context->GetOptionalInputShape(GMM_INDEX_IN_PERTOKEN_SCALE) != nullptr; + if (isPerTokenQuant) { + OP_CHECK_IF(CheckPerTokenScale(context, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Check perTokenScale failed!"), + return GRAPH_FAILED); + } + OP_CHECK_IF(IsGmmAntiQuantEmpty(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Detected quant, but antiquant inputs is not empty!"), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} +static bool isA8W4AsymmetricQuant(const gert::InferShapeContext* context) { + auto offsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_OFFSET, 0); + if (offsetShape == nullptr) { + return false; + } + size_t offsetDimNum = offsetShape->GetDimNum(); + if (offsetDimNum != GMM_A8W4_OFFSET_DIM_NUM) { + return false; + } + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + if (offsetShape->GetDim(0) == weightShape->GetDim(0) && offsetShape->GetDim(1) == 1 + && offsetShape->GetDim(GMM_A8W4_OFFSET_DIM_NUM - 1) == weightShape->GetDim(GMM_A8W4_OFFSET_DIM_NUM - 1)) { + return true; + } + return false; +} +static ge::graphStatus CheckA8W4AsymQuantParams(gert::InferShapeContext* context, const GMMParamsInfo& paramsInfo) { + OP_CHECK_IF(paramsInfo.platform == PlatformID::ASCEND310P, + OP_LOGE(context->GetNodeName(), "quant cases do not support on Ascend310P."), + return GRAPH_FAILED); + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weightShape); + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, biasShape); + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, scaleShape); + size_t biasDimNum = biasShape->GetDimNum(); + size_t scaleDimNum = scaleShape->GetDimNum(); + int64_t e = weightShape->GetDim(0); + int64_t n = weightShape->GetDim(GMM_A8W4_OFFSET_DIM_NUM - 1); + OP_CHECK_IF(IsGmmAntiQuantEmpty(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "antiquant inputs is not empty!"), + return GRAPH_FAILED); + OP_CHECK_IF(biasDimNum != GMM_A8W4_BIAS_DIM_NUM || biasShape->GetDim(0) != e || biasShape->GetDim(1) != n, + OP_LOGE(context->GetNodeName(), "bias shape is invalid, must be (e,n)."), + return GRAPH_FAILED); + auto isScaleInvalid = !(scaleDimNum == GMM_A8W4_OFFSET_DIM_NUM && scaleShape->GetDim(0) == e + && scaleShape->GetDim(1) == 1 && scaleShape->GetDim(GMM_A8W4_OFFSET_DIM_NUM - 1) == n); + OP_CHECK_IF(isScaleInvalid, OP_LOGE(context->GetNodeName(), "scale shape is invalid, must be (e,1,n)."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static int64_t GetPergroupSize(const GMMAttrs& gmmAttrs, bool isSingleWeight, + const gert::Shape* wShape, const gert::Shape* shape) { + int64_t pergroupSize = 0; + size_t shapeDimNum = shape->GetDimNum(); + if (isSingleWeight) { // antiquant param shape (E, N), (E, G, N) + if (shapeDimNum > GMM_SEPARATED_WEIGHT_DIM) { + int64_t k = gmmAttrs.transposeWeight ? wShape->GetDim(2) : wShape->GetDim(1); // 2: the k axis index + pergroupSize = k / shape->GetDim(shapeDimNum - 2); // 2: the last 2-th index + } + } else { // antiquant param shape (N), (G, N) + if (shapeDimNum > 1UL) { + int64_t k = gmmAttrs.transposeWeight ? wShape->GetDim(1): wShape->GetDim(0); + pergroupSize = k / shape->GetDim(shapeDimNum - 2); // 2: the last 2-th index + } + } + return pergroupSize; +} + +static ge::graphStatus CheckDlinferGroupedMatmulDirectAntiQuantForShape(gert::InferShapeContext* context, const GMMAttrs& gmmAttrs, const GMMParamsInfo& paramsInfo) { + OP_CHECK_IF(paramsInfo.platform == PlatformID::ASCEND310P, OP_LOGE(context->GetNodeName(), + "antiquant cases do not support on Ascend310P."), return GRAPH_FAILED); + OP_CHECK_IF(gmmAttrs.groupType == GMM_SPLIT_K, OP_LOGE(context->GetNodeName(), "antiquant cases do not support splited axis is K."), + return GRAPH_FAILED); + OP_CHECK_IF(IsTensorListNullOrEmpty(context, GMM_INDEX_IN_ANTIQUANT_SCALE), + OP_LOGE(context->GetNodeName(), "antiquantScale must not be nullptr in antiquant, but now is nullptr or empty."), + return GRAPH_FAILED); + OP_CHECK_IF(IsTensorListNullOrEmpty(context, GMM_INDEX_IN_ANTIQUANT_OFFSET), + OP_LOGE(context->GetNodeName(), "antiquantOffset must not be nullptr in antiquant, but now is nullptr or empty."), + return GRAPH_FAILED); + // check antiquantScale and antiquantOffset's tensor shape + OP_CHECK_IF(CheckOptionalTensorList(context, "antiquantScale", paramsInfo, gmmAttrs, GMM_INDEX_IN_ANTIQUANT_SCALE) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid antiquantScale"), + return GRAPH_FAILED); + OP_CHECK_IF(CheckOptionalTensorList(context, "antiquantOffset", paramsInfo, gmmAttrs, GMM_INDEX_IN_ANTIQUANT_OFFSET) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid antiquantOffset"), + return GRAPH_FAILED); + // check perGroupSize + auto w0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + if (w0Desc->GetDataType() == DT_INT4) { + auto antiquantScale0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + auto dimNum = antiquantScale0Shape->GetDimNum(); + bool isSingleWeight = ((paramsInfo.numWeight == 1UL) && (gmmAttrs.groupType != GMM_NO_SPLIT)); + int64_t pergroupSize = GetPergroupSize(gmmAttrs, isSingleWeight, context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0), antiquantScale0Shape); + OP_CHECK_IF(gmmAttrs.transposeWeight && pergroupSize % 2 != 0, // 2: a factor + OP_LOGE(context->GetNodeName(), "pergroupSize should be even when weight is transposed" + "in A16W4-pergroup case, but now is %ld", pergroupSize), return GRAPH_FAILED); + for (size_t i = 0; ; ++i) { + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, i); + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, i); + if (antiquantScaleShape == nullptr || antiquantOffsetShape == nullptr) { + break; + } + size_t antiquantScaleDimNum = antiquantScaleShape->GetDimNum(); + size_t antiquantOffsetDimNum = antiquantOffsetShape->GetDimNum(); + OP_CHECK_IF(antiquantScaleDimNum != dimNum || antiquantOffsetDimNum != dimNum, + OP_LOGE(context->GetNodeName(), "antiquantScale[%zu] dim num[%zu] or antiquantOffset[%zu] dim num[%zu] is not equal with %zu", + i, antiquantScaleDimNum, i, antiquantOffsetDimNum, dimNum), return GRAPH_FAILED); + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + int64_t pergroupSizeOfScale = GetPergroupSize(gmmAttrs, isSingleWeight, wShape, antiquantScaleShape); + int64_t pergroupSizeOfOffset = GetPergroupSize(gmmAttrs, isSingleWeight, wShape, antiquantOffsetShape); + OP_CHECK_IF(pergroupSizeOfScale != pergroupSize || pergroupSizeOfOffset != pergroupSize, + OP_LOGE(context->GetNodeName(), "antiquantScale[%zu]'s pergroup size[%ld] or antiquantOffset[%zu]'s pergroup size[%ld]" + "is not the required value[%ld]", i, pergroupSizeOfScale, i, pergroupSizeOfOffset, pergroupSize), + return GRAPH_FAILED); + } + } + OP_CHECK_IF(IsGmmQuantEmpty(context) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Detected antiquant, but quant inputs is not empty!"), return GRAPH_FAILED); + return GRAPH_SUCCESS; +} +static ge::graphStatus CheckQuantParams(gert::InferShapeContext* context, const GMMAttrs& gmmAttrs, GMMParamsInfo& paramsInfo) { + auto x0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, x0Desc); + DataType xDtype = x0Desc->GetDataType(); + auto w0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + DataType weightDtype = w0Desc->GetDataType(); + if (xDtype == DataType::DT_INT8 && weightDtype == DataType::DT_INT4) { + if (!isA8W4AsymmetricQuant(context)) { + return GRAPH_SUCCESS; + } + return CheckA8W4AsymQuantParams(context, paramsInfo); + } + if ((xDtype == DataType::DT_BF16 || xDtype == DataType::DT_FLOAT16 || + xDtype == DataType::DT_FLOAT) && xDtype == weightDtype) { + // nonquant + return CheckNonQuant(context); + } + if (xDtype == DataType::DT_INT8 && weightDtype == DataType::DT_INT8) { + // quant + return CheckDlinferGroupedMatmulDirectQuant(context, gmmAttrs, paramsInfo); + } + if ((xDtype == DataType::DT_BF16 || xDtype == DataType::DT_FLOAT16) && + (weightDtype == DataType::DT_INT8 || weightDtype == DataType::DT_INT4)) { + // antiquant + return CheckDlinferGroupedMatmulDirectAntiQuantForShape(context, gmmAttrs, paramsInfo); + } + return GRAPH_SUCCESS; +} +static ge::graphStatus CheckFunctionParamsForShape(gert::InferShapeContext* context, const GMMAttrs& gmmAttrs, + GMMParamsInfo& paramsInfo) { + if (context == nullptr) { + return GRAPH_FAILED; + } + fe::PlatformInfo platformInfo; + fe::OptionalInfo optionalInfo; + auto ret = fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo); + if (ret != ge::GRAPH_SUCCESS) { + paramsInfo.platform = PlatformID::UNKNOWN; + OP_LOGW(context->GetNodeName(), "Cannot get platform info!"); + return GRAPH_SUCCESS; + } else { + paramsInfo.platform = (optionalInfo.soc_version.find("310P") != std::string::npos) ? + PlatformID::ASCEND310P : (optionalInfo.soc_version.find("910_95") != std::string::npos) ? + PlatformID::ASCEND910_95 : PlatformID::ASCEND910B; + } + OP_CHECK_IF(CheckQuantParams(context, gmmAttrs, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckQuantParams failed!"), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} +static ge::graphStatus CheckDimNumAndGroupListNoSplitAndFormat(const gert::InferShapeContext* context, + uint64_t tensorListLength, const size_t numWeight) { + // when groupList is not empty, check its size equal with the length of x. + auto groupTensorOptionalShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + if (groupTensorOptionalShape != nullptr) { + OP_CHECK_IF(groupTensorOptionalShape->GetDim(0) != static_cast(tensorListLength), + OP_LOGE(context->GetNodeName(), "Size of groupList(tensor) %ld should be equal to size of x %lu.", + groupTensorOptionalShape->GetDim(0), tensorListLength), + return GRAPH_FAILED); + } + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + // check dimension + for (size_t i = 0; i < tensorListLength; ++i) { + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, i); + OP_CHECK_IF(xShape == nullptr, + OP_LOGE(context->GetNodeName(), "x[%lu] is null, which is not supported.", i), + return GRAPH_FAILED); + if (numWeight > 1) { + wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + size_t weightDimNum = wShape->GetDimNum(); + OP_CHECK_IF(weightDimNum != GMM_SEPARATED_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "weight[%lu] dimNum is %lu , but only support 2 when weight separated.", + i, weightDimNum), + return GRAPH_FAILED); + } + size_t xDimNum = xShape->GetDimNum(); + OP_CHECK_IF(xDimNum > GMM_MAX_FM_DIM || xDimNum < GMM_MIN_FM_DIM, + OP_LOGE(context->GetNodeName(), "x[%lu] dimNum is %lu , but only support 2-6.", i, xDimNum), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus TensorType2NodeId(const std::vector& tensorType, std::vector& nodeIdx) { + if (nodeIdx.size() > tensorType.size()) { + return GRAPH_FAILED; + } + for (size_t i(0); i < nodeIdx.size(); ++i) { + if (tensorType[i] == "x") { + nodeIdx[i] = GMM_INDEX_IN_X; + } else if (tensorType[i] == "weight") { + nodeIdx[i] = GMM_INDEX_IN_WEIGHT; + } else if (tensorType[i] == "y") { + nodeIdx[i] = GMM_INDEX_OUT_Y; + } else { + return GRAPH_FAILED; + } + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckDimNum(gert::InferShapeContext* context, uint64_t tensorListLength, + const size_t expectedDimNum, const std::string tensorType) { + int64_t nodeIdx = 0; + if (tensorType == "x") { + nodeIdx = static_cast(GMM_INDEX_IN_X); + } else if (tensorType == "weight") { + nodeIdx = static_cast(GMM_INDEX_IN_WEIGHT); + } else if (tensorType == "y") { + nodeIdx = static_cast(GMM_INDEX_OUT_Y); + } else { + return GRAPH_FAILED; + } + const gert::Shape* shape; + for (size_t i = 0; i < tensorListLength; ++i) { + if (tensorType == "y") { + shape = context->GetOutputShape(nodeIdx + i); + } else { + shape = context->GetDynamicInputShape(nodeIdx, i); + } + OP_CHECK_IF(shape == nullptr, + OP_LOGE(context->GetNodeName(), "%s[%lu] is null, which is not supported.", tensorType.c_str(), i), + return GRAPH_FAILED); + size_t dimNum = shape->GetDimNum(); + OP_CHECK_IF(dimNum != expectedDimNum, + OP_LOGE(context->GetNodeName(), "%s[%lu] dim num should be %lu in this case, but now is %lu.", + tensorType.c_str(), i, expectedDimNum, dimNum), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckWeightShapeInnerAxisEven(const gert::InferShapeContext* context, const size_t weightSize, + const int64_t innerAxisDimId) { + auto w0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + DataType wDtype = w0Desc->GetDataType(); + if (wDtype == DataType::DT_INT4) { + for (size_t i = 0; i < weightSize; ++i) { + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + int64_t n = wShape->GetDim(innerAxisDimId); + OP_CHECK_IF(n % 2 != 0, + OP_LOGE(context->GetNodeName(), "w[%zu] dim %ld value %ld should be even when weight is int4 dtype.", + i, innerAxisDimId, n), + return GRAPH_FAILED); + } + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus IsxSizeEqualWithWeightKAxis(const gert::InferShapeContext* context, + const GMMParamsInfo& paramsInfo, const gert::Shape* wShape, size_t& wKDimIdx, size_t& wNDimIdx) { + if (paramsInfo.numWeight == 1 && wShape->GetDimNum() > 2) { // 2: separated tensor's dim + wKDimIdx += 1UL; + wNDimIdx += 1UL; + OP_CHECK_IF(paramsInfo.numX != static_cast(wShape->GetDim(0)), + OP_LOGE(context->GetNodeName(), "When x and y are separated, and weight is not separated, size of x " + "%zu should equal to the first dim of weight tensor %ld.", paramsInfo.numX, wShape->GetDim(0)), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckCaseNoSplit(gert::InferShapeContext* context, bool transposeWeight, + const GMMParamsInfo& paramsInfo) { + const size_t& xSize = paramsInfo.numX; + const size_t& weightSize = paramsInfo.numWeight; + // check group num + OP_CHECK_IF(xSize != paramsInfo.numY, OP_LOGE(context->GetNodeName(), + "When y is separated, size of x %lu should equal to size of y %lu.", xSize, paramsInfo.numY), return GRAPH_FAILED); + OP_CHECK_IF(weightSize != 1 && xSize != weightSize, OP_LOGE(context->GetNodeName(), "When x and weight are separated, " + "size of x %lu should equal to size of weight %lu.", xSize, weightSize), return GRAPH_FAILED); + // check dimension + OP_CHECK_IF(CheckDimNumAndGroupListNoSplitAndFormat(context, xSize, weightSize) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor lists or grouplist is invalid."), + return GRAPH_FAILED); + // check shape + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + size_t wKDimIdx = transposeWeight ? 1UL : 0UL; + size_t wNDimIdx = transposeWeight ? 0UL : 1UL; + OP_CHECK_IF(IsxSizeEqualWithWeightKAxis(context, paramsInfo, wShape, wKDimIdx, wNDimIdx) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "IsxSizeEqualWithWeightKAxis failed."), return GRAPH_FAILED); + int64_t weightKDimValue = wShape->GetDim(wKDimIdx); + int64_t weightNDimValue = wShape->GetDim(wNDimIdx); + auto w0Desc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + DataType wDtype = w0Desc->GetDataType(); + // 2: an even factor + OP_CHECK_IF(wDtype == DataType::DT_INT4 && weightNDimValue % 2 != 0, OP_LOGE(context->GetNodeName(), + "w[0] dim %lu value %ld should be even when weight is int4 dtype.", wNDimIdx, weightNDimValue), + return GRAPH_FAILED); + for (size_t i = 0; i < xSize; i++) { + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, i); + size_t xDimNum = xShape->GetDimNum(); + // check inner axis of x, which should not be larger than 65535 + int64_t xKDimValue = xShape->GetDim(xDimNum - 1); // x always is not transposed + OP_CHECK_IF(xKDimValue > GMM_MAX_INNER_AXIS, + OP_LOGE(context->GetNodeName(), "x[%lu] dim %lu value %ld should less or equal to %ld.", + i, xDimNum - 1, xKDimValue, GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + if (weightSize > 1UL) { + wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + weightKDimValue = wShape->GetDim(wKDimIdx); + weightNDimValue = wShape->GetDim(wNDimIdx); + // 2: an even factor + OP_CHECK_IF(i > 0 && wDtype == DataType::DT_INT4 && weightNDimValue % 2 != 0, OP_LOGE(context->GetNodeName(), + "w[%lu] dim %lu value %ld should be even when weight is int4 dtype.", i, wNDimIdx, weightNDimValue), + return GRAPH_FAILED); + } + OP_CHECK_IF(xKDimValue != weightKDimValue, + OP_LOGE(context->GetNodeName(), "x[%lu] dim %lu value %ld should equal to weight[%lu] dim 0 value %ld.", + i, xDimNum - 1, xKDimValue, i, weightKDimValue), + return GRAPH_FAILED); + // if weight is not transposed, check N aisx; otherwise, check K axis, which can be skiped + OP_CHECK_IF(!transposeWeight && weightNDimValue > GMM_MAX_INNER_AXIS, + OP_LOGE(context->GetNodeName(), "w[%zu] dim %zu value %ld should less or equal to %ld.", + i, wNDimIdx, weightNDimValue, GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckInnerAxisOfTensorList(const gert::InferShapeContext* context, size_t nodeId, + int64_t innerAxisDimId, size_t checkNum) { + for (size_t i = 0; i < checkNum; i++) { + auto shape = context->GetDynamicInputShape(nodeId, i); + OP_CHECK_NULL_WITH_CONTEXT(context, shape); + int64_t innerAxisValue = shape->GetDim(innerAxisDimId); + OP_CHECK_IF(innerAxisValue > GMM_MAX_INNER_AXIS, + OP_LOGE(context->GetNodeName(), "Dim %ld value of %zu-th shape should less or equal to %ld, " + "but now is %ld.", innerAxisDimId, i, GMM_MAX_INNER_AXIS, innerAxisValue), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckShapeSameLengthTensorList(gert::InferShapeContext* context, + const std::vector& dimIds, const int64_t innerAxisDimId, + const std::vector tensorType, uint64_t groupNum) { + std::vector nodeIdx = {0, 0}; + OP_CHECK_IF(TensorType2NodeId(tensorType, nodeIdx) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "TensorType2NodeId failed."), + return GRAPH_FAILED); + // check two tensorlist's size to be the same, and tensors to have consistant dimension. + const gert::Shape* shape; + for (uint64_t i = 0; i < groupNum; i++) { + shape = context->GetDynamicInputShape(nodeIdx[0], i); + OP_CHECK_NULL_WITH_CONTEXT(context, shape); + int64_t dimValue1 = shape->GetDim(dimIds[0]); + // tensorType[2] indicates whether check tensorList0's inner axis(innerAxisDimId) + if (tensorType[2] == "true" && innerAxisDimId > -1) { + auto shape0 = context->GetDynamicInputShape(nodeIdx[0], i); + OP_CHECK_NULL_WITH_CONTEXT(context, shape0); + int64_t innerAxisValue = shape0->GetDim(innerAxisDimId); + if(innerAxisValue > GMM_MAX_INNER_AXIS){ + OP_LOGW(context->GetNodeName(), "Dim %lu value of %s[%lu] should less or equal to %ld," + "but now is %ld.", dimIds[0], tensorType[0].c_str(), i, GMM_MAX_INNER_AXIS, innerAxisValue); + } + } + if (tensorType[1] == "y") { + shape = context->GetOutputShape(nodeIdx[1] + i); + } else { + shape = context->GetDynamicInputShape(nodeIdx[1], i); + } + OP_CHECK_NULL_WITH_CONTEXT(context, shape); + int64_t dimValue2 = shape->GetDim(dimIds[1]); + if(dimValue1 != dimValue2){ + OP_LOGW(context->GetNodeName(), + "Dim %lu value of %s[%lu] should be equal with dim %lu value of %s[%lu]," + "but now is %ld and %ld respectively.", dimIds[0], tensorType[0].c_str(), + i, dimIds[1], tensorType[1].c_str(), i, dimValue1, dimValue2); + } + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckShapeDiffLengthTensorList(gert::InferShapeContext* context, + const std::vector& dimIds, + const int64_t innerAxisdimId, + const std::vector tensorType, + uint64_t groupNum) { + std::vector nodeIdx = {0, 0}; + OP_CHECK_IF(TensorType2NodeId(tensorType, nodeIdx) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "TensorType2NodeId failed."), + return GRAPH_FAILED); + // check each tensor's selected dimension size in a multi-tensor tensorlist's to equal with + // the tensor selected dimension in single-tensor tensorlist. + // the selected axis is not the split-axis. + const gert::Shape* singleTensor0; + if (tensorType[1] == "y") { + singleTensor0 = context->GetOutputShape(nodeIdx[1]); + } else { + singleTensor0 = context->GetDynamicInputShape(nodeIdx[1], 0); + } + OP_CHECK_NULL_WITH_CONTEXT(context, singleTensor0); + int64_t dimValueSingle = singleTensor0->GetDim(dimIds[1]); + // tensorType[2] indicates whether check single tensorList's inner axis(innerAxisDimId) + if (tensorType[2] == "true" && innerAxisdimId > -1) { + int64_t dimValue = singleTensor0->GetDim(innerAxisdimId); + OP_CHECK_IF(dimValue > GMM_MAX_INNER_AXIS, + OP_LOGE(context->GetNodeName(), + "Dim %ld value of %s[0] should less or equal to %ld, but now is %ld.", + innerAxisdimId, tensorType[1].c_str(), GMM_MAX_INNER_AXIS, dimValue), + return GRAPH_FAILED); + } + const gert::Shape* longTensor; + for (uint64_t i = 0; i < groupNum; i++) { + if (tensorType[0] == "y") { + longTensor = context->GetOutputShape(nodeIdx[0] + i); + } else { + longTensor = context->GetDynamicInputShape(nodeIdx[0], i); + } + OP_CHECK_NULL_WITH_CONTEXT(context, longTensor); + int64_t dimValueLong = longTensor->GetDim(dimIds[0]); + OP_CHECK_IF(dimValueLong != dimValueSingle, + OP_LOGE(context->GetNodeName(), + "Dim %lu value of %s[%lu] %ld should be equal with dim %lu value of %s[0] %ld.", + dimIds[0], tensorType[0].c_str(), i, dimValueLong, + dimIds[1], tensorType[1].c_str(), dimValueSingle), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckGroupListCommonTensor(const gert::InferShapeContext* context, + const bool isRequiredGroupList, const int64_t groupNum) { + auto groupTensorOptionalShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + bool isNull = groupTensorOptionalShape == nullptr; + OP_CHECK_IF(isNull && isRequiredGroupList, + OP_LOGE(context->GetNodeName(), "groupListOptional(tensor) is required in this case, but get nullptr."), + return GRAPH_FAILED); + if (isNull) { + return GRAPH_SUCCESS; + } + int64_t groupListSize = groupTensorOptionalShape->GetDim(0); + OP_CHECK_IF(groupListSize > GMM_MAX_GROUP_LIST_SIZE_TENSOR, + OP_LOGE(context->GetNodeName(), + "When groupList type is tenosr, size of groupList %ld should be less than or equal to %ld.", + groupListSize, GMM_MAX_GROUP_LIST_SIZE_TENSOR), + return GRAPH_FAILED); + const gert::RuntimeAttrs* attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* groupListType = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_LIST_TYPE); + bool validGroupListSize = (groupListSize == groupNum && groupNum > 1) || groupNum == 1; + if (groupListType != nullptr && *groupListType == 2L) { + validGroupListSize = groupListSize > 0 && groupListSize <= groupNum; + } + OP_CHECK_IF(!validGroupListSize, + OP_LOGE(context->GetNodeName(), + "When groupList is not null, size of groupList(tensor) %ld should be equal to groupNum %ld, " + "or in groupListType 2 be in range (0, %ld].", groupListSize, groupNum, groupNum), + return GRAPH_FAILED); + auto groupListDesc = context->GetOptionalInputDesc(GMM_INDEX_IN_GROUP_LIST); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListDesc); + OP_CHECK_IF(groupListDesc->GetDataType() != DataType::DT_INT64, + OP_LOGE(context->GetNodeName(), "Invalid dtype: Only int64 is supported for groupList, but now is %s.", + TypeUtils::DataTypeToAscendString(groupListDesc->GetDataType()).GetString()), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus SplitMSingleXSingleWeightSingleY(gert::InferShapeContext* context, bool transposeWeight, + const GMMParamsInfo& paramsInfo) { + std::vector tenorXAndWeight{"x", "weight", "true"}; + // check dimension + OP_CHECK_IF(CheckDimNum(context, paramsInfo.numX, GMM_MIN_FM_DIM, "x") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list x is invalid."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckDimNum(context, paramsInfo.numWeight, GMM_SPLIT_M_SINGLE_WEIGHT_DIM, "weight") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list weight is invalid."), + return GRAPH_FAILED); + // check shape, x(m,k), weight(b,k,n), y(m,n) + int64_t innerAxisDimId = 1; // x always is not transposed, check K axis + size_t kAxisOfWeight = transposeWeight ? 2UL : 1UL; // if weight is transposed, 2 is the k axis idx of the weight, otherwise is 1 + OP_CHECK_IF(CheckShapeSameLengthTensorList(context, {1, kAxisOfWeight}, innerAxisDimId, tenorXAndWeight, paramsInfo.numX) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "k dim value of x and weight is not matched."), + return GRAPH_FAILED); + innerAxisDimId = !transposeWeight ? 2 : -1; // If w is not transposed, check N(2) asix; otherwise, check k axis, which can be skiped + OP_CHECK_IF(CheckInnerAxisOfTensorList(context, GMM_INDEX_IN_WEIGHT, innerAxisDimId, paramsInfo.numWeight) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "inner axis size of weight is larger than %ld!", GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + OP_CHECK_IF(CheckWeightShapeInnerAxisEven(context, paramsInfo.numWeight, 2) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight's N axis size should be even when it is int4 dtype."), + return GRAPH_FAILED); + // check groupList + OP_CHECK_IF(CheckGroupListCommonTensor(context, true, context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0)->GetDim(0)) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid groupList."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus SplitMSingleXSeparatedWeightSingleY(gert::InferShapeContext* context, bool transposeWeight, + const GMMParamsInfo& paramsInfo) { + std::vector tenorWeightAndX{"weight", "x", "true"}; + // check dimension + OP_CHECK_IF(CheckDimNum(context, paramsInfo.numX, GMM_MIN_FM_DIM, "x") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list x is invalid."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckDimNum(context, paramsInfo.numWeight, GMM_SEPARATED_WEIGHT_DIM, "weight") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list weight is invalid."), + return GRAPH_FAILED); + // check shape, x(m,k), weight(k,n), y(m,n) + int64_t innerAxisDimId = 1; // x always is not transposed, check K axis + size_t kAxisOfWeight = transposeWeight ? 1UL : 0UL; + OP_CHECK_IF(CheckShapeDiffLengthTensorList(context, {kAxisOfWeight, 1}, innerAxisDimId, tenorWeightAndX, paramsInfo.numWeight) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "k dim value of x and weight is not matched."), + return GRAPH_FAILED); + innerAxisDimId = !transposeWeight ? 1 : -1; // if w is not transposed, check N asix; otherwise, check k axis, which can be skiped + OP_CHECK_IF(CheckInnerAxisOfTensorList(context, GMM_INDEX_IN_WEIGHT, innerAxisDimId, 1) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "inner axis size of weight is larger than %ld!", GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + OP_CHECK_IF(CheckWeightShapeInnerAxisEven(context, paramsInfo.numWeight, 1) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight's N axis size should be even when it is int4 dtype."), + return GRAPH_FAILED); + // check groupList + OP_CHECK_IF(CheckGroupListCommonTensor(context, true, paramsInfo.numWeight) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid groupList."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus SplitMSeparatedXSeparatedWeightSingleY(gert::InferShapeContext* context, + bool transposeWeight, const GMMParamsInfo& paramsInfo) { + const size_t& xSize = paramsInfo.numX; + const size_t& weightSize = paramsInfo.numWeight; + std::vector tenorWeightAndX{"weight", "x", "true"}; + OP_CHECK_IF(xSize != weightSize, + OP_LOGE(context->GetNodeName(), + "When x and weight are separated, size of x %lu should equal to size of weight %lu.", + xSize, weightSize), + return GRAPH_FAILED); + // check dimension + OP_CHECK_IF(CheckDimNum(context, xSize, GMM_MIN_FM_DIM, "x") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list x is invalid."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckDimNum(context, weightSize, GMM_SEPARATED_WEIGHT_DIM, "weight") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list weight is invalid."), + return GRAPH_FAILED); + // check shape, x(m,k), weight(k,n), y(m,n) + int64_t innerAxisDimId = 1; // originalShape's inner axis of weight + size_t kAxisOfWeight = transposeWeight ? 1UL : 0UL; + OP_CHECK_IF(CheckShapeSameLengthTensorList(context, {kAxisOfWeight, 1}, innerAxisDimId, tenorWeightAndX, weightSize) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "k dim value of x and weight is not matched."), + return GRAPH_FAILED); + innerAxisDimId = !transposeWeight ? 1 : -1; // if w is not transposed, N asix has been checked, need to check x's inner axis(K, when x is always not transposed) + OP_CHECK_IF(CheckInnerAxisOfTensorList(context, GMM_INDEX_IN_X, innerAxisDimId, 1) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "inner axis size of x is larger than %ld!", GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + OP_CHECK_IF(CheckWeightShapeInnerAxisEven(context, weightSize, 1) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight's N axis size should be even when it is int4 dtype."), + return GRAPH_FAILED); + // check groupList + OP_CHECK_IF(CheckGroupListCommonTensor(context, false, xSize) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid groupList."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus CheckCaseSplitM(gert::InferShapeContext* context, bool transposeWeight, + const GMMParamsInfo& paramsInfo) { + const size_t& xSize = paramsInfo.numX; + const size_t& weightSize = paramsInfo.numWeight; + const size_t& ySize = paramsInfo.numY; + if ((xSize == 1UL) && (weightSize == 1UL) && (ySize == 1UL)) { + OP_CHECK_IF(SplitMSingleXSingleWeightSingleY(context, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Split m, single x, single weight, single y case failed."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + if ((xSize == 1UL) && (weightSize > 1UL) && (ySize == 1UL)) { + OP_CHECK_IF(weightSize != paramsInfo.groupNum, OP_LOGE(context->GetNodeName(), + "weight Size [%zu] does not equal with groupNum %zu", weightSize, paramsInfo.groupNum), + return GRAPH_FAILED); + OP_CHECK_IF(SplitMSingleXSeparatedWeightSingleY(context, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Split m, single x, separated weight, single y case failed."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + if ((xSize == 1UL) && (weightSize > 1UL) && (ySize > 1UL)) { + const gert::Tensor* groupListTensor = context->GetOptionalInputTensor(GMM_INDEX_IN_GROUP_LIST); + OP_CHECK_IF(groupListTensor == nullptr || groupListTensor->GetData() == nullptr, + OP_LOGE(context->GetNodeName(), "Failed to obtain necessary data from groupListTensor. " + "When grouplist is an invalid tensor, split m, single x, separated weight, separated y cases do not support."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; // skip the check + } + if ((xSize > 1UL) && (weightSize > 1UL) && (ySize == 1UL)) { + OP_CHECK_IF(weightSize != paramsInfo.groupNum, OP_LOGE(context->GetNodeName(), + "weight Size [%zu] does not equal with groupNum %zu", weightSize, paramsInfo.groupNum), + return GRAPH_FAILED); + OP_CHECK_IF(SplitMSeparatedXSeparatedWeightSingleY(context, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Split m, separated x, separated weight, single y case failed."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + OP_LOGE(context->GetNodeName(), "When groupType is 0, current case with x %zu, weight %zu, y %zu is not supported.", + xSize, weightSize, ySize); + return GRAPH_FAILED; +} + +static ge::graphStatus CheckCaseSplitK(gert::InferShapeContext* context, bool transposeX, bool transposeWeight, + const GMMParamsInfo& paramsInfo) { + std::vector tenorXAndWeight{"x", "weight", "true"}; + const size_t& xSize = paramsInfo.numX; + const size_t& weightSize = paramsInfo.numWeight; + const size_t& ySize = paramsInfo.numY; + if (xSize == 1UL) { + if (paramsInfo.platform == PlatformID::ASCEND910_95) { + return GRAPH_SUCCESS; + } + OP_CHECK_IF(!transposeX, + OP_LOGE(context->GetNodeName(), + "When groupType is 2 and x is not separated, tensor in x should be transposed."), + return GRAPH_FAILED); + // check dimension + OP_CHECK_IF(CheckDimNum(context, xSize, GMM_MIN_FM_DIM, "x") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list x is invalid."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckDimNum(context, weightSize, GMM_SPLIT_K_SINGLE_WEIGHT_DIM, "weight") != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num or format of tensor in tensor list weight is invalid."), + return GRAPH_FAILED); + // check shape, x(m,k), weight(k,n), y(b,m,n) + int64_t innerAxisDimId = 1; // x always is transposed, and the inner axis is always the last axis, M axis. + size_t kAxisOfWeight = transposeWeight ? 1UL : 0UL; + if((weightSize == 1UL) && (ySize == 1UL)) { + OP_CHECK_IF(CheckShapeSameLengthTensorList(context, {0, kAxisOfWeight}, innerAxisDimId, tenorXAndWeight, xSize) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "k dim value of x and weight is not matched."), + return GRAPH_FAILED); + innerAxisDimId = 1; // w always is not transposed, and the inner axis is always the last axis, N axis. + // check groupList + OP_CHECK_IF(CheckGroupListCommonTensor(context, true, 1) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid groupList."), + return GRAPH_FAILED); + } + OP_CHECK_IF(CheckInnerAxisOfTensorList(context, GMM_INDEX_IN_WEIGHT, innerAxisDimId, weightSize) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "inner axis size of weight is larger than %ld!", GMM_MAX_INNER_AXIS), + return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + OP_LOGE(context->GetNodeName(), + "When groupType is 2, only support case with unseparated x, weight and y, " + "but now x size is %lu, weight size is %lu, y size is %lu.", xSize, weightSize, ySize); + return GRAPH_FAILED; +} + +static ge::graphStatus CheckParamDifferentGroupType(gert::InferShapeContext* context, const GMMAttrs& gmmAttrs, + const GMMParamsInfo& paramsInfo) { + OP_CHECK_IF(paramsInfo.platform == PlatformID::UNKNOWN, OP_LOGW(context->GetNodeName(), "Cannot get platform info!"), return GRAPH_SUCCESS); + const int64_t& groupType = gmmAttrs.groupType; + const bool& transposeX = gmmAttrs.transposeX; + const bool& transposeWeight = gmmAttrs.transposeWeight; + OP_CHECK_IF(transposeX && transposeWeight, OP_LOGE(context->GetNodeName(), + "x and weight can not be transposed at the same time."), return GRAPH_FAILED); + auto groupTensorOptionalShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* groupListTypePtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_GROUP_LIST_TYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListTypePtr); + size_t validGroupTensorDimNum = (*groupListTypePtr == 2L) ? 2UL: 1UL; // 2: split M sparse, group list shape [e, 2] + OP_CHECK_IF(groupTensorOptionalShape != nullptr && (groupTensorOptionalShape->GetDimNum() > validGroupTensorDimNum || + groupTensorOptionalShape->GetDim(0) < 1), + OP_LOGE(context->GetNodeName(), + "When groupList is a tensor, its dim only supports 1 or 2(only when groupListType is 2) and " + "size of elements should be larger than 0, but now are %zu and %ld, respectively.", + groupTensorOptionalShape->GetDimNum(), groupTensorOptionalShape->GetDim(0)), + return GRAPH_FAILED); + OP_CHECK_IF(paramsInfo.platform == PlatformID::ASCEND310P && !(groupType == GMM_SPLIT_M && paramsInfo.numX == 1 && + paramsInfo.numWeight == 1 && paramsInfo.numY == 1), + OP_LOGE(context->GetNodeName(), + "When on ASCEND310P, it only supports split m, single x, single weight, single y."), + return GRAPH_FAILED); + + if (groupType == GMM_NO_SPLIT) { + OP_CHECK_IF(transposeX, OP_LOGE(context->GetNodeName(), + "When x, weight and y are all separated, x can not be transposed."), return GRAPH_FAILED); + OP_CHECK_IF(CheckCaseNoSplit(context, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid inputs!"), return GRAPH_FAILED); + } else if (groupType == GMM_SPLIT_M) { + OP_CHECK_IF(transposeX, + OP_LOGE(context->GetNodeName(), "When groupType is 0, x can not be transposed."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckCaseSplitM(context, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid inputs!"), return GRAPH_FAILED); + } else if (groupType == GMM_SPLIT_K) { + OP_CHECK_IF(!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_BIAS), + OP_LOGE(context->GetNodeName(), "When groupType is 2, bias must be empty."), return GRAPH_FAILED); + OP_CHECK_IF(CheckCaseSplitK(context, transposeX, transposeWeight, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid inputs!"), return GRAPH_FAILED); + } + if (!IsTensorListNullOrEmpty(context, GMM_INDEX_IN_BIAS)) { + OP_CHECK_IF(CheckOptionalTensorList(context, "bias", paramsInfo, gmmAttrs, GMM_INDEX_IN_BIAS) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Invalid bias!"), return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus XNotSingleYSeparated(gert::InferShapeContext* context, + size_t weightDimN, bool isXTransposed, size_t xDimM) { + const gert::Tensor* groupListTensor = context->GetOptionalInputTensor(GMM_INDEX_IN_GROUP_LIST); + if (groupListTensor != nullptr) { + OP_CHECK_IF(UpdateMultipleShapeY(context, groupListTensor, weightDimN, isXTransposed, xDimM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } else { + OP_CHECK_IF(MultiInMultiOutWithoutGroupList(context)!= GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Failed to process multi-in-multi-out case without GroupList."), return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static ge::graphStatus XSingleYSeparated(gert::InferShapeContext* context, + size_t weightDimN, bool isXTransposed, size_t xDimM) { + const gert::Tensor* groupListTensor = context->GetOptionalInputTensor(GMM_INDEX_IN_GROUP_LIST); + OP_CHECK_IF(groupListTensor == nullptr, + OP_LOGE(context->GetNodeName(), "GroupList is required when x is single tensor while y is not."), + return GRAPH_FAILED); + OP_CHECK_IF(UpdateMultipleShapeY(context, groupListTensor, weightDimN, isXTransposed, xDimM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static ge::graphStatus GMMSetOutputShape(gert::InferShapeContext* context, GMMAttrs& gmmAttrs, + const GMMSetOutputParams& outputParams, const gert::Shape* x0Shape, + const gert::Shape* w0Shape) { + bool isSingleX = outputParams.isSingleX; + bool isSingleY = outputParams.isSingleY; + size_t xDimM = outputParams.xDimM; + size_t weightDimN = outputParams.weightDimN; + size_t numX = outputParams.numX; + size_t numWeight = outputParams.numWeight; + int64_t lenGroupList = outputParams.lenGroupList; + // X单 Y多 + if (isSingleX && !isSingleY) { + if(gmmAttrs.groupType != GMM_SPLIT_K) { + OP_CHECK_IF(XSingleYSeparated(context, weightDimN, gmmAttrs.transposeX, xDimM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } else { + OP_CHECK_IF(MultiWeightMultiOutWithoutGroupList(context)!= GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), + "Failed to process multi-in-multi-out case without GroupList."), return GRAPH_FAILED); + } + // X单 Y单 + } else if (isSingleX && isSingleY) { + OP_CHECK_IF(gmmAttrs.groupType != GMM_SPLIT_M && gmmAttrs.groupType != GMM_SPLIT_K, + OP_LOGE(context->GetNodeName(), + "When x is single tensor, input tensors can only be split along M or K axis."), return GRAPH_FAILED); + std::vector yDims = {x0Shape->GetDim(xDimM), w0Shape->GetDim(weightDimN)}; + if (gmmAttrs.groupType == GMM_SPLIT_K) { + yDims.insert(yDims.begin(), numWeight == 1 ? lenGroupList : numWeight); + } + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y, yDims) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update y shape."), return GRAPH_FAILED); + } + // X多 Y多 + else if (!isSingleX && !isSingleY) { + OP_CHECK_IF(XNotSingleYSeparated(context, weightDimN, gmmAttrs.transposeX, xDimM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } + // X多 Y单 + else if (!isSingleX && isSingleY) { + std::vector yDims = {GetDim0(context, gmmAttrs.transposeX, numX, xDimM), w0Shape->GetDim(weightDimN)}; + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y, yDims) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return GRAPH_FAILED); + } + + return GRAPH_SUCCESS; +} + +static graphStatus InferShape4DavidWeightQuantGMM(gert::InferShapeContext *context) +{ + DlinferGroupedMatmulDirectWeightQuantChecker davidWeightQuantGMMChecker; + DlinferGroupedMatmulDirectCommonUtil utilForDavidWeightQuantGMM; + OP_CHECK_IF(GetAttrsValue(context, utilForDavidWeightQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetAttrsValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidWeightQuantGMMChecker.GetXAndWeightDimValue(context, utilForDavidWeightQuantGMM.attrsInfo) != + GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetXAndWeightDimValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidWeightQuantGMMChecker.CheckShape(context, utilForDavidWeightQuantGMM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShape failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidWeightQuantGMMChecker.InferOutShape(context, utilForDavidWeightQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "InferOutShape failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus InferShape4DavidQuantGMM(gert::InferShapeContext* context) { + DlinferGroupedMatmulDirectQuantChecker davidQuantGMMChecker; + DlinferGroupedMatmulDirectCommonUtil utilForDavidQuantGMM; + OP_CHECK_IF(GetAttrsValue(context, utilForDavidQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetAttrsValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.GetXAndWeightDimValue(context, utilForDavidQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetXAndWeightDimValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.GetGroupNumValue(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetGroupNumValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.CheckShape(context, utilForDavidQuantGMM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShape failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.InferOutShape(context, utilForDavidQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "InferOutShape failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +template +static graphStatus IsDavidWeightQuantGMMByShape(T context) +{ + auto xDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_X, 0); + auto weightDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc); + DataType xDtype = xDesc->GetDataType(); + DataType weightDtype = weightDesc->GetDataType(); + return GetSizeByDataType(xDtype) != GetSizeByDataType(weightDtype) ? GRAPH_SUCCESS : GRAPH_FAILED; +} + +template +static graphStatus IsDavidQuantGMMByShape(T context) { + auto xDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_X, 0); + auto weightDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + auto scaleDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_SCALE, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, scaleDesc); + DataType xDtype = xDesc->GetDataType(); + DataType weightDtype = weightDesc->GetDataType(); + if (xDtype == ge::DT_FLOAT4_E1M2 || xDtype == ge::DT_FLOAT4_E2M1 || xDtype == ge::DT_INT4) { + return GRAPH_SUCCESS; + } + return (GetSizeByDataType(xDtype) == 1 && GetSizeByDataType(weightDtype) == 1) ? GRAPH_SUCCESS : GRAPH_FAILED; +} + +static ge::graphStatus InferShape4DlinferGroupedMatmulDirect(gert::InferShapeContext* context) { + OP_CHECK_NULL_WITH_CONTEXT(context, context); + fe::PlatformInfo platformInfo; + fe::OptionalInfo optionalInfo; + auto ret = fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo); + if (ret == GRAPH_SUCCESS && GmmDavidSupportSoc.count(platformInfo.str_info.short_soc_version) > 0) { + if (IsDavidQuantGMMByShape(context) == GRAPH_SUCCESS) { + OP_CHECK_IF(InferShape4DavidQuantGMM(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Check params failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; + } else if (IsDavidWeightQuantGMMByShape(context) == GRAPH_SUCCESS) { + OP_CHECK_IF(InferShape4DavidWeightQuantGMM(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Check params failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + } + GMMAttrs gmmAttrs{GMM_X_Y_SEPARATED, 0, GMM_NO_SPLIT, false, false, 0, 0}; + OP_CHECK_IF(GetAttrsValue(context, gmmAttrs) != GRAPH_SUCCESS || CheckAttrs(context, gmmAttrs) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to get attrs."), return GRAPH_FAILED); + + size_t numX = 0; // init numX + size_t numWeight = 0; // init numWeight + int64_t lenGroupList = 0; // init lenGroupList + size_t numY = context->GetComputeNodeOutputNum(); + if (GetNumOfInputs(context, numX, numWeight, lenGroupList) == GRAPH_SUCCESS) { // check input shape value inside + GMMParamsInfo paramsInfo{numX, numWeight, numY, lenGroupList, 0, 0, 0, 0, 0, PlatformID::UNKNOWN}; + OP_CHECK_IF(GetGroupSize(context, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "check groupNum failed"), return GRAPH_FAILED); + OP_CHECK_IF(CheckFunctionParamsForShape(context, gmmAttrs, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckFunctionParamsForShape failed."), return GRAPH_FAILED); + OP_CHECK_IF(CheckParamDifferentGroupType(context, gmmAttrs, paramsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckParamDifferentGroupType failed."), return GRAPH_FAILED); + } else { + OP_CHECK_IF(CheckDimNum(context, numX, GMM_MIN_FM_DIM, "x") != GRAPH_SUCCESS, // check dim number of tensors + OP_LOGE(context->GetNodeName(), "Dim num of tensor in tensorList x is invalid."), + return GRAPH_FAILED); + } + + const gert::Shape* x0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, x0Shape); + size_t xDimNum = x0Shape->GetDimNum(); + const gert::Shape* w0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Shape); + size_t weightDimNum = w0Shape->GetDimNum(); + bool isSingleX = (numX == 1UL) && (gmmAttrs.groupType != GMM_NO_SPLIT); + bool isSingleY = (numY == 1UL) && (gmmAttrs.groupType != GMM_NO_SPLIT); + size_t xDimM = gmmAttrs.transposeX ? xDimNum - 1UL : xDimNum - 2UL; + size_t weightDimN = gmmAttrs.transposeWeight ? weightDimNum - 2UL : weightDimNum - 1UL; + + GMMSetOutputParams outputParams; + outputParams.isSingleX = isSingleX; + outputParams.isSingleY = isSingleY; + outputParams.xDimM = xDimM; + outputParams.numX = numX; + outputParams.weightDimN = weightDimN; + outputParams.lenGroupList = lenGroupList; + outputParams.numWeight = numWeight; + OP_CHECK_IF(GMMSetOutputShape(context, gmmAttrs, outputParams, x0Shape, w0Shape) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GMMSetOutputShape failed"), return GRAPH_FAILED); + + return GRAPH_SUCCESS; +} + +// ========================================================================================= +// ========================================================================================= +static graphStatus CheckTensorListDataType(const gert::InferDataTypeContext* context, uint32_t index, + const DataType dtype) { + size_t inIdx = 0; + while (true) { + auto iDtype = context->GetDynamicInputDataType(index, inIdx); + if (iDtype == DT_UNDEFINED) { + break; + } + OP_CHECK_IF(iDtype != dtype, + OP_LOGE(context->GetNodeName(), "data type of tensors in a tensorList should all be the same!"), + return GRAPH_FAILED); + ++inIdx; + } + return GRAPH_SUCCESS; +} + +static graphStatus CheckMatmulDataType(gert::InferDataTypeContext* context, const DataType xDtype, + const DataType weightDtype, const DataType biasDtype) { + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_X, xDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "x dtype does not match with required dtype[%s].", + TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_WEIGHT, weightDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight dtype does not match with required dtype[%s].", + TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_BIAS, biasDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "bias dtype does not match with required dtype[%s].", + TypeUtils::DataTypeToAscendString(biasDtype).GetString()), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus CheckNonQuantMatmulParams(fe::PlatformInfo& platformInfo, gert::InferDataTypeContext* context, + const DataType xDtype, const DataType weightDtype) +{ + DataType biasDtype = xDtype == DataType::DT_BF16 ? DataType::DT_FLOAT: xDtype; + if (GmmDavidSupportSoc.count(platformInfo.str_info.short_soc_version) > 0) { + biasDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_BIAS, 0); + if (biasDtype != DT_UNDEFINED) { + OP_CHECK_IF(std::find(BIAS_DTYPE_SUPPORT_LIST.begin(), BIAS_DTYPE_SUPPORT_LIST.end(), biasDtype) == BIAS_DTYPE_SUPPORT_LIST.end(), + OP_LOGE(context->GetNodeName(),"non quant case bias only support dtype float16, bfloat16 and float32"), + return GRAPH_FAILED); + } + } + OP_CHECK_IF(CheckMatmulDataType(context, xDtype, weightDtype, biasDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "case with x dtype %s and weight dtype %s is not supported!", + TypeUtils::DataTypeToAscendString(xDtype).GetString(), TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus CheckFunctionQuantParams(gert::InferDataTypeContext* context) { + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_X, DataType::DT_INT8) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "x dtype does not match with required dtype[INT8]."), + return GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_WEIGHT, DataType::DT_INT8) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight dtype does not match with required dtype[INT8]."), + return GRAPH_FAILED); + OP_CHECK_IF((CheckTensorListDataType(context, GMM_INDEX_IN_BIAS, DataType::DT_INT32) != GRAPH_SUCCESS) && + (CheckTensorListDataType(context, GMM_INDEX_IN_BIAS, DataType::DT_BF16) != GRAPH_SUCCESS), + OP_LOGE(context->GetNodeName(), "bias dtype does not match with required dtype int32 or bfloat16."), + return GRAPH_FAILED); + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* outputDtype = attrs->GetInt(GMM_INDEX_ATTR_OUTPUT_DTYPE); + if (*outputDtype == GMM_OUT_DTYPE_INT32) { // output dtype is int32, this scene does not need scale + return GRAPH_SUCCESS; + } + auto scale0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_SCALE, 0); + // Now we cannot make sure if is pertoken quant case, so scale/offset dtype check is remained to the InferShape stage. + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_SCALE, scale0Dtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "dtypes of scales in the tensorList should all be the same."), + return GRAPH_FAILED); + auto offset0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_OFFSET, 0); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_OFFSET, offset0Dtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "dtypes of offsets in the tensorList should all be the same."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus CheckDlinferGroupedMatmulDirectAntiQuantForDtype(gert::InferDataTypeContext* context) { + auto xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_ANTIQUANT_SCALE, xDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "antiquantScale dtype does not match with x dtype[%s].", TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_ANTIQUANT_OFFSET, xDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "antiquantOffset dtype does not match with x dtype[%s].", TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus CheckFunctionParamsForDtype(gert::InferDataTypeContext* context) { + fe::PlatformInfo platformInfo; + fe::OptionalInfo optionalInfo; + graphStatus ret = fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo); + PlatformID platform = PlatformID::UNKNOWN; + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGW(context->GetNodeName(), "Cannot get platform info."); + return GRAPH_SUCCESS; + } else { + platform = (optionalInfo.soc_version.find("310P") != std::string::npos) ? + PlatformID::ASCEND310P : (optionalInfo.soc_version.find("910_95") != std::string::npos) ? + PlatformID::ASCEND910_95 : PlatformID::ASCEND910B; + } + DataType xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + DataType weightDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_WEIGHT, 0); + if (platform == PlatformID::ASCEND310P) { + bool isAllInputFP16 = xDtype == DataType::DT_FLOAT16 && weightDtype == DataType::DT_FLOAT16; + OP_CHECK_IF(!isAllInputFP16, OP_LOGE(context->GetNodeName(), + "Only float16 is supported on Ascend310P platforms."), return GRAPH_FAILED); + auto biasDtype = context->GetOptionalInputDataType(GMM_INDEX_IN_BIAS); + OP_CHECK_IF(biasDtype != ge::DT_UNDEFINED && biasDtype != DataType::DT_FLOAT16, OP_LOGE(context->GetNodeName(), + "only bias float16 is supported on Ascend310P platforms."), return GRAPH_FAILED); + } + if (xDtype == DataType::DT_INT8 && weightDtype == DataType::DT_INT4) { return GRAPH_SUCCESS; } + if ((xDtype == DataType::DT_BF16 || xDtype == DataType::DT_FLOAT16 || xDtype == DataType::DT_FLOAT) && + xDtype == weightDtype) { // nonquant + return CheckNonQuantMatmulParams(platformInfo, context, xDtype, weightDtype); + } + if (xDtype == DataType::DT_INT8 && weightDtype == DataType::DT_INT8) { + // quant + OP_CHECK_IF(CheckFunctionQuantParams(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckFunctionQuantParams failed."), + return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + if ((xDtype == DataType::DT_BF16 || xDtype == DataType::DT_FLOAT16) && + (weightDtype == DataType::DT_INT8 || weightDtype == DataType::DT_INT4)) { + // antiquant + DataType biasDtype = xDtype == DataType::DT_BF16 ? DataType::DT_FLOAT: DataType::DT_FLOAT16; + OP_CHECK_IF(CheckMatmulDataType(context, xDtype, weightDtype, biasDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "case with x dtype %s and weight dtype %s is not supported!", + TypeUtils::DataTypeToAscendString(xDtype).GetString(), TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return GRAPH_FAILED); + return CheckDlinferGroupedMatmulDirectAntiQuantForDtype(context); + } + OP_LOGE(context->GetNodeName(), "GMM: there is no matching xDtype and weightDtype pattern. " + "case with x dtype %s and weight dtype %s is not supported.", + TypeUtils::DataTypeToAscendString(xDtype).GetString(), TypeUtils::DataTypeToAscendString(weightDtype).GetString()); + return GRAPH_FAILED; +} + +static graphStatus CheckQuantParamsDtype(const gert::InferDataTypeContext* context, const int64_t outputDtype, + const DataType yDtype) { + size_t i = 0; + auto scale0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_SCALE, 0); + OP_CHECK_IF(scale0Dtype == ge::DT_UNDEFINED, OP_LOGE(context->GetNodeName(), "scale is undefined!"), + return GRAPH_FAILED); + auto perTokenScale0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + bool isPerTokenQuant = perTokenScale0Dtype != ge::DT_UNDEFINED; + if (isPerTokenQuant) { + bool isOutputBF16 = scale0Dtype == DataType::DT_BF16 && outputDtype == 1; + bool isOutputFloat16 = scale0Dtype == DataType::DT_FLOAT && outputDtype == 0; + OP_CHECK_IF(!isOutputBF16 && !isOutputFloat16, + OP_LOGE(context->GetNodeName(), "per-token quant case only supports scale data type bfloat16 with " + "output data type bfloat16, or scale with data type float32 when output is float16, but " + "now scale[%zu] has data type %s and output has data type %s!", + i, TypeUtils::DataTypeToAscendString(scale0Dtype).GetString(), TypeUtils::DataTypeToAscendString(yDtype).GetString()), + return GRAPH_FAILED); + } else { + bool isOutputInt8 = scale0Dtype == DataType::DT_UINT64 && outputDtype == -1; + bool isOutputBF16 = scale0Dtype == DataType::DT_BF16 && outputDtype == 1; + bool isOutputFP16 = scale0Dtype == DataType::DT_FLOAT && outputDtype == 0; + OP_CHECK_IF(!isOutputInt8 && !isOutputBF16 && !isOutputFP16, + OP_LOGE(context->GetNodeName(), "per-channel quant case only supports scale with data type uint64 " + "when output is int8, or data type bfloat16 when output is bfloat16, or data type float32 " + "when output is float16, but scale[%zu] has data type %s and output has data type %s!", + i, TypeUtils::DataTypeToAscendString(scale0Dtype).GetString(), TypeUtils::DataTypeToAscendString(yDtype).GetString()), + return GRAPH_FAILED); + } + if (isPerTokenQuant) { + OP_CHECK_IF(perTokenScale0Dtype != DataType::DT_FLOAT, + OP_LOGE(context->GetNodeName(), "pertoken quant case only support perTokenScale with dtype float32," + "but perTokenScale[%zu] has data type %s!", i, TypeUtils::DataTypeToAscendString(perTokenScale0Dtype).GetString()), + return GRAPH_FAILED); + } + return GRAPH_SUCCESS; +} + +static graphStatus InferDtype4DavidWeightQuantGMM(gert::InferDataTypeContext *context) +{ + DlinferGroupedMatmulDirectWeightQuantChecker davidWeightQuantGMMChecker; + OP_CHECK_IF(davidWeightQuantGMMChecker.CheckDtype(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckDtype failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidWeightQuantGMMChecker.InferOutDtype(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "SetYDtype failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus InferDtype4DavidQuantGMM(gert::InferDataTypeContext* context) { + DlinferGroupedMatmulDirectQuantChecker davidQuantGMMChecker; + DlinferGroupedMatmulDirectCommonUtil utilForDavidQuantGMM; + OP_CHECK_IF(GetAttrsValue(context, utilForDavidQuantGMM.attrsInfo) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetAttrsValue failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.CheckDtype(context, utilForDavidQuantGMM) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckDtype failed"), return GRAPH_FAILED); + OP_CHECK_IF(davidQuantGMMChecker.InferOutDtype(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "SetYDtype failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4DlinferGroupedMatmulDirect(gert::InferDataTypeContext *context){ + OP_CHECK_NULL_WITH_CONTEXT(context, context); + fe::PlatformInfo platformInfo; + fe::OptionalInfo optionalInfo; + auto ret = fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo); + if (ret == GRAPH_SUCCESS && GmmDavidSupportSoc.count(platformInfo.str_info.short_soc_version) > 0) { + if (IsDavidQuantGMMByShape(context) == GRAPH_SUCCESS) { + OP_CHECK_IF(InferDtype4DavidQuantGMM(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "InferDtype4DavidQuantGMM failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; + } else if (IsDavidWeightQuantGMMByShape(context) == GRAPH_SUCCESS) { + OP_CHECK_IF(InferDtype4DavidWeightQuantGMM(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "InferDtype4DavidWeightQuantGMM failed"), return GRAPH_FAILED); + return GRAPH_SUCCESS; + } + } + OP_CHECK_IF(CheckFunctionParamsForDtype(context) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckFunctionParamsForDtype failed!"), return GRAPH_FAILED); + + auto x0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto weight0Dtype = context->GetDynamicInputDataType(GMM_INDEX_IN_WEIGHT, 0); + size_t numY = context->GetComputeNodeOutputNum(); + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + bool isQuantCase = x0Dtype == ge::DT_INT8 && weight0Dtype == ge::DT_INT8; + bool isA8W4 = x0Dtype == ge::DT_INT8 && weight0Dtype == ge::DT_INT4; + const int64_t* outputDtype = attrs->GetInt(GMM_INDEX_ATTR_OUTPUT_DTYPE); + DataType yDtype = x0Dtype; + if (isQuantCase && outputDtype != nullptr) { + auto it = GMM_OUTPUT_DTYPE_MAP.find(*outputDtype); + OP_CHECK_IF(it == GMM_OUTPUT_DTYPE_MAP.end(), + OP_LOGE(context->GetNodeName(), + "value of attr dtype only supports -1/0/1/2, but now is %ld.", *outputDtype), + return GRAPH_FAILED); + yDtype = it->second; + if (*outputDtype != GMM_OUT_DTYPE_INT32) { // output dtype is int32, this scene does not need scale + OP_CHECK_IF(CheckQuantParamsDtype(context, *outputDtype, yDtype) != GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Check quant params data type failed!"), return GRAPH_FAILED); + } + } + if (isA8W4 && outputDtype != nullptr) { + auto it = GMM_OUTPUT_DTYPE_MAP.find(*outputDtype); + OP_CHECK_IF(it == GMM_OUTPUT_DTYPE_MAP.end(), + OP_LOGE(context->GetNodeName(), + "value of attr dtype only supports -1/0/1/2, but now is %ld.", *outputDtype), + return GRAPH_FAILED); + yDtype = it->second; + } + for (size_t k = 0; k < numY; k++) {context->SetOutputDataType(GMM_INDEX_OUT_Y + k, yDtype);} + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(DlinferGroupedMatmulDirect) + .InferShape(InferShape4DlinferGroupedMatmulDirect) + .InferDataType(InferDataType4DlinferGroupedMatmulDirect); +} // namespace ops diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_common_util.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_common_util.h new file mode 100644 index 00000000..a330a57d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_common_util.h @@ -0,0 +1,120 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef GROUPED_MATMUL_INFERSHAPE_COMMON_H_ +#define GROUPED_MATMUL_INFERSHAPE_COMMON_H_ + +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "platform/platform_info.h" + +namespace ops { + +constexpr size_t GMM_INDEX_IN_X = 0UL; +constexpr size_t GMM_INDEX_IN_WEIGHT = 1UL; +constexpr size_t GMM_INDEX_IN_BIAS = 2UL; +constexpr size_t GMM_INDEX_IN_SCALE = 3UL; +constexpr size_t GMM_INDEX_IN_OFFSET = 4UL; +constexpr size_t GMM_INDEX_IN_ANTIQUANT_SCALE = 5UL; +constexpr size_t GMM_INDEX_IN_ANTIQUANT_OFFSET = 6UL; +constexpr size_t GMM_INDEX_IN_GROUP_LIST = 7UL; +constexpr size_t GMM_INDEX_IN_PERTOKEN_SCALE = 8UL; + +constexpr size_t GMM_INDEX_OUT_Y = 0UL; +constexpr size_t GMM_INDEX_ATTR_SPLIT_ITEM = 0UL; +constexpr size_t GMM_INDEX_ATTR_OUTPUT_DTYPE = 1UL; +constexpr size_t GMM_INDEX_ATTR_TRANSPOSE_W = 2UL; +constexpr size_t GMM_INDEX_ATTR_TRANSPOSE_X = 3UL; +constexpr size_t GMM_INDEX_ATTR_GROUP_TYPE = 4UL; +constexpr size_t GMM_INDEX_ATTR_GROUP_LIST_TYPE = 5UL; +constexpr size_t GMM_INDEX_ATTR_ACT_TYPE = 6UL; +constexpr size_t GMM_INDEX_ATTR_TUNING_CONFIG = 7UL; + +constexpr int64_t GMM_X_Y_SEPARATED = 0; // x,y have been separated +constexpr int64_t GMM_Y_SEPARATED = 1; // y has been separated +constexpr int64_t GMM_X_SEPARATED = 2; // x has been separated +constexpr int64_t GMM_NO_SEPARATED = 3; // x,y have not been separated +constexpr int64_t GMM_NO_SPLIT = -1L; +constexpr int64_t GMM_SPLIT_M = 0L; +constexpr int64_t GMM_SPLIT_N = 1L; +constexpr int64_t GMM_SPLIT_K = 2L; + +constexpr int64_t GMM_OUT_DTYPE_INT32 = 2L; + +constexpr int64_t GMM_MAX_GROUP_LIST_SIZE_ARRAY = 128L; +constexpr int64_t GMM_MAX_GROUP_LIST_SIZE_TENSOR = 2560L; +constexpr int64_t GMM_MAX_INNER_AXIS = 65535L; +constexpr int64_t GMM_N_K_ALIGN_VALUE_WEIGHT_QUANT = 32L; +constexpr int64_t GMM_N_K_ALIGN_VALUE_WEIGHT_QUANT_4BIT = 64L; +constexpr size_t GMM_MAX_FM_DIM = 6UL; +constexpr size_t GMM_MIN_FM_DIM = 2UL; +constexpr size_t GMM_MIN_WEIGHT_DIM = 2UL; +constexpr size_t GMM_SEPARATED_WEIGHT_DIM = 2UL; +constexpr size_t GMM_SPLIT_M_SINGLE_WEIGHT_DIM = 3UL; +constexpr size_t GMM_SPLIT_K_SINGLE_WEIGHT_DIM = 2UL; +constexpr size_t PENULTIMATE_DIM = 2UL; +constexpr int64_t PERTILE_GROUP_SIZE = 128UL; +constexpr size_t GMM_A8W4_OFFSET_DIM_NUM = 3UL; +constexpr size_t GMM_A8W4_BIAS_DIM_NUM = 2UL; + +constexpr int64_t MXFP_DIVISOR_SIZE = 64; +constexpr int64_t MXFP_MULTI_BASE_SIZE = 2; +constexpr int64_t MXFP_TYPEM_SCALE_DIM_NUM = 4; +constexpr int64_t MXFP_TYPEK_SCALE_DIM_NUM = 3; +struct GMMAttrs { + int64_t splitItem; + int64_t outputDtype; + int64_t groupType; + bool transposeX; + bool transposeWeight; + int64_t activeType; + int64_t tuningConfig; +}; + +struct GMMInputParamsInfo { + size_t numX; + size_t numWeight; + size_t numBias; + size_t numScale; + size_t numOffset; + size_t numAntiquantScale; + size_t numAntiquantOffset; +}; + +const std::map GMM_OUTPUT_DTYPE_MAP = {{0, ge::DataType::DT_FLOAT16}, + {1, ge::DataType::DT_BF16}, + {2, ge::DataType::DT_INT32}, + {3, ge::DataType::DT_FLOAT}, + {-1, ge::DataType::DT_INT8}}; + +enum class GMMActType : int64_t { + GMM_ACT_TYPE_NONE, + GMM_ACT_TYPE_RELU, + GMM_ACT_TYPE_GELU_TANH, + GMM_ACT_TYPE_GELU_ERR_FUNC, + GMM_ACT_TYPE_FAST_GELU, + GMM_ACT_TYPE_SILU, + END_ACT_TYPE_ENUM +}; + +const std::initializer_list BIAS_DTYPE_SUPPORT_LIST = {ge::DataType::DT_FLOAT, ge::DataType::DT_FLOAT16, + ge::DataType::DT_BF16}; + +class DlinferGroupedMatmulDirectCommonUtil { +public: + GMMAttrs attrsInfo; + DlinferGroupedMatmulDirectCommonUtil(){}; + ~DlinferGroupedMatmulDirectCommonUtil(){}; + +private: +}; +} // namespace ops + +#endif // GROUPED_MATMUL_INFERSHAPE_COMMON_H_ diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.cpp new file mode 100644 index 00000000..6c4b7429 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.cpp @@ -0,0 +1,687 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_infershape_quant_checker.cpp + * \brief + */ +#include "grouped_matmul_infershape_common_util.h" +#include "grouped_matmul_infershape_quant_checker.h" + +namespace ops { + +static const std::unordered_set X_TYPE_SUPPORT_SET = {ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_HIFLOAT8, ge::DT_FLOAT4_E1M2, ge::DT_FLOAT4_E2M1}; +static const std::unordered_set WEIGHT_TYPE_SUPPORT_SET = {ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, ge::DT_HIFLOAT8, ge::DT_FLOAT4_E1M2, ge::DT_FLOAT4_E2M1}; +static const std::unordered_set BIAS_TYPE_SUPPORT_SET = {ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, + ge::DT_INT32}; +static const std::unordered_set SCALE_TYPE_SUPPORT_SET = {ge::DT_UINT64, ge::DT_INT64, ge::DT_FLOAT, + ge::DT_BF16, ge::DT_FLOAT8_E8M0}; +static const std::unordered_set PERTOEKN_SCALE_TYPE_SUPPORT_SET = {ge::DT_FLOAT, ge::DT_FLOAT8_E8M0}; + +static bool inline IsNonEmpty(const gert::Shape *shape) +{ + return (shape != nullptr && !(shape->GetDimNum() == 1 && shape->GetDim(0) == 0)); +} + +bool DlinferGroupedMatmulDirectQuantChecker::LogicXOR(bool cond1, bool cond2) const +{ + uint64_t result = static_cast(cond1) ^ static_cast(cond2); + return static_cast(result); +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::GetXAndWeightDimValue(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) +{ + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_IF(!(IsNonEmpty(xShape) && IsNonEmpty(weightShape)), + OP_LOGE(context->GetNodeName(), "The 1st tensor of tensor list x and weight cannot be empty."), + return ge::GRAPH_FAILED); + xdimNum_ = xShape->GetDimNum(); + weightdimNum_ = weightShape->GetDimNum(); + OP_CHECK_IF(xdimNum_ < GMM_MIN_FM_DIM, + OP_LOGE(context->GetNodeName(), "The dim num of x's 1st tensor should be greater than 1, but it is \ +[%zu].", + xdimNum_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightdimNum_ < GMM_MIN_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), "The dim num of weight's 1st tensor should be greater than 1, \ +but it is [%zu].", + weightdimNum_), + return ge::GRAPH_FAILED); + xMDim_ = gmmAttrs.transposeX ? xShape->GetDim(xdimNum_ - 1) : xShape->GetDim(xdimNum_ - PENULTIMATE_DIM); + xKDim_ = gmmAttrs.transposeX ? xShape->GetDim(xdimNum_ - PENULTIMATE_DIM) : xShape->GetDim(xdimNum_ - 1); + weightKDim_ = gmmAttrs.transposeWeight ? weightShape->GetDim(weightdimNum_ - 1) : + weightShape->GetDim(weightdimNum_ - PENULTIMATE_DIM); + weightNDim_ = gmmAttrs.transposeWeight ? weightShape->GetDim(weightdimNum_ - PENULTIMATE_DIM) : + weightShape->GetDim(weightdimNum_ - 1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckDtypeValid(const gert::InferDataTypeContext *context) const +{ + auto xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto weightDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_WEIGHT, 0); + // mandory param dtype check + OP_CHECK_IF(X_TYPE_SUPPORT_SET.find(xDtype) == X_TYPE_SUPPORT_SET.end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for x's 1st tensor.", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(WEIGHT_TYPE_SUPPORT_SET.find(weightDtype) == WEIGHT_TYPE_SUPPORT_SET.end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for weight.", + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + if (xDtype == ge::DataType::DT_INT8 || weightDtype == ge::DataType::DT_INT8) { + OP_CHECK_IF( + !(xDtype == ge::DataType::DT_INT8 && weightDtype == ge::DataType::DT_INT8), + OP_LOGE( + context->GetNodeName(), + "When data type of x is int8, data type of weight should also be int8, vice versa. But data type of \ +x is [%s] and the data type of weight is [%s].", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF( + LogicXOR((xDtype == ge::DataType::DT_HIFLOAT8), (weightDtype == ge::DataType::DT_HIFLOAT8)), + OP_LOGE(context->GetNodeName(), + "When one input dtype is HIFLOAT8, then the other input dtype must be HIFLOAT8, vice versa, actual \ +x is %s, weight is %s.", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + LogicXOR((xDtype == ge::DataType::DT_FLOAT8_E4M3FN || xDtype == ge::DataType::DT_FLOAT8_E5M2), + (weightDtype == ge::DataType::DT_FLOAT8_E4M3FN || weightDtype == ge::DataType::DT_FLOAT8_E5M2)), + OP_LOGE( + context->GetNodeName(), + "When x input dtype is FLOAT8, then the weight input dtype must be FLOAT8, vice versa, actual x is %s, \ +weight is %s.", ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(LogicXOR((xDtype == ge::DataType::DT_FLOAT4_E1M2 || xDtype == ge::DataType::DT_FLOAT4_E2M1), + (weightDtype == ge::DataType::DT_FLOAT4_E1M2 || weightDtype == ge::DataType::DT_FLOAT4_E2M1)), + OP_LOGE(context->GetNodeName(), + "When x input dtype is FLOAT4, then the weight input dtype must be FLOAT4, vice versa, actual x is %s, \ +weight is %s.", ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + auto ScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_SCALE, 0); + OP_CHECK_IF(SCALE_TYPE_SUPPORT_SET.find(ScaleDtype) == SCALE_TYPE_SUPPORT_SET.end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for scale.", + ge::TypeUtils::DataTypeToAscendString(ScaleDtype).GetString()), + return ge::GRAPH_FAILED); + + auto biasDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_BIAS, 0); + OP_CHECK_IF(BIAS_TYPE_SUPPORT_SET.find(biasDtype) == BIAS_TYPE_SUPPORT_SET.end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for bias.", + ge::TypeUtils::DataTypeToAscendString(biasDtype).GetString()), + return ge::GRAPH_FAILED); + + auto pertokenScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + if (pertokenScaleDtype != ge::DT_UNDEFINED) { + OP_CHECK_IF(PERTOEKN_SCALE_TYPE_SUPPORT_SET.find(pertokenScaleDtype) == PERTOEKN_SCALE_TYPE_SUPPORT_SET.end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for pertokenScale.", + ge::TypeUtils::DataTypeToAscendString(pertokenScaleDtype).GetString()), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckNotZeroValueForNoneSplitAxis(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + if (gmmAttrs.groupType == GMM_SPLIT_M) { + auto eDimValue = weightShape->GetDim(0); + OP_CHECK_IF(xKDim_ == 0 || eDimValue == 0 || weightNDim_ == 0, + OP_LOGE(context->GetNodeName(), + "Non split axis' dim value cannot be zero, but kDimValue is [%ld], eDimValue is [%ld], \ +nDimValue is [%ld].", + xKDim_, eDimValue, weightNDim_), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(xMDim_ == 0 || weightNDim_ == 0, + OP_LOGE(context->GetNodeName(), + "Non split axis' dim value cannot be zero but mDimValue is [%ld], nDimValue is [%ld].", + xMDim_, weightNDim_), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckScenarioValidForShape(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + // only support for single/single/single Scenario + auto xSecondTensorShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 1); + auto weightSecondTensorShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 1); + OP_CHECK_IF(IsNonEmpty(xSecondTensorShape), + OP_LOGE(context->GetNodeName(), "Only support single/single/single scenario for now, but the second \ +tensor of tensor list x is not empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(IsNonEmpty(weightSecondTensorShape), + OP_LOGE(context->GetNodeName(), "Only support single/single/single scenario for now, but the second \ +tensor of tensor list weight is not empty."), + return ge::GRAPH_FAILED); + // check split item value valid + OP_CHECK_IF(gmmAttrs.splitItem != GMM_X_SEPARATED && gmmAttrs.splitItem != GMM_NO_SEPARATED, + OP_LOGE(context->GetNodeName(), "Invalid splitItem, which can only be one of 2/3, but it is [%ld].", + gmmAttrs.splitItem), + return ge::GRAPH_FAILED); + OP_CHECK_IF(gmmAttrs.groupType != GMM_SPLIT_M && gmmAttrs.groupType != GMM_SPLIT_K, + OP_LOGE(context->GetNodeName(), "Invalid groupType, which can only be one of 0/2, but it is [%ld].", + gmmAttrs.groupType), + return ge::GRAPH_FAILED); + // check activation is null + OP_CHECK_IF(gmmAttrs.activeType != static_cast(GMMActType::GMM_ACT_TYPE_NONE), + OP_LOGE(context->GetNodeName(), "Activation function is not supported in quant mode now."), + return ge::GRAPH_FAILED); + // non split axis cannot be empty tensor + OP_CHECK_IF(CheckNotZeroValueForNoneSplitAxis(context, gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckNotZeroValueForNoneSplitAxis Failed."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckFormatValid(const gert::InferShapeContext *context) const +{ + const auto xDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + const auto xFormat = xDesc->GetOriginFormat(); + OP_CHECK_IF(xFormat != ge::FORMAT_ND && xFormat != ge::FORMAT_NCL && xFormat != ge::FORMAT_NCHW, + OP_LOGE(context->GetNodeName(), "Format of x only supports ND, NCL or NCHW for now, but it is [%s].", + ge::TypeUtils::FormatToAscendString(xFormat).GetString()), + return ge::GRAPH_FAILED); + const auto weightDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc); + const auto weightFormat = weightDesc->GetOriginFormat(); + OP_CHECK_IF(weightFormat != ge::FORMAT_ND && weightFormat != ge::FORMAT_NCL && weightFormat != ge::FORMAT_NCHW, + OP_LOGE(context->GetNodeName(), + "Format of weight only supports ND, NCL or NCHW for now, but it is [%s].", + ge::TypeUtils::FormatToAscendString(weightFormat).GetString()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShapeForBias(const gert::InferShapeContext *context) const +{ + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, 0); + if (IsNonEmpty(biasShape)) { + size_t biasDimNum = biasShape->GetDimNum(); + OP_CHECK_IF(biasDimNum != 2, + OP_LOGE(context->GetNodeName(), + "When bias is not null, its dim should be 2, but the actual is [%zu].", biasDimNum), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + biasShape->GetDim(0) != groupNum_ || biasShape->GetDim(1) != static_cast(weightNDim_), + OP_LOGE(context->GetNodeName(), + "The shape of bias should be (g, n), which is ([%ld], [%ld]), but the actual is ([%ld], [%ld]).", + groupNum_, weightNDim_, biasShape->GetDim(0), biasShape->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +bool DlinferGroupedMatmulDirectQuantChecker::IsDoubleScaleScenario(const gert::Shape *scaleShape, + const gert::Shape *perTokenScaleShape) const +{ // pertoken dim num support 2 or 1 for now + if (perTokenScaleShape->GetDimNum() != 2 && perTokenScaleShape->GetDimNum() != 1) { + return false; + } + bool IsScaleMatchedDoubleScale = false; + bool IsPertokenScaleMatchedDoubleScale = false; + if (scaleShape->GetDimNum() == 2) { // dim num 2 with shape (e, 1) + IsScaleMatchedDoubleScale = (scaleShape->GetDim(0) == groupNum_) && (scaleShape->GetDim(1) == 1); + } else { + IsScaleMatchedDoubleScale = scaleShape->GetDim(0) == groupNum_; + } + + if (perTokenScaleShape->GetDimNum() == 2) { // dim num 2 with shape (e, 1) + IsPertokenScaleMatchedDoubleScale = + (perTokenScaleShape->GetDim(0) == groupNum_) && (perTokenScaleShape->GetDim(1) == 1); + } else { + IsPertokenScaleMatchedDoubleScale = perTokenScaleShape->GetDim(0) == groupNum_; + } + return IsScaleMatchedDoubleScale && IsPertokenScaleMatchedDoubleScale; +} + +bool DlinferGroupedMatmulDirectQuantChecker::IsPerTileQuantMode(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto perTokenScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + bool isPerTileQuantMode = false; + if (!IsNonEmpty(perTokenScaleShape)) { + return isPerTileQuantMode; + } + auto perTokenScaleDimNum = perTokenScaleShape->GetDimNum(); + // 2 is the minimum dimension for perTokenScale in perTile case + if (perTokenScaleDimNum < 2 || perTokenScaleDimNum != xdimNum_) { + return false; + } + auto perTokenMDim = gmmAttrs.transposeX ? perTokenScaleShape->GetDim(xdimNum_ - 1) : + perTokenScaleShape->GetDim(xdimNum_ - PENULTIMATE_DIM); + if (perTokenMDim != xMDim_) { + return false; + } + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + auto scaleDimNum = scaleShape->GetDimNum(); + // 2 is the minimum dimension for scale in perTile case + if (scaleDimNum < 2 || scaleDimNum != weightdimNum_) { + return false; + } + auto scaleKDim = gmmAttrs.transposeWeight ? scaleShape->GetDim(scaleDimNum - 1) : + scaleShape->GetDim(scaleDimNum - PENULTIMATE_DIM); + auto scaleNDim = gmmAttrs.transposeWeight ? scaleShape->GetDim(scaleDimNum - PENULTIMATE_DIM) : + scaleShape->GetDim(scaleDimNum - 1); + bool isKdimValid = + ((gmmAttrs.groupType == GMM_SPLIT_K && scaleKDim == (weightKDim_ / PERTILE_GROUP_SIZE) + groupNum_) || + (gmmAttrs.groupType == GMM_SPLIT_M && + scaleKDim == (weightKDim_ + PERTILE_GROUP_SIZE - 1) / PERTILE_GROUP_SIZE)); + if (scaleNDim == (weightNDim_ + PERTILE_GROUP_SIZE - 1) / PERTILE_GROUP_SIZE && isKdimValid) { + isPerTileQuantMode = true; + } + return isPerTileQuantMode; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckDlinferGroupedMatmulDirectPerGroupDim(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto scaleDimNum = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0)->GetDimNum(); + auto perTokenDimNum = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0)->GetDimNum(); + OP_CHECK_IF(xdimNum_ != GMM_MIN_FM_DIM, + OP_LOGE(context->GetNodeName(), "The dim num of x should be 2, but actual dim num is %zu.", xdimNum_), + return ge::GRAPH_FAILED); + if (gmmAttrs.groupType == GMM_SPLIT_M) { + OP_CHECK_IF(weightdimNum_ != GMM_SPLIT_M_SINGLE_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "The dim num of weight should be 3 when groupType is 0 (split M), but \ +actual dim num is %zu.", + weightdimNum_), + return ge::GRAPH_FAILED); + } else if (gmmAttrs.groupType == GMM_SPLIT_K) { + OP_CHECK_IF(weightdimNum_ != GMM_SPLIT_K_SINGLE_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "The dim num of weight should be 2 when groupType is 2 (split K), but \ +actual dim num is %zu.", + weightdimNum_), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(scaleDimNum != weightdimNum_, + OP_LOGE(context->GetNodeName(), "The dim num of scale[%zu] should be equal to that of weight[%zu] when \ +groupType is %ld.", + scaleDimNum, weightdimNum_, gmmAttrs.groupType), + return ge::GRAPH_FAILED); + OP_CHECK_IF(perTokenDimNum != xdimNum_, + OP_LOGE(context->GetNodeName(), + "The dim num of per_token_scale[%zu] should be equal to that of x[%zu].", perTokenDimNum, + xdimNum_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckDlinferGroupedMatmulDirectPerTileShape(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + auto perTokenShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + auto perTokenMDim = + gmmAttrs.transposeX ? perTokenShape->GetDim(xdimNum_ - 1) : perTokenShape->GetDim(xdimNum_ - PENULTIMATE_DIM); + auto perTokenKDim = + gmmAttrs.transposeX ? perTokenShape->GetDim(xdimNum_ - PENULTIMATE_DIM) : perTokenShape->GetDim(xdimNum_ - 1); + auto scaleKDim = gmmAttrs.transposeWeight ? scaleShape->GetDim(weightdimNum_ - 1) : + scaleShape->GetDim(weightdimNum_ - PENULTIMATE_DIM); + auto scaleNDim = gmmAttrs.transposeWeight ? scaleShape->GetDim(weightdimNum_ - PENULTIMATE_DIM) : + scaleShape->GetDim(weightdimNum_ - 1); + OP_CHECK_IF(perTokenMDim != xMDim_, + OP_LOGE(context->GetNodeName(), + "When quantification mode is G-B quantification, the M value in x[%ld] and \ +per_token_scale[%ld] should be consistent.", + xMDim_, perTokenMDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(scaleNDim != (weightNDim_ + PERTILE_GROUP_SIZE - 1) / PERTILE_GROUP_SIZE, + OP_LOGE(context->GetNodeName(), + "When quantification mode is G-B quantification, the N value in scale [%ld] \ +must be equal to the N value in weight [%ld] divided by 128.", + scaleNDim, weightNDim_), + return ge::GRAPH_FAILED); + if (gmmAttrs.groupType == GMM_SPLIT_M) { + int64_t expectScaleKValue = (weightKDim_ + PERTILE_GROUP_SIZE - 1) / PERTILE_GROUP_SIZE; + OP_CHECK_IF(perTokenKDim != scaleKDim || scaleKDim != expectScaleKValue, + OP_LOGE(context->GetNodeName(), + "When quantification mode is G-B quantification, and groupType is 0 (split \ +M), the K dim of per_token_scale [%ld] should equal the K dim of scalel [%ld], and its value should be equal to the K \ +dim of weight [%ld] divided by 128, rounded up to the next integer.", + perTokenKDim, scaleNDim, weightNDim_), + return ge::GRAPH_FAILED); + } else { + int64_t expectScaleKValue = (weightKDim_ / PERTILE_GROUP_SIZE) + groupNum_; + OP_CHECK_IF(perTokenKDim != scaleKDim || scaleKDim != expectScaleKValue, + OP_LOGE(context->GetNodeName(), + "When quantification mode is G-B quantification, and groupType is 2 (split \ +K), the K dim of per_token_scale [%ld] should equal the K dim of scale [%ld], its value must be equal to the K dim of \ +weight [%ld] divided by 128, plus the groupSize [%ld].", + perTokenKDim, scaleKDim, weightKDim_, groupNum_), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShapeForPerGroupQuantParam(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + if (IsPerTileQuantMode(context, gmmAttrs)) { + OP_CHECK_IF(CheckDlinferGroupedMatmulDirectPerGroupDim(context, gmmAttrs) != ge::GRAPH_SUCCESS || + CheckDlinferGroupedMatmulDirectPerTileShape(context, gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForPerGroupQuantParam failed"), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckPertokenShapeInNormalQuantMode( + const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs, const gert::Shape *perTokenScaleShape) const +{ + OP_CHECK_IF(perTokenScaleShape->GetDimNum() != 1 && perTokenScaleShape->GetDimNum() != 2, + OP_LOGE(context->GetNodeName(), + "In T-C or K-C quant mode, the dim num of perTokenScale should be 1 or 2, \ +but the actual is [%zu].", + perTokenScaleShape->GetDimNum()), + return ge::GRAPH_FAILED); + if (gmmAttrs.groupType == GMM_SPLIT_M) { + if (perTokenScaleShape->GetDimNum() == 1) { + OP_CHECK_IF(perTokenScaleShape->GetDim(0) != xMDim_ && perTokenScaleShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "When perTokenScale dim num is 1 in split m scenario, the expected shape of \ +perTokenScale is (m,) or (g,), which is (%ld,) or (%ld,), but the actual shape is (%ld,).", + xMDim_, groupNum_, perTokenScaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(perTokenScaleShape->GetDim(0) != groupNum_ || perTokenScaleShape->GetDim(1) != 1, + OP_LOGE(context->GetNodeName(), + "When perTokenScale dim num is 2 in split m scenario, the expected shape of \ +perTokenScale is (g,1), which is (%ld,1), but the actual shape is (%ld, %ld).", + groupNum_, perTokenScaleShape->GetDim(0), perTokenScaleShape->GetDim(1)), + return ge::GRAPH_FAILED); + } + } else { + // split k + if (perTokenScaleShape->GetDimNum() == 1) { + OP_CHECK_IF(perTokenScaleShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "When perTokenScale dim num is 1 in split k scenario, the expected shape of \ +perTokenScale is (g,), which is (%ld,), but the actual shape is (%ld,).", + groupNum_, perTokenScaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF( + perTokenScaleShape->GetDim(0) != groupNum_ || + (perTokenScaleShape->GetDim(1) != xMDim_ && perTokenScaleShape->GetDim(1) != 1), + OP_LOGE( + context->GetNodeName(), + "When perTokenScale dim num is 2 in split k scenario, the expected shape of perTokenScale is (g,m) \ + or (g,1), which is (%ld,%ld) or (%ld,1), but the actual shape is (%ld,%ld).", + groupNum_, xMDim_, groupNum_, perTokenScaleShape->GetDim(0), perTokenScaleShape->GetDim(1)), + return ge::GRAPH_FAILED); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShapeForQuantParam(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + auto perTokenScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + if (IsPerTileQuantMode(context, gmmAttrs)) { + return CheckShapeForPerGroupQuantParam(context, gmmAttrs); + } + // mx在tiling中校验异常场景 + if (scaleShape->GetDimNum() == MXFP_TYPEM_SCALE_DIM_NUM || scaleShape->GetDimNum() == MXFP_TYPEK_SCALE_DIM_NUM) { + return ge::GRAPH_SUCCESS; + } + + OP_CHECK_IF(scaleShape->GetDimNum() != 2 && scaleShape->GetDimNum() != 1, + OP_LOGE(context->GetNodeName(), "The dim num of scale should be 1 or 2, but the actual is [%zu].", + scaleShape->GetDimNum()), + return ge::GRAPH_FAILED); + if (scaleShape->GetDimNum() == 1) { + OP_CHECK_IF(scaleShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "The 1st dim value of scale should be g[%ld], but the actual 1st dim value is [%ld].", + groupNum_, scaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF( + scaleShape->GetDim(0) != groupNum_ || (scaleShape->GetDim(1) != 1 && scaleShape->GetDim(1) != weightNDim_), + OP_LOGE(context->GetNodeName(), + "The 1st dim value of scale should be g[%ld] and 2nd dim value of scale should be 1 or n[%ld] \ +, but the actual 1st dim value is [%ld], and the 2nd dim value is [%ld].", + groupNum_, weightNDim_, scaleShape->GetDim(0), scaleShape->GetDim(1)), + return ge::GRAPH_FAILED); + } + if (IsNonEmpty(perTokenScaleShape)) { + if (IsDoubleScaleScenario(scaleShape, perTokenScaleShape)) { + return ge::GRAPH_SUCCESS; + } + if (CheckPertokenShapeInNormalQuantMode(context, gmmAttrs, perTokenScaleShape) == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::GetGroupNumValue(const gert::InferShapeContext *context) +{ + auto groupListShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListShape); + OP_CHECK_IF(groupListShape->GetDimNum() != 1, + OP_LOGE(context->GetNodeName(), + "The groupList only support 1 dim num for now, but the actual is [%zu].", + groupListShape->GetDimNum()), + return ge::GRAPH_FAILED); + groupNum_ = groupListShape->GetDim(0); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShapeForGrouplist(const gert::InferShapeContext *context) const +{ + OP_CHECK_IF(groupNum_ <= 0, + OP_LOGE(context->GetNodeName(), + "The groupList 1st dim value should be greater than 0, but the actual is [%ld].", groupNum_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupNum_ > 1024, + OP_LOGE(context->GetNodeName(), + "Only support 1024 groups at MAX for now, but the actual group number is [%ld].", groupNum_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShapeValid(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_IF( + xKDim_ != weightKDim_, + OP_LOGE(context->GetNodeName(), + "The k dim of x should be equal to the k dim of weight, but x's k is [%ld] and weight's k is [%ld].", + xKDim_, weightKDim_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeForGrouplist(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForGrouplist failed."), return ge::GRAPH_FAILED); + if (gmmAttrs.groupType == GMM_SPLIT_M) { + OP_CHECK_IF(xdimNum_ != 2 && weightdimNum_ != 3, + OP_LOGE(context->GetNodeName(), + "When split m, x dim num should be 2, weight dim num should be 3, y dim num should be 2 \ +but the actual x dim num is [%zu], actual weight dim num is [%zu].", + xdimNum_, weightdimNum_), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + weightShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "When split m, 1st dim value of weight should be g, which is [%ld], but actual value is [%ld].", + groupNum_, weightShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(xdimNum_ != 2 && weightdimNum_ != 2, + OP_LOGE(context->GetNodeName(), + "When split k, x dim num should be 2, weight dim num should be 2, y dim num should be 3 \ +but the actual x dim num is [%zu], actual weight dim num is [%zu].", + xdimNum_, weightdimNum_), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(CheckShapeForBias(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForBias failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(!IsNonEmpty(context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0)), + OP_LOGE(context->GetNodeName(), "The 1st tensor of scale cannot be empty."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeForQuantParam(context, gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForQuantParam failed."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +bool DlinferGroupedMatmulDirectQuantChecker::CheckDataContainsNegativeValue() const +{ + return xMDim_ < 0L || xKDim_ < 0L || weightNDim_ < 0L || weightKDim_ < 0L; +} + +bool DlinferGroupedMatmulDirectQuantChecker::CheckShapeContainsNegativeValue(const gert::Shape *shape) const +{ + if (shape == nullptr) { + return false; + } + size_t dimNum = shape->GetDimNum(); + for (size_t i = 0; i < dimNum; i++) { + if (shape->GetDim(i) < 0L) { + return true; + } + } + return false; +} + +bool DlinferGroupedMatmulDirectQuantChecker::CheckNonDataInputContainsNegativeValue(const gert::InferShapeContext *context) const +{ + if (groupNum_ < 0L) { + return true; + } + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, 0); + if (IsNonEmpty(biasShape) && CheckShapeContainsNegativeValue(biasShape)) { + return true; + } + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + if (IsNonEmpty(scaleShape) && CheckShapeContainsNegativeValue(scaleShape)) { + return true; + } + auto pertokenScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + if (IsNonEmpty(pertokenScaleShape) && CheckShapeContainsNegativeValue(pertokenScaleShape)) { + return true; + } + return false; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckShape(const gert::InferShapeContext *context, + const DlinferGroupedMatmulDirectCommonUtil &commonUtil) const +{ + if (CheckDataContainsNegativeValue() || CheckNonDataInputContainsNegativeValue(context)) { + return ge::GRAPH_SUCCESS; + } + OP_CHECK_IF(CheckScenarioValidForShape(context, commonUtil.attrsInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckScenarioValidForShape failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeValid(context, commonUtil.attrsInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeValid failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckFormatValid(context) != ge::GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), "CheckFormatValid failed."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::UpdateShapeY(gert::InferShapeContext *context, size_t idxY, + std::vector &yDims) +{ + gert::Shape *yShape = context->GetOutputShape(idxY); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + yShape->SetDimNum(yDims.size()); + for (size_t dim = 0; dim < yDims.size(); ++dim) { + yShape->SetDim(dim, yDims[dim]); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::InferOutShape(gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) +{ + std::vector yDims = {xMDim_, weightNDim_}; + if (gmmAttrs.groupType == GMM_SPLIT_K) { + yDims.insert(yDims.begin(), groupNum_); + } + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y, yDims) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update y shape."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckScenarioValidForDtype(const gert::InferDataTypeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto scaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_SCALE, 0); + auto perTokenScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + + if (gmmAttrs.groupType == GMM_SPLIT_K) { + OP_CHECK_IF(xDtype == ge::DT_INT8 || scaleDtype == ge::DT_INT64 || scaleDtype == ge::DT_UINT64, + OP_LOGE(context->GetNodeName(), "When split k, x with int8 dtype or scale with int64/uint64 dtype \ +is not supported, but the acutal x dtype is [%s] and actual scale dtype is [%s].", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(scaleDtype).GetString()), + return ge::GRAPH_FAILED); + } + if ((xDtype == ge::DT_HIFLOAT8 || xDtype == ge::DT_FLOAT8_E5M2 || + xDtype == ge::DT_FLOAT8_E4M3FN) && perTokenScaleDtype != ge::DT_UNDEFINED) { + OP_CHECK_IF( + scaleDtype != perTokenScaleDtype || (scaleDtype != ge::DataType::DT_FLOAT && scaleDtype != ge::DataType::DT_FLOAT8_E8M0), + OP_LOGE(context->GetNodeName(), "When data type of x is float8/hifloat8, data type of scale [%s] should be \ +equal to per_token_scale's dtype [%s], and be float32/float8_e8m0.", + ge::TypeUtils::DataTypeToAscendString(scaleDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(perTokenScaleDtype).GetString()), return ge::GRAPH_FAILED); + } else if (xDtype == ge::DT_FLOAT4_E1M2 || xDtype == ge::DT_FLOAT4_E2M1) { + OP_CHECK_IF( + scaleDtype != ge::DataType::DT_FLOAT8_E8M0, + OP_LOGE(context->GetNodeName(), "When data type of x is float4, data type of scale [%s] should be \ +equal to float8_e8m0.", + ge::TypeUtils::DataTypeToAscendString(scaleDtype).GetString()), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::CheckDtype(const gert::InferDataTypeContext *context, + const DlinferGroupedMatmulDirectCommonUtil &commonUtil) const +{ + OP_CHECK_IF(CheckDtypeValid(context) != ge::GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), "CheckDtypeValid failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckScenarioValidForDtype(context, commonUtil.attrsInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckScenarioValidForDtype failed."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectQuantChecker::InferOutDtype(gert::InferDataTypeContext *context) +{ + auto attrs = context->GetAttrs(); + const int64_t *outputDtype = attrs->GetInt(GMM_INDEX_ATTR_OUTPUT_DTYPE); + ge::DataType yDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto it = GMM_OUTPUT_DTYPE_MAP.find(*outputDtype); + OP_CHECK_IF(it == GMM_OUTPUT_DTYPE_MAP.end(), + OP_LOGE(context->GetNodeName(), "The output dtype should be int8, bfloat16, int32, float16 or float32 \ +, but the actual output dtype is [%s].", + ge::TypeUtils::DataTypeToAscendString(yDtype).GetString()), + return ge::GRAPH_FAILED); + yDtype = it->second; + context->SetOutputDataType(GMM_INDEX_OUT_Y, yDtype); + return ge::GRAPH_SUCCESS; +} +} // namespace ops \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.h new file mode 100644 index 00000000..d0e3c358 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_quant_checker.h @@ -0,0 +1,74 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef GROUPED_MATMUL_INFERSHAPE_QUANT_CHECKER_H_ +#define GROUPED_MATMUL_INFERSHAPE_QUANT_CHECKER_H_ + +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "platform/platform_info.h" +#include "graph/utils/type_utils.h" +#include "grouped_matmul_infershape_common_util.h" + +namespace ops { + +class DlinferGroupedMatmulDirectQuantChecker { +public: + DlinferGroupedMatmulDirectQuantChecker(){}; + ~DlinferGroupedMatmulDirectQuantChecker(){}; + ge::graphStatus GetXAndWeightDimValue(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs); + ge::graphStatus CheckShape(const gert::InferShapeContext *context, const DlinferGroupedMatmulDirectCommonUtil &commonUtil) const; + ge::graphStatus CheckDtype(const gert::InferDataTypeContext *context, + const DlinferGroupedMatmulDirectCommonUtil &commonUtil) const; + ge::graphStatus InferOutShape(gert::InferShapeContext *context, const GMMAttrs &gmmAttrs); + ge::graphStatus InferOutDtype(gert::InferDataTypeContext *context); + ge::graphStatus GetGroupNumValue(const gert::InferShapeContext *context); +private: + ge::graphStatus CheckFormatValid(const gert::InferShapeContext *context) const; + ge::graphStatus CheckDtypeValid(const gert::InferDataTypeContext *context) const; + ge::graphStatus CheckScenarioValidForShape(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckScenarioValidForDtype(const gert::InferDataTypeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckShapeValid(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckShapeForBias(const gert::InferShapeContext *context) const; + ge::graphStatus CheckShapeForQuantParam(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckNotZeroValueForNoneSplitAxis(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckShapeForGrouplist(const gert::InferShapeContext *context) const; + ge::graphStatus UpdateShapeY(gert::InferShapeContext *context, size_t idxY, std::vector &yDims); + ge::graphStatus CheckDlinferGroupedMatmulDirectPerGroupDim(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckShapeForPerGroupQuantParam(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckDlinferGroupedMatmulDirectPerTileShape(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckPertokenShapeInNormalQuantMode(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs, + const gert::Shape *perTokenScaleShape) const; + bool IsDoubleScaleScenario(const gert::Shape *scaleShape, const gert::Shape *perTokenScaleShape) const; + bool IsPerTileQuantMode(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + bool IsDoubleScaleScenario(const gert::Shape *scaleShape, const gert::Shape *perTokenScaleShape); + bool IsMxQuantMode(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + bool LogicXOR(bool cond1, bool cond2) const; + bool CheckShapeContainsNegativeValue(const gert::Shape *shape) const; + bool CheckDataContainsNegativeValue() const; + bool CheckNonDataInputContainsNegativeValue(const gert::InferShapeContext *context) const; +private: + int64_t groupNum_ = 0L; + int64_t xKDim_ = 0L; + int64_t xMDim_ = 0L; + int64_t weightKDim_ = 0L; + int64_t weightNDim_ = 0L; + size_t xdimNum_ = 0; + size_t weightdimNum_ = 0; +}; +} // namespace ops + +#endif // GROUPED_MATMUL_INFERSHAPE_DAVID_QUANT_CHECKER_H_ \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.cpp new file mode 100644 index 00000000..50dc5f7b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.cpp @@ -0,0 +1,956 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_infershape_weight_quant_checker.cpp + * \brief + */ +#include "grouped_matmul_infershape_weight_quant_checker.h" +#include "grouped_matmul_infershape_common_util.h" + +namespace ops { +static const std::map> BIAS_TYPE_SUPPORT_MAP = { + {ge::DT_FLOAT16, {ge::DT_FLOAT16}}, + {ge::DT_BF16, {ge::DT_BF16, ge::DT_FLOAT}}, + {ge::DT_INT8, {ge::DT_FLOAT}}, + {ge::DT_FLOAT8_E4M3FN, {ge::DT_BF16, ge::DT_FLOAT16}}}; +static const std::unordered_set FP8_SUPPORT_SET = {ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_HIFLOAT8}; +const int64_t UNKNOWN_SHAPE_VALUE = -1; +const int64_t SHAPE_UNKNOWN_DIM_NUM = -2; +const size_t ANTIQUANT_PARAM_DIM_NUM_PER_GROUP_SINGLE = 3; +const size_t OPTIONAL_PARAM_DIM_NUM_DEFAULT_SINGLE = 2; +const int64_t B4_NUMS_IN_B32 = 8; +const int64_t MX_GROUP_SIZE = 32; + +static bool inline IsNonEmpty(const gert::Shape *shape) +{ + return (shape != nullptr && !(shape->GetDimNum() == 1 && shape->GetDim(0) == 0)); +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsA16MxFp4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return (xDtype == ge::DT_FLOAT16 || xDtype == ge::DT_BF16) && + (weightDtype == ge::DT_FLOAT4_E2M1 || weightDtype == ge::DT_FLOAT); +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsMxA8W4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return xDtype == ge::DT_FLOAT8_E4M3FN && (weightDtype == ge::DT_FLOAT4_E2M1 || weightDtype == ge::DT_FLOAT); +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsS8S4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return xDtype == ge::DT_INT8 && (weightDtype == ge::DT_INT4 || weightDtype == ge::DT_INT32); +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsA16W8(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return (xDtype == ge::DT_FLOAT16 || xDtype == ge::DT_BF16) && weightDtype == ge::DT_INT8; +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsA16F8(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return (xDtype == ge::DT_FLOAT16 || xDtype == ge::DT_BF16) && + FP8_SUPPORT_SET.find(weightDtype) != FP8_SUPPORT_SET.end(); +} + +bool DlinferGroupedMatmulDirectWeightQuantChecker::IsA16W4(const ge::DataType xDtype, const ge::DataType weightDtype) const +{ + return (xDtype == ge::DT_FLOAT16 || xDtype == ge::DT_BF16) && + (weightDtype == ge::DT_INT4 || weightDtype == ge::DT_INT32); +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::GetXandWeightDtype(const gert::InferShapeContext *context) +{ + auto xDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_X, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + xDtype_ = xDesc->GetDataType(); + + auto weightDesc = context->GetDynamicInputDesc(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc); + weightDtype_ = weightDesc->GetDataType(); + + OP_LOGD(context->GetNodeName(), "Weight quant case xDtype is [%s], weightDtype is [%s]. ", + ge::TypeUtils::DataTypeToAscendString(xDtype_).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype_).GetString()); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::GetXAndWeightDimValue(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) +{ + OP_CHECK_IF(GetXandWeightDtype(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetXandWeightDtype failed."), return ge::GRAPH_FAILED); + + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_IF(!(IsNonEmpty(xShape) && IsNonEmpty(weightShape)), + OP_LOGE(context->GetNodeName(), "The 1st tensor of tensor list x and weight cannot be empty."), + return ge::GRAPH_FAILED); + + xdimNum_ = xShape->GetDimNum(); + weightdimNum_ = weightShape->GetDimNum(); + OP_CHECK_IF(xdimNum_ < GMM_MIN_FM_DIM, + OP_LOGE(context->GetNodeName(), + "The dim num of x's 1st tensor should be greater than 1, but it is [%zu].", xdimNum_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightdimNum_ < GMM_MIN_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "The dim num of weight's 1st tensor should be greater than 1, but it is [%zu].", weightdimNum_), + return ge::GRAPH_FAILED); + + if (gmmAttrs.groupType == GMM_NO_SPLIT) { + return ge::GRAPH_SUCCESS; + } + + xMDim_ = gmmAttrs.transposeX ? xShape->GetDim(xdimNum_ - 1) : xShape->GetDim(xdimNum_ - PENULTIMATE_DIM); + xKDim_ = gmmAttrs.transposeX ? xShape->GetDim(xdimNum_ - PENULTIMATE_DIM) : xShape->GetDim(xdimNum_ - 1); + weightKDim_ = gmmAttrs.transposeWeight ? weightShape->GetDim(weightdimNum_ - 1) + : weightShape->GetDim(weightdimNum_ - PENULTIMATE_DIM); + weightNDim_ = gmmAttrs.transposeWeight ? weightShape->GetDim(weightdimNum_ - PENULTIMATE_DIM) + : weightShape->GetDim(weightdimNum_ - 1); + + // 1个float32/int32表示8个float4_e2m1/int4，推导shape时，尾轴扩大8倍 + if ((weightDtype_ == ge::DT_FLOAT || weightDtype_ == ge::DT_INT32) && weightKDim_ > 0 && xKDim_ > 0) { + bool transWeightB32 = false; + if (!gmmAttrs.transposeWeight) { + transWeightB32 = weightKDim_ * B4_NUMS_IN_B32 == xKDim_; + } + if (gmmAttrs.transposeWeight || transWeightB32) { + weightKDim_ = weightKDim_ * B4_NUMS_IN_B32; + } else { + weightNDim_ = weightNDim_ * B4_NUMS_IN_B32; + } + } + + OP_LOGD(context->GetNodeName(), "xMDim = [%lld], xKDim = [%lld], weightKDim = [%lld], weightNDim = [%lld]. ", + xMDim_, xKDim_, weightKDim_, weightNDim_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShapeForXAndWeight(const gert::InferShapeContext *context) const +{ + OP_CHECK_IF(xKDim_ != weightKDim_, + OP_LOGE(context->GetNodeName(), + "The k dim of x should be equal to the k dim of weight, " + "but x's k is [%ld] and weight's k is [%ld].", + xKDim_, weightKDim_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightNDim_ <= 0, + OP_LOGE(context->GetNodeName(), "The n dim value should be positive, but the actual value is [%ld].", + weightNDim_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightKDim_ <= 0, + OP_LOGE(context->GetNodeName(), "The k dim value should be positive, but the actual value is [%ld].", + weightKDim_), + return ge::GRAPH_FAILED); + if (weightDtype_ == ge::DT_FLOAT4_E2M1 || weightDtype_ == ge::DT_FLOAT || IsS8S4NZ(xDtype_, weightDtype_)) { + // A16MxF4/MxA8W4/S8S4校验32B对齐 + OP_CHECK_IF( + !((weightNDim_ % GMM_N_K_ALIGN_VALUE_WEIGHT_QUANT_4BIT == 0) && + (weightKDim_ % GMM_N_K_ALIGN_VALUE_WEIGHT_QUANT_4BIT == 0)), + OP_LOGE(context->GetNodeName(), + "The value of dim n, k should be an integer multiple of [%ld], but actual n is [%ld], k is [%ld].", + GMM_N_K_ALIGN_VALUE_WEIGHT_QUANT_4BIT, weightNDim_, weightKDim_), + return ge::GRAPH_FAILED); + } + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_IF(xdimNum_ != GMM_MIN_FM_DIM || weightdimNum_ != GMM_SPLIT_M_SINGLE_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "When split m, x dim num should be 2, weight dim num should be 3, " + "but the actual x dim num is [%zu], actual weight dim num is [%zu].", + xdimNum_, weightdimNum_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "When split m, 1st dim value of weight should be g, " + "which is [%ld], but the actual value is [%ld].", + groupNum_, weightShape->GetDim(0)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckTensorDimEqualOne(const gert::InferShapeContext *context, + const gert::Shape *shape, + const std::string paramName, + const size_t index) const +{ + OP_CHECK_IF( + shape == nullptr, + OP_LOGE(context->GetNodeName(), "%s Shape[%zu] is null, which is not supported.", paramName.c_str(), index), + return ge::GRAPH_FAILED); + size_t dimNum = shape->GetDimNum(); + OP_CHECK_IF( + dimNum != 1, + OP_LOGE(context->GetNodeName(), "%s[%zu] dimNum is %zu, but only support 1.", paramName.c_str(), index, dimNum), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckDimNumNoSplit(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo) const +{ + const size_t &tensorListLength = paramsInputInfo.numX; + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + // check dimension + for (size_t i = 0; i < tensorListLength; ++i) { + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, i); + OP_CHECK_IF(xShape == nullptr, OP_LOGE(context->GetNodeName(), "x[%zu] is null, which is not supported.", i), + return ge::GRAPH_FAILED); + wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + OP_CHECK_NULL_WITH_CONTEXT(context, wShape); + size_t weightDimNum = wShape->GetDimNum(); + OP_CHECK_IF(weightDimNum != GMM_SEPARATED_WEIGHT_DIM, + OP_LOGE(context->GetNodeName(), + "weight[%zu] dimNum is %zu, but only support 2 when weight separated.", i, weightDimNum), + return ge::GRAPH_FAILED); + // 校验 x 中每个tensor的维度必须在[2,6]之间 + size_t xDimNum = xShape->GetDimNum(); + OP_CHECK_IF(xDimNum > GMM_MAX_FM_DIM || xDimNum < GMM_MIN_FM_DIM, + OP_LOGE(context->GetNodeName(), "x[%zu] dimNum is %zu, but only support 2-6.", i, xDimNum), + return ge::GRAPH_FAILED); + // 检测 bias antiquantScale antiquantOffset 的每个tensor的dim都需要为1 + if (paramsInputInfo.numBias != 0) { + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, i); + OP_CHECK_IF(CheckTensorDimEqualOne(context, biasShape, "bias", i) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckTensorDimEqualOne is failed."), return ge::GRAPH_FAILED); + } + if (paramsInputInfo.numAntiquantOffset != 0) { + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, i); + OP_CHECK_IF( + CheckTensorDimEqualOne(context, antiquantOffsetShape, "antiquantOffset", i) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckTensorDimEqualOne is failed."), return ge::GRAPH_FAILED); + } + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, i); + OP_CHECK_IF(CheckTensorDimEqualOne(context, antiquantScaleShape, "antiquantScale", i) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckTensorDimEqualOne is failed."), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckXWeightYGroupSizeMultiScenario( + const gert::InferShapeContext *context, const GMMInputParamsInfo ¶msInputInfo) const +{ + const size_t &xSize = paramsInputInfo.numX; + const size_t &weightSize = paramsInputInfo.numWeight; + size_t numY = context->GetComputeNodeOutputNum(); + // check group size + OP_CHECK_IF(xSize != numY, + OP_LOGE(context->GetNodeName(), "When y is separated, size of x %zu should equal to size of y %zu.", + xSize, numY), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xSize != weightSize, + OP_LOGE(context->GetNodeName(), + "When weight is separated, size of w %zu should equal to size of x %zu.", weightSize, xSize), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckTensorNDimMultiScenario(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo, + const size_t wNDimIdx, + const int64_t weightNDimValue, + const size_t index) const +{ + // 检验weight的n轴和antiquantScale的n轴一致 + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, index); + OP_CHECK_NULL_WITH_CONTEXT(context, antiquantScaleShape); + int64_t antiquantScaleNDim = antiquantScaleShape->GetDim(0); + OP_CHECK_IF(antiquantScaleNDim != weightNDimValue, + OP_LOGE(context->GetNodeName(), + "weight[%zu] dim %zu value %ld should equal to antiquantScale[%zu] dim 0 value %ld.", index, + wNDimIdx, weightNDimValue, index, antiquantScaleNDim), + return ge::GRAPH_FAILED); + + if (paramsInputInfo.numBias != 0) { + // 检验weigh的n轴和bias的n轴一致 + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, index); + OP_CHECK_NULL_WITH_CONTEXT(context, biasShape); + int64_t biasNDim = biasShape->GetDim(0); + OP_CHECK_IF( + biasNDim != weightNDimValue, + OP_LOGE(context->GetNodeName(), "weight[%zu] dim %zu value %ld should equal to bias[%zu] dim 0 value %ld.", + index, wNDimIdx, weightNDimValue, index, biasNDim), + return ge::GRAPH_FAILED); + } + if (paramsInputInfo.numAntiquantOffset != 0) { + // 检验weight的n轴和antiquantOffset的n轴一致 + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, index); + OP_CHECK_NULL_WITH_CONTEXT(context, antiquantOffsetShape); + int64_t antiquantOffsetNDim = antiquantOffsetShape->GetDim(0); + OP_CHECK_IF(antiquantOffsetNDim != weightNDimValue, + OP_LOGE(context->GetNodeName(), + "weight[%zu] dim %zu value %ld should equal to antiquantOffset[%zu] dim 0 value %ld.", + index, wNDimIdx, weightNDimValue, index, antiquantOffsetNDim), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckCaseMultiScenario(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs, + const GMMInputParamsInfo ¶msInputInfo) const +{ + const size_t &xSize = paramsInputInfo.numX; + // check group size + OP_CHECK_IF(CheckXWeightYGroupSizeMultiScenario(context, paramsInputInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "The size of X, Y and weight are not equal."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListSizeMultiScenario(context, paramsInputInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckTensorListSizeMultiScenario failed."), return ge::GRAPH_FAILED); + // check dimension + OP_CHECK_IF(CheckDimNumNoSplit(context, paramsInputInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Dim num of tensor in tensor lists or grouplist is invalid."), + return ge::GRAPH_FAILED); + + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + auto wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + size_t wKDimIdx = gmmAttrs.transposeWeight ? 1UL : 0UL; + size_t wNDimIdx = gmmAttrs.transposeWeight ? 0UL : 1UL; + + int64_t weightKDimValue = wShape->GetDim(wKDimIdx); + int64_t weightNDimValue = wShape->GetDim(wNDimIdx); + + for (size_t i = 0; i < xSize; i++) { + xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, i); + size_t xDimNum = xShape->GetDimNum(); + int64_t xKDimValue = xShape->GetDim(xDimNum - 1); // x always is not transposed + + wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, i); + weightKDimValue = wShape->GetDim(wKDimIdx); + weightNDimValue = wShape->GetDim(wNDimIdx); + // 校验M轴和batch轴大于等于0 + for (size_t j = 0; j < xDimNum - 1; j++) { + int64_t xNDimValue = xShape->GetDim(j); + OP_CHECK_IF(xNDimValue < 0, + OP_LOGE(context->GetNodeName(), "x[%zu] dim %zu value %ld should be more than or equal to 0.", + i, j, xNDimValue), + return ge::GRAPH_FAILED); + } + // 校验K轴和N轴大于0 + OP_CHECK_IF( + xKDimValue <= 0, + OP_LOGE(context->GetNodeName(), "x[%zu] dim %zu value %ld should more than 0.", i, xDimNum - 1, xKDimValue), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightNDimValue <= 0, + OP_LOGE(context->GetNodeName(), "w[%zu] dim %zu value %ld should more than 0.", i, wNDimIdx, + weightNDimValue), + return ge::GRAPH_FAILED); + // 校验X和weight矩阵的K轴 + OP_CHECK_IF( + xKDimValue != weightKDimValue, + OP_LOGE(context->GetNodeName(), "x[%zu] dim %zu value %ld should equal to weight[%zu] dim 0 value %ld.", i, + xDimNum - 1, xKDimValue, i, weightKDimValue), + return ge::GRAPH_FAILED); + // 校验 antiquantScale bias antiquantOffset 每个tensor的中n轴的大小 + OP_CHECK_IF(CheckTensorNDimMultiScenario(context, paramsInputInfo, wNDimIdx, weightNDimValue, i), + OP_LOGE(context->GetNodeName(), "CheckTensorNDimMultiScenario is failed."), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckScenarioValid(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto xSecondTensorShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 1); + auto weightSecondTensorShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 1); + // 检验groupType + OP_CHECK_IF((gmmAttrs.groupType != GMM_SPLIT_M) && (gmmAttrs.groupType != GMM_NO_SPLIT), + OP_LOGE(context->GetNodeName(), "Invalid groupType, which can only be 0 or -1, but it is [%ld].", + gmmAttrs.groupType), + return ge::GRAPH_FAILED); + + if (gmmAttrs.groupType == GMM_SPLIT_M) { // single/single/single Scenario + OP_CHECK_IF(IsNonEmpty(xSecondTensorShape), + OP_LOGE(context->GetNodeName(), "The second tensor of tensor list x is not empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(IsNonEmpty(weightSecondTensorShape), + OP_LOGE(context->GetNodeName(), "The second tensor of tensor list weight is not empty."), + return ge::GRAPH_FAILED); + + // check split item value valid + OP_CHECK_IF(gmmAttrs.splitItem != GMM_X_SEPARATED && gmmAttrs.splitItem != GMM_NO_SEPARATED, + OP_LOGE(context->GetNodeName(), + "Invalid splitItem, which can only be one of 2 or 3, but it is [%ld].", gmmAttrs.splitItem), + return ge::GRAPH_FAILED); + } else { // multi/multi/multi Scenario + // check split item value valid + OP_CHECK_IF(!IsA16W8(xDtype_, weightDtype_), + OP_LOGE(context->GetNodeName(), + "Multi/Multi/Multi Scenario is supported only when xDtype-weightDtype is fp16/bf16-int8."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(gmmAttrs.splitItem != GMM_X_Y_SEPARATED && gmmAttrs.splitItem != GMM_Y_SEPARATED, + OP_LOGE(context->GetNodeName(), + "Invalid splitItem, which can only be one of 0 or 1, but it is [%ld].", gmmAttrs.splitItem), + return ge::GRAPH_FAILED); + } + // check activation is null + OP_CHECK_IF(gmmAttrs.activeType != static_cast(GMMActType::GMM_ACT_TYPE_NONE), + OP_LOGE(context->GetNodeName(), "Activation function is not supported in weight quant mode now."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShapeForGrouplist(const gert::InferShapeContext *context, + const gert::Shape *groupListShape) const +{ + OP_CHECK_IF(groupListShape->GetDimNum() != 1, + OP_LOGE(context->GetNodeName(), + "In single-single-single scenario, " + "the groupList only support 1 dim num for now, but the actual dim num is [%zu].", + groupListShape->GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupListShape->GetDim(0) <= 0, + OP_LOGE(context->GetNodeName(), + "The groupList 1st dim value should be greater than 0, but the actual value is [%ld].", + groupListShape->GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupListShape->GetDim(0) > GMM_MAX_GROUP_LIST_SIZE_TENSOR, + OP_LOGE(context->GetNodeName(), + "Only support [%ld] groups at MAX for now, but the actual group number is [%ld].", + GMM_MAX_GROUP_LIST_SIZE_TENSOR, groupListShape->GetDim(0)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckPertokenScaleForA8W4(const gert::InferShapeContext *context) const +{ + auto tensorShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + size_t tensorDimNum = tensorShape->GetDimNum(); + // S8S4的PerTokenScale维度为1，shape为(m) + // MxA8W4的PerTokenScale维度为2，shape为(m, k/32) + size_t tensorDimNumExp = IsS8S4NZ(xDtype_, weightDtype_) ? 1 : 2; + OP_CHECK_IF(tensorDimNum != tensorDimNumExp, + OP_LOGE(context->GetNodeName(), "PertokenScale dim should be [%zu], but the actual dim num is [%zu].", + tensorDimNumExp, tensorDimNum), + return ge::GRAPH_FAILED); + OP_CHECK_IF(tensorShape->GetDim(0) != xMDim_, + OP_LOGE(context->GetNodeName(), + "The shape of PertokenScale should be (m), which is (%ld), but the actual shape is (%ld).", + xMDim_, tensorShape->GetDim(0)), + return ge::GRAPH_FAILED); + if (IsMxA8W4NZ(xDtype_, weightDtype_)) { + OP_CHECK_IF( + tensorShape->GetDim(1) != weightKDim_ / MX_GROUP_SIZE, + OP_LOGE(context->GetNodeName(), + "PerTokenScale shape should be (m, k/%ld) when xDtype-weightDtype is fp8_e4m3-fp4_e2m1, which is " + "(%ld, %ld), but the actual shape is (%ld, %ld).", + MX_GROUP_SIZE, xMDim_, weightKDim_ / MX_GROUP_SIZE, tensorShape->GetDim(0), tensorShape->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShapeForTensorList(const gert::InferShapeContext *context, + size_t gmm_index, + const std::string &tensorType, + const GMMAttrs &gmmAttrs) const +{ + // 校验单单单场景antiquantScale/antiquantOffset/bias/scale的shape + // A16MxF4/MxA8W4/S8S4校验antiquant params的Shape为(g, k/groupSize, n)/(g, n, k/groupsize)，维度数为3 + // 其他场景及参数校验shape为(g, n) + auto tensorShape = context->GetDynamicInputShape(gmm_index, 0); + if (IsNonEmpty(tensorShape)) { + size_t tensorDimNum = tensorShape->GetDimNum(); + size_t expectedDimNum = OPTIONAL_PARAM_DIM_NUM_DEFAULT_SINGLE; + if ((gmm_index == GMM_INDEX_IN_ANTIQUANT_SCALE || gmm_index == GMM_INDEX_IN_ANTIQUANT_OFFSET) && + (IsA16MxFp4NZ(xDtype_, weightDtype_) || IsMxA8W4NZ(xDtype_, weightDtype_) || + IsS8S4NZ(xDtype_, weightDtype_))) { + expectedDimNum = ANTIQUANT_PARAM_DIM_NUM_PER_GROUP_SINGLE; + } + // check dim num + OP_CHECK_IF(tensorDimNum != expectedDimNum, + OP_LOGE(context->GetNodeName(), "%s dim num should be [%zu], but the actual dim num is [%zu].", + tensorType.c_str(), expectedDimNum, tensorDimNum), + return ge::GRAPH_FAILED); + + // check shape + OP_CHECK_IF(tensorShape->GetDim(0) != groupNum_, + OP_LOGE(context->GetNodeName(), + "The first dim of %s should be g, which is [%ld], but the actual value is [%ld].", + tensorType.c_str(), groupNum_, tensorShape->GetDim(0)), + return ge::GRAPH_FAILED); + + size_t tensorNDimIdx = tensorDimNum - 1; + if (expectedDimNum == ANTIQUANT_PARAM_DIM_NUM_PER_GROUP_SINGLE && gmmAttrs.transposeWeight) { + // mx/per_group量化weight转置时antiquant params同步转置，shape为(g, n, k/groupsize)，n轴的索引为-2 + tensorNDimIdx = tensorDimNum - 2; + } + + OP_CHECK_IF(tensorShape->GetDim(tensorNDimIdx) != weightNDim_, + OP_LOGE(context->GetNodeName(), + "The n dim of %s should be equal to the weightNDim[%ld], but the actual value is [%ld].", + tensorType.c_str(), weightNDim_, tensorShape->GetDim(tensorNDimIdx)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShapeForWeightQuantParam(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + auto scaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_SCALE, 0); + auto offsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_OFFSET, 0); + auto perTokenScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, 0); + OP_CHECK_IF(IsNonEmpty(antiquantOffsetShape) && (FP8_SUPPORT_SET.find(weightDtype_) != FP8_SUPPORT_SET.end() || + weightDtype_ == ge::DT_FLOAT4_E2M1 || + weightDtype_ == ge::DT_FLOAT || IsS8S4NZ(xDtype_, weightDtype_)), + OP_LOGE(context->GetNodeName(), + "In weight quant case, only support antiquantOffset is none when weightDtype is fp8/hif8/fp4 " + "or xDtype-weightDtype is int8-int4."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(IsNonEmpty(scaleShape) && !IsS8S4NZ(xDtype_, weightDtype_), + OP_LOGE(context->GetNodeName(), "In weight quant case, scale must be empty."), return ge::GRAPH_FAILED); + OP_CHECK_IF(!IsNonEmpty(scaleShape) && IsS8S4NZ(xDtype_, weightDtype_), + OP_LOGE(context->GetNodeName(), "When xDtype-weightDtype is int8-int4, scale must not be empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(IsNonEmpty(offsetShape), OP_LOGE(context->GetNodeName(), "In weight quant case, offset must be empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + IsNonEmpty(perTokenScaleShape) && (!IsMxA8W4NZ(xDtype_, weightDtype_) && !IsS8S4NZ(xDtype_, weightDtype_)), + OP_LOGE(context->GetNodeName(), + "If xDtype-weightDtype is not fp8_e4m3-fp4_e2m1 or int8-int4, pertokenscale must be empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + !IsNonEmpty(perTokenScaleShape) && (IsMxA8W4NZ(xDtype_, weightDtype_) || IsS8S4NZ(xDtype_, weightDtype_)), + OP_LOGE(context->GetNodeName(), + "When xDtype-weightDtype is fp8_e4m3-fp4_e2m1 or int8-int4, pertokenscale must be not empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(!IsNonEmpty(antiquantScaleShape), + OP_LOGE(context->GetNodeName(), "In weight quant case, antiquantScale must be not empty."), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + CheckShapeForTensorList(context, GMM_INDEX_IN_ANTIQUANT_SCALE, "antiquantScale", gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForAntiquantScale failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeForTensorList(context, GMM_INDEX_IN_ANTIQUANT_OFFSET, "antiquantOffset", gmmAttrs) != + ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForAntiquantOffset failed."), return ge::GRAPH_FAILED); + if (IsMxA8W4NZ(xDtype_, weightDtype_) || IsS8S4NZ(xDtype_, weightDtype_)) { + OP_CHECK_IF(CheckPertokenScaleForA8W4(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckPertokenScale failed."), return ge::GRAPH_FAILED); + } + if (IsS8S4NZ(xDtype_, weightDtype_)) { + OP_CHECK_IF(CheckShapeForTensorList(context, GMM_INDEX_IN_SCALE, "scale", gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForscale failed."), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckGroupSize(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + if (!(IsA16MxFp4NZ(xDtype_, weightDtype_) || IsMxA8W4NZ(xDtype_, weightDtype_) || + IsS8S4NZ(xDtype_, weightDtype_))) { + return ge::GRAPH_SUCCESS; + } + + int64_t groupSize = 0; + + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + auto antiquantScaleDimNum = antiquantScaleShape->GetDimNum(); + // 2含义: (g, k/groupSize, n)的k轴索引，此处groupNum是K轴上量化分组的groupNum，与groupNum_含义不同 + int64_t groupNum = gmmAttrs.transposeWeight ? antiquantScaleShape->GetDim(antiquantScaleDimNum - 1) + : antiquantScaleShape->GetDim(antiquantScaleDimNum - 2); + OP_CHECK_IF(groupNum <= 0, OP_LOGE(context->GetNodeName(), "GroupNum must be greater than 0."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(weightKDim_ % groupNum != 0, + OP_LOGE(context->GetNodeName(), "GroupNum must be multiple of the k axis of weight."), + return ge::GRAPH_FAILED); + groupSize = weightKDim_ / groupNum; + + if (IsS8S4NZ(xDtype_, weightDtype_)) { + // 伪量化S8S4场景支持groupsize为128/192/256/512 + OP_CHECK_IF(groupSize != 128 && groupSize != 256 && groupSize != 512 && groupSize != 192, + OP_LOGE(context->GetNodeName(), + "groupSize must be 128/192/256/512, but current groupSize is (%ld).", groupSize), + return ge::GRAPH_FAILED); + } else { + // Mx量化的groupSize为32 + OP_CHECK_IF(groupSize != MX_GROUP_SIZE, + OP_LOGE(context->GetNodeName(), "groupSize must be [%ld], but current groupSize is (%ld).", + MX_GROUP_SIZE, groupSize), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::GetNumOfInputs(const gert::InferShapeContext *context, + GMMInputParamsInfo ¶msInputInfo) const +{ + ge::graphStatus res = ge::GRAPH_SUCCESS; + const gert::Shape *shape = nullptr; + struct ParamInfoTmp { + int index; + size_t &count; + const char *name; + }; + ParamInfoTmp params[] = {{GMM_INDEX_IN_X, paramsInputInfo.numX, "numX"}, + {GMM_INDEX_IN_WEIGHT, paramsInputInfo.numWeight, "numWeight"}, + {GMM_INDEX_IN_BIAS, paramsInputInfo.numBias, "numBias"}, + {GMM_INDEX_IN_ANTIQUANT_SCALE, paramsInputInfo.numAntiquantScale, "numAntiquantScale"}, + {GMM_INDEX_IN_ANTIQUANT_OFFSET, paramsInputInfo.numAntiquantOffset, "numAntiquantOffset"}}; + + for (auto ¶m : params) { + param.count = 0; + for (int i = 0; i < GMM_MAX_GROUP_LIST_SIZE_ARRAY; i++) { + shape = context->GetDynamicInputShape(param.index, param.count); + if (!IsNonEmpty(shape)) { + break; + } + ++param.count; + } + OP_CHECK_IF(param.count >= GMM_MAX_GROUP_LIST_SIZE_ARRAY, + OP_LOGE(context->GetNodeName(), + "In multi/multi/multi Scenario, each tensorlist's length cannot exceed 128"), + return ge::GRAPH_FAILED); + OP_LOGI(context->GetNodeName(), "%s = %zu", param.name, param.count); + } + return res; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckTensorListSizeMultiScenario( + const gert::InferShapeContext *context, const GMMInputParamsInfo ¶msInputInfo) const +{ + // 检测 bias antiquantScale antiquantOffset的 tensorListsize 需要等于 weightSize + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, 0); + if (IsNonEmpty(biasShape)) { + OP_CHECK_IF( + paramsInputInfo.numBias != paramsInputInfo.numWeight, + OP_LOGE(context->GetNodeName(), "Bias size should be equal to weight size, actual size are [%zu] and [%zu]", + paramsInputInfo.numBias, paramsInputInfo.numWeight), + return ge::GRAPH_FAILED); + } + + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, 0); + if (IsNonEmpty(antiquantOffsetShape)) { + OP_CHECK_IF(paramsInputInfo.numAntiquantOffset != paramsInputInfo.numWeight, + OP_LOGE(context->GetNodeName(), + "AntiquantOffset size should be equal to weight size, actual size are [%zu] and [%zu]", + paramsInputInfo.numAntiquantOffset, paramsInputInfo.numWeight), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(paramsInputInfo.numAntiquantScale != paramsInputInfo.numWeight, + OP_LOGE(context->GetNodeName(), + "AntiquantScale size should be equal to weight size, actual size are [%zu] and [%zu]", + paramsInputInfo.numAntiquantScale, paramsInputInfo.numWeight), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShapeValid(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) +{ + if (gmmAttrs.groupType == GMM_NO_SPLIT) { + // 多多多场景校验 + GMMInputParamsInfo paramsInputInfo{0, 0, 0, 0, 0, 0, 0}; + OP_CHECK_IF(GetNumOfInputs(context, paramsInputInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "GetNumOfInputs failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckCaseMultiScenario(context, gmmAttrs, paramsInputInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckCaseMultiScenario failed."), return ge::GRAPH_FAILED); + } else { + // 单单单场景校验 + auto groupListShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListShape); + OP_CHECK_IF(CheckShapeForGrouplist(context, groupListShape) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForGrouplist failed."), return ge::GRAPH_FAILED); + groupNum_ = groupListShape->GetDim(0); + OP_CHECK_IF(CheckShapeForXAndWeight(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForXAndWeight failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeForTensorList(context, GMM_INDEX_IN_BIAS, "bias", gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForBias failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeForWeightQuantParam(context, gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeForWeightQuantParam failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckGroupSize(context, gmmAttrs) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckGroupSize failed."), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +bool IsUnknownShape(const gert::Shape *shape) +{ + if (IsNonEmpty(shape)) { + size_t size = shape->GetDimNum(); + for (size_t i = 0; i < size; i++) { + if (shape->GetDim(i) == UNKNOWN_SHAPE_VALUE || shape->GetDim(i) == SHAPE_UNKNOWN_DIM_NUM) { + return true; + } + } + return false; + } + return false; +} + +bool CheckUnknownShape(const gert::InferShapeContext *context) +{ + bool hasUnknownShape = false; + auto xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, 0); + auto weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + auto antiquantScaleShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + auto antiquantOffsetShape = context->GetDynamicInputShape(GMM_INDEX_IN_ANTIQUANT_OFFSET, 0); + auto biasShape = context->GetDynamicInputShape(GMM_INDEX_IN_BIAS, 0); + auto groupListShape = context->GetOptionalInputShape(GMM_INDEX_IN_GROUP_LIST); + hasUnknownShape = IsUnknownShape(xShape) || IsUnknownShape(weightShape) || IsUnknownShape(antiquantScaleShape) || + IsUnknownShape(groupListShape); + if (!hasUnknownShape && IsNonEmpty(antiquantOffsetShape)) { + hasUnknownShape |= IsUnknownShape(antiquantOffsetShape); + } + if (!hasUnknownShape && IsNonEmpty(biasShape)) { + hasUnknownShape |= IsUnknownShape(biasShape); + } + return hasUnknownShape; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckShape(const gert::InferShapeContext *context, + const DlinferGroupedMatmulDirectCommonUtil &commonUtil) +{ + if (CheckUnknownShape(context)) { + return ge::GRAPH_SUCCESS; + } + + OP_CHECK_IF(CheckScenarioValid(context, commonUtil.attrsInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckScenarioValid failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckShapeValid(context, commonUtil.attrsInfo) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckShapeValid failed."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::UpdateShapeY(gert::InferShapeContext *context, size_t idxY, + std::vector &yDims) const +{ + gert::Shape *yShape = context->GetOutputShape(idxY); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + yShape->SetDimNum(yDims.size()); + for (size_t dim = 0; dim < yDims.size(); ++dim) { + yShape->SetDim(dim, yDims[dim]); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::UpdateShapeYMultiDim(gert::InferShapeContext *context, size_t idxY, + const gert::Shape *xShape, + const gert::Shape *weightShape) const +{ + gert::Shape *yShape = context->GetOutputShape(idxY); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + *yShape = *xShape; + size_t dimY = yShape->GetDimNum(); + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const bool *transposeWPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_W); + const bool *transposeXPtr = attrs->GetAttrPointer(GMM_INDEX_ATTR_TRANSPOSE_X); + + OP_CHECK_NULL_WITH_CONTEXT(context, weightShape); + if (transposeWPtr != nullptr && *transposeWPtr) { + yShape->SetDim(dimY - 1, weightShape->GetDim(weightShape->GetDimNum() - 2)); // -2: transpose weight + } else { + yShape->SetDim(dimY - 1, weightShape->GetDim(weightShape->GetDimNum() - 1)); + } + if (transposeXPtr != nullptr && *transposeXPtr) { + yShape->SetDim(dimY - 2, xShape->GetDim(xShape->GetDimNum() - 1)); // -2: last two dim of Y + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::InferOutShape(gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const +{ + if (gmmAttrs.groupType == GMM_NO_SPLIT) { + size_t idx = 0; + size_t idw = 0; + const gert::Shape *w0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Shape); + for (int i = 0; i < GMM_MAX_GROUP_LIST_SIZE_ARRAY; i++) { + const gert::Shape *xShape = context->GetDynamicInputShape(GMM_INDEX_IN_X, idx); + if (xShape == nullptr) { + break; + } + ++idx; + const gert::Shape *wShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idw); + if (wShape) { + ++idw; + } else { + wShape = w0Shape; + } + OP_CHECK_IF(UpdateShapeYMultiDim(context, GMM_INDEX_OUT_Y + idx - 1, xShape, wShape) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update shape of y."), return ge::GRAPH_FAILED); + } + } else { + std::vector yDims = {xMDim_, weightNDim_}; + OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y, yDims) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "Failed to update y shape."), return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckScaleDtypeForS8S4(const gert::InferDataTypeContext *context) const +{ + auto perTokenScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + auto scaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_SCALE, 0); + auto antiquantScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + OP_CHECK_IF(scaleDtype != ge::DT_FLOAT, + OP_LOGE(context->GetNodeName(), "scaleDtype datatype [%s] does not match float32.", + ge::TypeUtils::DataTypeToAscendString(scaleDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(antiquantScaleDtype != ge::DT_FLOAT16, + OP_LOGE(context->GetNodeName(), "antiquantScaleDtype datatype [%s] does not match float16.", + ge::TypeUtils::DataTypeToAscendString(antiquantScaleDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(perTokenScaleDtype != ge::DT_FLOAT, + OP_LOGE(context->GetNodeName(), "perTokenScaleDtype datatype [%s] does not match float32.", + ge::TypeUtils::DataTypeToAscendString(perTokenScaleDtype).GetString()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckBiasDtype(const gert::InferDataTypeContext *context) const +{ + auto xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto biasDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_BIAS, 0); + OP_CHECK_IF(BIAS_TYPE_SUPPORT_MAP.find(xDtype) == BIAS_TYPE_SUPPORT_MAP.end(), + OP_LOGE(context->GetNodeName(), "Cannot find bias dtype match with xDtype [%s].", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(BIAS_TYPE_SUPPORT_MAP.at(xDtype).find(biasDtype) == BIAS_TYPE_SUPPORT_MAP.at(xDtype).end(), + OP_LOGE(context->GetNodeName(), "Data type [%s] is not supported for bias, when xDtype is [%s].", + ge::TypeUtils::DataTypeToAscendString(biasDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckTensorListDataType(const gert::InferDataTypeContext *context, + uint32_t index, const ge::DataType dtype) const +{ + size_t inIdx = 0; + for (int i = 0; i < GMM_MAX_GROUP_LIST_SIZE_ARRAY; i++) { + auto iDtype = context->GetDynamicInputDataType(index, inIdx); + if (iDtype == ge::DT_UNDEFINED) { + break; + } + OP_CHECK_IF(iDtype != dtype, + OP_LOGE(context->GetNodeName(), "data type of tensors in a tensorList should all be the same!"), + return ge::GRAPH_FAILED); + ++inIdx; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckMatmulDataType( + const gert::InferDataTypeContext *context, const ge::DataType xDtype, const ge::DataType weightDtype, + const ge::DataType biasDtype, const ge::DataType antiquantScaleDtype, const ge::DataType antiquantOffsetDtype) const +{ + // 单单单/多多多常规参数通用校验，PerTokenScale和scale当前仅支持单单单，暂不在此处校验 + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_X, xDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "x dtype does not match with required dtype[%s].", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_WEIGHT, weightDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "weight dtype does not match with required dtype[%s].", + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckTensorListDataType(context, GMM_INDEX_IN_BIAS, biasDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "bias dtype does not match with required dtype[%s].", + ge::TypeUtils::DataTypeToAscendString(biasDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + CheckTensorListDataType(context, GMM_INDEX_IN_ANTIQUANT_SCALE, antiquantScaleDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "antiquantScaleDtype dtype does not match with required dtype[%s].", + ge::TypeUtils::DataTypeToAscendString(antiquantScaleDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + CheckTensorListDataType(context, GMM_INDEX_IN_ANTIQUANT_OFFSET, antiquantOffsetDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "antiquantOffsetDtype dtype does not match with required dtype[%s].", + ge::TypeUtils::DataTypeToAscendString(antiquantOffsetDtype).GetString()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::CheckDtype(const gert::InferDataTypeContext *context) const +{ + auto xDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto weightDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_WEIGHT, 0); + auto biasDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_BIAS, 0); + auto antiquantScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_ANTIQUANT_SCALE, 0); + auto antiquantOffsetDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_ANTIQUANT_OFFSET, 0); + auto perTokenScaleDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_PERTOKEN_SCALE, 0); + + // 必选参数校验，同时校验数据流是否支持 + if (IsA16W8(xDtype, weightDtype) || IsA16F8(xDtype, weightDtype) || IsA16W4(xDtype, weightDtype)) { + OP_CHECK_IF(antiquantScaleDtype != xDtype, + OP_LOGE(context->GetNodeName(), "AntiquantScale datatype [%s] does not match xDtype [%s].", + ge::TypeUtils::DataTypeToAscendString(antiquantScaleDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + } else if (IsA16MxFp4NZ(xDtype, weightDtype) || IsMxA8W4NZ(xDtype, weightDtype)) { + OP_CHECK_IF(antiquantScaleDtype != ge::DT_FLOAT8_E8M0, + OP_LOGE(context->GetNodeName(), + "Only support float8_e8m0 for antiquantScaleDataType when xDtype-weightDtype is " + "fp16/bf16-fp4_e2m1 or fp8_e4m3-fp4_e2m1, but got [%s].", + ge::TypeUtils::DataTypeToAscendString(antiquantScaleDtype).GetString()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(perTokenScaleDtype != ge::DT_FLOAT8_E8M0 && IsMxA8W4NZ(xDtype, weightDtype), + OP_LOGE(context->GetNodeName(), + "Only support float8_e8m0 for perTokenScaleDtype when xDtype-weightDtype is " + "fp8_e4m3-fp4_e2m1, but got [%s].", + ge::TypeUtils::DataTypeToAscendString(perTokenScaleDtype).GetString()), + return ge::GRAPH_FAILED); + } else if (IsS8S4NZ(xDtype, weightDtype)) { + OP_CHECK_IF(CheckScaleDtypeForS8S4(context) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckScaleDtype failed."), return ge::GRAPH_FAILED); + } else { + OP_LOGE(context->GetNodeName(), "Weight quant case does not support xDtype [%s] and weightDtype [%s].", + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(weightDtype).GetString()); + return ge::GRAPH_FAILED; + } + + if (IsA16W8(xDtype, weightDtype) || IsA16W4(xDtype, weightDtype)) { + OP_CHECK_IF(antiquantOffsetDtype != xDtype, + OP_LOGE(context->GetNodeName(), "AntiquantOffset datatype [%s] does not match xDtype [%s].", + ge::TypeUtils::DataTypeToAscendString(antiquantOffsetDtype).GetString(), + ge::TypeUtils::DataTypeToAscendString(xDtype).GetString()), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(CheckBiasDtype(context) != ge::GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(), "CheckBiasDtype failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckMatmulDataType(context, xDtype, weightDtype, biasDtype, antiquantScaleDtype, + antiquantOffsetDtype) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "CheckMatmulDataType is failed!"), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DlinferGroupedMatmulDirectWeightQuantChecker::InferOutDtype(gert::InferDataTypeContext *context) const +{ + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t *outputDtype = attrs->GetInt(GMM_INDEX_ATTR_OUTPUT_DTYPE); + ge::DataType yDtype = context->GetDynamicInputDataType(GMM_INDEX_IN_X, 0); + auto it = GMM_OUTPUT_DTYPE_MAP.find(*outputDtype); + OP_CHECK_IF(it == GMM_OUTPUT_DTYPE_MAP.end(), + OP_LOGE(context->GetNodeName(), + "The output dtype should be bfloat16 or float16 , but the actual output dtype is [%s].", + ge::TypeUtils::DataTypeToAscendString(yDtype).GetString()), + return ge::GRAPH_FAILED); + yDtype = it->second; + context->SetOutputDataType(GMM_INDEX_OUT_Y, yDtype); + return ge::GRAPH_SUCCESS; +} + +} // namespace ops \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.h new file mode 100644 index 00000000..54339fa3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/grouped_matmul_infershape_weight_quant_checker.h @@ -0,0 +1,94 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef GROUPED_MATMUL_INFERSHAPE_WEIGHT_QUANT_CHECKER_H_ +#define GROUPED_MATMUL_INFERSHAPE_WEIGHT_QUANT_CHECKER_H_ + +#include "graph/utils/type_utils.h" +#include "grouped_matmul_infershape_common_util.h" +#include "log/log.h" +#include "platform/platform_info.h" +#include "register/op_impl_registry.h" + +namespace ops { + +class DlinferGroupedMatmulDirectWeightQuantChecker { +public: + DlinferGroupedMatmulDirectWeightQuantChecker(){}; + ~DlinferGroupedMatmulDirectWeightQuantChecker(){}; + ge::graphStatus GetXAndWeightDimValue(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs); + ge::graphStatus CheckShape(const gert::InferShapeContext *context, const DlinferGroupedMatmulDirectCommonUtil &commonUtil); + ge::graphStatus InferOutShape(gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckScaleDtypeForS8S4(const gert::InferDataTypeContext *context) const; + ge::graphStatus CheckBiasDtype(const gert::InferDataTypeContext *context) const; + ge::graphStatus CheckDtype(const gert::InferDataTypeContext *context) const; + ge::graphStatus InferOutDtype(gert::InferDataTypeContext *context) const; + +private: + ge::graphStatus GetXandWeightDtype(const gert::InferShapeContext *context); + ge::graphStatus CheckTensorDimEqualOne(const gert::InferShapeContext *context, const gert::Shape *shape, + const std::string paramName, const size_t index) const; + ge::graphStatus UpdateShapeYMultiDim(gert::InferShapeContext *context, size_t idxY, const gert::Shape *xShape, + const gert::Shape *weightShape) const; + ge::graphStatus CheckMatmulDataType(const gert::InferDataTypeContext *context, const ge::DataType xDtype, + const ge::DataType weightDtype, const ge::DataType biasDtype, + const ge::DataType antiquantScaleDtype, + const ge::DataType antiquantOffsetDtype) const; + ge::graphStatus CheckTensorListDataType(const gert::InferDataTypeContext *context, uint32_t index, + const ge::DataType dtype) const; + ge::graphStatus CheckShapeForXAndWeight(const gert::InferShapeContext *context) const; + ge::graphStatus CheckDimNumNoSplit(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo) const; + ge::graphStatus CheckXWeightYGroupSizeMultiScenario(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo) const; + ge::graphStatus CheckTensorNDimMultiScenario(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo, const size_t wNDimIdx, + const int64_t weightNDimValue, const size_t index) const; + ge::graphStatus CheckCaseMultiScenario(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs, + const GMMInputParamsInfo ¶msInputInfo) const; + ge::graphStatus CheckShapeForTensorList(const gert::InferShapeContext *context, size_t gmm_index, + const std::string &tensorType, const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckScenarioValid(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + ge::graphStatus GetNumOfInputs(const gert::InferShapeContext *context, GMMInputParamsInfo ¶msInputInfo) const; + ge::graphStatus CheckTensorListSizeMultiScenario(const gert::InferShapeContext *context, + const GMMInputParamsInfo ¶msInputInfo) const; + ge::graphStatus CheckShapeValid(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs); + + ge::graphStatus CheckShapeForWeightQuantParam(const gert::InferShapeContext *context, + const GMMAttrs &gmmAttrs) const; + ge::graphStatus CheckShapeForGrouplist(const gert::InferShapeContext *context, + const gert::Shape *groupListShape) const; + ge::graphStatus UpdateShapeY(gert::InferShapeContext *context, size_t idxY, std::vector &yDims) const; + ge::graphStatus CheckGroupAntiS(const gert::Shape *tensorShape, const gert::InferShapeContext *context, + const std::string &tensorType) const; + ge::graphStatus CheckPertokenScaleForA8W4(const gert::InferShapeContext *context) const; + ge::graphStatus CheckGroupSize(const gert::InferShapeContext *context, const GMMAttrs &gmmAttrs) const; + bool IsA16MxFp4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const; + bool IsMxA8W4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const; + bool IsS8S4NZ(const ge::DataType xDtype, const ge::DataType weightDtype) const; + bool IsA16W8(const ge::DataType xDtype, const ge::DataType weightDtype) const; + bool IsA16F8(const ge::DataType xDtype, const ge::DataType weightDtype) const; + bool IsA16W4(const ge::DataType xDtype, const ge::DataType weightDtype) const; + +private: + int64_t groupNum_; //当前含义为M分组数g + int64_t xKDim_; + int64_t xMDim_; + int64_t weightKDim_; + int64_t weightNDim_; + size_t xdimNum_; + size_t weightdimNum_; + ge::DataType xDtype_ = ge::DT_UNDEFINED; + ge::DataType weightDtype_ = ge::DT_UNDEFINED; +}; + +} // namespace ops + +#endif // GROUPED_MATMUL_INFERSHAPE_DAVID_QUANT_CHECKER_H_ \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.cpp new file mode 100644 index 00000000..19033716 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.cpp @@ -0,0 +1,565 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_no_quant_matmul_tiling.cpp + * \brief + */ +#include "grouped_no_quant_matmul_tiling.h" +#include "../../../op_kernel/arch35/non_quant/grouped_matmul_tiling_key.h" + +enum class GmmTrans { + NoTrans = 0, + ATrans = 1, + BTrans = 2, + ABTrans = 3 +}; +namespace optiling { +bool GroupedNoQuantMatmulTiling::SetTiling(gert::TilingContext *context) +{ + auto compileInfoPtr = context->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context->GetNodeName(), "compileInfoPtr is nullptr."), return false); + usedCoreNum_ = compileInfoPtr->aicNum; + OP_CHECK_IF(!Init(context), OP_LOGE(context->GetNodeName(), "Init failed"), return false); + OP_CHECK_IF(!CalMatMulTiling(context, compileInfoPtr), + OP_LOGE(context->GetNodeName(), "Unable to calculate matmul-tiling"), return false); + auto ret = SetGMMTiling(); + if (!ret) { + OP_LOGE(context->GetNodeName(), "Unable to set GMM tiling data"); + return false; + } + SetMatMulTiling(); + SetTilingKey(context); + OP_CHECK_IF(!SetCustomParam(context), OP_LOGE(context->GetNodeName(), "Unable to set custom param"), return false); + PrintTilingResult(context); + return true; +} + +bool GroupedNoQuantMatmulTiling::CalBaseMMTiling(const gert::TilingContext *context, + const GMMCompileInfo *compileInfoPtr) +{ + // according to the double buffer enabled L0B, compute baseK + baseK_ = (compileInfoPtr->l0BSize / DB_SIZE) / (baseN_ * GetSizeByDataType(xDType_)); + baseK_ = baseK_ & ~ALIGN_DOWN_16; // 16 bytes down-align + OP_CHECK_IF(baseK_ == 0, OP_LOGE(context->GetNodeName(), "baseK_ cannot be 0."), return false); + // according to the double buffer enabled L0A/L0C, compute baseM(cube) + uint32_t maxBaseM = static_cast(compileInfoPtr->l0CSize / (baseN_ * FP32_DTYPE_SIZE)); + baseM_ = std::min((compileInfoPtr->l0ASize / DB_SIZE) / (baseK_ * GetSizeByDataType(xDType_)), maxBaseM); + if (baseM_ > BASE_M_DEFAULT) { + baseM_ = BASE_M_DEFAULT; + } + OP_CHECK_IF(baseM_ == 0, OP_LOGE(context->GetNodeName(), "baseM_ cannot be 0."), return false); + return CalL1Tiling(context, compileInfoPtr); +} + +void GroupedNoQuantMatmulTiling::FormulateBasicBlock(const GMMCompileInfo *compileInfoPtr, uint32_t remainCoreNum) +{ + uint64_t mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + uint64_t nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + if (mCnt * nCnt >= remainCoreNum) { + return; + } + if (mCnt <= nCnt) { + baseM_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(m_, mCnt), BASIC_BLOCK_SIZE_16); + mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + nCnt = remainCoreNum / mCnt; + baseN_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(n_, nCnt), BASIC_BLOCK_SIZE_16); + } else { + baseN_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(n_, nCnt), BASIC_BLOCK_SIZE_16); + nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + mCnt = remainCoreNum / nCnt; + baseM_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(m_, mCnt), BASIC_BLOCK_SIZE_16); + } + + while (baseN_ >= baseM_ * NUM_TWO && nCnt < remainCoreNum / NUM_TWO) { + nCnt = nCnt * NUM_TWO; + mCnt = remainCoreNum / nCnt; + baseM_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(m_, mCnt), BASIC_BLOCK_SIZE_16); + baseN_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(n_, nCnt), BASIC_BLOCK_SIZE_16); + mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + } + + while (baseM_ >= baseN_ * NUM_TWO && mCnt < remainCoreNum / NUM_TWO) { + mCnt = mCnt * NUM_TWO; + nCnt = remainCoreNum / mCnt; + baseM_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(m_, mCnt), BASIC_BLOCK_SIZE_16); + baseN_ = DlinferGroupedMatmulDirect::CeilAlign(DlinferGroupedMatmulDirect::CeilDiv(n_, nCnt), BASIC_BLOCK_SIZE_16); + mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + } + mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + uint64_t kValueAlign = DlinferGroupedMatmulDirect::CeilAlign(k_, BASIC_BLOCK_SIZE_16); + uint64_t kValueMax = + DlinferGroupedMatmulDirect::FloorAlign(static_cast(compileInfoPtr->l0ASize / DB_SIZE / GetSizeByDataType(xDType_) / + std::max(baseM_, baseN_)), + BASIC_BLOCK_SIZE_16); + baseK_ = std::min(kValueAlign, kValueMax); + usedCoreNum_ = std::min(static_cast(mCnt * nCnt * groupNum_), compileInfoPtr->aicNum); +} + +void GroupedNoQuantMatmulTiling::CalAswtL1Tiling(const GMMCompileInfo *compileInfoPtr) +{ + uint64_t totalL1Size = compileInfoPtr->l1Size + 256UL; // 256B为预留给rpc使用，单算子不涉及 + uint64_t reserveBTSize = hasBias_ ? BIAS_TABLE_NUM * DATA_SIZE_FP32 : 0UL; + depthA1_ = totalL1Size / NUM_TWO / baseM_ / baseK_ / GetSizeByDataType(xDType_); // 2: half of l1 + depthB1_ = totalL1Size / NUM_TWO / baseN_ / baseK_ / GetSizeByDataType(weightDtype_); // 2: half of l1 + + uint64_t depthASize = depthA1_ * baseM_ * baseK_ * GetSizeByDataType(xDType_); + uint64_t depthBSize = depthB1_ * baseN_ * baseK_ * GetSizeByDataType(weightDtype_); + if (depthASize + depthBSize > totalL1Size - reserveBTSize) { + if (baseM_ <= baseN_) { + depthA1_ = std::max(depthA1_ / NUM_TWO, 1UL); // 2: adjust deptch for l1 buffer + } else { + depthB1_ = std::max(depthB1_ / NUM_TWO, 1UL); // 2: adjust deptch for l1 buffer + } + } + stepKa_ = std::max(depthA1_ / DB_SIZE, 1UL); + stepKb_ = std::max(depthB1_ / DB_SIZE, 1UL); + if (stepKa_ >= stepKb_) { + stepKa_ = stepKa_ / stepKb_ * stepKb_; + } else { + stepKb_ = stepKb_ / stepKa_ * stepKa_; + } + depthA1_ = stepKa_ * DB_SIZE; // depth % (stepKa * stepM) == 0 + depthB1_ = stepKb_ * DB_SIZE; // depth % (stepKb * stepN) == 0 + return; +} + +bool GroupedNoQuantMatmulTiling::CalL1Tiling(const gert::TilingContext *context, const GMMCompileInfo *compileInfoPtr) +{ + uint64_t reserveBTSize = hasBias_ ? BIAS_TABLE_NUM * DATA_SIZE_FP32 : 0UL; + uint64_t totalL1Size = hasBias_ ? compileInfoPtr->l1Size - reserveBTSize : compileInfoPtr->l1Size; + uint64_t l1ASize = baseM_ > baseN_ ? PARTA_L1_SIZE : totalL1Size - PARTA_L1_SIZE; + uint64_t l1BSize = totalL1Size - l1ASize; + stepKa_ = l1ASize / NUM_TWO / baseM_ / baseK_ / GetSizeByDataType(xDType_); // 2: half of l1 + stepKb_ = l1BSize / NUM_TWO / baseN_ / baseK_ / GetSizeByDataType(weightDtype_); // 2: half of l1 + + OP_CHECK_IF(stepKa_ == 0 || stepKb_ == 0, OP_LOGE(context->GetNodeName(), "stepka or stepkb cannot be 0"), + return false); + if (stepKa_ >= stepKb_) { + stepKa_ = stepKa_ / stepKb_ * stepKb_; + } else { + stepKb_ = stepKb_ / stepKa_ * stepKa_; + } + depthA1_ = stepKa_ * DB_SIZE; + depthB1_ = stepKb_ * DB_SIZE; + return true; +} + +void GroupedNoQuantMatmulTiling::CalcTailBasicBlock(const GMMCompileInfo *compileInfoPtr) +{ + uint64_t mCnt = DlinferGroupedMatmulDirect::CeilDiv(m_, baseM_); + uint64_t nCnt = DlinferGroupedMatmulDirect::CeilDiv(n_, baseN_); + uint64_t mnCnt = mCnt * nCnt; + uint64_t tailCnt = mnCnt <= compileInfoPtr->aicNum ? 0UL : mnCnt % compileInfoPtr->aicNum; + + if (tailCnt != 0UL) { + while ((mTailCnt_ + 1UL) * nTailCnt_ * tailCnt <= compileInfoPtr->aicNum) { + mTailCnt_ += 1UL; + if (mTailCnt_ * (nTailCnt_ + 1UL) * tailCnt <= compileInfoPtr->aicNum) { + nTailCnt_ += 1UL; + } + } + } +} + +bool GroupedNoQuantMatmulTiling::CheckWeightNzShape(const gert::TilingContext *context, int64_t c0) +{ + int i = 0; + while (true) { + auto wTensor = context->GetDynamicInputTensor(INDEX_WEIGHT, i++); + if (wTensor == nullptr) { + break; + } + gert::Shape wOriginShape = wTensor->GetOriginShape(); + int64_t lastDimValue = wOriginShape.GetDim(wOriginShape.GetDimNum() - 1); + OP_CHECK_IF(lastDimValue % c0 != 0, + OP_LOGE(context->GetNodeName(), + "the inner axis size of nz weight is expected to be a multiple of 32B, but now is %ld", + lastDimValue), + return false); + } + return true; +} + +bool GroupedNoQuantMatmulTiling::Init(const gert::TilingContext *context) +{ + OP_CHECK_IF(!GetAttrs(context), OP_LOGE(context->GetNodeName(), "Unable to calculate matmul-tiling"), return false); + + auto xTensor = context->GetDynamicInputTensor(INDEX_X, 0); + OP_CHECK_IF(xTensor == nullptr, OP_LOGE(context->GetNodeName(), "xTensor is nullptr."), return false); + gert::Shape xShape = xTensor->GetStorageShape(); + xDimNum_ = static_cast(xShape.GetDimNum()); + + auto wTensor = context->GetDynamicInputTensor(INDEX_WEIGHT, 0); + OP_CHECK_IF(wTensor == nullptr, OP_LOGE(context->GetNodeName(), "wTensor is nullptr."), return false); + gert::Shape wShape = wTensor->GetOriginShape(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + + OP_CHECK_IF(wDimNum < MIN_DIM || xDimNum_ < MIN_DIM, + OP_LOGE(context->GetNodeName(), "The dimension of x or weight should be at least 2"), return false); + xKDim_ = transposeX_ ? 0U : xDimNum_ - 1U; + weightNDim_ = transposeWeight_ ? wDimNum - DIM_TWO : wDimNum - DIM_ONE; + + auto biasPtr = context->GetDynamicInputTensor(INDEX_BIAS, 0); // 0: obtain the first tensor of the tensorList + hasBias_ = !(biasPtr == nullptr || biasPtr->GetStorageShape().GetShapeSize() == 0); + + isSingleWeight_ = (context->GetDynamicInputTensor(INDEX_WEIGHT, 1) == nullptr); + isSingleX_ = (context->GetDynamicInputTensor(INDEX_X, 1) == nullptr); + isSingleY_ = (splitItem_ == 2 || splitItem_ == 3); // 2&3: output tensor is single + + if (weightNzFlag_) { + uint64_t c0 = BLOCK_SIZE / std::max(1, GetSizeByDataType(weightDtype_)); + if (wDimNum > NZ_DIM_NUM) { + weightNDim_ = transposeWeight_ ? wDimNum - DIM_THREE : wDimNum - DIM_FOUR; + nzFactor_ = transposeWeight_ ? BASIC_BLOCK_SIZE_16 : static_cast(c0); + } else { + OP_CHECK_IF(!CheckWeightNzShape(context, static_cast(c0)), + OP_LOGE(context->GetNodeName(), "the size of nz weight is invaild."), return false); + } + } + + if (groupType_ == SPLIT_K) { + return GMMGetTensorShapeSplitK(context, xShape, wShape); + } + if (groupType_ == SPLIT_M) { + return GMMGetTensorShapeSplitM(context, xShape, wShape); + } + if (groupType_ == NO_SPLIT) { // not split any axis + if (isSingleWeight_ && wDimNum > 2U) { // 2: dim of splited weight tensor + return SeparatedXSingleWeight(context, wShape); + } + return SeparatedXSeparatedWeight(context); + } + OP_LOGE(context->GetNodeName(), + "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", groupType_, + isSingleWeight_, isSingleX_, isSingleY_); + return false; +} + +bool GroupedNoQuantMatmulTiling::GetAttrs(const gert::TilingContext *context) +{ + auto attr = context->GetAttrs(); + OP_CHECK_IF(attr == nullptr, OP_LOGE(context->GetNodeName(), "attr is nullptr."), + return false); // check attr is not null + const bool *transposeWeightPtr = attr->GetAttrPointer(ATTR_IDX_TRANS_W); + const bool *transposeXPtr = attr->GetAttrPointer(ATTR_IDX_TRANS_X); + const int32_t *groupTypePtr = attr->GetAttrPointer(ATTR_IDX_GROUPTYPE); + const int64_t *splitItemPtr = attr->GetAttrPointer(ATTR_IDX_SPLIT_ITEM); + const uint32_t *groupListTypePtr = attr->GetAttrPointer(ATTR_IDX_GROUP_LIST_TYPE); + transposeWeight_ = transposeWeightPtr != nullptr ? *transposeWeightPtr : false; + transposeX_ = transposeXPtr != nullptr ? *transposeXPtr : false; + groupType_ = groupTypePtr != nullptr ? *groupTypePtr : NO_SPLIT; + splitItem_ = splitItemPtr != nullptr ? *splitItemPtr : 0; + groupListType_ = groupListTypePtr != nullptr ? *groupListTypePtr : 0; + + auto xDesc = context->GetDynamicInputDesc(INDEX_X, 0); + OP_CHECK_IF(xDesc == nullptr, OP_LOGE(context->GetNodeName(), "xDesc is nullptr."), return false); + xDType_ = xDesc->GetDataType(); + + auto w0Desc = context->GetDynamicInputDesc(INDEX_WEIGHT, 0); + OP_CHECK_IF(w0Desc == nullptr, OP_LOGE(context->GetNodeName(), "w0Desc is nullptr."), return false); + weightDtype_ = w0Desc->GetDataType(); + auto wFormat0 = static_cast(ge::GetPrimaryFormat(w0Desc->GetStorageFormat())); + weightNzFlag_ = wFormat0 == ge::FORMAT_FRACTAL_NZ; + return true; +} + +bool GroupedNoQuantMatmulTiling::CalMatMulTiling(const gert::TilingContext *context, + const GMMCompileInfo *compileInfoPtr) +{ + if (groupNum_ == 0U) { + OP_LOGE(context->GetNodeName(), "gmm no quant groupNum_ cannot be 0"); + return false; + } + if (groupNum_ == 1U) { + FormulateBasicBlock(compileInfoPtr, usedCoreNum_); + CalcTailBasicBlock(compileInfoPtr); + CalAswtL1Tiling(compileInfoPtr); + return true; + } + if (groupType_ == SPLIT_M || groupType_ == NO_SPLIT) { + return CalBaseMMTiling(context, compileInfoPtr); + } else if (groupType_ == SPLIT_K) { + uint32_t remainCoreNum = std::max(1U, compileInfoPtr->aicNum / groupNum_); + FormulateBasicBlock(compileInfoPtr, remainCoreNum); + return true; + } + return false; +} + +bool GroupedNoQuantMatmulTiling::SetGMMTiling() +{ + errno_t retM = memcpy_s(tilingData_.gmmArray.mList, sizeof(tilingData_.gmmArray.mList), mList_, sizeof(mList_)); + if (retM != EOK) { + return false; + } + errno_t retK = memcpy_s(tilingData_.gmmArray.kList, sizeof(tilingData_.gmmArray.kList), kList_, sizeof(kList_)); + if (retK != EOK) { + return false; + } + errno_t retN = memcpy_s(tilingData_.gmmArray.nList, sizeof(tilingData_.gmmArray.nList), nList_, sizeof(nList_)); + if (retN != EOK) { + return false; + } + tilingData_.gmmNoQuantParam.groupNum = groupNum_; + tilingData_.gmmNoQuantParam.hasBias = static_cast(hasBias_); + tilingData_.gmmNoQuantParam.groupType = groupType_; + tilingData_.gmmNoQuantParam.groupListType = groupListType_; + tilingData_.gmmNoQuantParam.singleWeight = static_cast(isSingleWeight_); + tilingData_.gmmNoQuantParam.singleX = static_cast(isSingleX_); + tilingData_.gmmNoQuantParam.singleY = static_cast(isSingleY_); + tilingData_.gmmNoQuantParam.coreNum = usedCoreNum_; + tilingData_.gmmNoQuantParam.mTailCnt = static_cast(mTailCnt_); + tilingData_.gmmNoQuantParam.nTailCnt = static_cast(nTailCnt_); + return true; +} + +void GroupedNoQuantMatmulTiling::SetMatMulTiling() +{ + tilingData_.mmTilingData.isBias = static_cast(hasBias_); + tilingData_.mmTilingData.M = m_; + tilingData_.mmTilingData.N = n_; + tilingData_.mmTilingData.Ka = k_; + tilingData_.mmTilingData.Kb = k_; + tilingData_.mmTilingData.singleCoreM = m_; + tilingData_.mmTilingData.singleCoreN = baseN_; + tilingData_.mmTilingData.singleCoreK = k_; + tilingData_.mmTilingData.dbL0A = DB_SIZE; + tilingData_.mmTilingData.dbL0B = DB_SIZE; + tilingData_.mmTilingData.dbL0C = 1; + tilingData_.mmTilingData.baseM = baseM_; + tilingData_.mmTilingData.baseN = baseN_; + tilingData_.mmTilingData.baseK = baseK_; + tilingData_.mmTilingData.stepKa = stepKa_; + tilingData_.mmTilingData.stepKb = stepKb_; + tilingData_.mmTilingData.depthA1 = depthA1_; + tilingData_.mmTilingData.depthB1 = depthB1_; + tilingData_.mmTilingData.stepM = 1; + tilingData_.mmTilingData.stepN = 1; + tilingData_.mmTilingData.usedCoreNum = usedCoreNum_; +} + +bool GroupedNoQuantMatmulTiling::SetCustomParam(gert::TilingContext *context) +{ + size_t *workspaces = context->GetWorkspaceSizes(1); // get second variable + OP_CHECK_IF(workspaces == nullptr, OP_LOGE(context->GetNodeName(), "workspaces is nullptr."), + return false); // check workspaces is not null + workspaces[0] = 32U; // 32: default workspace size + + context->SetBlockDim(usedCoreNum_); + OP_CHECK_IF(context->GetRawTilingData() == nullptr, OP_LOGE(context->GetNodeName(), "RawTilingData is nullptr."), + return false); + errno_t ret = memcpy_s(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity(), + reinterpret_cast(&tilingData_), sizeof(tilingData_)); + if (ret != EOK) { + OP_LOGE(context->GetNodeName(), "memcpy_s failed, ret = %d", ret); + return false; + } + context->GetRawTilingData()->SetDataSize(sizeof(tilingData_)); + return true; +} + +void GroupedNoQuantMatmulTiling::SetTilingKey(gert::TilingContext *context) +{ + tilingKeyBuilder_.gmmTrans = static_cast(transposeX_) | (static_cast(transposeWeight_) << 1); + context->SetTilingKey(tilingKeyBuilder_.GenTilingKey()); +} + +bool GroupedNoQuantMatmulTiling::GMMGetTensorShapeSplitM(const gert::TilingContext *context, const gert::Shape xShape, + const gert::Shape wShape) +{ + if (isSingleX_ && isSingleWeight_ && isSingleY_) { // split M, s-s-s + return SplitMSingleXSingleWeightSingleY(xShape, wShape); + } + if (!isSingleX_ && !isSingleWeight_ && isSingleY_) { // split M, m-m-s + return SeparatedXSeparatedWeight(context); + } + if (isSingleX_ && !isSingleWeight_ && !isSingleY_) { // splitM, s-m-m + return SplitMSingleXSeparatedWeight(context, xShape); + } + if (isSingleX_ && !isSingleWeight_ && isSingleY_) { // split M, s-m-s + return SplitMSingleXSeparatedWeight(context, xShape); + } + if (!isSingleX_ && !isSingleWeight_ && !isSingleY_) { // split M, m-m-m + return SeparatedXSeparatedWeight(context); + } + if (!isSingleX_ && isSingleWeight_) { // split M, m-s-m/m-s-s + return SeparatedXSingleWeight(context, wShape); + } + OP_LOGE(context->GetNodeName(), + "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", groupType_, + isSingleWeight_, isSingleX_, isSingleY_); + return false; +} + +bool GroupedNoQuantMatmulTiling::GMMGetTensorShapeSplitK(const gert::TilingContext *context, const gert::Shape xShape, + const gert::Shape wShape) +{ + if (isSingleX_ && isSingleWeight_ && isSingleY_) { // splitK, s-s-s + return SplitKSingleXSingleWeightSingleY(context, xShape, wShape); + } + OP_LOGE(context->GetNodeName(), + "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", groupType_, + isSingleWeight_, isSingleX_, isSingleY_); + return false; +} + +/** @brief split M：single-single-single(s-s-s) + */ +bool GroupedNoQuantMatmulTiling::SplitMSingleXSingleWeightSingleY(const gert::Shape xShape, const gert::Shape wShape) +{ + groupNum_ = static_cast(wShape.GetDim(0)); + int64_t m = xShape.GetDim(0); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + kList_[0] = static_cast(k); // if split M axis, the K axis values of x tensorList are all the same. + nList_[0] = static_cast(n); + mList_[0] = -1; + m_ = static_cast(m); + k_ = static_cast(k); + n_ = static_cast(n); + return true; +} + +/** @brief split M：single-multi-single(s-m-s)/single-multi-multi(s-m-m), share the same function. + */ +bool GroupedNoQuantMatmulTiling::SplitMSingleXSeparatedWeight(const gert::TilingContext *context, + const gert::Shape xShape) +{ + int64_t m = xShape.GetDim(0); + int64_t k = xShape.GetDim(xKDim_); + for (uint32_t i = 0; i < MAX_TENSOR; i++) { + auto wTensor = context->GetDynamicInputTensor(INDEX_WEIGHT, i); + if (wTensor == nullptr) { + break; + } // when x has multi tensors, xTensor is allowed to be empty + auto wShape = wTensor->GetOriginShape(); + + groupNum_ += 1U; + kList_[i] = static_cast(k); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + nList_[i] = static_cast(n); + n_ = std::max(n_, static_cast(n)); + } + mList_[0] = -1; // mList is unknown right now + m_ = static_cast(m); + k_ = static_cast(k); + return true; +} + +/** @brief split M：multi-multi-single(m-m-s); no split: multi-multi-multi(m-m-m), share the same function + */ +bool GroupedNoQuantMatmulTiling::SeparatedXSeparatedWeight(const gert::TilingContext *context) +{ + for (uint32_t i = 0; i < MAX_TENSOR; i++) { + auto wTensor = context->GetDynamicInputTensor(INDEX_WEIGHT, i); + auto xTensor = context->GetDynamicInputTensor(INDEX_X, i); + if (wTensor == nullptr || xTensor == nullptr) { + break; + } + auto wShape = wTensor->GetOriginShape(); + auto xShape = xTensor->GetStorageShape(); + groupNum_ += 1U; + int64_t m = xShape.GetDim(0); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + mList_[i] = static_cast(m); + kList_[i] = static_cast(k); + nList_[i] = static_cast(n); + m_ = std::max(m_, static_cast(m)); + k_ = std::max(k_, static_cast(k)); + n_ = std::max(n_, static_cast(n)); + } + groupType_ = NO_SPLIT; + return true; +} + +/** @brief split M : multi-single-multi(m-s-m), split K : multi-single-multi(m-s-m), share the same function + */ +bool GroupedNoQuantMatmulTiling::SeparatedXSingleWeight(const gert::TilingContext *context, const gert::Shape wShape) +{ + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + for (uint32_t i = 0; i < MAX_TENSOR; i++) { + auto xTensor = context->GetDynamicInputTensor(INDEX_X, i); + if (xTensor == nullptr) { + break; + } // when x has multi tensors, xTensor is allowed to be empty + auto xShape = xTensor->GetStorageShape(); + groupNum_ += 1U; + int64_t m = xShape.GetDim(0); + int64_t k = xShape.GetDim(xKDim_); + mList_[i] = static_cast(m); + kList_[i] = static_cast(k); + nList_[i] = static_cast(n); + m_ = std::max(m_, static_cast(m)); + k_ = std::max(k_, static_cast(k)); + } + n_ = static_cast(n); + groupType_ = NO_SPLIT; + return true; +} + +/** @brief split K single-single-single + */ +bool GroupedNoQuantMatmulTiling::SplitKSingleXSingleWeightSingleY(const gert::TilingContext *context, + const gert::Shape xShape, const gert::Shape wShape) +{ + int64_t m = xShape.GetDim(1); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + + auto groupListTensor = context->GetDynamicInputTensor(INDEX_GROUPLIST, 0); + if (groupListTensor == nullptr) { + OP_LOGE(context->GetNodeName(), "groupListTensor is nullptr"); + return false; + } + gert::Shape groupListShape = groupListTensor->GetStorageShape(); + groupNum_ = static_cast(groupListShape.GetDim(0)); // 0: the first dim of groupList is groupNum + mList_[0] = static_cast(m); + nList_[0] = static_cast(n); + kList_[0] = -1; + m_ = static_cast(m); + n_ = static_cast(n); + k_ = static_cast(k); + return true; +} + +void GroupedNoQuantMatmulTiling::PrintTilingResult(const gert::TilingContext *context) +{ + OP_LOGI(context->GetNodeName(), + "GMM Tiling result: groupNum: %u, singleX: %u, singleWeight: %u, singleY: %u," + "groupType: %d, groupListType: %u, hasBias: %u, mTailCnt: %u, nTailCnt: %u", + tilingData_.gmmNoQuantParam.groupNum, tilingData_.gmmNoQuantParam.singleX, + tilingData_.gmmNoQuantParam.singleWeight, tilingData_.gmmNoQuantParam.singleY, + tilingData_.gmmNoQuantParam.groupType, tilingData_.gmmNoQuantParam.groupListType, + tilingData_.gmmNoQuantParam.hasBias, tilingData_.gmmNoQuantParam.mTailCnt, + tilingData_.gmmNoQuantParam.nTailCnt); + + OP_LOGI(context->GetNodeName(), + "GMM MatMul Tiling result: usedCoreNum: %d, baseM: %d, baseN: %d, baseK: %d, stepKa: %d," + "stepKb: %d, depthA1: %d, depthB1: %d", + tilingData_.mmTilingData.usedCoreNum, tilingData_.mmTilingData.baseM, tilingData_.mmTilingData.baseN, + tilingData_.mmTilingData.baseK, tilingData_.mmTilingData.stepKa, tilingData_.mmTilingData.stepKb, + tilingData_.mmTilingData.depthA1, tilingData_.mmTilingData.depthB1); +} + +uint64_t TilingKeyBuilder::GenTilingKey() +{ + uint64_t transInfo = static_cast(this->gmmTrans); + bool atrans_ = (transInfo == static_cast(GmmTrans::ATrans)) || + (transInfo == static_cast(GmmTrans::ABTrans)); + bool btrans_ = (transInfo == static_cast(GmmTrans::BTrans)) || + (transInfo == static_cast(GmmTrans::ABTrans)); + return GET_TPL_TILING_KEY(static_cast(btrans_), static_cast(atrans_)); +} +} // namespace optiling diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.h new file mode 100644 index 00000000..cf4fcbdc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_no_quant_matmul_tiling.h @@ -0,0 +1,155 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_no_quant_matmul_tiling.h + * \brief + */ +#ifndef GROUPED_NO_QUANT_MATMUL_TILING_H +#define GROUPED_NO_QUANT_MATMUL_TILING_H + +#include "../grouped_matmul_tiling.h" +#include "../../../op_kernel/arch35/grouped_matmul_tiling_data_apt.h" +#include "log/log.h" +#include "log/error_code.h" +#include "register/op_impl_registry.h" + +namespace optiling { +constexpr uint64_t TILING_KEY_DELIMITER = 9UL; +constexpr uint64_t BASIC_BLOCK_SIZE_16 = 16UL; +constexpr uint64_t STEP_K_DEFAULT = 4UL; +constexpr uint64_t DEPTH_DEFAULT = 8UL; +constexpr uint64_t DB_SIZE = 2UL; +constexpr uint32_t NZ_DIM_NUM = 4U; +constexpr uint32_t MIN_DIM = 2U; +constexpr uint64_t BASE_M_DEFAULT = 256UL; +constexpr uint64_t BASE_N_DEFAULT = 256UL; +constexpr uint64_t BASE_K_DEFAULT = 64UL; +constexpr uint64_t NUM_TWO = 2UL; +constexpr uint64_t BIAS_TABLE_NUM = 256UL; +constexpr uint64_t DATA_SIZE_FP32 = 4UL; +constexpr uint32_t INDEX_X = 0U; +constexpr uint32_t INDEX_WEIGHT = 1U; +constexpr uint32_t INDEX_BIAS = 2U; +constexpr uint32_t MAX_TENSOR = 128U; +constexpr uint32_t INDEX_GROUPLIST = 7U; +constexpr uint32_t DIM_ONE = 1U; +constexpr uint32_t DIM_TWO = 2U; +constexpr uint32_t DIM_THREE = 3U; +constexpr uint32_t DIM_FOUR = 4U; +constexpr uint32_t FP32_DTYPE_SIZE = 4U; +constexpr uint64_t ATTR_IDX_SPLIT_ITEM = 0UL; +constexpr uint64_t ATTR_IDX_TRANS_W = 2UL; +constexpr uint64_t ATTR_IDX_TRANS_X = 3UL; +constexpr uint64_t ATTR_IDX_GROUPTYPE = 4UL; +constexpr uint64_t ATTR_IDX_GROUP_LIST_TYPE = 5UL; +constexpr int32_t NO_SPLIT = -1; +constexpr int32_t SPLIT_M = 0; +constexpr int32_t SPLIT_K = 2; +constexpr uint64_t BLOCK_SIZE = 32UL; +constexpr uint64_t ALIGN_DOWN_16 = 15UL; +constexpr uint64_t PARTA_L1_SIZE = 256UL * 1024UL; + +enum class GMMTrans : std::uint8_t { + NO_TRANS = 0, + A_TRANS = 1, + B_TRANS = 2, + AB_TRANS = 3 +}; + +enum class Model : std::uint8_t { + BASIC = 0 +}; + +class TilingKeyBuilder { +public: + // 对应0-1位 平台大类，平台小类 + uint8_t model = 0; + + // 对应20位，转置场景 + uint8_t gmmTrans = 0; + +public: + uint64_t GenTilingKey(); +}; + +class GroupedNoQuantMatmulTiling { +public: + bool SetTiling(gert::TilingContext *context); + +protected: + bool Init(const gert::TilingContext* context); + bool CalBaseMMTiling(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr); + void FormulateBasicBlock(const GMMCompileInfo* compileInfoPtr, uint32_t remainCoreNum); + void CalAswtL1Tiling(const GMMCompileInfo* compileInfoPtr); + bool CalL1Tiling(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr); + void CalcTailBasicBlock(const GMMCompileInfo* compileInfoPtr); + bool SetCustomParam(gert::TilingContext *context); + bool GetAttrs(const gert::TilingContext* context); + bool CalMatMulTiling(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr); + bool SetGMMTiling(); + void SetMatMulTiling(); + void SetTilingKey(gert::TilingContext *context); + bool GMMGetTensorShapeSplitM(const gert::TilingContext* context, const gert::Shape xShape, const gert::Shape wShape); + bool GMMGetTensorShapeSplitK(const gert::TilingContext* context, const gert::Shape xShape, const gert::Shape wShape); + bool SplitMSingleXSingleWeightSingleY(const gert::Shape xShape, const gert::Shape wShape); + bool SplitMSingleXSeparatedWeight(const gert::TilingContext* context, const gert::Shape xShape); + bool SeparatedXSeparatedWeight(const gert::TilingContext* context); + bool SeparatedXSingleWeight(const gert::TilingContext* context, const gert::Shape wShape); + bool SplitKSingleXSingleWeightSingleY(const gert::TilingContext* context, const gert::Shape xShape, const gert::Shape wShape); + bool CheckWeightNzShape(const gert::TilingContext* context, int64_t c0); + void PrintTilingResult(const gert::TilingContext *context); + +private: + int32_t mList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t kList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t nList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + + bool transposeX_ = false; + bool transposeWeight_ = false; + bool isSingleX_ = true; + bool isSingleWeight_ = true; + bool isSingleY_ = true; + bool hasBias_ = false; + bool weightNzFlag_ = false; + + uint64_t m_ = 0; + uint64_t k_ = 0; + uint64_t n_ = 0; + uint64_t nSizeOri_ = 0; + int32_t groupType_ = 0; + int64_t splitItem_ = 0; + uint32_t groupNum_ = 0; + uint32_t groupListType_ = 0; + uint32_t groupSize_ = 0; + uint32_t xKDim_ = 0; + uint32_t weightNDim_ = 0; + uint32_t xDimNum_ = 0; + int64_t nzFactor_ = 1; // for weight nz format + uint64_t baseM_ = BASE_M_DEFAULT; + uint64_t baseN_ = BASE_N_DEFAULT; + uint64_t baseK_ = BASE_K_DEFAULT; + uint32_t usedCoreNum_ = 0; + uint64_t stepKa_ = STEP_K_DEFAULT; + uint64_t stepKb_ = STEP_K_DEFAULT; + uint64_t depthA1_ = DEPTH_DEFAULT; + uint64_t depthB1_ = DEPTH_DEFAULT; + uint64_t mTailCnt_ = 1UL; + uint64_t nTailCnt_ = 1UL; + + ge::DataType xDType_ = ge::DT_UNDEFINED; + ge::DataType weightDtype_ = ge::DT_UNDEFINED; + + TilingKeyBuilder tilingKeyBuilder_; + DlinferGroupedMatmulDirectTilingData::GMMNoQuantTilingData tilingData_; +}; +} // namespace optiling + +#endif // GROUPED_WEIGHT_QUANT_BATCH_MATMUL_TILING_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.cpp new file mode 100644 index 00000000..061d8a65 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.cpp @@ -0,0 +1,985 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include +#include "grouped_quant_matmul_tiling.h" + +#include "log/log.h" +#include "log/error_code.h" +#include "tiling_base/tiling_templates_registry.h" +#include "tiling_base/tiling_type.h" +#include "../../../op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_tiling_key.h" +using namespace Ops::Transformer::OpTiling; +using namespace DlinferGroupedMatmulDirect; +using namespace optiling::GmmConstant; +using GMMQuantTilingData = DlinferGroupedMatmulDirectTilingData::GMMQuantTilingData; +using GMMQuantParams = DlinferGroupedMatmulDirectTilingData::GMMQuantParams; +namespace optiling { + +bool GroupedQbmmTiling::IsCapable() +{ + return true; +} + +void GroupedQbmmTiling::Reset() +{ + tilingData_ = GMMQuantTilingData(); +} + +ge::graphStatus GroupedQbmmTiling::GetPlatformInfo() +{ + auto platformInfoPtr = context_->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, + OP_LOGE(context_->GetNodeName(), "CompileInfoPtr is null."), + return ge::GRAPH_FAILED); + + aicoreParams_.aicNum = compileInfoPtr->aicNum; + aicoreParams_.ubSize = compileInfoPtr->ubSize; + aicoreParams_.l1Size = compileInfoPtr->l1Size; + aicoreParams_.l0aSize = compileInfoPtr->l0ASize; + aicoreParams_.l0bSize = compileInfoPtr->l0BSize; + aicoreParams_.l0cSize = compileInfoPtr->l0CSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + aicoreParams_.aicNum = ascendcPlatform.GetCoreNumAic(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, aicoreParams_.ubSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, aicoreParams_.l1Size); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_A, aicoreParams_.l0aSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, aicoreParams_.l0bSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, aicoreParams_.l0cSize); + } + + OP_LOGI(context_, "Platform info: aicNum(%lu) ubSize(%lu) l1Size(%lu) l0aSize(%lu) l0bSize(%lu) l0cSize(%lu).", + aicoreParams_.aicNum, aicoreParams_.ubSize, aicoreParams_.l1Size, aicoreParams_.l0aSize, + aicoreParams_.l0bSize, aicoreParams_.l0cSize); + return ge::GRAPH_SUCCESS; +} + +bool GroupedQbmmTiling::IsMicroScaling() const +{ + return inputParams_.scaleDtype == ge::DT_FLOAT8_E8M0; +} + +bool GroupedQbmmTiling::AnalyzeAttrs() +{ + auto attrs = context_->GetAttrs(); + if (attrs) { + OP_CHECK_IF(attrs->GetAttrNum() < ATTR_INDEX_ACT_TYPE + 1, + OP_LOGE(inputParams_.opName, + "The num of attrs should be greater than %lu, actual is %zu", + ATTR_INDEX_ACT_TYPE + 1, attrs->GetAttrNum()), + return false); + const int64_t *splitItemPtr = attrs->GetAttrPointer(ATTR_INDEX_SPLIT_ITEM); + const bool *transposeWeightPtr = attrs->GetAttrPointer(ATTR_INDEX_TRANS_W); + const bool *transposeXPtr = attrs->GetAttrPointer(ATTR_INDEX_TRANS_X); + const int64_t *groupTypePtr = attrs->GetAttrPointer(ATTR_INDEX_GROUPTYPE); + const int64_t *groupListTypePtr = attrs->GetAttrPointer(ATTR_INDEX_GROUP_LIST_TYPE); // 通路保证非负数 + const int64_t *actTypePtr = attrs->GetAttrPointer(ATTR_INDEX_ACT_TYPE); + + inputParams_.transB = transposeWeightPtr != nullptr ? *transposeWeightPtr : false; + inputParams_.transA = transposeXPtr != nullptr ? *transposeXPtr : false; + inputParams_.groupType = groupTypePtr != nullptr ? *groupTypePtr : inputParams_.groupType; + inputParams_.splitItem = splitItemPtr != nullptr ? *splitItemPtr : inputParams_.splitItem; + inputParams_.actType = actTypePtr != nullptr ? *actTypePtr : inputParams_.actType; + inputParams_.groupListType = groupListTypePtr != nullptr ? *groupListTypePtr : inputParams_.groupListType; + } + OP_CHECK_IF( + inputParams_.groupType != SPLIT_M && inputParams_.groupType != SPLIT_K, + OP_LOGE(inputParams_.opName, "Only support group type is 0 or 2 when the dtype of x is %s, actual is %d", + ge::TypeUtils::DataTypeToSerialString(inputParams_.aDtype).c_str(), inputParams_.groupType), + return false); + OP_CHECK_IF( + (inputParams_.aDtype == ge::DT_FLOAT4_E2M1 || inputParams_.aDtype == ge::DT_FLOAT4_E1M2) && + inputParams_.groupType != SPLIT_M, + OP_LOGE(inputParams_.opName, "Only support group type to be 0 when the dtype of x is FLOAT4, actual is %d.", + inputParams_.groupType), + return false); + if (inputParams_.groupType == SPLIT_M) { + OP_CHECK_IF(inputParams_.transA, + OP_LOGE(inputParams_.opName, "When group type is 0, transA can only be false."), + return false); + } else { + OP_CHECK_IF(!inputParams_.transA, + OP_LOGE(inputParams_.opName, "When group type is 2, transA can only be true."), + return false); + OP_CHECK_IF(inputParams_.transB, + OP_LOGE(inputParams_.opName, "When group type is 2, transB can only be false."), + return false); + } + + inputParams_.isSingleX = (context_->GetDynamicInputDesc(X_INDEX, 1) == nullptr); + inputParams_.isSingleW = (context_->GetDynamicInputDesc(WEIGHT_INDEX, 1) == nullptr); + // 2: when x is multi-tensor, y is single-tensor; 3: when x is single-tensor, y is single-tensor + inputParams_.isSingleY = (inputParams_.splitItem == 2 || inputParams_.splitItem == 3); + return true; +} + +bool GroupedQbmmTiling::CheckBiasDtype() const +{ + if ((inputParams_.aDtype == ge::DT_FLOAT4_E2M1 || inputParams_.aDtype == ge::DT_FLOAT4_E1M2)) { + OP_CHECK_IF(inputParams_.biasDtype != ge::DT_FLOAT, + OP_LOGE(inputParams_.opName, + "The dtype of bias should be FLOAT when the dtype of x is FLOAT4, actual is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.biasDtype).c_str()), + return false); + } else if (inputParams_.aDtype == ge::DT_INT8) { + if (inputParams_.cDtype == ge::DT_BF16) { + OP_CHECK_IF(inputParams_.biasDtype != ge::DT_INT32 && inputParams_.biasDtype != ge::DT_BF16 && + inputParams_.biasDtype != ge::DT_FLOAT, + OP_LOGE(inputParams_.opName, + "The dtype of bias should be INT32, BF16 or FLOAT when the dtype of x is INT8 and the \ +dtype of output is BF16, actual is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.biasDtype).c_str()), + return false); + } else if (inputParams_.cDtype == ge::DT_FLOAT16) { + OP_CHECK_IF(inputParams_.biasDtype != ge::DT_INT32 && inputParams_.biasDtype != ge::DT_FLOAT16 && + inputParams_.biasDtype != ge::DT_FLOAT, + OP_LOGE(inputParams_.opName, + "The dtype of bias should be INT32, FLOAT16 or FLOAT when the dtype of x is INT8 and \ +the dtype of output is FLOAT16, actual is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.biasDtype).c_str()), + return false); + } else { + OP_LOGE(inputParams_.opName, "Invalid dtype of output %s with the dtype of x being INT8", + ge::TypeUtils::DataTypeToSerialString(inputParams_.cDtype).c_str()); + return false; + } + } else { + OP_LOGE(inputParams_.opName, "Bias is not supported when the dtype of x is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.aDtype).c_str()); + return false; + } + return true; +} + +bool GroupedQbmmTiling::AnalyzeDtype() +{ + static const std::vector legalInputDtypes = { + ge::DT_INT8, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2}; + auto xDesc = context_->GetDynamicInputDesc(X_INDEX, 0); + OP_CHECK_IF(xDesc == nullptr, OP_LOGE(context_->GetNodeName(), "xDesc is nullptr."), return false); + inputParams_.aDtype = xDesc->GetDataType(); + OP_CHECK_IF( + std::find(legalInputDtypes.begin(), legalInputDtypes.end(), inputParams_.aDtype) == legalInputDtypes.end(), + OP_LOGE(inputParams_.opName, + "The dtype of x should be in {INT8, HIFLOAT8, FLOAT8_E4M3, FLOAT8_E5M2, FLOAT4_E2M1, FLOAT4_E1M2}, \ +actual is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.aDtype).c_str()), + return false); + auto wDesc = context_->GetDynamicInputDesc(WEIGHT_INDEX, 0); + OP_CHECK_IF(wDesc == nullptr, OP_LOGE(context_->GetNodeName(), "wDesc is nullptr."), return false); + inputParams_.bDtype = wDesc->GetDataType(); + OP_CHECK_IF( + std::find(legalInputDtypes.begin(), legalInputDtypes.end(), inputParams_.bDtype) == legalInputDtypes.end(), + OP_LOGE(inputParams_.opName, + "The dtype of weight should be in {INT8, HIFLOAT8, FLOAT8_E4M3, FLOAT8_E5M2, FLOAT4_E2M1, \ +FLOAT4_E1M2}, actual is %s.", + ge::TypeUtils::DataTypeToSerialString(inputParams_.bDtype).c_str()), + return false); + inputParams_.bFormat = static_cast(ge::GetPrimaryFormat(wDesc->GetStorageFormat())); + auto biasStorageShape = context_->GetDynamicInputShape(BIAS_INDEX, 0); + inputParams_.hasBias = !(biasStorageShape == nullptr || biasStorageShape->GetStorageShape().GetShapeSize() == 0); + auto biasDesc = context_->GetDynamicInputDesc(BIAS_INDEX, 0); + OP_CHECK_IF(inputParams_.hasBias && biasDesc == nullptr, + OP_LOGE(inputParams_.opName, + "Bias from tensor is not nullptr, but bias from desc is nullptr."), + return false); + inputParams_.biasDtype = inputParams_.hasBias ? biasDesc->GetDataType() : inputParams_.biasDtype; + auto scaleDesc = context_->GetDynamicInputDesc(SCALE_INDEX, 0); + inputParams_.scaleDtype = scaleDesc != nullptr ? scaleDesc->GetDataType() : inputParams_.scaleDtype; + auto pertokenScaleDesc = context_->GetOptionalInputDesc(PER_TOKEN_SCALE_INDEX); + inputParams_.perTokenScaleDtype = + pertokenScaleDesc != nullptr ? pertokenScaleDesc->GetDataType() : inputParams_.perTokenScaleDtype; + + auto yDesc = context_->GetOutputDesc(Y_INDEX); + OP_CHECK_IF(yDesc == nullptr, OP_LOGE(context_->GetNodeName(), "yDesc is nullptr."), return false); + inputParams_.cDtype = yDesc->GetDataType(); + if (inputParams_.hasBias) { + OP_CHECK_IF(!CheckBiasDtype(), OP_LOGE(inputParams_.opName, "CheckBiasDtype failed."), return false); + } + return true; +} + +bool GroupedQbmmTiling::CheckQuantParamsForMXTypeM(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape) const +{ + auto xScaleDimNum = xScaleShape.GetDimNum(); + auto wScaleDimNum = wScaleShape.GetDimNum(); + OP_CHECK_IF(wScaleDimNum != MXFP_TYPE_M_SCALE_DIM_NUM, + OP_LOGE(inputParams_.opName, + "When split m, the dim num of scale should be 4 in mx quant mode, but actual \ +is %zu", wScaleDimNum), return false); + OP_CHECK_IF(xScaleDimNum != MXFP_PER_TOKEN_SCALE_DIM_NUM, + OP_LOGE( + inputParams_.opName, "When split m, the dim num of pertokenScale should be 3 in mx quant mode, but \ +actual is %zu", xScaleDimNum), return false); + auto wScaleEDim = static_cast(wScaleShape.GetDim(0)); + auto wScaleNDim = + static_cast(inputParams_.transB ? wScaleShape.GetDim(1) : + wScaleShape.GetDim(2)); // 2 is index for the third dim + auto wScaleKDim = + static_cast(inputParams_.transB ? wScaleShape.GetDim(2) : // 2 is index for the third dim + wScaleShape.GetDim(1)); + auto xScaleMDim = + static_cast(inputParams_.transA ? xScaleShape.GetDim(1) : + xScaleShape.GetDim(0)); + auto xScaleKDim = + static_cast(inputParams_.transA ? xScaleShape.GetDim(0) : + xScaleShape.GetDim(1)); + auto wScaleLastDim = static_cast(wScaleShape.GetDim(wScaleDimNum - 1)); + auto xScaleLastDim = static_cast(xScaleShape.GetDim(xScaleDimNum - 1)); + auto expectedKDimValue = CeilDiv(inputParams_.kSize, MXFP_BASEK_FACTOR); + OP_CHECK_IF(wScaleEDim != inputParams_.groupNum || wScaleKDim != expectedKDimValue || + wScaleNDim != inputParams_.nSize || wScaleLastDim != MXFP_MULTI_BASE_SIZE, + OP_LOGE( + inputParams_.opName, + "When split m in mx quant mode, the expected shape of scale is (%lu,%lu,%lu,2), but the actual \ +is (%lu,%lu,%lu,%lu).", + inputParams_.groupNum, inputParams_.nSize, expectedKDimValue, wScaleEDim, wScaleNDim, wScaleKDim, + wScaleLastDim), return false); + OP_CHECK_IF(xScaleMDim != inputParams_.mSize || xScaleKDim != expectedKDimValue || + xScaleLastDim != MXFP_MULTI_BASE_SIZE, + OP_LOGE( + inputParams_.opName, + "When split m in mx quant mode, the expected shape of pertokenScale is (%lu,%lu,2), but the actual \ +is (%lu,%lu,%lu).", inputParams_.mSize, expectedKDimValue, xScaleMDim, xScaleKDim, xScaleLastDim), return false); + return true; +} + +bool GroupedQbmmTiling::CheckQuantParamsForMXTypeK(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape) const +{ + auto xScaleDimNum = xScaleShape.GetDimNum(); + auto wScaleDimNum = wScaleShape.GetDimNum(); + OP_CHECK_IF(wScaleDimNum != MXFP_TYPE_K_SCALE_DIM_NUM, + OP_LOGE(inputParams_.opName, + "When split k, the dim num of scale should be 3 in mx quant mode, but actual \ +is %zu", wScaleDimNum), return false); + OP_CHECK_IF(xScaleDimNum != MXFP_PER_TOKEN_SCALE_DIM_NUM, + OP_LOGE( + inputParams_.opName, "When split k, the dim num of pertokenScale should be 3 in mx quant mode, but \ +actual is %zu", xScaleDimNum), return false); + auto xScaleLastDim = static_cast(xScaleShape.GetDim(xScaleDimNum - 1)); + auto xScaleKDim = static_cast( + inputParams_.transA ? xScaleShape.GetDim(0) : xScaleShape.GetDim(xScaleDimNum - LAST_SECOND_DIM_INDEX)); + auto xScaleMDim = static_cast( + inputParams_.transA ? xScaleShape.GetDim(xScaleDimNum - LAST_SECOND_DIM_INDEX) : xScaleShape.GetDim(0)); + auto wScaleLastDim = static_cast(wScaleShape.GetDim(wScaleDimNum - 1)); + auto wScaleNDim = static_cast( + inputParams_.transB ? wScaleShape.GetDim(0) : wScaleShape.GetDim(wScaleDimNum - LAST_SECOND_DIM_INDEX)); + auto wScaleKDim = static_cast( + inputParams_.transB ? wScaleShape.GetDim(wScaleDimNum - LAST_SECOND_DIM_INDEX) : wScaleShape.GetDim(0)); + auto expectedKDimValue = inputParams_.kSize / MXFP_BASEK_FACTOR + inputParams_.groupNum; + OP_CHECK_IF(!inputParams_.transA || inputParams_.transB, + OP_LOGE(inputParams_.opName, + "When split m in mx quant mode, the expected transpose attrs of x and \ +weight are true and false, but the actual transpose attrs of x and weight are %d and %d.", + inputParams_.transA, inputParams_.transB), return false); + OP_CHECK_IF(xScaleLastDim != MXFP_MULTI_BASE_SIZE || xScaleKDim != expectedKDimValue || + xScaleMDim != inputParams_.mSize, + OP_LOGE( + inputParams_.opName, "When split k in mx quant mode, the expected shape of pertokenScale is \ +(%lu,%lu,%lu), but the actual is (%lu,%lu,%lu).", + expectedKDimValue, inputParams_.mSize, MXFP_MULTI_BASE_SIZE, xScaleKDim, xScaleMDim, xScaleLastDim), + return false); + OP_CHECK_IF(wScaleLastDim != MXFP_MULTI_BASE_SIZE || wScaleKDim != expectedKDimValue || + wScaleNDim != inputParams_.nSize, + OP_LOGE( + inputParams_.opName, "When split k in mx quant mode, the expected shape of scale is (%lu,%lu,%lu), \ +but the actual is (%lu,%lu,%lu).", + expectedKDimValue, inputParams_.nSize, MXFP_MULTI_BASE_SIZE, wScaleKDim, wScaleNDim, wScaleLastDim), + return false); + return true; +} + +bool GroupedQbmmTiling::CheckQuantParamsForMxQuantMode(const gert::StorageShape *xScaleStorageShape, + const gert::Shape &wScaleShape) const +{ + // 多数参数在CheckQuantParamsForMxQuantMode函数调用前已有非空校验 + OP_CHECK_IF(xScaleStorageShape == nullptr, OP_LOGE(context_->GetNodeName(), "xScaleStorageShape is nullptr."), return false); + auto &xScaleShape = xScaleStorageShape->GetStorageShape(); + if (inputParams_.groupType == SPLIT_M) { + OP_CHECK_IF(!CheckQuantParamsForMXTypeM(xScaleShape, wScaleShape), + OP_LOGE(inputParams_.opName, "CheckQuantParamsForMXTypeM failed."), return false); + } else { + OP_CHECK_IF(!CheckQuantParamsForMXTypeK(xScaleShape, wScaleShape), + OP_LOGE(inputParams_.opName, "CheckQuantParamsForMXTypeK failed."), return false); + } + return true; +} + + +bool GroupedQbmmTiling::CheckQuantParamsForNonKGroupQuantMode(const gert::Shape &wScaleShape) const +{ + auto wScaleDimNum = wScaleShape.GetDimNum(); + // dim num 1 for the shape (g,), dim num 2 for the shape (g,1) or (g,n) + OP_CHECK_IF(wScaleDimNum != 1 && wScaleDimNum != 2, + OP_LOGE(inputParams_.opName, "In non k axis group quant mode, the dim num of scale \ +should be 1 or 2, but the actual dim num is %zu.", wScaleDimNum), return false); + return true; +} + +bool GroupedQbmmTiling::CheckFp4Shape() const +{ + OP_CHECK_IF(inputParams_.kSize % EVEN_FACTOR != 0, + OP_LOGE(inputParams_.opName, + "When the dtype of x is FLOAT4, the k size should be even number, but actual k size is %lu", + inputParams_.kSize), + return false); + // 2: mxfp4场景下不支持K轴为2 + OP_CHECK_IF(inputParams_.kSize == 2, + OP_LOGE(inputParams_.opName, "When the dtype of x is FLOAT4, the k size should not be 2"), + return false); + if (!inputParams_.transB) { + OP_CHECK_IF( + inputParams_.nSize % EVEN_FACTOR != 0, + OP_LOGE(inputParams_.opName, + "When the dtype of x is FLOAT4 and weight is not transposed, the n size should be even number, \ +but actual n size is %lu", + inputParams_.nSize), + return false); + } + return true; +} + +bool GroupedQbmmTiling::CheckBiasShape(const gert::StorageShape *biasStorageShape) const +{ + auto &biasShape = biasStorageShape->GetStorageShape(); + OP_CHECK_IF(biasStorageShape->GetStorageShape().GetDimNum() != BIAS_DIMS, + OP_LOGE(inputParams_.opName, "The dim num of bias should be 2, but actual is %zu.", + biasStorageShape->GetStorageShape().GetDimNum()), + return false); + auto biasEDim = static_cast(biasShape.GetDim(0)); + auto biasNDim = static_cast(biasShape.GetDim(1)); + OP_CHECK_IF(biasEDim != inputParams_.groupNum || biasNDim != inputParams_.nSize, + OP_LOGE(inputParams_.opName, "The expected shape of bias is (%lu, %lu), but the actual is (%lu, %lu).", + inputParams_.groupNum, inputParams_.nSize, biasEDim, biasNDim), + return false); + return true; +} + +bool GroupedQbmmTiling::CheckQuantParams(const gert::StorageShape *xScaleStorageShape, + const gert::Shape &wScaleShape) const +{ + // 非k分组量化校验 + if (inputParams_.bQuantMode != optiling::QuantMode::MX_PERGROUP_MODE && + inputParams_.bQuantMode != optiling::QuantMode::PERGROUP_MODE && + inputParams_.bQuantMode != optiling::QuantMode::PERBLOCK_MODE) { + OP_CHECK_IF(!CheckQuantParamsForNonKGroupQuantMode(wScaleShape), + OP_LOGE(inputParams_.opName, "CheckQuantParamsForNonKGroupQuantMode failed."), + return false); + } + // mx量化校验 + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + OP_CHECK_IF(!CheckQuantParamsForMxQuantMode(xScaleStorageShape, wScaleShape), + OP_LOGE(inputParams_.opName, "CheckParamsForMxQuantMode failed."), return false); + } + + return true; +} + +bool GroupedQbmmTiling::AnalyzeInputs() +{ + auto xStorageShape = context_->GetDynamicInputShape(X_INDEX, 0); + + OP_CHECK_IF(xStorageShape == nullptr, OP_LOGE(context_->GetNodeName(), "xStorageShape is nullptr."), return false); + const gert::Shape &xShape = xStorageShape->GetOriginShape(); + + auto wStorageShape = context_->GetDynamicInputShape(WEIGHT_INDEX, 0); + OP_CHECK_IF(wStorageShape == nullptr, OP_LOGE(context_->GetNodeName(), "wStorageShape is nullptr."), return false); + const gert::Shape &wShape = wStorageShape->GetOriginShape(); + + // 全量化scale必须有值，目前无输出int32等不需要scale的场景 + auto scaleStorageShape = context_->GetDynamicInputShape(SCALE_INDEX, 0); + OP_CHECK_IF(scaleStorageShape == nullptr, OP_LOGE(context_->GetNodeName(), "scaleStorageShape is nullptr."), return false); + const gert::Shape &wScaleShape = scaleStorageShape->GetOriginShape(); + auto scaleDimNum = wScaleShape.GetDimNum(); + OP_CHECK_IF(scaleDimNum < 1, + OP_LOGE(inputParams_.opName, + "The dimension of scale should be positive integer, actual is %zu", + scaleDimNum), + return false); + auto xScaleStorageShape = context_->GetOptionalInputShape(PER_TOKEN_SCALE_INDEX); + OP_CHECK_IF(!SetGroupNum(GROUPLIST_INDEX), OP_LOGE(inputParams_.opName, "SetGroupNum failed."), + return false); + OP_CHECK_IF(!SetMKN(xShape, wShape), OP_LOGE(inputParams_.opName, "SetMKN failed."), return false); + OP_CHECK_IF(!SetMKNList(), OP_LOGE(inputParams_.opName, "SetMKNList failed."), return false); + OP_CHECK_IF(!SetQuantMode(wScaleShape, xScaleStorageShape, wShape), + OP_LOGE(inputParams_.opName, "SetQuantMode failed."), return false); + OP_CHECK_IF(!CheckQuantParams(xScaleStorageShape, wScaleShape), + OP_LOGE(inputParams_.opName, "CheckQuantParams failed."), return false); + if (inputParams_.aDtype == ge::DT_FLOAT4_E2M1 || inputParams_.aDtype == ge::DT_FLOAT4_E1M2) { + OP_CHECK_IF(!CheckFp4Shape(), OP_LOGE(inputParams_.opName, "CheckFp4Shape failed."), return false); + if (inputParams_.hasBias) { + auto biasStorageShape = context_->GetDynamicInputShape(BIAS_INDEX, 0); + OP_CHECK_IF(!CheckBiasShape(biasStorageShape), + OP_LOGE(inputParams_.opName, "CheckBiasShape failed."), return false); + } + } + SetKernelType(); + return true; +} + +bool GroupedQbmmTiling::SetQuantMode(const gert::Shape &wScaleShape, const gert::StorageShape *xScaleStorageShape, + const gert::Shape &wShape) +{ + auto wScaleDims = wScaleShape.GetDimNum(); + if (IsMicroScaling()) { + inputParams_.bQuantMode = optiling::QuantMode::MX_PERGROUP_MODE; + inputParams_.aQuantMode = optiling::QuantMode::MX_PERGROUP_MODE; + return true; + } + // scale pertensor: (g,1) 2维或（g,）1维, perchannel:（g, N), 2维 + if (wScaleDims == 2 && static_cast(wScaleShape.GetDim(wScaleDims - 1)) == inputParams_.nSize && + inputParams_.nSize != 1UL) { + inputParams_.bQuantMode = optiling::QuantMode::PERCHANNEL_MODE; + } else if ((wScaleDims == 2 && wScaleShape[wScaleDims - 1] == 1) || // 2:（g,1) 2维 + (wScaleDims == 1 && static_cast(wScaleShape[0]) == inputParams_.groupNum)) { + inputParams_.bQuantMode = optiling::QuantMode::PERTENSOR_MODE; + } + if (xScaleStorageShape != nullptr) { + // split_m: pertoken (M,), pertensor（g,1) 2维或（g,）1维; + // split_k: pertoken (g, M), pertensor（g,1) 2维或（g,) 1维 + auto &xScaleShape = xScaleStorageShape->GetStorageShape(); + auto xScaleDims = xScaleShape.GetDimNum(); + if (inputParams_.aDtype != ge::DT_INT8 && + ((xScaleDims == 2 && xScaleShape[xScaleDims - 1] == 1) || // 2:（g,1) 2维 + (xScaleDims == 1 && static_cast(xScaleShape[0]) == inputParams_.groupNum && + inputParams_.groupNum != inputParams_.mSize))) { + inputParams_.aQuantMode = optiling::QuantMode::PERTENSOR_MODE; + } else { + inputParams_.aQuantMode = optiling::QuantMode::PERTOKEN_MODE; + } + SetPerGroupQuantMode(xScaleShape, wScaleShape, wShape); + } + return true; +} + +void GroupedQbmmTiling::SetPerGroupQuantMode(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape, + const gert::Shape &wShape) +{ + if (inputParams_.aDtype == ge::DT_INT8) { + return; + } + auto xScaleDims = xScaleShape.GetDimNum(); + if (xScaleDims < X_DIMS) { + return; + } + auto wScaleDims = wScaleShape.GetDimNum(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + if (wDimNum != wScaleDims || (inputParams_.groupType == SPLIT_M && wScaleDims < SPLIT_M_W_DIMS) || + (inputParams_.groupType == SPLIT_K && wScaleDims < SPLIT_K_W_DIMS)) { + return; + } + optiling::QuantMode aQuantMode = optiling::QuantMode::DEFAULT; + optiling::QuantMode bQuantMode = optiling::QuantMode::DEFAULT; + if (inputParams_.groupType == SPLIT_M) { + for (uint64_t i = 1; i < wScaleDims; ++i) { + if (wScaleShape.GetDim(i) != CeilDiv(wShape.GetDim(i), PER_BLOCK_GROUP_SIZE)) { + return; + } + } + bQuantMode = optiling::QuantMode::PERBLOCK_MODE; + + uint64_t scaleKPerBlock = CeilDiv(inputParams_.kSize, PER_BLOCK_GROUP_SIZE); + if (static_cast(xScaleShape.GetDim(xScaleDims - LAST_FIRST_DIM_INDEX)) == scaleKPerBlock && + static_cast(xScaleShape.GetDim(xScaleDims - LAST_SECOND_DIM_INDEX)) == inputParams_.mSize) { + aQuantMode = optiling::QuantMode::PERGROUP_MODE; + } + } + if (inputParams_.groupType == SPLIT_K) { + uint64_t scaleKPerBlock = inputParams_.kSize / PER_BLOCK_GROUP_SIZE + inputParams_.groupNum; + uint64_t scaleNPerBlock = CeilDiv(inputParams_.nSize, PER_BLOCK_GROUP_SIZE); + if (static_cast(xScaleShape.GetDim(xScaleDims - LAST_SECOND_DIM_INDEX)) == scaleKPerBlock && + static_cast(xScaleShape.GetDim(xScaleDims - LAST_FIRST_DIM_INDEX)) == inputParams_.mSize) { + aQuantMode = optiling::QuantMode::PERGROUP_MODE; + } + if (static_cast(wScaleShape.GetDim(wScaleDims - LAST_SECOND_DIM_INDEX)) == scaleKPerBlock && + static_cast(wScaleShape.GetDim(wScaleDims - LAST_FIRST_DIM_INDEX)) == scaleNPerBlock) { + bQuantMode = optiling::QuantMode::PERBLOCK_MODE; + } + } + if (aQuantMode == optiling::QuantMode::PERGROUP_MODE && bQuantMode == optiling::QuantMode::PERBLOCK_MODE) { + inputParams_.aQuantMode = optiling::QuantMode::PERGROUP_MODE; + inputParams_.bQuantMode = optiling::QuantMode::PERBLOCK_MODE; + } +} + +bool GroupedQbmmTiling::SetGroupNum(uint32_t groupListIndex) +{ + auto groupListStorageShape = context_->GetOptionalInputShape(groupListIndex); + OP_CHECK_IF(groupListStorageShape == nullptr, OP_LOGE(context_->GetNodeName(), "groupListStorageShape is nullptr."), return false); + const gert::Shape &groupListShape = groupListStorageShape->GetStorageShape(); + OP_CHECK_IF(groupListShape.GetDimNum() != 1, + OP_LOGE(inputParams_.opName, "The dimension of groupList should be 1, actual is %zu", + groupListShape.GetDimNum()), + return false); + inputParams_.groupNum = static_cast(groupListShape.GetDim(0)); + return true; +} + +bool GroupedQbmmTiling::SetMKN(const gert::Shape &xShape, const gert::Shape &wShape) +{ + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + OP_CHECK_IF(wDimNum < MIN_ND_DIM, + OP_LOGE(inputParams_.opName, + "The dimension of weight should be at least 2, actual is %u", wDimNum), + return false); + uint32_t xDimNum = static_cast(xShape.GetDimNum()); + OP_CHECK_IF(xDimNum < MIN_ND_DIM, + OP_LOGE(inputParams_.opName, + "Invalid x dimension for format ND, expect at least 2, actual is %u", xDimNum), + return false); + auto mSize = inputParams_.transA ? xShape.GetDim(xDimNum - LAST_FIRST_DIM_INDEX) : + xShape.GetDim(xDimNum - LAST_SECOND_DIM_INDEX); + auto kSize = inputParams_.transA ? xShape.GetDim(xDimNum - LAST_SECOND_DIM_INDEX) : + xShape.GetDim(xDimNum - LAST_FIRST_DIM_INDEX); + auto nSize = inputParams_.transB ? wShape.GetDim(wDimNum - LAST_SECOND_DIM_INDEX) : + wShape.GetDim(wDimNum - LAST_FIRST_DIM_INDEX); + OP_CHECK_IF(mSize <= 0 || kSize <= 0 || nSize <= 0, + OP_LOGE(inputParams_.opName, + "Invalid mSize[%ld] kSize[%ld] or nSize[%ld], expect all greater than 0", + mSize, kSize, nSize), + return false); + inputParams_.mSize = mSize; + inputParams_.kSize = kSize; + inputParams_.nSize = nSize; + return true; +} + +bool GroupedQbmmTiling::SetMKNList() +{ + if (inputParams_.groupType == SPLIT_M) { + mList_[0] = -1; + kList_[0] = static_cast(inputParams_.kSize); + nList_[0] = static_cast(inputParams_.nSize); + } else { + mList_[0] = static_cast(inputParams_.mSize); + kList_[0] = -1; + nList_[0] = static_cast(inputParams_.nSize); + } + return true; +} + +ge::graphStatus GroupedQbmmTiling::GetShapeAttrsInfo() +{ + inputParams_.opName = context_->GetNodeName(); + OP_CHECK_IF(!AnalyzeDtype() || !AnalyzeAttrs() || !AnalyzeInputs(), + OP_LOGE(inputParams_.opName, "Failed to analyze context_ info."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GroupedQbmmTiling::DoOpTiling() +{ + tilingData_.gmmQuantParams.groupNum = inputParams_.groupNum; + tilingData_.gmmQuantParams.activeType = inputParams_.actType; + tilingData_.gmmQuantParams.aQuantMode = static_cast(inputParams_.aQuantMode); + tilingData_.gmmQuantParams.bQuantMode = static_cast(inputParams_.bQuantMode); + tilingData_.gmmQuantParams.singleX = static_cast(inputParams_.isSingleX); + tilingData_.gmmQuantParams.singleW = static_cast(inputParams_.isSingleW); + tilingData_.gmmQuantParams.singleY = static_cast(inputParams_.isSingleY); + tilingData_.gmmQuantParams.groupType = static_cast(inputParams_.groupType); + tilingData_.gmmQuantParams.groupListType = static_cast(inputParams_.groupListType); + tilingData_.gmmQuantParams.hasBias = static_cast(inputParams_.hasBias); + errno_t retM = memcpy_s(tilingData_.gmmArray.mList, sizeof(tilingData_.gmmArray.mList), mList_, sizeof(mList_)); + if (retM != EOK) { + OP_LOGE(context_->GetNodeName(), "memcpy_s failed, ret = %d", retM); + return ge::GRAPH_FAILED; + } + errno_t retK = memcpy_s(tilingData_.gmmArray.kList, sizeof(tilingData_.gmmArray.kList), kList_, sizeof(kList_)); + if (retK!= EOK) { + OP_LOGE(context_->GetNodeName(), "memcpy_s failed, ret = %d", retK); + return ge::GRAPH_FAILED; + } + errno_t retN = memcpy_s(tilingData_.gmmArray.nList, sizeof(tilingData_.gmmArray.nList), nList_, sizeof(nList_)); + if (retN != EOK) { + OP_LOGE(context_->GetNodeName(), "memcpy_s failed, ret = %d", retN); + return ge::GRAPH_FAILED; + } + PrintQuantParams(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GroupedQbmmTiling::DoLibApiTiling() +{ + CalBasicBlock(); + OP_CHECK_IF(CalL1Tiling() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "CalL1Tiling failed"), return ge::GRAPH_FAILED); + tilingData_.mmTilingData.M = inputParams_.mSize; + tilingData_.mmTilingData.N = inputParams_.nSize; + tilingData_.mmTilingData.Ka = inputParams_.kSize; + tilingData_.mmTilingData.Kb = inputParams_.kSize; + tilingData_.mmTilingData.usedCoreNum = aicoreParams_.aicNum; + tilingData_.mmTilingData.baseM = basicTiling_.baseM; + tilingData_.mmTilingData.baseN = basicTiling_.baseN; + tilingData_.mmTilingData.baseK = basicTiling_.baseK; + tilingData_.mmTilingData.singleCoreM = basicTiling_.singleCoreM; + tilingData_.mmTilingData.singleCoreN = basicTiling_.singleCoreN; + tilingData_.mmTilingData.singleCoreK = basicTiling_.singleCoreK; + tilingData_.mmTilingData.depthA1 = basicTiling_.depthA1; + tilingData_.mmTilingData.depthB1 = basicTiling_.depthB1; + tilingData_.mmTilingData.stepM = basicTiling_.stepM; + tilingData_.mmTilingData.stepN = basicTiling_.stepN; + tilingData_.mmTilingData.stepKa = basicTiling_.stepKa; + tilingData_.mmTilingData.stepKb = basicTiling_.stepKb; + tilingData_.mmTilingData.isBias = inputParams_.hasBias ? 1 : 0; + tilingData_.mmTilingData.iterateOrder = basicTiling_.iterateOrder; + tilingData_.mmTilingData.dbL0A = 2; // db switch, 1: off, 2: on + tilingData_.mmTilingData.dbL0B = 2; // db switch, 1: off, 2: on + tilingData_.mmTilingData.dbL0C = basicTiling_.dbL0c; + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + if (basicTiling_.scaleFactorA >= SCALER_FACTOR_MIN && basicTiling_.scaleFactorA <= SCALER_FACTOR_MAX && + basicTiling_.scaleFactorB >= SCALER_FACTOR_MIN && basicTiling_.scaleFactorB <= SCALER_FACTOR_MAX) { + tilingData_.mmTilingData.mxTypePara = (SCALER_FACTOR_DEFAULT << SCALER_FACTOR_N_BIT) + (SCALER_FACTOR_DEFAULT << SCALER_FACTOR_M_BIT) + + (basicTiling_.scaleFactorB << SCALER_FACTOR_B_BIT) + basicTiling_.scaleFactorA; + } else { + tilingData_.mmTilingData.mxTypePara = (SCALER_FACTOR_DEFAULT << SCALER_FACTOR_N_BIT) + (SCALER_FACTOR_DEFAULT << SCALER_FACTOR_M_BIT) + + (SCALER_FACTOR_DEFAULT << SCALER_FACTOR_B_BIT) + SCALER_FACTOR_DEFAULT; + } + } + + return ge::GRAPH_SUCCESS; +} + +void GroupedQbmmTiling::SetKernelType() +{ + // 以选择主模板设置kernelType, 0: dequant fixp随路（包含K轴分组）；1：dequant vector计算；2：perGroup-perBlock + inputParams_.kernelType = 0UL; + // mx K轴分组当前是独立的模板，后续归一 + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + return; + } + // perGroup-perBlock(GB)有独立pertile模板 + if (inputParams_.bQuantMode == optiling::QuantMode::PERBLOCK_MODE) { + inputParams_.kernelType = 2UL; + return; + } + // pertensor-pertensor且没有后处理的bias，都可以走dequant fixp随路 + bool isPertensorCube = inputParams_.aQuantMode <= optiling::QuantMode::PERTENSOR_MODE && + inputParams_.bQuantMode == optiling::QuantMode::PERTENSOR_MODE; + bool isBiasEpilogue = + inputParams_.aDtype == ge::DT_INT8 && inputParams_.hasBias && inputParams_.biasDtype != ge::DT_INT32; + // 如果bias bf16/fp16/fp32，需mix模板进行后处理 + if (isPertensorCube && !isBiasEpilogue) { + return; + } + bool isScaleEpilogue = (inputParams_.scaleDtype != ge::DT_UINT64 && inputParams_.scaleDtype != ge::DT_INT64); + // 后处理的bias和（scale非64bits && ！isPertensorCube）需要走dequant vec模板 + if (isBiasEpilogue || isScaleEpilogue) { + inputParams_.kernelType = 1UL; + } +} + +uint64_t GroupedQbmmTiling::GetTilingKey() const +{ + return GET_TPL_TILING_KEY(static_cast(inputParams_.transB), static_cast(inputParams_.transA), + static_cast(inputParams_.kernelType)); +} + +ge::graphStatus GroupedQbmmTiling::GetWorkspaceSize() +{ + size_t *workspaces = context_->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context_, workspaces); + workspaces[0] = SYS_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GroupedQbmmTiling::PostTiling() +{ + context_->SetBlockDim(aicoreParams_.aicNum); + OP_CHECK_IF(sizeof(tilingData_) % sizeof(uint64_t) != 0, + OP_LOGE(context_->GetNodeName(), "Tiling data size[%zu] is not aligned to 8", + sizeof(tilingData_)), + return ge::GRAPH_FAILED); + errno_t ret = memcpy_s(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity(), reinterpret_cast(&tilingData_), sizeof(tilingData_)); + if (ret != EOK) { + OP_LOGE(context_->GetNodeName(), "memcpy_s failed, ret = %d", ret); + return ge::GRAPH_FAILED; + } + context_->GetRawTilingData()->SetDataSize(sizeof(tilingData_)); + return ge::GRAPH_SUCCESS; +} + +void GroupedQbmmTiling::PrintQuantParams() +{ + int32_t enable = AlogCheckDebugLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + GMMQuantParams ¶ms = tilingData_.gmmQuantParams; + std::ostringstream oss; + oss << "GMMQuantParams: groupNum = " << params.groupNum << ", activeType = " << params.activeType + << ", aQuantMode = " << params.aQuantMode << ", bQuantMode = " << params.bQuantMode + << ", singleX=" << static_cast(params.singleX) + << ", singleW = " << static_cast(params.singleW) + << ", singleY = " << static_cast(params.singleY) + << ", groupType = " << static_cast(params.groupType) + << ", groupListType = " << static_cast(params.groupListType) + << ", hasBias = " << static_cast(params.hasBias); + OP_LOGD(inputParams_.opName, "%s", oss.str().c_str()); +} + +void GroupedQbmmTiling::CalBasicBlock() +{ + bool isGBQuantMode = inputParams_.aQuantMode == optiling::QuantMode::PERGROUP_MODE && + inputParams_.bQuantMode == optiling::QuantMode::PERBLOCK_MODE; + basicTiling_.baseM = std::min(inputParams_.mSize, static_cast(GmmConstant::BASIC_BLOCK_SIZE_256)); + basicTiling_.baseM = !inputParams_.transA ? + CeilAlign(basicTiling_.baseM, CUBE_BLOCK) : + CeilAlign(basicTiling_.baseM, GetShapeWithDataType(L1_ALIGN_SIZE, inputParams_.aDtype)); + if (isGBQuantMode) { + // 不管M/K轴分组，单单单场景下，N不变，可以确定baseN + if (inputParams_.nSize <= PER_BLOCK_GROUP_SIZE || basicTiling_.baseM > PER_BLOCK_GROUP_SIZE) { + basicTiling_.baseN = PER_BLOCK_GROUP_SIZE; + } else { + basicTiling_.baseN = GmmConstant::BASIC_BLOCK_SIZE_256; + } + basicTiling_.baseK = PER_BLOCK_GROUP_SIZE; + return; + } + basicTiling_.baseN = std::min(inputParams_.nSize, static_cast(GmmConstant::BASIC_BLOCK_SIZE_256)); + basicTiling_.baseN = inputParams_.transB ? + CeilAlign(basicTiling_.baseN, CUBE_BLOCK) : + CeilAlign(basicTiling_.baseN, GetShapeWithDataType(L1_ALIGN_SIZE, inputParams_.bDtype)); + basicTiling_.baseK = CeilAlign( + std::min(GetShapeWithDataType(GmmConstant::BASIC_BLOCK_SIZE_128, inputParams_.aDtype), inputParams_.kSize), + GetShapeWithDataType(CUBE_REDUCE_BLOCK, inputParams_.aDtype)); + + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + basicTiling_.baseK = CeilAlign(basicTiling_.baseK, MXFP_BASEK_FACTOR); // mx_mmad requires basek align to 64 + bool isFp4Input = inputParams_.aDtype == ge::DT_FLOAT4_E2M1 || inputParams_.aDtype == ge::DT_FLOAT4_E1M2; + if (isFp4Input && !inputParams_.transB) { + // 64: mx_mmad requires the inner axis to align to 64 + basicTiling_.baseN = CeilAlign(basicTiling_.baseN, static_cast(64)); + } + } +} + +bool GroupedQbmmTiling::IsBiasInL1() const +{ + // 目前仅int8进bias int32需要进L1 + return inputParams_.hasBias && inputParams_.biasDtype == ge::DT_INT32; +} + +ge::graphStatus GroupedQbmmTiling::CalL1Tiling() +{ + basicTiling_.stepM = 1UL; + basicTiling_.stepN = 1UL; + basicTiling_.singleCoreM = std::min(inputParams_.mSize, basicTiling_.baseM); + basicTiling_.singleCoreN = std::min(inputParams_.nSize, basicTiling_.baseN); + basicTiling_.singleCoreK = inputParams_.kSize; + + uint64_t biasDtypeSize = ge::GetSizeByDataType(inputParams_.biasDtype); + uint64_t scaleDtypeSize = ge::GetSizeByDataType(inputParams_.scaleDtype); + uint64_t totalL1Size = aicoreParams_.l1Size; + + basicTiling_.iterateOrder = 0U; + basicTiling_.dbL0c = + (basicTiling_.baseM * basicTiling_.baseN * DATA_SIZE_L0C * DB_SIZE <= aicoreParams_.l0cSize) ? DB_SIZE : 1; + uint64_t singleCoreBiasSize = IsBiasInL1() ? basicTiling_.baseN * biasDtypeSize : 0; + uint64_t singleCoreScaleSize = + inputParams_.bQuantMode == optiling::QuantMode::PERCHANNEL_MODE && inputParams_.kernelType == 0 ? + basicTiling_.baseN * scaleDtypeSize : + 0; + uint64_t usedSize = singleCoreBiasSize + singleCoreScaleSize; + OP_CHECK_IF(totalL1Size <= usedSize, + OP_LOGE(context_->GetNodeName(), "L1 space overflow. L1Size: %lu, used space: %lu", + totalL1Size, usedSize), + return ge::GRAPH_FAILED); + uint64_t leftL1Size = totalL1Size - usedSize; + return CalL1Depth(leftL1Size); +} + +ge::graphStatus GroupedQbmmTiling::CalL1Depth(uint64_t leftL1Size) +{ + uint64_t baseASize = GetSizeWithDataType(basicTiling_.baseM * basicTiling_.baseK, inputParams_.aDtype); + uint64_t baseBSize = GetSizeWithDataType(basicTiling_.baseN * basicTiling_.baseK, inputParams_.bDtype); + + uint64_t baseScaleASize = 0; + uint64_t baseScaleBSize = 0; + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + if (inputParams_.groupType == SPLIT_M) { + baseScaleASize = + GetSizeWithDataType(CeilAlign(CeilDiv(basicTiling_.baseK, MX_GROUP_SIZE), 2UL) * basicTiling_.baseM, + inputParams_.perTokenScaleDtype); + baseScaleBSize = + GetSizeWithDataType(CeilAlign(CeilDiv(basicTiling_.baseK, MX_GROUP_SIZE), 2UL) * basicTiling_.baseN, + inputParams_.scaleDtype); + } else { + baseScaleASize = GetSizeWithDataType( + (basicTiling_.baseK / (MX_GROUP_SIZE * MXFP_MULTI_BASE_SIZE) + inputParams_.groupNum) * + MXFP_MULTI_BASE_SIZE * basicTiling_.baseM, // 2 is dim value of last scale dim + inputParams_.perTokenScaleDtype); + baseScaleBSize = GetSizeWithDataType( + (basicTiling_.baseK / (MX_GROUP_SIZE * MXFP_MULTI_BASE_SIZE) + inputParams_.groupNum) * + MXFP_MULTI_BASE_SIZE * basicTiling_.baseN, // 2 is dim value of last pertokenScale dim + inputParams_.scaleDtype); + } + } + uint64_t baseL1Size = baseASize + baseBSize + baseScaleASize + baseScaleBSize; + OP_CHECK_IF(leftL1Size < baseL1Size, + OP_LOGE(context_->GetNodeName(), + "L1 space overflow. Free L1Size : %lu, used space: %lu", leftL1Size, + baseL1Size), + return ge::GRAPH_FAILED); + uint64_t depthInit = GetDepthA1B1(leftL1Size, baseL1Size, 1UL); + uint64_t leftL1SizeByDepthInit = leftL1Size - depthInit * (baseL1Size); + uint64_t depthASec = GetDepthA1B1(leftL1SizeByDepthInit, (baseASize + baseScaleASize) * depthInit, depthInit); + uint64_t depthBSec = GetDepthA1B1(leftL1SizeByDepthInit, (baseBSize + baseScaleBSize) * depthInit, depthInit); + basicTiling_.depthA1 = std::max(depthASec, depthBSec); + basicTiling_.depthB1 = basicTiling_.depthA1; + if (basicTiling_.depthA1 * baseL1Size > leftL1Size) { + basicTiling_.depthA1 = depthASec >= depthBSec ? depthASec : depthInit; + basicTiling_.depthB1 = depthASec < depthBSec ? depthBSec : depthInit; + } + CalStepKs(); + if (inputParams_.bQuantMode == optiling::QuantMode::MX_PERGROUP_MODE) { + CalScaleFactors(); + } + return ge::GRAPH_SUCCESS; +} + +uint64_t GroupedQbmmTiling::GetDepthA1B1(uint64_t leftSize, uint64_t perDepthSize, uint64_t depthInit) +{ + if (depthInit > 1UL && perDepthSize > DB_SIZE * MTE2_MIN_LOAD_SIZE_V120) { + return depthInit; + } + uint64_t depthScale = leftSize / perDepthSize; + if (depthInit > 1UL) { + uint64_t baseKSize = GetSizeWithDataType(basicTiling_.baseK, inputParams_.aDtype); + while ((depthScale * baseKSize) % GmmConstant::BASIC_BLOCK_SIZE_512 != 0 && + (depthScale * baseKSize) > GmmConstant::BASIC_BLOCK_SIZE_512) { + depthScale -= 1UL; + } + if ((depthScale * baseKSize) % GmmConstant::BASIC_BLOCK_SIZE_512 != 0 && + (depthScale * baseKSize) >= GmmConstant::BASIC_BLOCK_SIZE_256) { + depthScale = GmmConstant::BASIC_BLOCK_SIZE_256 / baseKSize; + } + depthScale = std::max(depthScale, static_cast(1)); + } else { + constexpr uint64_t index = 2; // 2: depth的值是2的幂 + depthScale = 1UL; + while (depthScale * (perDepthSize) < leftSize) { + depthScale *= index; + } + depthScale = depthScale == 1UL ? depthScale : depthScale / index; + } + return depthInit * depthScale; +} + +void GroupedQbmmTiling::CalStepKs() +{ + // depthA,depthB 为1时，stepka, stepkb 只能是1. + basicTiling_.stepKa = basicTiling_.depthA1 == 1UL ? 1UL : basicTiling_.depthA1 / DB_SIZE; + basicTiling_.stepKb = basicTiling_.depthB1 == 1UL ? 1UL : basicTiling_.depthB1 / DB_SIZE; + + if (basicTiling_.stepKa * basicTiling_.baseK > inputParams_.kSize) { + basicTiling_.stepKa = CeilDiv(inputParams_.kSize, basicTiling_.baseK); + } + + if (basicTiling_.stepKb * basicTiling_.baseK >= inputParams_.kSize) { + basicTiling_.stepKb = CeilDiv(inputParams_.kSize, basicTiling_.baseK); + } + // G-B量化场景下，限制stepK最大为4, 防止issue queue阻塞 + if (inputParams_.aQuantMode == optiling::QuantMode::PERGROUP_MODE && + inputParams_.bQuantMode == optiling::QuantMode::PERBLOCK_MODE) { + basicTiling_.stepKa = std::min(basicTiling_.stepKa, static_cast(4)); // 4: G-B最大stepk值 + basicTiling_.stepKb = std::min(basicTiling_.stepKb, static_cast(4)); // 4: G-B最大stepk值 + } + if (basicTiling_.stepKa >= basicTiling_.stepKb && basicTiling_.stepKa * basicTiling_.baseK < inputParams_.kSize) { + basicTiling_.stepKa = basicTiling_.stepKa / basicTiling_.stepKb * basicTiling_.stepKb; + } + if (basicTiling_.stepKb > basicTiling_.stepKa && basicTiling_.stepKb * basicTiling_.baseK < inputParams_.kSize) { + basicTiling_.stepKb = basicTiling_.stepKb / basicTiling_.stepKa * basicTiling_.stepKa; + } + + basicTiling_.depthA1 = basicTiling_.stepKa * DB_SIZE; + basicTiling_.depthB1 = basicTiling_.stepKb * DB_SIZE; +} + +void GroupedQbmmTiling::CalScaleFactors() +{ + uint64_t baseASize = GetSizeWithDataType(basicTiling_.baseM * basicTiling_.baseK, inputParams_.aDtype); + uint64_t baseBSize = GetSizeWithDataType(basicTiling_.baseN * basicTiling_.baseK, inputParams_.bDtype); + uint64_t baseScaleASize = GetSizeWithDataType(CeilDiv(basicTiling_.baseK, MX_GROUP_SIZE) * basicTiling_.baseM, + inputParams_.perTokenScaleDtype); + uint64_t baseScaleBSize = + GetSizeWithDataType(CeilDiv(basicTiling_.baseK, MX_GROUP_SIZE) * basicTiling_.baseN, inputParams_.scaleDtype); + uint64_t biasDtypeSize = ge::GetSizeByDataType(inputParams_.biasDtype); + uint64_t baseBiasSize = inputParams_.hasBias ? basicTiling_.baseN * biasDtypeSize : 0; + uint64_t leftL1Size = + aicoreParams_.l1Size - (basicTiling_.depthA1 * baseASize + basicTiling_.depthB1 * baseBSize + baseBiasSize); + uint32_t scaleInit = static_cast(leftL1Size / (basicTiling_.depthA1 * baseScaleASize + + basicTiling_.depthB1 * baseScaleBSize)); + + // 计算scaleFactorA, scaleFactorB + // 来自K轴的约束 + uint32_t scaleFactorAMax = + std::min(static_cast(MTE2_MIN_LOAD_SIZE_V120 / baseScaleASize), SCALER_FACTOR_MAX); + uint32_t scaleFactorBMax = + std::min(static_cast(MTE2_MIN_LOAD_SIZE_V120 / baseScaleBSize), SCALER_FACTOR_MAX); + uint32_t scaleFactorA = static_cast(inputParams_.kSize / (basicTiling_.stepKa * basicTiling_.baseK)); + uint32_t scaleFactorB = static_cast(inputParams_.kSize / (basicTiling_.stepKb * basicTiling_.baseK)); + basicTiling_.scaleFactorA = std::max(SCALER_FACTOR_MIN, scaleFactorA); + basicTiling_.scaleFactorB = std::max(SCALER_FACTOR_MIN, scaleFactorB); + basicTiling_.scaleFactorA = std::min(scaleFactorAMax, basicTiling_.scaleFactorA); + basicTiling_.scaleFactorB = std::min(scaleFactorBMax, basicTiling_.scaleFactorB); + + // 来自L1 size 的约束 + if (basicTiling_.scaleFactorA <= scaleInit && basicTiling_.scaleFactorB > scaleInit) { + leftL1Size -= (basicTiling_.scaleFactorA * basicTiling_.depthA1 * baseScaleASize); + basicTiling_.scaleFactorB = std::min(static_cast(leftL1Size / (basicTiling_.depthB1 * baseScaleBSize)), + basicTiling_.scaleFactorB); + } else if (basicTiling_.scaleFactorB <= scaleInit && basicTiling_.scaleFactorA > scaleInit) { + leftL1Size -= (basicTiling_.scaleFactorB * basicTiling_.depthB1 * baseScaleBSize); + basicTiling_.scaleFactorA = std::min(static_cast(leftL1Size / (basicTiling_.depthA1 * baseScaleASize)), + basicTiling_.scaleFactorA); + } else if (basicTiling_.scaleFactorA > scaleInit && basicTiling_.scaleFactorB > scaleInit) { + leftL1Size -= + (scaleInit * basicTiling_.depthB1 * baseScaleBSize + scaleInit * basicTiling_.depthA1 * baseScaleASize); + uint32_t scaleASec = std::min(static_cast(leftL1Size / (basicTiling_.depthA1 * baseScaleASize)), + basicTiling_.scaleFactorA - scaleInit); + uint32_t scaleBSec = std::min(static_cast(leftL1Size / (basicTiling_.depthB1 * baseScaleBSize)), + basicTiling_.scaleFactorB - scaleInit); + basicTiling_.scaleFactorA = scaleASec >= scaleBSec ? (scaleASec + scaleInit) : scaleInit; + basicTiling_.scaleFactorB = scaleASec < scaleBSec ? (scaleBSec + scaleInit) : scaleInit; + } +} + +uint64_t GroupedQbmmTiling::GetSizeWithDataType(uint64_t shapeSize, ge::DataType dtype) const +{ + // shapeSize应该是偶数 + bool is4BitInput = (dtype == ge::DT_FLOAT4_E2M1 || dtype == ge::DT_FLOAT4_E1M2 || dtype == ge::DT_INT4); + if (is4BitInput) { + // 2: 判断是否是偶数 + OP_CHECK_IF(shapeSize % 2 != 0, + OP_LOGE( + context_->GetNodeName(), + "To get size of matrix/array, the number of elements must be even when dtype is FLOAT4/INT4"), + return 0); + // 1/2: 这几种数据类型的dsize=1/2 + return shapeSize / 2UL; + } else { + return shapeSize * static_cast(ge::GetSizeByDataType(dtype)); + } +} + +uint64_t GroupedQbmmTiling::GetShapeWithDataType(uint64_t shapeSize, ge::DataType dtype) const +{ + bool is4BitInput = (dtype == ge::DT_FLOAT4_E2M1 || dtype == ge::DT_FLOAT4_E1M2 || dtype == ge::DT_INT4); + if (is4BitInput) { + return shapeSize + shapeSize; + } else { + return shapeSize / static_cast(ge::GetSizeByDataType(dtype)); + } +} + +REGISTER_OPS_TILING_TEMPLATE(DlinferGroupedMatmulDirect, GroupedQbmmTiling, 0); +} // namespace optiling \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.h new file mode 100644 index 00000000..d95baf18 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_quant_matmul_tiling.h @@ -0,0 +1,198 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_quant_matmul_tiling.h + * \brief + */ +#ifndef GROUPED_QUANT_MATMUL_TILING_H +#define GROUPED_QUANT_MATMUL_TILING_H + +#include "../grouped_matmul_tiling.h" +#include "../../../op_kernel/arch35/grouped_matmul_tiling_data_apt.h" +#include "tiling_base/tiling_base.h" +namespace optiling { +namespace GmmConstant { +constexpr uint64_t MX_GROUP_SIZE = 32; +constexpr uint64_t NUM_HALF = 2; +constexpr uint64_t EVEN_FACTOR = 2; +constexpr uint32_t DB_SIZE = 2; +constexpr uint32_t BASIC_BLOCK_SIZE_512 = 512; +constexpr uint32_t BASIC_BLOCK_SIZE_256 = 256; +constexpr uint32_t BASIC_BLOCK_SIZE_128 = 128; +constexpr uint64_t CUBE_BLOCK = 16; +constexpr uint64_t L1_ALIGN_SIZE = 32; +constexpr uint64_t UB_ALIGN_SIZE = 32; +constexpr uint64_t CUBE_REDUCE_BLOCK = 32; +constexpr uint32_t DATA_SIZE_L0C = 4; +constexpr uint32_t SCALER_FACTOR_MAX = 127; +constexpr uint32_t SCALER_FACTOR_MIN = 1; +constexpr uint32_t SCALER_FACTOR_DEFAULT = 1; +constexpr uint32_t SCALER_FACTOR_B_BIT = 8; +constexpr uint32_t SCALER_FACTOR_M_BIT = 16; +constexpr uint32_t SCALER_FACTOR_N_BIT = 24; +constexpr uint64_t MTE2_MIN_LOAD_SIZE_V120 = 64 * 1024UL; +constexpr uint64_t MAX_REPEAT_TIMES = 255; // InitOutput接口取值 +constexpr size_t LAST_FIRST_DIM_INDEX = 1; +constexpr size_t LAST_SECOND_DIM_INDEX = 2; +constexpr uint64_t PER_BLOCK_GROUP_SIZE = 128; +constexpr uint64_t SPLIT_M_W_DIMS = 3; +constexpr uint64_t SPLIT_K_W_DIMS = 2; +constexpr uint64_t X_DIMS = 2; +constexpr uint64_t BIAS_DIMS = 2; +constexpr uint64_t MXFP_MULTI_BASE_SIZE = 2; +constexpr uint64_t MXFP_BASEK_FACTOR = 64; +constexpr size_t MXFP_TYPE_K_SCALE_DIM_NUM = 3; +constexpr size_t MXFP_TYPE_M_SCALE_DIM_NUM = 4; +constexpr size_t MXFP_PER_TOKEN_SCALE_DIM_NUM = 3; +} // namespace GmmConstant + +enum class QuantMode : uint32_t { + DEFAULT = 0x0U, + PERTENSOR_MODE = 0x1U, + PERCHANNEL_MODE = 0x1U << 1, + PERTOKEN_MODE = 0x1U << 2, + MX_PERGROUP_MODE = 0x1U << 3, + PERGROUP_MODE = 0x1U << 4, + PERBLOCK_MODE = 0x1U << 5, +}; + +struct GQmmBasicTiling { + uint32_t usedCoreNum = 1; + uint32_t tilingMode = 0; + uint64_t singleCoreM = 1; + uint64_t singleCoreN = 1; + uint64_t singleCoreK = 1; + uint64_t baseM = 1; + uint64_t baseN = 1; + uint64_t baseK = 1; + uint64_t stepKa = 1; + uint64_t stepKb = 1; + uint64_t depthA1 = 1; + uint64_t depthB1 = 1; + uint64_t stepM = 1; + uint64_t stepN = 1; + uint32_t iterateOrder = 0; + uint32_t dbL0c = 1; + uint32_t calOrder = 0; + uint32_t scaleFactorA = 1; + uint32_t scaleFactorB = 1; +}; + +struct GQmmInputInfo { + uint64_t mSize = 0UL; + uint64_t kSize = 0UL; + uint64_t nSize = 0UL; + uint64_t groupNum = 0UL; + int64_t outDtype = 0L; + uint64_t kernelType = 0UL; + QuantMode aQuantMode = QuantMode::DEFAULT; + QuantMode bQuantMode = QuantMode::DEFAULT; + int8_t groupType = DlinferGroupedMatmulDirect::NO_SPLIT; + int8_t groupListType = 0; + int8_t splitItem = 0; + int8_t actType = 0; + const char *opName = nullptr; + ge::DataType aDtype = ge::DT_INT8; + ge::DataType bDtype = ge::DT_INT8; + ge::DataType cDtype = ge::DT_FLOAT16; + ge::DataType biasDtype = ge::DT_INT32; + ge::DataType scaleDtype = ge::DT_UINT64; + ge::DataType perTokenScaleDtype = ge::DT_FLOAT; + ge::DataType outDataDtype = ge::DT_FLOAT16; + ge::DataType outScaleDtype = ge::DT_FLOAT; + + ge::Format aFormat = ge::FORMAT_ND; + ge::Format bFormat = ge::FORMAT_ND; + ge::Format cFormat = ge::FORMAT_ND; + bool transA = false; + bool transB = false; + bool hasBias = false; + bool isSingleX = false; + bool isSingleW = false; + bool isSingleY = false; +}; + +class GroupedQbmmTiling : public Ops::Transformer::OpTiling::TilingBaseClass { +public: + explicit GroupedQbmmTiling(gert::TilingContext *context) : Ops::Transformer::OpTiling::TilingBaseClass(context) + { + Reset(); + } + ~GroupedQbmmTiling() override = default; + + void Reset(gert::TilingContext *context) override + { + Ops::Transformer::OpTiling::TilingBaseClass::Reset(context); + Reset(); + } + +protected: + bool IsCapable() override; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo() override; + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo() override; + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() override; + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling() override; + // 5、计算TilingKey + uint64_t GetTilingKey() const override; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize() override; + // 7、保存Tiling数据 + ge::graphStatus PostTiling() override; + virtual void Reset(); + void CalBasicBlock(); + ge::graphStatus CalL1Tiling(); + ge::graphStatus CalL1Depth(uint64_t leftL1Size); + bool SetGroupNum(uint32_t groupListIndex); + bool SetMKN(const gert::Shape &xShape, const gert::Shape &wShape); + void SetKernelType(); + virtual bool AnalyzeAttrs(); + virtual bool AnalyzeDtype(); + virtual bool AnalyzeInputs(); + virtual void PrintQuantParams(); + bool IsMicroScaling() const; + GQmmBasicTiling basicTiling_; + GQmmInputInfo inputParams_; + +private: + uint64_t GetDepthA1B1(uint64_t leftSize, uint64_t perDepthSize, uint64_t depthInit); + void CalStepKs(); + void CalScaleFactors(); + uint64_t GetSizeWithDataType(uint64_t shapeSize, ge::DataType dtype) const; + uint64_t GetShapeWithDataType(uint64_t shapeSize, ge::DataType dtype) const; + bool SetQuantMode(const gert::Shape &wScaleShape, const gert::StorageShape *xScaleStorageShape, + const gert::Shape &wShape); + void SetPerGroupQuantMode(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape, + const gert::Shape &wShape); + bool CheckQuantParamsForMXTypeM(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape) const; + bool CheckQuantParamsForMXTypeK(const gert::Shape &xScaleShape, const gert::Shape &wScaleShape) const; + bool CheckFp4Shape() const; + bool CheckBiasDtype() const; + bool CheckBiasShape(const gert::StorageShape *biasStorageShape) const; + bool CheckQuantParamsForMxQuantMode(const gert::StorageShape *xScaleStorageShape, + const gert::Shape &wScaleShape) const; + bool CheckQuantParams(const gert::StorageShape *xScaleStorageShape, const gert::Shape &wScaleShape) const; + bool CheckQuantParamsForNonKGroupQuantMode(const gert::Shape &wScaleShape) const; + bool SetMKNList(); + bool IsBiasInL1() const; + + DlinferGroupedMatmulDirectTilingData::GMMQuantTilingData tilingData_; + + int32_t mList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t kList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t nList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; +}; +} // namespace optiling + +#endif // GROUPED_QUANT_MATMUL_TILING_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.cpp new file mode 100644 index 00000000..bf3746f7 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.cpp @@ -0,0 +1,1180 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_weight_quant_batch_matmul_tiling.cpp + * \brief + */ +#include "grouped_weight_quant_batch_matmul_tiling.h" +#include "../../../op_kernel/arch35/weight_quant_basic_block/weight_quant_tiling_key.h" + +enum class GmmTrans { + NoTrans = 0, + BTrans = 1, + ATrans = 2, + ABTrans = 3 +}; +namespace optiling { + +static const std::map> BIAS_TYPE_SUPPORT_MAP = { + {ge::DT_FLOAT16, {ge::DT_FLOAT16}}, + {ge::DT_BF16, {ge::DT_BF16, ge::DT_FLOAT}}}; + +static bool inline IsNonEmpty(const gert::StorageShape *shapePtr) +{ + return (shapePtr != nullptr && + !(shapePtr->GetStorageShape().GetDimNum() == 1 && shapePtr->GetStorageShape().GetDim(0) == 0)); +} + +bool GroupedWeightQuantBatchMatmulTiling::SetTiling(gert::TilingContext *context) +{ + OP_CHECK_IF(!AnalyzeAttr(context), OP_LOGE(context->GetNodeName(), "Invalid attr param"), + return false); + OP_CHECK_IF(!CalcResplitTiling(context), + OP_LOGE(context->GetNodeName(), "Unable to calculate resplit-tiling"), return false); + auto ret = SetBaseTiling(); + if (!ret) { + OP_LOGE(context->GetNodeName(), "Set Base Tiling Failed"); + return false; + } + SetMatMulTiling(); + SetTilingKey(context); + OP_CHECK_IF(!SetCustomParam(context), + OP_LOGE(context->GetNodeName(), "Unable to set custom param"), return false); + PrintTilingResult(context); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling:: SetShapeList(const gert::TilingContext *context) +{ + if (isSingleX_ && isSingleWeight_ && isSingleY_) { + OP_CHECK_IF(!SetShapeListSplitMSingleXSingleWeightSingleY(context), + OP_LOGE(context->GetNodeName(), "Unable to get shape list"), return false); + } else if (!isSingleX_ && !isSingleWeight_ && !isSingleY_) { + OP_CHECK_IF(!SetShapeListMultiXMultiWeightMultiY(context), + OP_LOGE(context->GetNodeName(), "Unable to get MMM shape list"), return false); + } else { + OP_LOGE(context->GetNodeName(), + "Only support single-single-single or multi-multi-multi mode, actual " + "groupType: %d, singleX: %s, singleW: %s, singleY: %s", + static_cast(groupType_), isSingleX_ ? "true" : "false", isSingleWeight_ ? "true" : "false", + isSingleY_ ? "true" : "false"); + return false; + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorListSize(const gert::TilingContext *context) +{ + GetNumOfInputs(context); + OP_CHECK_IF( + numX_ >= DlinferGroupedMatmulDirect::MAX_TENSOR_CONT, + OP_LOGE(context->GetNodeName(), + "In multi/multi/multi Scenario, tensorlist's length cannot exceed 128, but it is more than 128"), + return false); + if (groupType_ == GroupType::NO_SPLIT) { + OP_CHECK_IF( + numX_ != numWeight_, + OP_LOGE(context->GetNodeName(), + "When groupType is -1 (no split), the sizes of x and weight should be all the same, but the " + "actual sizes are [%hu] and [%hu].", + numX_, numWeight_), + return false); + } else { + OP_CHECK_IF(numX_ != 1 || numWeight_ != 1, + OP_LOGE(context->GetNodeName(), + "When groupType is 0 (split M), the sizes of x and weight should all be 1, but the actual " + "sizes are [%hu] and [%hu].", + numX_, numWeight_), + return false); + } + OP_CHECK_IF(numAntiquantScale_ != numWeight_, + OP_LOGE(context->GetNodeName(), + "antiquantScaleOptional size should be equal to weight size, actual sizes are [%hu], [%hu]", + numAntiquantScale_, numWeight_), + return false); + if (hasAntiquantOffset_) { + OP_CHECK_IF( + numAntiquantOffset_ != numWeight_, + OP_LOGE(context->GetNodeName(), + "antiquantOffsetOptional size should be equal to weight size, actual sizes are [%hu], [%hu]", + numAntiquantOffset_, numWeight_), + return false); + } + if (hasBias_) { + OP_CHECK_IF(numBias_ != numWeight_, + OP_LOGE(context->GetNodeName(), + "biasOptional size should be equal to weight size, actual sizes are [%hu], [%hu]", numBias_, + numWeight_), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorDtype(const gert::TilingContext *context, uint32_t attrIdx, + size_t idx, const ge::DataType &tensorDtype, + const std::string &tensorType) const +{ + if (!IsA16W4ND()) { + return true; + } + auto tmpDesc = context->GetDynamicInputDesc(attrIdx, idx); + auto tmpDType = tmpDesc->GetDataType(); + OP_CHECK_IF(tmpDType != tensorDtype, + OP_LOGE(context->GetNodeName(), + "The dtype of each tensor in %s tensor list must be consistent. %s[%zu]'s dtype " + "is different from the first tensor's dtype. ", + tensorType.c_str(), tensorType.c_str(), idx), + return false); + return true; +} +bool GroupedWeightQuantBatchMatmulTiling::IsNzFormat(const gert::TilingContext *context, uint32_t attrIdx, + size_t idx) const +{ + auto tmpDesc = context->GetDynamicInputDesc(attrIdx, idx); + auto tmpFormat = static_cast(ge::GetPrimaryFormat(tmpDesc->GetStorageFormat())); + if (tmpFormat == ge::FORMAT_FRACTAL_NZ_C0_16 || tmpFormat == ge::FORMAT_FRACTAL_NZ_C0_32 || + tmpFormat == ge::FORMAT_FRACTAL_NZ_C0_4) { + tmpFormat = ge::FORMAT_FRACTAL_NZ; + } + return tmpFormat == ge::FORMAT_FRACTAL_NZ; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckXAndWeightFormat(const gert::TilingContext *context, size_t idx) const +{ + bool xNzFlag = IsNzFormat(context, X_IDX, idx); + OP_CHECK_IF(xNzFlag, OP_LOGE(context->GetNodeName(), "The format of x[%zu] is NZ. It should only be ND.", idx), + return false); + if (IsA16W4ND()) { + bool weightNzFlag = IsNzFormat(context, WEIGHT_IDX, idx); + OP_CHECK_IF(weightNzFlag, + OP_LOGE(context->GetNodeName(), "The format of weight[%zu] is NZ. It should only be ND.", idx), + return false); + return true; + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckNotNullPtr(const gert::TilingContext *context, uint32_t attrIdx, + size_t idx) const +{ + auto tmpIDesc = context->GetDynamicInputDesc(attrIdx, idx); + auto tmpITensor = context->GetDynamicInputTensor(attrIdx, idx); + auto tmpIShape = context->GetDynamicInputShape(attrIdx, idx); + OP_CHECK_IF(tmpIDesc == nullptr, OP_LOGE(context->GetNodeName(), "Desc[%zu] is nullptr.", idx), return false); + OP_CHECK_IF(tmpITensor == nullptr, OP_LOGE(context->GetNodeName(), "Tensor[%zu] is nullptr.", idx), return false); + OP_CHECK_IF(tmpIShape == nullptr, OP_LOGE(context->GetNodeName(), "Shape[%zu] is nullptr.", idx), return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckNotNull(const gert::TilingContext *context, size_t idx) const +{ + OP_CHECK_IF(!CheckNotNullPtr(context, X_IDX, idx), + OP_LOGE(context->GetNodeName(), "x's tenosr or Desc is nullptr."), return false); + OP_CHECK_IF(!CheckNotNullPtr(context, WEIGHT_IDX, idx), + OP_LOGE(context->GetNodeName(), "weight's tenosr or Desc is nullptr."), return false); + OP_CHECK_IF(!CheckNotNullPtr(context, ANTIQUANT_SCALE_IDX, idx), + OP_LOGE(context->GetNodeName(), "antiquantscale's tenosr or Desc is nullptr."), return false); + if (hasBias_) { + OP_CHECK_IF(!CheckNotNullPtr(context, BIAS_IDX, idx), + OP_LOGE(context->GetNodeName(), "bias's tenosr or Desc is nullptr."), return false); + } + if (hasAntiquantOffset_) { + OP_CHECK_IF(!CheckNotNullPtr(context, ANTIQUANT_OFFSET_IDX, idx), + OP_LOGE(context->GetNodeName(), "antiquantoffset's tenosr or Desc is nullptr."), return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorDimEqualTarget(const gert::TilingContext *context, + uint32_t attrIdx, size_t idx, uint32_t targetDim, + const std::string &tensorType) const +{ + auto tmpShapePtr = context->GetDynamicInputShape(attrIdx, idx); + OP_CHECK_IF(tmpShapePtr == nullptr, + OP_LOGE(context->GetNodeName(), "%s[%zu] shape is nullptr.", tensorType.c_str(), idx), return false); + auto tmpShape = tmpShapePtr->GetStorageShape(); + uint32_t tmpDimNum = static_cast(tmpShape.GetDimNum()); + OP_CHECK_IF(tmpDimNum != targetDim, + OP_LOGE(context->GetNodeName(), "The expected dimension of %s[%zu] is [%u], actual %u.", + tensorType.c_str(), idx, targetDim, tmpDimNum), + return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorDimSingleXSingleWeightSingleY(const gert::TilingContext *context, + size_t idx) const +{ + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, X_IDX, idx, 2, "x"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, " + "the dimension of x does not match the expected value."), + return false); + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, WEIGHT_IDX, idx, 3, "weight"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, the " + "dimension of weight does not match the expected value."), + return false); + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, ANTIQUANT_SCALE_IDX, idx, 2, "antiquantScale"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, the " + "dimension of antiquantScale does not match the expected value."), + return false); + if (hasBias_) { + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, BIAS_IDX, idx, 2, "bias"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, " + "the dimension of bias does not match the expected value."), + return false); + } + if (hasAntiquantOffset_) { + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, ANTIQUANT_OFFSET_IDX, idx, 2, "antiquantOffset"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, the " + "dimension of antiquantOffset does not match the expected value."), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorDimMultiXMultiWeightMultiY(const gert::TilingContext *context, + size_t idx) const +{ + // 多多多场景x和weight的维度已经在SetShapeListMultiXMultiWeightMultiY函数中校验 + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, ANTIQUANT_SCALE_IDX, idx, 1, "antiquantScale"), + OP_LOGE(context->GetNodeName(), + "When x-weight is bf16/fp16-int4/int32 and grouptype is -1, antiquantScale " + "dimension mismatch."), + return false); + if (hasBias_) { + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, BIAS_IDX, idx, 1, "bias"), + OP_LOGE(context->GetNodeName(), + "When x-weight is bf16/fp16-int4/int32 and grouptype -1, bias dimension mismatch"), + return false); + } + if (hasAntiquantOffset_) { + OP_CHECK_IF(!CheckTensorDimEqualTarget(context, ANTIQUANT_OFFSET_IDX, idx, 1, "antiquantOffset"), + OP_LOGE(context->GetNodeName(), "When x-weight is bf16/fp16-int4/int32 and grouptype is -1, the " + "dimension of antiquantOffset does not match the expected value."), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorDim(const gert::TilingContext *context, size_t idx) const +{ + if (!IsA16W4ND()) { + return true; + } + if (groupType_ == GroupType::SPLIT_M) { + OP_CHECK_IF(!CheckTensorDimSingleXSingleWeightSingleY(context, idx), + OP_LOGE(context->GetNodeName(), "CheckTensorDimSingleXSingleWeightSingleY failed"), return false); + } else { + OP_CHECK_IF(!CheckTensorDimMultiXMultiWeightMultiY(context, idx), + OP_LOGE(context->GetNodeName(), "CheckTensorDimMultiXMultiWeightMultiY failed"), return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTensorShape(const gert::TilingContext *context, uint32_t attrIdx, + size_t idx, const std::string &tensorType) const +{ + if (!IsA16W4ND()) { + return true; + } + // 校验bias、antiquantScale和antiquantOffset的shape + auto tensorShapePtr = context->GetDynamicInputShape(attrIdx, idx); + auto wShapePtr = context->GetDynamicInputShape(WEIGHT_IDX, idx); + auto tensorShape = tensorShapePtr->GetStorageShape(); + auto wShape = wShapePtr->GetStorageShape(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + uint32_t tensorDimNum = static_cast(tensorShape.GetDimNum()); + + if (groupType_ == GroupType::SPLIT_M) { + // 校验bias、antiquantScale和antiquantOffset的第一根轴 + uint64_t groupNum = static_cast(wShape.GetDim(0)); + uint64_t batchSize = static_cast(tensorShape.GetDim(0)); + OP_CHECK_IF(groupNum != batchSize, + OP_LOGE(context->GetNodeName(), "%s batch size %lu should be euqal to groupList length %lu.", + tensorType.c_str(), batchSize, groupNum), + return false); + } + // 校验bias、antiquantScale和antiquantOffset的N轴 + int64_t weightNDimValue = wShape.GetDim(wDimNum - (transB_ ? 2 : 1)); + int64_t tensorNDimValue = tensorShape.GetDim(tensorDimNum - 1); + OP_CHECK_IF(weightNDimValue != tensorNDimValue, + OP_LOGE(context->GetNodeName(), + "NDim of %s should be equal to NDim of weight, but actual is %ld and %ld.", + tensorType.c_str(), tensorNDimValue, weightNDimValue), + return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckDimValue(const gert::TilingContext *context, size_t idx) const +{ + auto xDimNum = context->GetDynamicInputTensor(X_IDX, idx)->GetStorageShape().GetDimNum(); + // 校验batch轴和M轴大于等于0，A16W4ND场景下不考虑X转置的情况 + for (size_t dimIdx = 0; dimIdx < xDimNum - 1; dimIdx++) { + int64_t xDimValue = context->GetDynamicInputTensor(X_IDX, idx)->GetStorageShape().GetDim(dimIdx); + OP_CHECK_IF(xDimValue < 0, + OP_LOGE(context->GetNodeName(), + "The dimension[%zu] of the tensor[%zu] of x should be >= 0, but actual value is %ld", + dimIdx, idx, xDimValue), + return false); + } + int64_t xKDimValue = context->GetDynamicInputTensor(X_IDX, idx)->GetStorageShape().GetDim(xDimNum - 1); + OP_CHECK_IF(xKDimValue <= 0, + OP_LOGE(context->GetNodeName(), + "The K dimension of the tensor[%zu] should be positive, but actual is %ld ", idx, xKDimValue), + return false); + if (!IsA16W4ND()) { + return true; + } + auto wShape = context->GetDynamicInputShape(WEIGHT_IDX, idx)->GetStorageShape(); + auto wDimNum = wShape.GetDimNum(); + auto wNDimNum = transB_ ? (wDimNum - 2) : (wDimNum - 1); + auto wKDimNum = transB_ ? (wDimNum - 1) : (wDimNum - 2); + int64_t weightNDimValue = wShape.GetDim(wNDimNum); + int64_t weightKDimValue = wShape.GetDim(wKDimNum); + OP_CHECK_IF(weightNDimValue <= 0, + OP_LOGE(context->GetNodeName(), + "The N dimensions of the tensor[%zu] should be positive, but actual is %ld.", idx, + weightKDimValue), + return false); + OP_CHECK_IF( + xKDimValue != weightKDimValue, + OP_LOGE(context->GetNodeName(), + "The k dimension of the tensor[%zu] of x and weight should be equal, but actual are %ld and %ld.", + idx, xKDimValue, weightKDimValue), + return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckWeightInnerAxisEven(const gert::TilingContext *context, size_t idx) const +{ + if (weightDtype_ == ge::DT_INT4) { + auto wShapePtr = context->GetDynamicInputShape(WEIGHT_IDX, idx); + auto wShape = wShapePtr->GetOriginShape(); + auto wDimNum = wShape.GetDimNum() - 1; + int64_t wInnerDimvalue = wShape.GetDim(wDimNum); + OP_CHECK_IF((wInnerDimvalue % 2) != 0, + OP_LOGE(context->GetNodeName(), + "The last dimension of weight tensor[%zu] must be even, but got %ld.", idx, wInnerDimvalue), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckXAndWeightShape(const gert::TilingContext *context) const +{ + // 多多多场景在SetShapeListMultiXMultiWeightMultiY中校验过了，这里只校验单单单的场景 + auto xTensor = context->GetDynamicInputTensor(X_IDX, 0); // 0: get first tensor + auto wTensor = context->GetDynamicInputTensor(WEIGHT_IDX, 0); // 0: get first tensor + auto xShape = xTensor->GetStorageShape(); + auto wShape = wTensor->GetStorageShape(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + uint32_t xDimNum = static_cast(xShape.GetDimNum()); + OP_CHECK_IF(weightNzFlag_ && wDimNum < DlinferGroupedMatmulDirect::MIN_NZ_DIM, + OP_LOGE(context->GetNodeName(), + "Invalid weight dimension for format FRACTAL_NZ, expect at least 4, actual %u", wDimNum), + return false); + OP_CHECK_IF(!weightNzFlag_ && wDimNum < DlinferGroupedMatmulDirect::MIN_ND_DIM, + OP_LOGE(context->GetNodeName(), "Invalid weight dimension for format ND, expect at least 2, actual %u", + wDimNum), + return false); + OP_CHECK_IF( + xDimNum < DlinferGroupedMatmulDirect::MIN_ND_DIM, + OP_LOGE(context->GetNodeName(), "Invalid x dimension for format ND, expect at least 2, actual %u", xDimNum), + return false); + if (IsA16W4ND()) { + OP_CHECK_IF(xDimNum != 2, OP_LOGE(context->GetNodeName(), "Invalid x dimension, expect 2, actual %u", xDimNum), + return false); + OP_CHECK_IF(wDimNum != 3, + OP_LOGE(context->GetNodeName(), "Invalid weight dimension, expect 3, actual %u", xDimNum), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckEveryTensor(const gert::TilingContext *context) const +{ + for (size_t i = 0; i < numX_; i++) { + OP_CHECK_IF(!CheckNotNull(context, i), OP_LOGE(context->GetNodeName(), "CheckTensorNotNull failed"), + return false); + OP_CHECK_IF(!CheckTensorDim(context, i), OP_LOGE(context->GetNodeName(), "CheckTensorDim failed"), + return false); + OP_CHECK_IF(!CheckDimValue(context, i), OP_LOGE(context->GetNodeName(), "CheckDimValue failed"), return false); + OP_CHECK_IF(!CheckXAndWeightFormat(context, i), + OP_LOGE(context->GetNodeName(), "Check X And Weight Format failed"), return false); + OP_CHECK_IF(!CheckTensorDtype(context, X_IDX, i, xDType_, "x"), + OP_LOGE(context->GetNodeName(), "Check Tensor x Dtype failed"), return false); + OP_CHECK_IF(!CheckTensorDtype(context, WEIGHT_IDX, i, weightDtype_, "weight"), + OP_LOGE(context->GetNodeName(), "Check Tensor weight Dtype failed"), return false); + OP_CHECK_IF(!CheckTensorDtype(context, ANTIQUANT_SCALE_IDX, i, antiquantScaleDtype_, "antiquantScale"), + OP_LOGE(context->GetNodeName(), "Check Tensor antiquantScale Dtype failed"), return false); + OP_CHECK_IF(!CheckTensorShape(context, ANTIQUANT_SCALE_IDX, i, "antiquantScale"), + OP_LOGE(context->GetNodeName(), "Check antiquantScale tensor shape failed"), return false); + if (hasAntiquantOffset_) { + OP_CHECK_IF(!CheckTensorDtype(context, ANTIQUANT_OFFSET_IDX, i, antiquantOffsetDtype_, "antiquantOffset"), + OP_LOGE(context->GetNodeName(), "Check Tensor antiquantOffset Dtype failed"), return false); + OP_CHECK_IF(!CheckTensorShape(context, ANTIQUANT_OFFSET_IDX, i, "antiquantOffset"), + OP_LOGE(context->GetNodeName(), "Check antiquantOffset tensor shape failed"), return false); + } + if (hasBias_) { + OP_CHECK_IF(!CheckTensorDtype(context, BIAS_IDX, i, biasDtype_, "bias"), + OP_LOGE(context->GetNodeName(), "Check Tensor bias Dtype failed"), return false); + OP_CHECK_IF(!CheckTensorShape(context, BIAS_IDX, i, "bias"), + OP_LOGE(context->GetNodeName(), "Check bias tensor shape failed"), return false); + } + OP_CHECK_IF(!CheckWeightInnerAxisEven(context, i), + OP_LOGE(context->GetNodeName(), "CheckWeightInnerAxisEven failed"), return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckGroupList(const gert::TilingContext *context) const +{ + if (groupType_ == GroupType::SPLIT_M) { + OP_CHECK_IF(groupListType_ != 0 && groupListType_ != 1, + OP_LOGE(context->GetNodeName(), + "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, " + "groupListType only supports values 0 or 1, but actual is %u.", + groupListType_), + return false); + auto groupListPtr = context->GetOptionalInputShape(DlinferGroupedMatmulDirect::GROUPLIST_INDEX); + OP_CHECK_IF(groupListPtr == nullptr, OP_LOGE(context->GetNodeName(), "GroupList is nullptr."), return false); + int32_t groupListSize = static_cast(groupListPtr->GetOriginShape().GetDim(0)); + + auto wTensor = context->GetDynamicInputTensor(WEIGHT_IDX, 0); + OP_CHECK_IF(wTensor == nullptr, OP_LOGE(context->GetNodeName(), "wTensor is nullptr."), return false); + gert::Shape wShape = wTensor->GetStorageShape(); + int32_t groupNum = static_cast(wShape.GetDim(0)); + + OP_LOGD(context->GetNodeName(), "groupNum is %d, groupListSize is %d.", groupNum, groupListSize); + OP_CHECK_IF(groupListSize != groupNum, + OP_LOGE(context->GetNodeName(), + "When x-weight is bf16/fp16-int4/int32 and grouptype is 0, the length of " + "groupList must match the value of first dimension of weight, but actual is %d and %d.", + groupListSize, groupNum), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckRequiredInputs(const gert::TilingContext *context) const +{ + auto xShape = context->GetDynamicInputShape(X_IDX, 0); + OP_CHECK_IF(!IsNonEmpty(xShape), OP_LOGE(context->GetNodeName(), "x should not be null, but is null."), + return false); + auto weightShape = context->GetDynamicInputShape(WEIGHT_IDX, 0); + OP_CHECK_IF(!IsNonEmpty(weightShape), OP_LOGE(context->GetNodeName(), "weight should not be null, but is null."), + return false); + auto antiquantScaleShape = context->GetDynamicInputShape(ANTIQUANT_SCALE_IDX, 0); + OP_CHECK_IF(!IsNonEmpty(antiquantScaleShape), + OP_LOGE(context->GetNodeName(), "antiquantScale should not be null, but is null."), return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::AnalyzeAttr(const gert::TilingContext *context) +{ + auto compileInfoPtr = context->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context->GetNodeName(), "compileInfoPtr is nullptr."), return false); + coreNum_ = compileInfoPtr->aicNum; + + OP_CHECK_IF(!AnalyzeInput(context), OP_LOGE(context->GetNodeName(), "Invalid Input param"), return false); + + auto attr = context->GetAttrs(); + OP_CHECK_IF(attr == nullptr, OP_LOGE(context->GetNodeName(), "attr is nullptr."), return false); + const bool *transposeWeightPtr = attr->GetAttrPointer(ATTR_TRANS_W_IDX); + const bool *transposeXPtr = attr->GetAttrPointer(ATTR_TRANS_X_IDX); + const int32_t *groupTypePtr = attr->GetAttrPointer(ATTR_GROUPTYPE_IDX); + const int64_t *splitItemPtr = attr->GetAttrPointer(ATTR_SPLIT_ITEM_IDX); + const uint32_t *groupListTypePtr = attr->GetAttrPointer(ATTR_GROUP_LIST_TYPE_IDX); + transA_ = transposeXPtr != nullptr ? *transposeXPtr : false; + transB_ = transposeWeightPtr != nullptr ? *transposeWeightPtr : false; + groupType_ = groupTypePtr != nullptr ? static_cast(*groupTypePtr) : GroupType::NO_SPLIT; + splitItem_ = splitItemPtr != nullptr ? *splitItemPtr : 0; // 0: 默认split_item + groupListType_ = groupListTypePtr != nullptr ? *groupListTypePtr : 0; + isSingleX_ = (groupType_ != GroupType::NO_SPLIT && context->GetDynamicInputTensor(X_IDX, 1) == nullptr); + isSingleWeight_ = (groupType_ != GroupType::NO_SPLIT && context->GetDynamicInputTensor(WEIGHT_IDX, 1) == nullptr); + // 2: when x is multi-tensor, y is single-tensor; 3: when x is single-tensor, y is single-tensor + isSingleY_ = (splitItem_ == 2 || splitItem_ == 3); + + // 参数的设置和校验 + OP_CHECK_IF(coreNum_ <= 0, OP_LOGE(context->GetNodeName(), "Invalid coreNum[%u], expect greater than 0", coreNum_), + return false); + OP_CHECK_IF(!CheckUnsupportDataFlow(context), + OP_LOGE(context->GetNodeName(), "Input data contains unsupported dtype or format."), return false); + OP_CHECK_IF(!CheckTransposeStatus(context), OP_LOGE(context->GetNodeName(), "CheckTransposeStatus failed."), + return false); + OP_CHECK_IF(!CheckGroupTypeAndSplitItem(context), OP_LOGE(context->GetNodeName(), "CheckParam failed."), + return false); + OP_CHECK_IF(!SetShapeList(context), OP_LOGE(context->GetNodeName(), "SetShapeList failed."), return false); + OP_CHECK_IF(!CheckAntiQuantDtype(context), OP_LOGE(context->GetNodeName(), "CheckAntiQuantDtype failed."), + return false); + OP_CHECK_IF(!CheckBiasDtype(context), OP_LOGE(context->GetNodeName(), "CheckBiasDtype failed."), return false); + OP_CHECK_IF(!CheckGroupList(context), OP_LOGE(context->GetNodeName(), "CheckGroupList failed."), return false); + OP_CHECK_IF(!CheckRequiredInputs(context), OP_LOGE(context->GetNodeName(), "CheckRequiredInputs failed."), + return false); + OP_CHECK_IF(!CheckTensorListSize(context), OP_LOGE(context->GetNodeName(), "CheckTensorListSize failed."), + return false); + OP_CHECK_IF(!CheckEveryTensor(context), OP_LOGE(context->GetNodeName(), "CheckEveryTensor failed."), return false); + OP_CHECK_IF(!SetAntiquantGroupSize(context), OP_LOGE(context->GetNodeName(), "Unable to get antiquant groupSize"), + return false); + OP_CHECK_IF(!CheckGroupSize(context), OP_LOGE(context->GetNodeName(), "CheckGroupSize failed."), return false); + PrintInputParam(context); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::AnalyzeInput(const gert::TilingContext *context) +{ + auto xDesc = context->GetDynamicInputDesc(X_IDX, 0); + OP_CHECK_IF(xDesc == nullptr, OP_LOGE(context->GetNodeName(), "xDesc is nullptr."), return false); + xDType_ = xDesc->GetDataType(); + + auto wDesc = context->GetDynamicInputDesc(WEIGHT_IDX, 0); + OP_CHECK_IF(wDesc == nullptr, OP_LOGE(context->GetNodeName(), "wDesc is nullptr."), return false); + weightDtype_ = wDesc->GetDataType(); + auto wFormat = static_cast(ge::GetPrimaryFormat(wDesc->GetStorageFormat())); + if (wFormat == ge::FORMAT_FRACTAL_NZ_C0_16 || wFormat == ge::FORMAT_FRACTAL_NZ_C0_32 || + wFormat == ge::FORMAT_FRACTAL_NZ_C0_4 || wFormat == ge::FORMAT_FRACTAL_NZ_C0_2) { + wFormat = ge::FORMAT_FRACTAL_NZ; + } + weightNzFlag_ = wFormat == ge::FORMAT_FRACTAL_NZ; + + auto biasShape = context->GetDynamicInputShape(BIAS_IDX, 0); + hasBias_ = !(biasShape == nullptr || biasShape->GetStorageShape().GetShapeSize() == 0); + if (hasBias_) { + auto biasDesc = context->GetDynamicInputDesc(BIAS_IDX, 0); + OP_CHECK_IF(biasDesc == nullptr, OP_LOGE(context->GetNodeName(), "biasDesc is nullptr."), + return false); + biasDtype_ = biasDesc->GetDataType(); + } + + auto antiquantScaleDesc = context->GetDynamicInputDesc(ANTIQUANT_SCALE_IDX, 0); + OP_CHECK_IF(antiquantScaleDesc == nullptr, OP_LOGE(context->GetNodeName(), "antiquantScaleDesc is nullptr."), + return false); + antiquantScaleDtype_ = antiquantScaleDesc->GetDataType(); + + auto antiquantOffsetShape = context->GetDynamicInputShape(ANTIQUANT_OFFSET_IDX, 0); + hasAntiquantOffset_ = + !(antiquantOffsetShape == nullptr || antiquantOffsetShape->GetStorageShape().GetShapeSize() == 0); + if (hasAntiquantOffset_) { + auto antiquantOffsetDesc = context->GetDynamicInputDesc(ANTIQUANT_OFFSET_IDX, 0); + OP_CHECK_IF(antiquantOffsetDesc == nullptr, OP_LOGE(context->GetNodeName(), "antiquantOffsetDesc is nullptr."), + return false); + antiquantOffsetDtype_ = antiquantOffsetDesc->GetDataType(); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::EnableTailResplit() const { + // 不使能尾块重切分的场景 + // 1. 多多多 + // 2. 单单单 weight ND B非转置 + if (!isSingleX_ && !isSingleWeight_ && !isSingleY_) { + return false; + } + + if (!weightNzFlag_ && !transB_) { + return false; + } + + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CalcResplitTiling(const gert::TilingContext *context) +{ + if (!EnableTailResplit()) { + return true; + } + + uint64_t c0Size = 0; + OP_CHECK_IF(!GetC0Size(context, xDType_, c0Size), OP_LOGE(context->GetNodeName(), "Get C0 size failed"), + return false); + OP_CHECK_IF( + (weightNzFlag_ && !transB_ && nSize_ % c0Size > 0), + OP_LOGE(context->GetNodeName(), + "Invalid C0 size[%lu], expect greater than 0 and divisible by N[%lu] when weight format is FRACTAL_NZ", + c0Size, nSize_), + return false); + OP_CHECK_IF(coreNum_ <= 0, OP_LOGE(context->GetNodeName(), "Invalid core num[%u], expect greater than 0", coreNum_), + return false); + + cubeBlockDimN_ = static_cast(coreNum_); + if (nSize_ % (coreNum_ * static_cast(BASIC_BLOCK_BASE_N)) == 0UL) { + resplitParam_.mainBlockSize = BASIC_BLOCK_BASE_N; + resplitParam_.mainBlockCount = nSize_ / (coreNum_ * static_cast(BASIC_BLOCK_BASE_N)); + } else if (nSize_ >= coreNum_ * static_cast(BASIC_BLOCK_BASE_N_MIN)) { + // 该场景下可以保证分满核且尾块在128~256之间 + CalcFullBlockDimResplitTiling(c0Size); + } else { + // N <= 4096场景，优先保证单核尾块大于128，可能无法分满核 + CalcNoFullBlockDimResplitTiling(c0Size); + } + OP_CHECK_IF(!CheckResplitTilingResult(context), OP_LOGE(context->GetNodeName(), "Invalid resplit tiling result"), + return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::SetBaseTiling() +{ + tilingData_.gmmWeightQuantParam.groupNum = groupNum_; + tilingData_.gmmWeightQuantParam.coreNum = coreNum_; + tilingData_.gmmWeightQuantParam.kSize = kSize_; + tilingData_.gmmWeightQuantParam.nSize = nSizeOri_; + tilingData_.gmmWeightQuantParam.singleX = static_cast(isSingleX_); + tilingData_.gmmWeightQuantParam.singleWeight = static_cast(isSingleWeight_); + tilingData_.gmmWeightQuantParam.singleY = static_cast(isSingleY_); + tilingData_.gmmWeightQuantParam.groupType = static_cast(groupType_); + tilingData_.gmmWeightQuantParam.groupListType = static_cast(groupListType_); + tilingData_.gmmWeightQuantParam.hasBias = static_cast(hasBias_); + tilingData_.gmmWeightQuantParam.cubeBlockDimN = cubeBlockDimN_; + tilingData_.gmmWeightQuantParam.groupSize = groupSize_; + tilingData_.gmmWeightQuantParam.mainBlockSize = resplitParam_.mainBlockSize; + tilingData_.gmmWeightQuantParam.mainBlockCount = resplitParam_.mainBlockCount * coreNum_; + tilingData_.gmmWeightQuantParam.firstTailBlockSize = resplitParam_.firstTailBlockSize; + tilingData_.gmmWeightQuantParam.secondTailBlockSize = resplitParam_.secondTailBlockSize; + tilingData_.gmmWeightQuantParam.firstTailBlockCount = resplitParam_.firstTailBlockCount; + tilingData_.gmmWeightQuantParam.secondTailBlockCount = resplitParam_.secondTailBlockCount; + errno_t retM = memcpy_s(tilingData_.gmmArray.mList, sizeof(tilingData_.gmmArray.mList), mList_, sizeof(mList_)); + if (retM != EOK) { + return false; + } + errno_t retK = memcpy_s(tilingData_.gmmArray.kList, sizeof(tilingData_.gmmArray.kList), kList_, sizeof(kList_)); + if (retK != EOK) { + return false; + } + errno_t retN = memcpy_s(tilingData_.gmmArray.nList, sizeof(tilingData_.gmmArray.nList), nList_, sizeof(nList_)); + if (retN != EOK) { + return false; + } + return true; +} + +void GroupedWeightQuantBatchMatmulTiling::SetMatMulTiling() +{ + tilingData_.mmTilingData.baseM = BASIC_BLOCK_BASE_M; + tilingData_.mmTilingData.singleCoreM = BASIC_BLOCK_BASE_M; + tilingData_.mmTilingData.isBias = static_cast(hasBias_); + tilingData_.mmTilingData.M = mSize_; + tilingData_.mmTilingData.N = nSize_; + tilingData_.mmTilingData.Ka = kSize_; + tilingData_.mmTilingData.Kb = kSize_; + tilingData_.mmTilingData.singleCoreN = BASIC_BLOCK_BASE_N; + tilingData_.mmTilingData.singleCoreK = kSize_; + tilingData_.mmTilingData.dbL0A = BUFFER_NUM_2; + tilingData_.mmTilingData.dbL0B = BUFFER_NUM_2; + tilingData_.mmTilingData.dbL0C = 1; + tilingData_.mmTilingData.shareL0CSize = BASIC_BLOCK_BASE_M * BASIC_BLOCK_BASE_N * sizeof(float); + + tilingData_.mmTilingData.baseN = BASIC_BLOCK_BASE_N; + tilingData_.mmTilingData.baseK = BASIC_BLOCK_BASE_K; + tilingData_.mmTilingData.stepKa = STEP_K_4; + tilingData_.mmTilingData.stepKb = STEP_K_4; + tilingData_.mmTilingData.depthA1 = DEPTH_8; + tilingData_.mmTilingData.depthB1 = DEPTH_8; + tilingData_.mmTilingData.stepM = 1; + tilingData_.mmTilingData.stepN = 1; + tilingData_.mmTilingData.usedCoreNum = coreNum_; + + if (xDType_ == ge::DT_INT8 && weightDtype_ == ge::DT_INT4) { + // 2含义：S8S4场景，MAD采用S8类型，baseK需要放大2倍 + tilingData_.mmTilingData.baseK = BASIC_BLOCK_BASE_K * 2; + // A8W4场景在UB中处理bias，mm api默认无bias + tilingData_.mmTilingData.isBias = 0; + // groupsize=192时, ubMte2InnerSize=384, stepKb=ubMte2InnerSize/baseK=3 + if (groupSize_ == 192u) { + tilingData_.mmTilingData.stepKa = STEP_K_3; + tilingData_.mmTilingData.stepKb = STEP_K_3; + } + OP_LOGI("SetMatMulTiling", "stepKb = %llu", tilingData_.mmTilingData.stepKb); + } else if (xDType_ == ge::DT_FLOAT8_E4M3FN && weightDtype_ == ge::DT_FLOAT4_E2M1 && + antiquantScaleDtype_ == ge::DT_FLOAT8_E8M0) { + // MxA8W4场景配置mxTypePara + tilingData_.mmTilingData.mxTypePara = (SCALE_FACTOR_DEFAULT << SCALE_FACTOR_N_BIT) + + (SCALE_FACTOR_DEFAULT << SCALE_FACTOR_M_BIT) + + (SCALE_FACTOR_MIN << SCALE_FACTOR_B_BIT) + SCALE_FACTOR_MIN; + } else if (hasBias_) { + tilingData_.mmTilingData.baseM = BASIC_BLOCK_BASE_M_WITH_BIAS; + tilingData_.mmTilingData.singleCoreM = BASIC_BLOCK_BASE_M_WITH_BIAS; + } +} + +void GroupedWeightQuantBatchMatmulTiling::SetTilingKey(gert::TilingContext *context) +{ + constexpr uint8_t DECIMAL = 10U; + // 平台类型占2位(平台大类， 平台小类)，平台大类在高位，需要乘10 + tilingKeyConfig_.socVersionType = static_cast(SocVersionType::SUPPORT_L1_TO_BT_BF16) * DECIMAL; + tilingKeyConfig_.quantizationScenario = static_cast(QuantizationScenario::DEFAULT); + // 算法类型占2位(算法大类，算法小类)，算法大类在高位，需要乘10 + if (EnableTailResplit()) { + tilingKeyConfig_.algorithm = static_cast(OptimizationAlgorithmCategory::VECTOR_ANTIQUANT) * DECIMAL + + static_cast(OptimizationAlgorithmSubCategory::N_FIRST_TAIL_RESPLIT); + } else { + tilingKeyConfig_.algorithm = static_cast(OptimizationAlgorithmCategory::VECTOR_ANTIQUANT) * DECIMAL + + static_cast(OptimizationAlgorithmSubCategory::N_FIRST_BASIC_BLOCK); + } + + tilingKeyConfig_.transposeSituation = (static_cast(transA_) << 1) | static_cast(transB_); + + if (antiquantScaleDtype_ == ge::DT_FLOAT8_E8M0) { + tilingKeyConfig_.antiquantType = static_cast(QuantType::MX); + } else { + tilingKeyConfig_.antiquantType = static_cast(QuantType::PER_CHANNEL); + } + + tilingKeyConfig_.quantType = static_cast(QuantType::NONE); + tilingKeyConfig_.optionInputSituation = static_cast(hasAntiquantOffset_) << 1; + tilingKeyConfig_.weightFormat = + weightNzFlag_ ? static_cast(WeightFormat::FRACTAL_NZ) : static_cast(WeightFormat::ND); + tilingKeyConfig_.templateCustom = static_cast(Mte2Configuration::MTE2_INNER_SIZE_256_BUF_NUM_4); + if (xDType_ == ge::DT_INT8 && weightDtype_ == ge::DT_INT4) { + tilingKeyConfig_.templateCustom = static_cast(Mte2Configuration::MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT); + if (groupSize_ == 192u) { + tilingKeyConfig_.templateCustom = static_cast(Mte2Configuration::MTE2_INNER_SIZE_384_BUF_NUM_3); + } + } + tilingKeyConfig_.apiConstexpr = 0U; + context->SetTilingKey(tilingKeyConfig_.GenTilingKey()); +} + +bool GroupedWeightQuantBatchMatmulTiling::SetCustomParam(gert::TilingContext *context) +{ + size_t *workspaces = context->GetWorkspaceSizes(1); // get second variable + OP_CHECK_IF(workspaces == nullptr, OP_LOGE(context->GetNodeName(), "workspaces is nullptr."), + return false); // check workspaces is not null + workspaces[0] = 16777216U; // 16 * 1024 * 1024: default workspace size + + context->SetBlockDim(coreNum_); + OP_CHECK_IF(context->GetRawTilingData() == nullptr, OP_LOGE(context->GetNodeName(), "RawTilingData is nullptr."), + return false); + errno_t ret = memcpy_s(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity(), reinterpret_cast(&tilingData_), sizeof(tilingData_)); + if (ret != EOK) { + OP_LOGE(context->GetNodeName(), "memcpy_s failed, ret = %d", ret); + return false; + } + context->GetRawTilingData()->SetDataSize(sizeof(tilingData_)); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::IsA16W4ND() const +{ + if (ge::GetSizeByDataType(xDType_) == B16_DATA_SIZE && + ((weightDtype_ == ge::DT_INT4) || (weightDtype_ == ge::DT_INT32)) && !weightNzFlag_) { + return true; + } + return false; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckUnsupportDataFlow(const gert::TilingContext *context) const +{ + if (ge::GetSizeByDataType(xDType_) == B16_DATA_SIZE && + (weightDtype_ == ge::DT_INT8 || weightDtype_ == ge::DT_INT32 || weightDtype_ == ge::DT_INT4)) { + OP_CHECK_IF(weightNzFlag_, + OP_LOGE(context->GetNodeName(), "In weight quant case, when x-weight is bf16/fp16-int8/int4/int32, " + "weight only support weight with ND-format . "), + return false); + } else if (ge::GetSizeByDataType(xDType_) == B16_DATA_SIZE && + (weightDtype_ == ge::DT_FLOAT8_E4M3FN || weightDtype_ == ge::DT_FLOAT8_E5M2 || + weightDtype_ == ge::DT_HIFLOAT8)) { + OP_CHECK_IF(!(!weightNzFlag_ && transB_), + OP_LOGE(context->GetNodeName(), "In weight quant case, when x-weight is bf16/fp16-fp8/hif8, " + "weight only supports transposed and ND-format."), + return false); + } else if (ge::GetSizeByDataType(xDType_) == B16_DATA_SIZE && + (weightDtype_ == ge::DT_FLOAT4_E2M1 || weightDtype_ == ge::DT_FLOAT4_E1M2 || + weightDtype_ == ge::DT_FLOAT) && + antiquantScaleDtype_ == ge::DT_FLOAT8_E8M0) { + OP_CHECK_IF(!(weightNzFlag_ && !transB_), + OP_LOGE(context->GetNodeName(), "In weight quant case, when x-weight is bf16/fp16-fp4/fp16, weight " + " only supports untransposed and FRACTAL_NZ-format "), + return false); + } else if (xDType_ == ge::DT_INT8 && weightDtype_ == ge::DT_INT4) { + OP_CHECK_IF(!(weightNzFlag_ && !transB_), + OP_LOGE(context->GetNodeName(), "In weight quant case, when x-weight is int8-int4, weight only " + "supports untransposed and FRACTAL_NZ-format "), + return false); + } else if (xDType_ == ge::DT_FLOAT8_E4M3FN && + (weightDtype_ == ge::DT_FLOAT4_E2M1 || weightDtype_ == ge::DT_FLOAT) && + antiquantScaleDtype_ == ge::DT_FLOAT8_E8M0) { + OP_CHECK_IF( + !(weightNzFlag_ && transB_), + OP_LOGE(context->GetNodeName(), "In weight quant case, when x-weight is float8_e4m3fn-float4_e2m1/fp32, " + "weight only supports transposed weight and FRACTAL_NZ-format"), + return false); + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckAntiQuantDtype(const gert::TilingContext *context) const +{ + if (IsA16W4ND()) { + OP_CHECK_IF(antiquantScaleDtype_ != xDType_, + OP_LOGE(context->GetNodeName(), "The dtype of antiquantscale should be same with xdtype."), + return false); + if (hasAntiquantOffset_) { + OP_CHECK_IF(antiquantOffsetDtype_ != xDType_, + OP_LOGE(context->GetNodeName(), "The dtype of antiquantOffset should be same with xdtype."), + return false); + } + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckBiasDtype(const gert::TilingContext *context) const +{ + if (IsA16W4ND()) { + if (hasBias_) { + OP_CHECK_IF(BIAS_TYPE_SUPPORT_MAP.find(xDType_) == BIAS_TYPE_SUPPORT_MAP.end(), + OP_LOGE(context->GetNodeName(), "Cannot find bias dtype match with xDtype."), return false); + OP_CHECK_IF(BIAS_TYPE_SUPPORT_MAP.at(xDType_).find(biasDtype_) == BIAS_TYPE_SUPPORT_MAP.at(xDType_).end(), + OP_LOGE(context->GetNodeName(), + "Data type [%s] is not supported for bias, when xDtype is [%s].", + ge::TypeUtils::DataTypeToSerialString(biasDtype_).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDType_).c_str()), + return false); + } + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckGroupTypeAndSplitItem(const gert::TilingContext *context) const +{ + if (IsA16W4ND()) { + OP_CHECK_IF( + (groupType_ != GroupType::NO_SPLIT) && (groupType_ != GroupType::SPLIT_M), + OP_LOGE(context->GetNodeName(), "when x-weight is bf16/fp16-int32/int4, grouptype only supports -1 or 0."), + return false); + if (groupType_ == GroupType::NO_SPLIT) { + OP_CHECK_IF((splitItem_ != 0 && splitItem_ != 1), + OP_LOGE(context->GetNodeName(), "When grouptype is -1. splititem can only be 0 or 1."), + return false); + } else { + OP_CHECK_IF((splitItem_ != 2 && splitItem_ != 3), + OP_LOGE(context->GetNodeName(), "When grouptype is 0. splititem can only be 2 or 3."), + return false); + } + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckTransposeStatus(const gert::TilingContext *context) const +{ + OP_CHECK_IF(transA_, OP_LOGE(context->GetNodeName(), "Transposed A is not supported. "), return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::SetShapeListSplitMSingleXSingleWeightSingleY( + const gert::TilingContext *context) +{ + auto xTensor = context->GetDynamicInputTensor(X_IDX, 0); // 0: get first tensor + OP_CHECK_IF(xTensor == nullptr, OP_LOGE(context->GetNodeName(), "xTensor is nullptr."), return false); + gert::Shape xShape = xTensor->GetStorageShape(); + + auto wTensor = context->GetDynamicInputTensor(WEIGHT_IDX, 0); // 0: get first tensor + OP_CHECK_IF(wTensor == nullptr, OP_LOGE(context->GetNodeName(), "wTensor is nullptr."), return false); + gert::Shape wShape = wTensor->GetStorageShape(); + + groupNum_ = static_cast(wShape.GetDim(0)); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + uint32_t xDimNum = static_cast(xShape.GetDimNum()); + OP_CHECK_IF(!CheckXAndWeightShape(context), OP_LOGE(context->GetNodeName(), "CheckXAndWeightShape failed"), + return false); + mSize_ = transA_ ? xShape.GetDim(1) : xShape.GetDim(0); + kSize_ = transA_ ? xShape.GetDim(0) : xShape.GetDim(xDimNum - 1); + // -1含义为(K, N)场景N索引，-2含义为(N, K)场景N索引 + nSize_ = transB_ ? wShape.GetDim(wDimNum - 2) : wShape.GetDim(wDimNum - 1); + if (weightNzFlag_) { + // 非转置NZ排布(N1, K1, K0, N0), 转置NZ排布(K1, N1, N0, K0) + // -3含义：转置N1索引；-2含义：转置N0索引 + nSize_ = transB_ ? wShape.GetDim(wDimNum - 3) * wShape.GetDim(wDimNum - 2) + // -4含义：非转置N1索引；-1含义：非转置N0索引 + : wShape.GetDim(wDimNum - 4) * wShape.GetDim(wDimNum - 1); + const gert::StorageShape *yShapePtr = context->GetOutputShape(0); + OP_CHECK_IF(yShapePtr == nullptr, OP_LOGE(context->GetNodeName(), "yShapePtr is nullptr."), return false); + const gert::Shape &yShape = yShapePtr->GetOriginShape(); + nSizeOri_ = yShape.GetDim(static_cast(yShape.GetDimNum()) - 1); + } + OP_CHECK_IF( + mSize_ <= 0 || kSize_ <= 0 || nSize_ <= 0, + OP_LOGE(context->GetNodeName(), "Invalid mSize[%lu] kSize[%lu] or nSize[%lu], expect all greater than 0", + mSize_, kSize_, nSize_), + return false); + + if (weightDtype_ == ge::DT_FLOAT || weightDtype_ == ge::DT_INT32) { + weightDtype_ = weightDtype_ == ge::DT_FLOAT ? ge::DT_FLOAT4_E2M1 : ge::DT_INT4; + if (!transB_) { + // 一个float32/int32表示8个fp4/int4，设置为正确shape；kSize来自x，不需要考虑转置场景k轴扩大 + nSize_ = static_cast(8) * nSize_; + } + } + if (!weightNzFlag_) { + nSizeOri_ = nSize_; + } + kList_[0] = static_cast(kSize_); + nList_[0] = static_cast(nSizeOri_); + mList_[0] = -1; + return true; +} + +uint16_t GroupedWeightQuantBatchMatmulTiling::GetTensorListSize(const gert::TilingContext *context, + uint32_t attrIdx) const +{ + uint16_t count = 0; + for (int i = 0; i <= DlinferGroupedMatmulDirect::MAX_TENSOR_CONT; i++) { + auto shapePtr = context->GetDynamicInputShape(attrIdx, count); + if (!IsNonEmpty(shapePtr)) { + break; + } + ++count; + } + return count; +} + +void GroupedWeightQuantBatchMatmulTiling::GetNumOfInputs(const gert::TilingContext *context) +{ + numX_ = GetTensorListSize(context, X_IDX); + numWeight_ = GetTensorListSize(context, WEIGHT_IDX); + numBias_ = GetTensorListSize(context, BIAS_IDX); + numAntiquantScale_ = GetTensorListSize(context, ANTIQUANT_SCALE_IDX); + numAntiquantOffset_ = GetTensorListSize(context, ANTIQUANT_OFFSET_IDX); +} + +bool GroupedWeightQuantBatchMatmulTiling::SetShapeListMultiXMultiWeightMultiY(const gert::TilingContext *context) +{ + for (uint16_t i = 0; i < DlinferGroupedMatmulDirect::MAX_TENSOR_CONT; i++) { + auto xShapePtr = context->GetDynamicInputShape(X_IDX, i); + auto wShapePtr = context->GetDynamicInputShape(WEIGHT_IDX, i); + if (xShapePtr == nullptr || wShapePtr == nullptr) { + break; + } + auto xShape = xShapePtr->GetStorageShape(); + auto wShape = wShapePtr->GetOriginShape(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + uint32_t xDimNum = static_cast(xShape.GetDimNum()); + OP_CHECK_IF(xDimNum < MIN_X_DIM || xDimNum > MAX_X_DIM, + OP_LOGE(context->GetNodeName(), "Invalid x dimension, expect 2-6, actual %u", xDimNum), + return false); + OP_CHECK_IF(wDimNum != DlinferGroupedMatmulDirect::MIN_ND_DIM, + OP_LOGE(context->GetNodeName(), "Invalid weight dimension, expect 2, actual %u", xDimNum), + return false); + groupNum_ += 1U; + // -1含义为(K, M)场景M索引，-2含义为(M, K)场景M索引 + int64_t m = transA_ ? xShape.GetDim(xDimNum - 1) : xShape.GetDim(xDimNum - 2); + // -2含义：x的最后2维为M和K，对除M, K的batch轴进行累乘 + for (uint16_t xDim = 0; xDim < static_cast(xDimNum) - 2; xDim++) { + m *= xShape.GetDim(xDim); + } + int64_t k = transB_ ? wShape.GetDim(1) : wShape.GetDim(0); + int64_t n = transB_ ? wShape.GetDim(0) : wShape.GetDim(1); + mList_[i] = static_cast(m); + kList_[i] = static_cast(k); + nList_[i] = static_cast(n); + mSize_ = std::max(mSize_, static_cast(m)); + kSize_ = std::max(kSize_, static_cast(k)); + nSize_ = std::max(nSize_, static_cast(n)); + } + nSizeOri_ = nSize_; + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckGroupSize(const gert::TilingContext *context) const +{ + if (xDType_ != ge::DT_INT8 || weightDtype_ != ge::DT_INT4) { + return true; + } + // 伪量化S8S4场景支持groupsize为128/192/256/512 + OP_CHECK_IF( + groupSize_ != 128u && groupSize_ != 256u && groupSize_ != 512u && groupSize_ != 192u, + OP_LOGE(context->GetNodeName(), "groupSize must be 128/192/256/512, but current groupSize is %u.", groupSize_), + return false); + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::SetAntiquantGroupSize(const gert::TilingContext *context) +{ + auto antiquantScale = context->GetDynamicInputTensor(ANTIQUANT_SCALE_IDX, 0); + OP_CHECK_IF(antiquantScale == nullptr, OP_LOGE(context->GetNodeName(), "antiquantScale is nullptr."), return false); + auto antiquantScaleShape = antiquantScale->GetStorageShape(); + int64_t antiquantScaleDimNum = antiquantScaleShape.GetDimNum(); + // 3含义：单场景antiquantScale维度在切M时是(g, N, K)或(g, K, N) + if (antiquantScaleDimNum == 3) { + // 2含义：(g, K, N)格式K轴索引 + int64_t groupNum = transB_ ? antiquantScaleShape.GetDim(antiquantScaleDimNum - 1) + : antiquantScaleShape.GetDim(antiquantScaleDimNum - 2); + OP_CHECK_IF(groupNum <= 0 || kSize_ % groupNum > 0, + OP_LOGE( + context->GetNodeName(), + "Invalid groupNum[%ld], expect greater than 0 and divisible by kSize[%lu]", groupNum, kSize_), + return false); + // GMM伪量化场景支持K=groupSize + groupSize_ = groupNum > 0 ? kSize_ / static_cast(groupNum) : 0; + } + return true; +} + +bool GroupedWeightQuantBatchMatmulTiling::GetC0Size(const gert::TilingContext *context, ge::DataType dtype, + uint64_t &c0Size) const +{ + if (dtype == ge::DT_INT4) { + c0Size = AscendC::ONE_BLK_SIZE + AscendC::ONE_BLK_SIZE; + } else { + int64_t dtypeSize = ge::GetSizeByDataType(dtype); + OP_CHECK_IF(dtypeSize <= 0, + OP_LOGE(context->GetNodeName(), "Invalid dtypeSize[%ld], expect greater than 0", + dtypeSize), + return false); + if (dtypeSize > 0) { + c0Size = AscendC::ONE_BLK_SIZE / dtypeSize; + } + } + return true; +} + +void GroupedWeightQuantBatchMatmulTiling::CalcFullBlockDimResplitTiling(uint64_t c0Size) +{ + resplitParam_.mainBlockSize = BASIC_BLOCK_BASE_N; + resplitParam_.mainBlockCount = 0UL; + if (nSize_ / (coreNum_ * static_cast(BASIC_BLOCK_BASE_N)) > 1UL) { + resplitParam_.mainBlockCount = nSize_ / (coreNum_ * static_cast(BASIC_BLOCK_BASE_N)) - 1UL; + } + uint64_t tailSizeOri = nSize_ - resplitParam_.mainBlockCount * resplitParam_.mainBlockSize * coreNum_; + uint64_t tailSize = tailSizeOri; + if (weightNzFlag_ && c0Size != 0UL) { + tailSize = tailSizeOri / c0Size; + } + if (tailSizeOri > coreNum_ * static_cast(BASIC_BLOCK_BASE_N)) { + // 2含义：当尾块大于核数*256，除以2倍核数以保证单核尾块大小小于256 + constexpr uint64_t resplitFactor = 2; + resplitParam_.firstTailBlockSize = static_cast(tailSize / (coreNum_ * resplitFactor)); + resplitParam_.secondTailBlockSize = static_cast(resplitParam_.firstTailBlockSize + 1U); + resplitParam_.secondTailBlockCount = static_cast(tailSize % (coreNum_ * resplitFactor)); + resplitParam_.firstTailBlockCount = + static_cast(coreNum_ * resplitFactor - resplitParam_.secondTailBlockCount); + } else { + resplitParam_.firstTailBlockSize = static_cast(tailSize / coreNum_); + resplitParam_.secondTailBlockSize = static_cast(resplitParam_.firstTailBlockSize + 1U); + resplitParam_.secondTailBlockCount = static_cast(tailSize % coreNum_); + resplitParam_.firstTailBlockCount = static_cast(coreNum_ - resplitParam_.secondTailBlockCount); + } + if (weightNzFlag_) { + resplitParam_.firstTailBlockSize *= static_cast(c0Size); + resplitParam_.secondTailBlockSize *= static_cast(c0Size); + } +} + +void GroupedWeightQuantBatchMatmulTiling::CalcNoFullBlockDimResplitTiling(uint64_t c0Size) +{ + resplitParam_.mainBlockSize = BASIC_BLOCK_BASE_N; + resplitParam_.mainBlockCount = 0UL; + uint64_t taskNum = std::max(1UL, nSize_ / BASIC_BLOCK_BASE_N_MIN); // 实际任务数，必然小于核数 + cubeBlockDimN_ = static_cast(taskNum); + if (weightNzFlag_ && c0Size != 0UL && taskNum != 0UL) { + resplitParam_.firstTailBlockSize = static_cast(nSize_ / c0Size / taskNum); + resplitParam_.secondTailBlockSize = static_cast(resplitParam_.firstTailBlockSize + 1U); + resplitParam_.secondTailBlockCount = static_cast(nSize_ / c0Size % taskNum); + resplitParam_.firstTailBlockCount = static_cast(taskNum - resplitParam_.secondTailBlockCount); + resplitParam_.firstTailBlockSize *= static_cast(c0Size); + resplitParam_.secondTailBlockSize *= static_cast(c0Size); + } else if (taskNum != 0UL) { + resplitParam_.firstTailBlockSize = static_cast(nSize_ / taskNum); + resplitParam_.secondTailBlockSize = static_cast(resplitParam_.firstTailBlockSize + 1U); + resplitParam_.secondTailBlockCount = static_cast(nSize_ % taskNum); + resplitParam_.firstTailBlockCount = static_cast(taskNum - resplitParam_.secondTailBlockCount); + } +} + +bool GroupedWeightQuantBatchMatmulTiling::CheckResplitTilingResult(const gert::TilingContext *context) const +{ + OP_CHECK_IF(nSize_ != static_cast(resplitParam_.mainBlockCount) * coreNum_ * resplitParam_.mainBlockSize + + resplitParam_.firstTailBlockCount * resplitParam_.firstTailBlockSize + + resplitParam_.secondTailBlockCount * resplitParam_.secondTailBlockSize, + OP_LOGE(context->GetNodeName(), + "Invalid resplit tiling result, expect nSize[%lu] == mainBlockCount[%lu] x coreNum[%u] x " + "mainBlockSize[%u] + firstTailBlockCount[%hu] x firstTailBlockSize[%hu] + " + "secondTailBlockCount[%hu] x secondTailBlockSize[%hu]", + nSize_, resplitParam_.mainBlockCount, coreNum_, resplitParam_.mainBlockSize, + resplitParam_.firstTailBlockCount, resplitParam_.firstTailBlockSize, + resplitParam_.secondTailBlockCount, resplitParam_.secondTailBlockSize), + return false); + OP_CHECK_IF( + nSize_ >= static_cast(coreNum_) * BASIC_BLOCK_BASE_N_MIN && + (resplitParam_.firstTailBlockCount + resplitParam_.secondTailBlockCount) % coreNum_ > 0, + OP_LOGE(context->GetNodeName(), + "Invalid resplit tiling result, expect core num [%u] is divisible by " + "(firstTailBlockCount[%hu] + secondTailBlockCount[%hu])", + coreNum_, resplitParam_.firstTailBlockCount, resplitParam_.secondTailBlockCount), + return false); + OP_CHECK_IF( + nSize_ >= static_cast(BASIC_BLOCK_BASE_N_MIN) && resplitParam_.firstTailBlockCount > 0 && + (resplitParam_.firstTailBlockSize < BASIC_BLOCK_BASE_N_MIN || + resplitParam_.firstTailBlockSize > BASIC_BLOCK_BASE_N), + OP_LOGE(context->GetNodeName(), + "Invalid resplit tiling result, expect [%u] <= firstTailBlockSize [%hu] <= [%u] ", + BASIC_BLOCK_BASE_N_MIN, resplitParam_.firstTailBlockSize, BASIC_BLOCK_BASE_N), + return false); + OP_CHECK_IF( + nSize_ >= static_cast(BASIC_BLOCK_BASE_N_MIN) && resplitParam_.secondTailBlockCount > 0 && + (resplitParam_.secondTailBlockSize < BASIC_BLOCK_BASE_N_MIN || + resplitParam_.secondTailBlockSize > BASIC_BLOCK_BASE_N), + OP_LOGE(context->GetNodeName(), + "Invalid resplit tiling result, expect [%u] <= secondTailBlockSize [%hu] <= [%u] ", + BASIC_BLOCK_BASE_N_MIN, resplitParam_.secondTailBlockSize, BASIC_BLOCK_BASE_N), + return false); + return true; +} + +void GroupedWeightQuantBatchMatmulTiling::PrintInputParam(const gert::TilingContext *context) const +{ + OP_LOGI(context->GetNodeName(), + "Input params: coreNum: %u, gmm groupNum: %u, groupType: %d, groupListType: %u, splitItem: %lld, " + "antiquant-groupSize: %u, mSize: %llu, kSize: %llu, nSize: %llu, nSizeOri: %llu, transA: %s, transB: %s, " + "isSingleX: %s, isSingleWeight: %s, isSingleY: %s, hasBias: %s, weightNzFlag: %s, hasAntiquantOffset: " + "%s, xDtype: %s, weightDtype: %s", + coreNum_, groupNum_, static_cast(groupType_), groupListType_, splitItem_, groupSize_, mSize_, + kSize_, nSize_, nSizeOri_, transA_ ? "true" : "false", transB_ ? "true" : "false", + isSingleX_ ? "true" : "false", isSingleWeight_ ? "true" : "false", isSingleY_ ? "true" : "false", + hasBias_ ? "true" : "false", weightNzFlag_ ? "true" : "false", hasAntiquantOffset_ ? "true" : "false", + ge::TypeUtils::DataTypeToSerialString(xDType_).c_str(), + ge::TypeUtils::DataTypeToSerialString(weightDtype_).c_str()); +} + +void GroupedWeightQuantBatchMatmulTiling::PrintTilingResult(const gert::TilingContext *context) +{ + OP_LOGI( + context->GetNodeName(), + "Tiling result: groupNum: %u, coreNum: %u, kSize: %lu, nSize: %lu, singleX: %u, singleWeight: %u, singleY: %u, " + "groupType: %d, groupListType: %u, hasBias: %u, groupSize: %u, mainBlockSize: %u, mainBlockCount: %lu, " + "firstTailBlockSize: %u, secondTailBlockSize: %u, firstTailBlockCount: %u, secondTailBlockCount: %u", + tilingData_.gmmWeightQuantParam.groupNum, tilingData_.gmmWeightQuantParam.coreNum, + tilingData_.gmmWeightQuantParam.kSize, tilingData_.gmmWeightQuantParam.nSize, + tilingData_.gmmWeightQuantParam.singleX, tilingData_.gmmWeightQuantParam.singleWeight, + tilingData_.gmmWeightQuantParam.singleY, tilingData_.gmmWeightQuantParam.groupType, + tilingData_.gmmWeightQuantParam.groupListType, tilingData_.gmmWeightQuantParam.hasBias, + tilingData_.gmmWeightQuantParam.groupSize, tilingData_.gmmWeightQuantParam.mainBlockSize, + tilingData_.gmmWeightQuantParam.mainBlockCount, tilingData_.gmmWeightQuantParam.firstTailBlockSize, + tilingData_.gmmWeightQuantParam.secondTailBlockSize, + tilingData_.gmmWeightQuantParam.firstTailBlockCount, + tilingData_.gmmWeightQuantParam.secondTailBlockCount); +} + +uint64_t TilingKeyConfigure::GenTilingKey() const +{ + PrintTilingKeyLog(); + constexpr uint8_t DECIMAL = 10U; + uint64_t transInfo = this->transposeSituation; + bool atrans_ = (transInfo == static_cast(GmmTrans::ATrans)) || + (transInfo == static_cast(GmmTrans::ABTrans)); + bool btrans_ = (transInfo == static_cast(GmmTrans::BTrans)) || + (transInfo == static_cast(GmmTrans::ABTrans)); + return GET_TPL_TILING_KEY( + static_cast(this->weightFormat), static_cast(this->optionInputSituation), + static_cast(this->quantType), static_cast(this->antiquantType), + static_cast(btrans_), static_cast(atrans_), static_cast(this->templateCustom), + static_cast(this->algorithm % DECIMAL), static_cast(this->algorithm / DECIMAL)); +} + +} // namespace optiling \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.h new file mode 100644 index 00000000..01a2b31d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/arch35/grouped_weight_quant_batch_matmul_tiling.h @@ -0,0 +1,311 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_weight_quant_batch_matmul_tiling.h + * \brief + */ +#ifndef GROUPED_WEIGHT_QUANT_BATCH_MATMUL_TILING_H +#define GROUPED_WEIGHT_QUANT_BATCH_MATMUL_TILING_H + +#include +#include +#include +#include + +#include "../grouped_matmul_tiling.h" +#include "../../../op_kernel/arch35/grouped_matmul_tiling_data_apt.h" +#include "log/log.h" +#include "register/op_impl_registry.h" + +namespace optiling { +constexpr uint32_t X_IDX = 0; +constexpr uint32_t WEIGHT_IDX = 1; +constexpr uint32_t BIAS_IDX = 2; +constexpr uint32_t ANTIQUANT_SCALE_IDX = 5; +constexpr uint32_t ANTIQUANT_OFFSET_IDX = 6; +constexpr uint64_t ATTR_SPLIT_ITEM_IDX = 0; +constexpr uint64_t ATTR_TRANS_W_IDX = 2; +constexpr uint64_t ATTR_TRANS_X_IDX = 3; +constexpr uint64_t ATTR_GROUPTYPE_IDX = 4; +constexpr uint32_t ATTR_GROUP_LIST_TYPE_IDX = 5; + +constexpr uint32_t MAX_X_DIM = 6UL; +constexpr uint32_t MIN_X_DIM = 2UL; + +constexpr uint32_t BASIC_BLOCK_BASE_M = 256; +constexpr uint32_t BASIC_BLOCK_BASE_M_WITH_BIAS = 240; +constexpr uint32_t BASIC_BLOCK_BASE_N = 256; +constexpr uint32_t BASIC_BLOCK_BASE_K = 64; +constexpr uint32_t BASIC_BLOCK_BASE_N_MIN = 128; +constexpr uint32_t STEP_K_4 = 4; +constexpr uint32_t STEP_K_3 = 3; +constexpr uint32_t DEPTH_8 = 8; +constexpr uint32_t BUFFER_NUM_2 = 2; + +constexpr uint32_t SCALE_FACTOR_DEFAULT = 1; +constexpr uint32_t SCALE_FACTOR_MIN = 1; +constexpr uint32_t SCALE_FACTOR_B_BIT = 8; +constexpr uint32_t SCALE_FACTOR_M_BIT = 16; +constexpr uint32_t SCALE_FACTOR_N_BIT = 24; + +constexpr int32_t B16_DATA_SIZE = 2; +constexpr int32_t B8_DATA_SIZE = 1; + +struct TailBlockResplitParam { + uint32_t mainBlockSize = 0; + uint64_t mainBlockCount = 0; + uint16_t firstTailBlockSize = 0; + uint16_t secondTailBlockSize = 0; + uint16_t firstTailBlockCount = 0; + uint16_t secondTailBlockCount = 0; +}; + +enum class GroupType : int8_t { + NO_SPLIT = -1, + SPLIT_M = 0, + SPLIT_N = 1, + SPLIT_K = 2, +}; + +enum class QuantType : uint8_t { + NONE = 0, + PER_TENSOR = 1, + PER_CHANNEL = 2, + PER_GROUP = 3, + MX = 4 +}; + +enum class WeightFormat : uint8_t { + ND = 0, + FRACTAL_NZ = 1, +}; + +// 对应0位 平台大类 +enum class SocVersionType : uint8_t { + RESERVERD = 0, + SUPPORT_L0C_TO_OUT = 1, + SUPPORT_L1_TO_BT_BF16 = 2, +}; + +// 对应1位 平台小类 +enum class SocVersionSubType : uint8_t { + RESERVERD = 0, +}; + +// 对应2-3位 伪量化场景 +enum class QuantizationScenario : uint8_t { + DEFAULT = 0, +}; + +// 对应4位 算法大类 +enum class OptimizationAlgorithmCategory : uint8_t { + VECTOR_ANTIQUANT = 0, + MULTI_SCALE_DEQUANT = 1, + FIXPIPE_ANTIQUANT = 2, +}; + +// 对应5位 算法小类 +enum class OptimizationAlgorithmSubCategory : uint8_t { + VDEFAULT = 0, + SPLIT_K = 1, + N_FIRST_TAIL_RESPLIT = 2, + N_FIRST_BASIC_BLOCK = 3, +}; + +// 对应6-9位 fixp模板自定义组合 +enum class FixpipeConfiguration : uint16_t { + A_NORMAL_LOAD = 0, + A_SINGLE_M_SINGLE_K_FULL_LOAD = 1, +}; + +enum class CustomSplitKConfiguration : uint8_t { + A_NORMAL_LOAD = 0, + A_MK_FULL_LOAD = 1, +}; + +// 对应14位 表示转置场景，transA/transB +enum class TransposeSituation : uint8_t { + A_NOT_TRANS_B_NOT_TRANS = 0, + A_NOT_TRANS_B_TRANS = 1, + A_TRANS_B_NOT_TRANS = 2, + A_TRANS_B_TRANS = 3, +}; + +// 对应17位 可选输入是否存在 hasAntiquantOffset/hasBias +enum class OptionInputSituation : uint8_t { + ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST = 0, + ANTIQUANT_OFFSET_NOT_EXIST_BIAS_EXIST = 1, + ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST = 2, + ANTIQUANT_OFFSET_EXIST_BIAS_EXIST = 3, + ANTIQUANT_OFFSET_NOT_EXIST_BIAS_FP32_EXIST = 4, + ANTIQUANT_OFFSET_EXIST_BIAS_FP32_EXIST = 6, +}; + +enum class Mte2Configuration : uint8_t { + MTE2_INNER_SIZE_512_BUF_NUM_2 = 0, + MTE2_INNER_SIZE_512_BUF_NUM_4 = 1, + MTE2_INNER_SIZE_1024_BUF_NUM_2 = 2, + MTE2_INNER_SIZE_256_BUF_NUM_4 = 3, + MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT = 4, // w8 w4在非性能场景下复用一组设置 + MTE2_INNER_SIZE_384_BUF_NUM_3 = 5, +}; + +class TilingKeyConfigure { +public: + // 对应0-1位 平台大类，平台小类 + uint8_t socVersionType = 0; + + // 对应2-3位 伪量化场景 + uint8_t quantizationScenario = 0; + + // 对应4-5位 算法大类、算法小类 + uint8_t algorithm = 0; + + // 对应14位 表示转置场景，transA/transB + uint8_t transposeSituation = 0; + + // 对应15位 表示反量化的类型 perchannel/pertensor/perGroup + uint8_t antiquantType = 0; + + // 对应16位 表示量化的类型 perchannel/perTensor/None + uint8_t quantType = 0; + + // 对应17位 可选输入是否存在 hasAntiquantOffset/hasBias + uint8_t optionInputSituation = 0; + + // 对应18位 weight的数据类型 weightNd/weightNz + uint8_t weightFormat = 0; + + // 对应6-9位 模板自定义组合 + uint16_t templateCustom = 0; + + // 对应10-13位 api常量化保留位 + uint16_t apiConstexpr = 0; + +public: + void PrintTilingKeyLog() const + { + std::stringstream ss; + ss << "socVersionType: " << static_cast(this->socVersionType) + << " quantizationScenario: " << static_cast(this->quantizationScenario) + << " algorithm: " << static_cast(this->algorithm) + << " transposeSituation: " << static_cast(this->transposeSituation) + << " antiquantType: " << static_cast(this->antiquantType) + << " quantType: " << static_cast(this->quantType) + << " optionInputSituation: " << static_cast(this->optionInputSituation) + << " weightFormat: " << static_cast(this->weightFormat) + << " templateCustom: " << static_cast(this->templateCustom) + << " apiConstexpr: " << static_cast(this->apiConstexpr); + OP_LOGI("GMMWeightQuantBatchMatmul", "tilingKeyConfigure: %s", ss.str().c_str()); + return; + } + + uint64_t GenTilingKey() const; +}; + +class GroupedWeightQuantBatchMatmulTiling { +public: + bool SetTiling(gert::TilingContext *context); + +protected: + bool SetShapeList(const gert::TilingContext *context); + bool CheckTensorListSize(const gert::TilingContext *context); + bool CheckTensorDtype(const gert::TilingContext *context, uint32_t attrIdx, size_t idx, + const ge::DataType &tensorDtype, const std::string &tensorType) const; + bool IsNzFormat(const gert::TilingContext *context, uint32_t attrIdx, size_t idx) const; + bool CheckXAndWeightFormat(const gert::TilingContext *context, size_t idx) const; + bool CheckNotNullPtr(const gert::TilingContext *context, uint32_t attrIdx, size_t idx) const; + bool CheckNotNull(const gert::TilingContext *context, size_t idx) const; + bool CheckTensorDimEqualTarget(const gert::TilingContext *context, uint32_t attrIdx, size_t idx, uint32_t targetDim, + const std::string &tensorType) const; + bool CheckTensorDimSingleXSingleWeightSingleY(const gert::TilingContext *context, size_t idx) const; + bool CheckTensorDimMultiXMultiWeightMultiY(const gert::TilingContext *context, size_t idx) const; + bool CheckTensorDim(const gert::TilingContext *context, size_t idx) const; + bool CheckTensorShape(const gert::TilingContext *context, uint32_t attrIdx, size_t idx, + const std::string &tensorType) const; + bool CheckDimValue(const gert::TilingContext *context, size_t idx) const; + bool CheckWeightInnerAxisEven(const gert::TilingContext *context, size_t idx) const; + bool CheckXAndWeightShape(const gert::TilingContext *context) const; + bool CheckEveryTensor(const gert::TilingContext *context) const; + bool CheckGroupList(const gert::TilingContext *context) const; + bool CheckRequiredInputs(const gert::TilingContext *context) const; + bool AnalyzeAttr(const gert::TilingContext *context); + bool AnalyzeInput(const gert::TilingContext *context); + bool CalcResplitTiling(const gert::TilingContext *context); + bool SetBaseTiling(); + void SetMatMulTiling(); + void SetTilingKey(gert::TilingContext *context); + bool SetCustomParam(gert::TilingContext *context); + bool IsA16W4ND() const; + bool CheckUnsupportDataFlow(const gert::TilingContext *context) const; + bool CheckAntiQuantDtype(const gert::TilingContext *context) const; + bool CheckBiasDtype(const gert::TilingContext *context) const; + bool CheckGroupTypeAndSplitItem(const gert::TilingContext *context) const; + bool CheckTransposeStatus(const gert::TilingContext *context) const; + bool SetShapeListSplitMSingleXSingleWeightSingleY(const gert::TilingContext *context); + bool SetShapeListMultiXMultiWeightMultiY(const gert::TilingContext *context); + uint16_t GetTensorListSize(const gert::TilingContext *context, uint32_t attrIdx) const; + void GetNumOfInputs(const gert::TilingContext *context); + bool SetAntiquantGroupSize(const gert::TilingContext *context); + bool CheckGroupSize(const gert::TilingContext *context) const; + bool GetC0Size(const gert::TilingContext *context, ge::DataType dtype, uint64_t &c0Size) const; + void CalcFullBlockDimResplitTiling(uint64_t c0Size); + void CalcNoFullBlockDimResplitTiling(uint64_t c0Size); + bool CheckResplitTilingResult(const gert::TilingContext *context) const; + void PrintInputParam(const gert::TilingContext *context) const; + void PrintTilingResult(const gert::TilingContext *context); + bool EnableTailResplit() const; + +private: + int32_t mList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t kList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t nList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + + bool transA_ = false; + bool transB_ = false; + bool isSingleX_ = true; + bool isSingleWeight_ = true; + bool isSingleY_ = true; + bool hasBias_ = false; + bool weightNzFlag_ = false; + bool hasAntiquantOffset_ = false; + + uint64_t mSize_ = 0; + uint64_t kSize_ = 0; + uint64_t nSize_ = 0; + uint64_t nSizeOri_ = 0; + GroupType groupType_ = GroupType::SPLIT_M; + int64_t splitItem_ = 0; + uint32_t groupNum_ = 0; + uint32_t groupListType_ = 0; + uint32_t coreNum_ = 0; + uint32_t groupSize_ = 0; + uint8_t cubeBlockDimN_ = 0; + + uint16_t numX_ = 0; + uint16_t numWeight_ = 0; + uint16_t numBias_ = 0; + uint16_t numAntiquantScale_ = 0; + uint16_t numAntiquantOffset_ = 0; + + ge::DataType xDType_ = ge::DT_UNDEFINED; + ge::DataType weightDtype_ = ge::DT_UNDEFINED; + ge::DataType biasDtype_ = ge::DT_UNDEFINED; + ge::DataType antiquantScaleDtype_ = ge::DT_UNDEFINED; + ge::DataType antiquantOffsetDtype_ = ge::DT_UNDEFINED; + + TailBlockResplitParam resplitParam_; + TilingKeyConfigure tilingKeyConfig_; + DlinferGroupedMatmulDirectTilingData::GMMWeightQuantTilingData tilingData_; +}; +} // namespace optiling + +#endif // GROUPED_WEIGHT_QUANT_BATCH_MATMUL_TILING_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.cpp new file mode 100644 index 00000000..fbe5a66d --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.cpp @@ -0,0 +1,1913 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_tiling.cpp + * \brief + */ +#include "grouped_matmul_tiling.h" + +#include +#include "register/op_impl_registry.h" +#include "arch35/grouped_weight_quant_batch_matmul_tiling.h" +#include "arch35/grouped_no_quant_matmul_tiling.h" +#include "tiling_base/tiling_templates_registry.h" +#include "err/ops_err.h" +#include "../../op_kernel/grouped_matmul_tiling_key.h" +using namespace Ops::Transformer::OpTiling; +using namespace ge; +using namespace AscendC; +using namespace DlinferGroupedMatmulDirect; + +namespace optiling { +static const GMMCompileInfo* GetGMMCompileInfo(const gert::TilingContext* context) { + auto compileInfoPtr = context->GetCompileInfo(); + if (compileInfoPtr != nullptr) { + return compileInfoPtr; + } + + static thread_local GMMCompileInfo fallbackCompileInfo{}; + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + return nullptr; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + fallbackCompileInfo.aicNum = ascendcPlatform.GetCoreNumAic(); + fallbackCompileInfo.aivNum = ascendcPlatform.GetCoreNumAiv(); + fallbackCompileInfo.socVersion = ascendcPlatform.GetSocVersion(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, fallbackCompileInfo.ubSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, fallbackCompileInfo.l1Size); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_A, fallbackCompileInfo.l0ASize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, fallbackCompileInfo.l0BSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, fallbackCompileInfo.l0CSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L2, fallbackCompileInfo.l2Size); + return &fallbackCompileInfo; +} + +static inline uint32_t SixteenAlign(uint32_t a, bool up = false) { + if (up) { + a += 15U; // 15: 16 bytes up-align + } + return a & ~15U; // ~15: 16 bytes down-align +} + +static inline int64_t SixteenAlign(int64_t a, bool up = false) { + if (up) { + a += 15; // 15: 16 bytes up-align + } + return a & ~15; // ~15: 16 bytes down-align +} + +template +static inline auto AlignUp(T num1, T num2) -> T +{ + if (num2 == 0) { + return 0; + } + if (num1 < 0) { + return -(-num1 / num2) * num2; + } + return (num1 + num2 - 1) / num2 * num2; +} + +constexpr uint32_t GROUP_LIST_SPARSE_M = 2U; +// EFFECTIVE_TASK_RATIO表示动态分块后的任务数量占可处理任务数量的比例 +// 通常情况下阈值越高，优化效果越好，但是阈值调高会导致可能找不到更优分块 +// 基于当前模型case 测试，整体有优化 +// DSKV3,k/n=7168/4096和2048/7168 +// A2双机e=16/17， m/e平均32/48/64/96 +// A2大EP/A3四机 e=4/5， m/e平均32/48/64/96 +// A3八机e=2/3，m/e平均192 +// Qwen3 k/n=2048/1536和768/2048 +// 30B单机 e=128，m/e平均 8/16 +constexpr float EFFECTIVE_TASK_RATIO = 0.95f; +// 小K时,由于vector bound,所以开启双vector会有优化, +// 中K时,开启两个vector时会导致MTE2带宽争抢严重,可能会劣化 +// 大K时,开启两个vector会提升部分带宽利用率，此时也会有提升 +// 实测数据表明,当K小于等于1024或者K大于2048时开启双Vector会有优化 +constexpr int64_t DOUBLE_VECTOT_THRESHOLD_K_LOWER = 1024L; +constexpr int64_t DOUBLE_VECTOT_THRESHOLD_K_UPPER = 2048L; +// 实测当单专家token数低于128时cube算力不能完全发挥，导致开启2个vector核可能会劣化 +constexpr int32_t SMALL_TUNING_CONFIG_THRESHOLD = 128; +constexpr int32_t BIAS_REMAIN_SPACE = 2 * 1024; +constexpr int32_t MIN_BASE_M = 16; +// 定轴搬移算法K的范围 +constexpr int64_t FIXAXISMOVE_K1 = 2048L; +constexpr int64_t FIXAXISMOVE_K2 = 7168L; +// 定轴搬移算法N的范围 +constexpr int64_t FIXAXISMOVE_N1 = 7168L; +constexpr int64_t FIXAXISMOVE_N2 = 4096L; +// 定轴搬移算法group_num的范围 +constexpr int32_t FIXAXISMOVE_GROUP_NUM = 4; +// 定轴搬移算法每个专家M的范围 +constexpr int64_t FIXAXISMOVE_PERM_LOWER = 128L; +constexpr int64_t FIXAXISMOVE_PERM_UPPER = 512L; +// 定轴搬移算法split_item的范围 +constexpr int64_t FIXAXISMOVE_SPLIT_ITEM2 = 2L; +constexpr int64_t FIXAXISMOVE_SPLIT_ITEM3 = 3L; +// 定轴搬移算法group_list_type的范围 +constexpr int64_t FIXAXISMOVE_GROUP_LIST_TYPE = 0L; +// 定轴搬移算法group_type的范围 +constexpr int32_t FIXAXISMOVE_GROUP_TYPE = 0; +constexpr size_t TUNING_CONFIG_TOKEN_PER_EXPECT_INDEX = 0; +constexpr size_t TUNING_CONFIG_A8W4_SPEC_SCENARIO_INDEX = 1; +constexpr size_t TUNING_CONFIG_ALLOW_WORKSPACE_INDEX = 2; + +ge::graphStatus GMMTiling::CheckWeightNZShape(const gert::TilingContext* context, int64_t numInOneBlk) const { + OP_CHECK_IF(numInOneBlk <= 0, OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "numInOneBlk, the " + "input of CheckWeightNZShape has an invaild value %ld", numInOneBlk), return ge::GRAPH_FAILED); + size_t i = 0; + while (true) { + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, i++); + if (wTensor == nullptr) { break; } + gert::Shape wOriginShape = wTensor->GetOriginShape(); + int64_t lastDimValue = wOriginShape.GetDim(wOriginShape.GetDimNum() - 1); // inner axis + OP_CHECK_IF(lastDimValue % numInOneBlk != 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "the inner axis size of nz weight is expected to be a multiple of 32B, " + "but now the inner axis size is %ld.", lastDimValue), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::CheckMKN(const gert::TilingContext* context) { + mmDataTypeSize_ = GetSizeByDataType(mmDType_); + OP_CHECK_IF(mmDataTypeSize_ == 0, OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "GMM get mm dtype[%s] size is 0.", TypeUtils::DataTypeToAscendString(mmDType_).GetString()), + return ge::GRAPH_FAILED); + uint32_t numInOneBlk = 0; + if (isA4W4_) { + numInOneBlk = static_cast(ONE_BLK_SIZE / INT4_DATA_TYPE_SIZE); + } else { + numInOneBlk = ONE_BLK_SIZE / mmDataTypeSize_; + } + OP_CHECK_IF(numInOneBlk == 0, OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "GMM numInOneBlk cannot be 0."), return ge::GRAPH_FAILED); + int64_t maxMKN = INT_MAX / numInOneBlk * numInOneBlk; + OP_CHECK_IF(maxM_ > maxMKN || maxN_ > maxMKN || maxK_ > maxMKN, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "32B-aligned m, n or k axis is out of range int32!"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +void GMMTiling::SetTilingDataIsSingleTensor() { + tilingData.gmmBaseParams.set_singleWeight(static_cast(isSingleWeight_)); + tilingData.gmmBaseParams.set_singleX(static_cast(isSingleX_)); + tilingData.gmmBaseParams.set_singleY(static_cast(isSingleY_)); +} + +ge::graphStatus GMMTiling::PrepareTilingData(const gert::TilingContext* context) { + // get transpose and groupType + OP_CHECK_IF(GMMGetAttrs(context) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMMGetAttrs failed"), + return ge::GRAPH_FAILED); + // get the first tensor's shape of weight and x + auto xTensor = context->GetDynamicInputTensor(X_INDEX, 0); // 0: get first tensor + OP_CHECK_NULL_WITH_CONTEXT(context, xTensor); + gert::Shape xShape = xTensor->GetStorageShape(); + xDimNum_ = static_cast(xShape.GetDimNum()); + xKDim_ = transposeX_ ? 0U : xDimNum_ - 1U; // 0: when x is transposed, the first dim is k; -1：otherwise, the last dim is k + + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, wTensor); + gert::Shape wShape = wTensor->GetOriginShape(); + uint32_t wDimNum = static_cast(wShape.GetDimNum()); + weightNDim_ = transposeWeight_ ? wDimNum - 2U : wDimNum - 1U; // -2: when w is transposed, the last 2 dim is n; -1: otherwise, the last dim is n + weightKDim_ = transposeWeight_ ? wDimNum - 1U : wDimNum - 2U; // -2: when w is transposed, the last 1 dim is k; -1: otherwise, the last 2 dim is k + nzFactor_ = 1; // init + if (wFormat_ == matmul_tiling::CubeFormat::NZ) { + uint32_t numInOneBlk = isA4W4_ ? static_cast(UB_BLOCK_UNIT_SIZE / INT4_DATA_TYPE_SIZE) : + UB_BLOCK_UNIT_SIZE / std::max(1, GetSizeByDataType(weightDtype_)); + if (isA8W4FakeA8W8_) { + numInOneBlk = UB_BLOCK_UNIT_SIZE; + } + if (wDimNum >= 4U) { // 4: least dim num of nz format tensor + weightNDim_ = transposeWeight_ ? wDimNum - 3U : wDimNum - 4U; // -3: when w is transposed, the last 3 dim is n/nzFactor; -4: when w has nz format, the last 4 dim is n/nzFactor + // nzFactor_ is a factor used to compute n axis size. If weight is transposed, nzFactor_ is 16; otherwise nzFactor_ is 16 for bf16, 32 for int8 + nzFactor_ = transposeWeight_ ? 16 : static_cast(numInOneBlk); + } else { + OP_CHECK_IF(CheckWeightNZShape(context, static_cast(numInOneBlk)) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "the shape of nz weight is invaild."), + return ge::GRAPH_FAILED); + } + } + isSingleWeight_ = (context->GetDynamicInputTensor(WEIGHT_INDEX, 1) == nullptr); + isSingleX_ = (context->GetDynamicInputTensor(X_INDEX, 1) == nullptr); + isSingleY_ = (splitItem_ == 2 || splitItem_ == 3); // 2: when x is multi-tensor, y is single-tensor; 3: when x is single-tensor, y is single-tensor + SetTilingDataIsSingleTensor(); + + ge::graphStatus shapeStatus = ge::GRAPH_FAILED; + if (groupType_ == SPLIT_M) { + shapeStatus = GMMGetTensorShapeSplitM(context, xShape, wShape); + } else if (groupType_ == SPLIT_K) { + shapeStatus = GMMGetTensorShapeSplitK(context, xShape, wShape); + } else if (groupType_ == NO_SPLIT) { // not split any axis + if (isSingleWeight_ && wDimNum > 2U) { // 2: dim of splited weight tensor + shapeStatus = SeparatedXSingleWeight(context, wShape); + } else { + shapeStatus = SeparatedXSeparatedWeight(context); + } + } else { + OP_LOGE(context->GetNodeName(), "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", + groupType_, isSingleWeight_, isSingleX_, isSingleY_); + return ge::GRAPH_FAILED; + } + if (shapeStatus == ge::GRAPH_SUCCESS && groupListType_ == GROUP_LIST_SPARSE_M) { + auto groupListTensor = context->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListTensor); + groupNum_ = static_cast(groupListTensor->GetStorageShape().GetDim(0)); + } + return shapeStatus; +} + +ge::graphStatus GMMTiling::GMMGetTensorShapeSplitM(const gert::TilingContext* context, const gert::Shape &xShape, + const gert::Shape &wShape) { + if (isSingleX_ && isSingleWeight_ && isSingleY_) { // split M, s-s-s + return SplitMSingleXSingleWeightSingleY(xShape, wShape); + } + if (isSingleX_ && !isSingleWeight_ && isSingleY_) { // split M, s-m-s + return SplitMSingleXSeparatedWeight(context, xShape); + } + if (isSingleX_ && !isSingleWeight_ && !isSingleY_) { // splitM, s-m-m + return SplitMSingleXSeparatedWeight(context, xShape); + } + if (!isSingleX_ && !isSingleWeight_ && isSingleY_) { // split M, m-m-s + return SeparatedXSeparatedWeight(context); + } + if (!isSingleX_ && isSingleWeight_) { // split M, m-s-m/m-s-s + return SeparatedXSingleWeight(context, wShape); + } + if (!isSingleX_ && !isSingleWeight_ && !isSingleY_) { // split M, m-m-m + return SeparatedXSeparatedWeight(context); + } + OP_LOGE(context->GetNodeName(), "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", + groupType_, isSingleWeight_, isSingleX_, isSingleY_); + return ge::GRAPH_FAILED; +} + +ge::graphStatus GMMTiling::GMMGetTensorShapeSplitK(const gert::TilingContext* context, const gert::Shape &xShape, + const gert::Shape &wShape) { + if (isSingleX_ && isSingleWeight_ && isSingleY_) { // splitK, s-s-s + return SplitKSingleXSingleWeightSingleY(context, xShape, wShape); + } + if (isSingleX_ && !isSingleWeight_ && !isSingleY_) { // splitK, s-m-s + return SplitKSingleXSeparatedWeight(context, xShape, wShape); + } + if (!isSingleX_ && isSingleWeight_) { // splitK, m-s-m/m-s-s + return SeparatedXSingleWeight(context, wShape); + } + OP_LOGE(context->GetNodeName(), "GMM_tiling: not support groupType_=%d, isSingleWeight_=%d, isSingleX_=%d, isSingleY_=%d", + groupType_, isSingleWeight_, isSingleX_, isSingleY_); + return ge::GRAPH_FAILED; +} + +/** @brief split M：single-single-single(s-s-s) +*/ +ge::graphStatus GMMTiling::SplitMSingleXSingleWeightSingleY(const gert::Shape &xShape, const gert::Shape &wShape) { + groupNum_ = static_cast(wShape.GetDim(0)); + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * static_cast(nzFactor_); + kList_[0] = static_cast(k); // if split M axis, the K axis values of x tensorList are all the same. + nList_[0] = static_cast(n); + mList_[0] = -1; + maxM_ = m; + maxK_ = k; + maxN_ = n; + totalM_ = static_cast(m); + FixedAxisMoveWorkspace_ = maxM_ * maxN_ * sizeof(int32_t); + return ge::GRAPH_SUCCESS; +} + +/** @brief split M：single-multi-single(s-m-s)/single-multi-multi(s-m-m), share the same function. +*/ +ge::graphStatus GMMTiling::SplitMSingleXSeparatedWeight(const gert::TilingContext* context, const gert::Shape &xShape) { + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + for (uint32_t i = 0; i < MAX_TENSOR_CONT; i++) { + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, i); + if (wTensor == nullptr) { break; } // when x has multi tensors, xTensor is allowed to be empty + auto wShape = wTensor->GetOriginShape(); + + groupNum_ += 1U; + kList_[i] = static_cast(k); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + nList_[i] = static_cast(n); + maxN_ = std::max(maxN_, n); + } + mList_[0] = -1; // mList is unknown right now + maxM_ = m; + maxK_ = k; + totalM_ = static_cast(m); + + return ge::GRAPH_SUCCESS; +} + +/** @brief split M：multi-multi-single(m-m-s); no split: multi-multi-multi(m-m-m), share the same function +*/ +ge::graphStatus GMMTiling::SeparatedXSeparatedWeight(const gert::TilingContext* context) { + for (uint32_t i = 0; i < MAX_TENSOR_CONT; i++) { + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, i); + auto xTensor = context->GetDynamicInputTensor(X_INDEX, i); + if (wTensor == nullptr || xTensor == nullptr) { break; } + auto wShape = wTensor->GetOriginShape(); + auto xShape = xTensor->GetStorageShape(); + groupNum_ += 1U; + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + mList_[i] = static_cast(m); + kList_[i] = static_cast(k); + nList_[i] = static_cast(n); + maxM_ = std::max(maxM_, m); + maxK_ = std::max(maxK_, k); + maxN_ = std::max(maxN_, n); + totalM_ += static_cast(m); + } + groupType_ = NO_SPLIT; + return ge::GRAPH_SUCCESS; +} + +/** @brief split M : multi-single-multi(m-s-m), split K : multi-single-multi(m-s-m), share the same function +*/ +ge::graphStatus GMMTiling::SeparatedXSingleWeight(const gert::TilingContext* context, const gert::Shape &wShape) { + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + for (uint32_t i = 0; i < MAX_TENSOR_CONT; i++) { + auto xTensor = context->GetDynamicInputTensor(X_INDEX, i); + if (xTensor == nullptr) { break; } // when x has multi tensors, xTensor is allowed to be empty + auto xShape = xTensor->GetStorageShape(); + groupNum_ += 1U; + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + mList_[i] = static_cast(m); + kList_[i] = static_cast(k); + nList_[i] = static_cast(n); + maxM_ = std::max(maxM_, m); + maxK_ = std::max(maxK_, k); + totalM_ += static_cast(m); + } + maxN_ = n; + groupType_ = NO_SPLIT; + return ge::GRAPH_SUCCESS; +} + +/** @brief split K single-single-single +*/ +ge::graphStatus GMMTiling::SplitKSingleXSingleWeightSingleY(const gert::TilingContext* context, + const gert::Shape &xShape, const gert::Shape &wShape) { + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + + auto groupListTensor = context->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + if (groupListTensor == nullptr) { + OP_LOGE(context->GetNodeName(), "groupListTensor is nullptr"); + return ge::GRAPH_FAILED; + } + gert::Shape groupListShape = groupListTensor->GetStorageShape(); + groupNum_ = static_cast(groupListShape.GetDim(0)); // 0: the first dim of groupList is groupNum + mList_[0] = static_cast(m); + nList_[0] = static_cast(n); + kList_[0] = -1; + maxM_ = m; + maxN_ = n; + maxK_ = k; + totalM_ = static_cast(m); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::SplitKSingleXSeparatedWeight(const gert::TilingContext* context, + const gert::Shape &xShape, const gert::Shape &wShape) { + int64_t m = GMMGetBS(xShape); + int64_t k = xShape.GetDim(xKDim_); + int64_t n = wShape.GetDim(weightNDim_) * nzFactor_; + for (uint32_t i = 0; i < MAX_TENSOR_CONT; i++) { + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, i); + if (wTensor == nullptr) { break; } + auto wTensorShape = wTensor->GetOriginShape(); + groupNum_ += 1; + mList_[i] = static_cast(m); + k = wTensorShape.GetDim(weightKDim_); + kList_[i] = static_cast(k); + maxK_ = std::max(maxK_, k); + n = wTensorShape.GetDim(weightNDim_); + nList_[i] = static_cast(n); + maxN_ = std::max(maxN_, n); + } + maxM_ = m; + totalM_ = static_cast(m); + groupType_ = NO_SPLIT; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::Init(const gert::TilingContext* context) { + OP_CHECK_IF(PrepareTilingData(context) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM PrepareTilingData failed."), + return ge::GRAPH_FAILED); + auto compileInfoPtr = GetGMMCompileInfo(context); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null + + // check tuningConfig_ + if (tuningConfig_ < 0 || tuningConfig_ > maxM_) { + OP_LOGE(context->GetNodeName(), "Invalid tuningConfig_: %ld. Valid range: [0, maxM]", tuningConfig_); + return ge::GRAPH_FAILED; + } + // check whether x, weight and y are all single tensor + isAllSingleTensor_ = isSingleX_ && isSingleWeight_ && isSingleY_; + bool isA16W8 = (xDType_ == ge::DT_FLOAT16 || xDType_ == ge::DT_BF16) && weightDtype_ == ge::DT_INT8; + // check whether k and n are supported in msd + bool isKNForA16W8MSD = maxN_ % static_cast(A16W8_MSD_KN_BASE_BLOCK) == 0L && + maxK_ % static_cast(A16W8_MSD_KN_BASE_BLOCK) == 0L && + maxK_ <= static_cast(A16W8_MSD_MAX_K) && + maxN_ >= static_cast(A16W8_MSD_MIN_N); + // check whether total token num and average token num are supported in msd + bool isMForA16W8MSD = totalM_ <= A16W8_MSD_AVERAGE_TOKEN_NUM * groupNum_; + isA16W8Msd_ = isAllSingleTensor_ && groupType_ == SPLIT_M && isA16W8 && isKNForA16W8MSD && isMForA16W8MSD; + mmDType_ = isA16W8Msd_ ? ge::DT_INT8 : xDType_; + OP_CHECK_IF(CheckMKN(context) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM CheckMKN failed."), + return ge::GRAPH_FAILED); + auto biasPtr = context->GetDynamicInputTensor(BIAS_INDEX, 0); // 0: obtain the first tensor of the tensorList + hasBias_ = !(biasPtr == nullptr || biasPtr->GetStorageShape().GetShapeSize() == 0); + if (isA4W4_) { + uint64_t quantGroupNum = 0; + uint32_t scaleDimNum = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDimNum(); + // 3: pergroup scale shape is [e,g,n] + if (scaleDimNum == 3U) { + quantGroupNum = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDim(1); + // 2: perchannel scale shape is [e, n] + } else if (scaleDimNum == 2U) { + quantGroupNum = 1UL; + } else { + OP_LOGE(context->GetNodeName(), "GMM A4W4: scale dim should be 2 or 3, but now is %u", scaleDimNum); + return ge::GRAPH_FAILED; + } + tilingData.gmmBaseParams.set_k(maxK_); + tilingData.gmmBaseParams.set_n(maxN_); + tilingData.gmmBaseParams.set_quantGroupNum(quantGroupNum); + } + if (isA8W4FakeA8W8_) { + hasBias_ = false; + } + tilingData.gmmArray.set_mList(mList_); + tilingData.gmmArray.set_kList(kList_); + tilingData.gmmArray.set_nList(nList_); + tilingData.gmmBaseParams.set_groupNum(groupNum_); + tilingData.gmmBaseParams.set_m(totalM_); + tilingData.gmmBaseParams.set_hasBias(static_cast(hasBias_)); + tilingData.gmmBaseParams.set_groupType(static_cast(groupType_)); + tilingData.gmmBaseParams.set_activeType(actType_); + tilingData.gmmBaseParams.set_quantParam(perTokenOrPerGroupSize_); + tilingData.gmmBaseParams.set_groupListType(groupListType_); + tilingData.gmmBaseParams.set_k(maxK_); + tilingData.gmmBaseParams.set_n(maxN_); + OP_LOGI(context->GetNodeName(), "GMM_tiling: groupNum_ is %u, maxM_ is %ld, maxK_ is %ld, maxN_ is %ld.", + groupNum_, maxM_, maxK_, maxN_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::GetPerGroupNum(const gert::TilingContext* context) { + auto antiquantScale = context->GetDynamicInputTensor(ANTIQUANT_SCALE_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, antiquantScale); + auto antiquantScaleShape = antiquantScale->GetStorageShape(); + int64_t dimNum = antiquantScaleShape.GetDimNum(); + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, wTensor); + gert::Shape wShape = wTensor->GetOriginShape(); + size_t wDimNum = wShape.GetDimNum(); + if ((isSingleWeight_ && wDimNum > 2UL && dimNum == 3L) || (!isSingleWeight_ && dimNum == 2L)) { // 2 and 3: dim threshold + int64_t g = antiquantScaleShape.GetDim(dimNum - 2L); + perTokenOrPerGroupSize_ = g > 1L ? static_cast(kList_[0] / g) : 0U; + tilingData.gmmBaseParams.set_quantParam(perTokenOrPerGroupSize_); + } + return ge::GRAPH_SUCCESS; +} + +void GMMTiling::DivideUbAndSetWorkspaceAntiquant(size_t* workspaces, const uint32_t& aicNum, uint32_t &ubSize) { + if (isA16W8Msd_) { + // whole workspace for a16w8 msd scene is combined by workspace for global max of each row (m * 8), + // workspace for local reduce sum of each row (m * aicNum), workspace for data after prerpocessing x (2 * m * k) + // and workspace for matmul output (2 * m * n) + // 32: need 32 byte to store global max + workspaces[0] += totalM_ * (aicNum * sizeof(float) + 32UL + + A16W8_MSD_STEP * (static_cast(maxK_) * sizeof(int8_t) + + static_cast(maxN_) * sizeof(int32_t))); + // 7: make aicnum align up to 8 + uint32_t alignedAicNum = (aicNum + 7U) & (~7U); + ubSize = static_cast(ubSize_ - (static_cast(baseM_) / A16W8_MSD_STEP) * + alignedAicNum * sizeof(float)); + // workspacesSize in GMMBaseParams is size of matmul left input matrix of matmul. + workspacesSize_ += A16W8_MSD_STEP * static_cast(maxK_) * static_cast(maxM_); + } else { + for (uint32_t i = 0; i < groupNum_; i++) { + bool isAllSingleTensor = isSingleX_ && isSingleWeight_ && isSingleY_; + int32_t kInList = isAllSingleTensor ? kList_[0] : kList_[i]; // in s-s-s case，k only exits in the first of the list + int32_t nInList = isAllSingleTensor ? nList_[0] : nList_[i]; // in s-s-s case，n only exits in the first of the list + int32_t k = kList_[0] == -1 ? static_cast(maxK_) : kInList; + int32_t n = nList_[0] == -1 ? static_cast(maxN_) : nInList; + minK_ = std::min(minK_, k); + workspacesSize_ += static_cast(k) * static_cast(n); + } + // when minK * baseN * coreNum * sizeof(float16) > 12M, it goes into antiquantPerformance branch (12M is obtained by test). + int32_t dimMN = + CeilDiv(CeilDiv(maxM_, groupNum_), baseM_) * CeilDiv(maxN_, baseN_); + bool goodCubeUtility = dimMN * (xDType_ == ge::DT_BF16 ? 2 : 1) >= static_cast(aicNum * 0.4); // 0.4: a factor, in practice. + antiquantPerformance_ = + goodCubeUtility && static_cast(minK_) * baseN_ * static_cast(aicNum) >= ANTIQUANT_PERFORMANCE_THRESHOLD && !transposeWeight_; + uint32_t maxUbBaseN = static_cast(BEST_UB_BASEN); + if (transposeWeight_) { + maxUbBaseN = baseN_; + } else if (antiquantPerformance_) { + // 2: use 2 pieces of workspace in antiquantPerformance branch + workspacesSize_ = static_cast(maxN_) * static_cast(maxK_) * 2UL; + } + // 2: 2 InQueue(antiquant_scale,antiquant_offset) + ubSize = static_cast(ubSize_ - 2U * maxUbBaseN * mmDataTypeSize_ * QUEUE_DOUBLE_BUFFER); + workspaces[0] += workspacesSize_ * mmDataTypeSize_; + } +} + +int32_t GMMTiling::FindBestSingleN(const uint32_t& aicNum) { + if(maxN_ < baseN_ || tuningConfig_ <= 0|| xDType_ != ge::DT_INT8 || weightDtype_ != ge::DT_INT8) { + return baseN_; + } + int32_t mDim = CeilDiv(tuningConfig_ , baseM_); + int32_t nDim = CeilDiv(maxN_, baseN_); + int32_t taskNum = mDim * nDim * static_cast(groupNum_); + int32_t taskNumPerCore = CeilDiv(taskNum, aicNum); + // 每个核只需要做1个基本块的时候，任务量太少，无需处理 + if(taskNumPerCore <= 1) { + return baseN_; + } + int32_t curNDim = 0; + int32_t curTaskNum = 0; + int32_t bestSingleN = baseN_; + float ratio = 0; + for (uint32_t i = 1; i <= aicNum; ++i) { + if(wFormat_ == matmul_tiling::CubeFormat::NZ) { + bestSingleN = CeilDiv(static_cast(maxN_), i); + if(bestSingleN != maxN_ && bestSingleN % baseN_ != 0) { + continue; + } + } else { + //暂时只NZ格式开启动态分块 + return baseN_; + } + curNDim = CeilDiv(maxN_, bestSingleN); + curTaskNum = mDim * curNDim * static_cast(groupNum_); + ratio = static_cast(curTaskNum) / AlignUp(static_cast(curTaskNum), aicNum); + if (ratio >= EFFECTIVE_TASK_RATIO) { + return bestSingleN; + } + } + return baseN_; +} + +bool GMMTiling::TryFullLoadA(int32_t baseM,const GMMCompileInfo *compileInfoPtr) { + auto l1Size = compileInfoPtr->l1Size; + //暂时只支持A8W8 + float sizeofweightDtype = 1.0f; + float sizeofxDtype = 1.0f; + auto matBl1Size = static_cast(tilingData.mmTilingData.get_depthB1() * baseN_ * baseK_ * sizeofweightDtype); + auto remainL1Size = l1Size - matBl1Size; + if(hasBias_) { + remainL1Size -= BIAS_REMAIN_SPACE; + } + int32_t newDepthA1 = CeilDiv(maxK_, baseK_); + if(static_cast(newDepthA1 * baseM * baseK_ * sizeofxDtype) < static_cast(remainL1Size)) { + tilingData.mmTilingData.set_stepKa(newDepthA1); + tilingData.mmTilingData.set_depthA1(newDepthA1); + return true; + } + return false; +} + +ge::graphStatus GMMTiling::DynamicTilingSingleN(gert::TilingContext* context, const uint32_t& aicNum, const GMMCompileInfo *compileInfoPtr) { + OP_CHECK_IF(compileInfoPtr == nullptr, OPS_REPORT_CUBE_INNER_ERR( + context->GetNodeName(), "compileInfoPtr is nullptr."), return ge::GRAPH_FAILED); + if (maxN_ < baseN_ || tuningConfig_ <= 0 || wFormat_ == matmul_tiling::CubeFormat::ND) { + return ge::GRAPH_SUCCESS; + } + int32_t bestSingleN = FindBestSingleN(aicNum); + if(bestSingleN == baseN_) {//没找到更优的singleN + return ge::GRAPH_SUCCESS; + } + tilingData.gmmBaseParams.set_singleN(bestSingleN); + //先不改看看baseM能否全载左矩阵 + if(TryFullLoadA(baseM_, compileInfoPtr)) { + return ge::GRAPH_SUCCESS; + } + //可以尝试减小baseM来全载左矩阵 + int32_t newBaseM = static_cast(SixteenAlign(tuningConfig_, true)); + //防止不均匀情况 + newBaseM += MIN_BASE_M; + //再看看能否全载左矩阵 + if(newBaseM < baseM_ && TryFullLoadA(newBaseM, compileInfoPtr)) { + tilingData.mmTilingData.set_baseM(newBaseM); + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_SUCCESS; +} +ge::graphStatus GMMTiling::DivideUbAndSetWorkspace(gert::TilingContext* context, const uint32_t& aicNum) { + size_t* workspaces = context->GetWorkspaceSizes(1); // get second variable + OP_CHECK_NULL_WITH_CONTEXT(context, workspaces); // check workspaces is not null + workspaces[0] = SYS_WORKSPACE_SIZE; // default size + if (weightDtype_ != ge::DT_INT8 && weightDtype_ != ge::DT_INT4) { + return ge::GRAPH_SUCCESS; + } + uint32_t ubSize = static_cast(ubSize_); + if ((xDType_ == ge::DT_BF16 || xDType_ == ge::DT_FLOAT16)) { + DivideUbAndSetWorkspaceAntiquant(workspaces, aicNum, ubSize); + OP_CHECK_IF(GetPerGroupNum(context) != ge::GRAPH_SUCCESS, OPS_REPORT_VECTOR_INNER_ERR( + context->GetNodeName(), "GetPerGroupNum failed."), return ge::GRAPH_FAILED); + } else if (xDType_ == ge::DT_INT8) { + isFixedAxisMove_ = IsFixedAxisMoveCondition(); + if (isFixedAxisMove_) { + workspaces[0] += FixedAxisMoveWorkspace_; + } else { + // if tuningConfig_ in [1,256], recompute coreNum + constexpr int32_t tuningConfigLowerLimit = 1; + constexpr int32_t tuningConfigUpperLimit = 256; + if (tuningConfig_ >= tuningConfigLowerLimit && tuningConfig_ <= tuningConfigUpperLimit) { + FindBestUsedCoreNumOneGroup(aicNum); + } + if (yDtype_ == ge::DT_INT32) { + return ge::GRAPH_SUCCESS; + } + uint32_t scaleDataTypeSize = GetSizeByDataType(scaleDtype_); + ubSize = perTokenOrPerGroupSize_ == 1U ? // is perToken + static_cast(ubSize_ - + (static_cast(baseN_) * scaleDataTypeSize + + static_cast(baseM_) * sizeof(float)) * QUEUE_DOUBLE_BUFFER) : + static_cast(ubSize_ - baseN_ * scaleDataTypeSize * QUEUE_DOUBLE_BUFFER); + OP_CHECK_IF(SetWorkspscesPerTokenQuant(aicNum, workspaces) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "SetWorkspscesPerTokenQuant failed."), + return ge::GRAPH_FAILED); + if (isA8W4FakeA8W8_) { + workspaces[0] += A8W4noMsdSpace_; + } + } + } else if (xDType_ == ge::DT_INT4) { + ubSize = perTokenOrPerGroupSize_ == 1U ? // is perToken + static_cast(ubSize_ - (static_cast(baseM_) * sizeof(float)) * QUEUE_DOUBLE_BUFFER) : ubSize_; + OP_CHECK_IF(SetWorkspscesPerTokenQuant(aicNum, workspaces) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "SetWorkspscesPerTokenQuant failed."), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(GMMSetUbDivideBlk() != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "GMMSetUbDivideBlk failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(GMMCalUbSize(context, ubSize) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "GMMCalUbSize failed."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +int32_t GMMTiling::FindBestSingleNPertoken(const uint32_t aicNum) const { + if (CeilDiv(maxN_, baseN_) * groupNum_ <= aicNum) { // if all matmuls only occupy a part of cores + return baseN_; + } + if (maxN_ >= 2048) { // 2048: a threshold + return 1024; // 1024: max singleN + } + int32_t bestSingleN = baseN_; // init bestSingleN + uint32_t bestLastCycleCoreNum = (groupNum_ * CeilDiv(maxN_, bestSingleN)) % aicNum; // init lastCycleCoreNum + // 1024: max singleN + for (int32_t tempSingleN = 1024 / baseN_ * baseN_; tempSingleN > baseN_; tempSingleN -= baseN_) { + uint32_t lastCycleCoreNum = (groupNum_ * CeilDiv(maxN_, tempSingleN)) % aicNum; + if (lastCycleCoreNum == 0U) { + bestSingleN = tempSingleN; + break; + } + if (lastCycleCoreNum > bestLastCycleCoreNum || + (lastCycleCoreNum == bestLastCycleCoreNum && maxN_ % tempSingleN == 0)) { + bestSingleN = tempSingleN; + bestLastCycleCoreNum = lastCycleCoreNum; + } + } + return bestSingleN; +} + +void GMMTiling::FindBestUsedCoreNumOneGroup(const uint32_t aicNum) { + if (groupNum_ > 1U) { + return; + } + uint32_t totalCoreNums = CeilDiv(maxN_, baseN_); + // 3: if cube iterNum less or equal to 3, and more than half cores are unused in last iter, use less cores each iter + if ((aicNum * 3U > totalCoreNums && totalCoreNums % aicNum <= aicNum / 2U) || totalCoreNums < aicNum) { // 2: half of aicNum + uint32_t cubeIterNum = CeilDiv(totalCoreNums, aicNum); + usedCoreNum_ = CeilDiv(totalCoreNums, cubeIterNum); + } +} + + +ge::graphStatus GMMTiling::SetWorkspscesPerTokenQuant(const uint32_t aicNum, size_t* workspaces) { + if (aicNum == 0U) { // invaild value + return ge::GRAPH_FAILED; + } + bool opt = (maxM_ <= 32 * groupNum_ && wFormat_ == matmul_tiling::CubeFormat::NZ) && + (!transposeWeight_ || maxN_ >= 2048); // 32: a factor, 2048: a threshold. + if (opt) { // non-basic strategy. matmul output in non-continugous mode with singleN >= baseN + int32_t bestSingleN = FindBestSingleNPertoken(aicNum); + tilingData.gmmBaseParams.set_singleN(bestSingleN); + } + if (isA4W4_) { + // 4： when do cv parallelism, four pieces of workspace are used for storing four cycles of matmul output + workspaces[0] += 4UL * baseM_ * baseN_ * usedCoreNum_ * sizeof(short); // a4w4 mmout dtype is half + } else { + // 4： when do cv parallelism, four pieces of workspace are used for storing four cycles of matmul output + workspaces[0] += 4UL * baseM_ * baseN_ * usedCoreNum_ * sizeof(int32_t); + } + + return ge::GRAPH_SUCCESS; +} + +bool GMMTiling::StaticTilingProcess(gert::TilingContext *context) { + // cond.1 A8W8 + // cond.2 singleX-singleW-singleY scenario + // cond.3 without bias + // cond.4 without activation + // cond.5 no pretiling + // cond.6 only support typeM + // cond.7 tilingdata corresponds to expected value + if ((xDType_ != ge::DT_INT8 || weightDtype_ != ge::DT_INT8) || + (isSingleX_ == 0 || isSingleWeight_ == 0 || isSingleY_ == 0) || + hasBias_ || + actType_ != 0U || + tilingData.gmmBaseParams.get_isPreTiling() != 0 || + tilingData.gmmBaseParams.get_groupType() != 0 || + !CheckTilingMatchStaticValue()) { + return false; + } + // cond.8 only support milan platform + auto compileInfoPtr = GetGMMCompileInfo(context); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context->GetNodeName(), "CompileInfoPtr is nullptr."), return false); + if (!(compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910B || + compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910_93)) { + return false; + } + + tilingData.gmmBaseParams.set_m(tilingData.mmTilingData.get_M()); + tilingData.gmmBaseParams.set_n(tilingData.mmTilingData.get_N()); + tilingData.gmmBaseParams.set_k(tilingData.mmTilingData.get_Ka()); + return true; +} + +bool GMMTiling::CheckTilingMatchStaticValue() { + if (tilingData.mmTilingData.get_depthA1() != STATIC_TILING_DEPTH_A1_B1 || + tilingData.mmTilingData.get_depthB1() != STATIC_TILING_DEPTH_A1_B1 || + tilingData.mmTilingData.get_stepM() != 1 || + tilingData.mmTilingData.get_stepN() != 1 || + tilingData.mmTilingData.get_stepKa() != STATIC_TILING_STEP_KA_KB || + tilingData.mmTilingData.get_stepKb() != STATIC_TILING_STEP_KA_KB || + tilingData.mmTilingData.get_dbL0A() != DOUBLE_BUFFER_L0A_L0B || + tilingData.mmTilingData.get_dbL0B() != DOUBLE_BUFFER_L0A_L0B || + tilingData.mmTilingData.get_dbL0C() != 1 || + maxK_ > STATIC_TILING_MAX_K) { + return false; + } + if (tilingData.mmTilingData.get_baseM() == BASIC_BLOCK_SIZE_128 && + tilingData.mmTilingData.get_baseN() == BASIC_BLOCK_SIZE_256 && + tilingData.mmTilingData.get_baseK() == BASIC_BLOCK_SIZE_128) { + return true; + } + return false; +} + +ge::graphStatus GMMTiling::RunFusionKernelTiling(gert::TilingContext* context) { + OP_LOGI(context->GetNodeName(), "Begin Run GMM Tiling"); + auto compileInfoPtr = GetGMMCompileInfo(context); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null + + ubSize_ = compileInfoPtr->ubSize; // get ubSize from compileInfo + const uint32_t& aicNum = compileInfoPtr->aicNum; // get aicNum from compileInfo + if (aicNum == 0U) { // invaild value + return ge::GRAPH_FAILED; + } + usedCoreNum_ = aicNum; + + OP_CHECK_IF(CalMMTiling(context, compileInfoPtr) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM CalMMTiling failed"), return ge::GRAPH_FAILED); + + OP_CHECK_IF(GMMSetMMTiling(context, compileInfoPtr) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM GMMSetMMTiling failed"), + return ge::GRAPH_FAILED); + tilingData.gmmBaseParams.set_singleN(0); // 0 is the default value + FullLoadK(compileInfoPtr); + OP_CHECK_IF(DivideUbAndSetWorkspace(context, aicNum) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM DivideUbAndSetWorkspace failed"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(DynamicTilingSingleN(context, usedCoreNum_, compileInfoPtr) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM DynamicTilingSingleN failed"), + return ge::GRAPH_FAILED); + tilingData.gmmBaseParams.set_workspaceSize(workspacesSize_); + tilingData.mmTilingData.set_usedCoreNum(usedCoreNum_); // usedCoreNum is ai_core num + tilingData.gmmBaseParams.set_coreNum(usedCoreNum_); // ai cube number + GMMSetTplTilingKey(context); + tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->SetBlockDim(usedCoreNum_); // block dim is the number of aicube + context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + PrintTilingInfo(context); + return ge::GRAPH_SUCCESS; +} + +void GMMTiling::PrintTilingInfo(gert::TilingContext *context) { + OP_LOGD(context->GetNodeName(), "End Run GMM Tiling"); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: usedCoreNum is %d.", tilingData.mmTilingData.get_usedCoreNum()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: bestSingleN is %u.", tilingData.gmmBaseParams.get_singleN()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: tuning_config is %ld.", tuningConfig_); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: Ka is %d.", tilingData.mmTilingData.get_Ka()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: Kb is %d.", tilingData.mmTilingData.get_Kb()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: baseM is %d.", tilingData.mmTilingData.get_baseM()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: baseN is %d.", tilingData.mmTilingData.get_baseN()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: baseK is %d.", tilingData.mmTilingData.get_baseK()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: depthA1 is %d.", tilingData.mmTilingData.get_depthA1()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: depthB1 is %d.", tilingData.mmTilingData.get_depthB1()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: stepKa is %d.", tilingData.mmTilingData.get_stepKa()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: stepKb is %d.", tilingData.mmTilingData.get_stepKb()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: stepM is %d.", tilingData.mmTilingData.get_stepM()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: stepN is %d.", tilingData.mmTilingData.get_stepN()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: isBias is %d.", tilingData.mmTilingData.get_isBias()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: transLength is %d.", tilingData.mmTilingData.get_transLength()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: iterateOrder is %d.", tilingData.mmTilingData.get_iterateOrder()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: dbL0A is %d.", tilingData.mmTilingData.get_dbL0A()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: dbL0B is %d.", tilingData.mmTilingData.get_dbL0B()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: dbL0C is %d.", tilingData.mmTilingData.get_dbL0C()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: usedL1Size is %d.", tilingData.mmTilingData.get_shareL1Size()); + OP_LOGD(context->GetNodeName(), "GMM_tiling_new: usedUBSize is %d.", tilingData.mmTilingData.get_shareUbSize()); + auto buf = (uint32_t *)context->GetRawTilingData()->GetData(); + auto bufLen = context->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context, "%s", oss.str().c_str()); +} + +ge::graphStatus GMMTiling::GMMCalUbSize(const gert::TilingContext* context, uint32_t ubSize) { + OP_CHECK_IF((ubDivideBlkNum_ == 0 || ubBlockAlign_ == 0), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "ubDivideBlkNum and ubBlockAlign cannot be 0"), + return ge::GRAPH_FAILED); + uint32_t ubCalSize = ubSize / ubDivideBlkNum_; // divide the UB into ubDivideBlkNum_ pieces + ubCalSize = ubCalSize / ubBlockAlign_ * ubBlockAlign_; // 16k/8k/4k align. + uint32_t ubRestBytes = ubSize - ubCalSize * ubIoBlkNum_; // compute the rest memory in UB space + ubRestBytes = ubRestBytes / UB_BLOCK_UNIT_SIZE * UB_BLOCK_UNIT_SIZE; // 32B align. + OP_CHECK_IF((ubCalSize == 0 || ubRestBytes == 0), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "ubCalSize and ubRestBytes cannot be 0"), + return ge::GRAPH_FAILED); + uint32_t ubBaseN = 0; // init + uint32_t ubBaseK = 0; // init + uint32_t ubBaseM = 0; // init + if (transposeWeight_) { + ubBaseK = static_cast(BEST_UB_BASEK); + ubBaseN = ubCalSize / ubBaseK; + uint32_t alignFactor = UB_BLOCK_UNIT_SIZE; + if (weightDtype_ == ge::DT_INT4) { + alignFactor <<= 1U; // int4 need 64 elements algin. + } + ubBaseN = ubBaseN / alignFactor * alignFactor; + } else { + if ((xDType_ == ge::DT_BF16 || xDType_ == ge::DT_FLOAT16) && + (weightDtype_ == ge::DT_INT8 || weightDtype_ == ge::DT_INT4)) { + if (perTokenOrPerGroupSize_ > 0U) { + ubBaseK = perTokenOrPerGroupSize_; + ubBaseN = std::min(BEST_UB_BASEN, std::max(MIN_UB_BASEN, (ubCalSize / ubBaseK + MIN_UB_BASEN - 1) / MIN_UB_BASEN * MIN_UB_BASEN)); + } else if (antiquantPerformance_) { + ubBaseN = static_cast(BEST_UB_BASEN); + } else { + ubBaseN = static_cast(baseN_); + } + } else { + ubBaseN = static_cast(baseN_); + } + ubBaseK = ubCalSize / ubBaseN; // ubCalSize is the number of elements, not in bytes unit. + ubBaseM = ubCalSize / ubBaseN; + } + if (xDType_ == ge::DT_BF16 && (weightDtype_ == ge::DT_INT8 || weightDtype_ == ge::DT_INT4) && !isA16W8Msd_) { + OP_CHECK_IF(ubBaseK == 0 || ubBaseN == 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "ubBaseK or ubBaseN cannot be 0"), + return ge::GRAPH_FAILED); + } + tilingData.gmmBaseParams.set_ubCalSize(ubCalSize); + tilingData.gmmBaseParams.set_ubRestBytes(ubRestBytes); // in byte unit + tilingData.gmmBaseParams.set_ubBaseK(ubBaseK); + tilingData.gmmBaseParams.set_ubBaseN(ubBaseN); + tilingData.gmmBaseParams.set_vBaseM(ubBaseM); + return ge::GRAPH_SUCCESS; +} + +int64_t GMMTiling::GMMGetBS(const gert::Shape &xShape) const { + int64_t bs = 0; // init bs + if (transposeX_) { + bs = xShape.GetDim(1); // x shape is [k, m] if x is transpose_ + } else { + if (groupType_ == -1) { // -1: no group case, may exits a situation that multi dims product equals to bs. + bs = xShape.GetDim(0); // 0: x first dim + size_t bsDimNum = xDimNum_ >= 1U ? xDimNum_ - 1UL : 0UL; // 1: x last dim k, the other dimensions are bs + for (size_t i = 1; i < bsDimNum; i++) { + bs *= xShape.GetDim(i); + } + } else { + bs = xShape.GetDim(0); // in group case，x's shapeis [m,k], 0 is the m axis. + } + } + return bs; +} + +uint32_t GMMTiling::GetTplDataType(const ge::DataType &dtype) { + static const std::map dtype_map = { + {ge::DT_INT4, GMM_TPL_INT4}, + {ge::DT_INT8, GMM_TPL_INT8}, + {ge::DT_FLOAT16, GMM_TPL_FLOAT16}, + {ge::DT_BF16, GMM_TPL_BF16}, + {ge::DT_INT32, GMM_TPL_INT32}, + {ge::DT_FLOAT, GMM_TPL_FLOAT} + }; + auto it = dtype_map.find(dtype); + if (it != dtype_map.end()) { + return it->second; + } + return GMM_TPL_INVALID; +} + +bool GMMTiling::IsAivAicRatioTwoRequired() { + // Must be valid group type (not 2) + if (groupListType_ == GROUP_LIST_SPARSE_M) { + return false; + } + + // Condition 1: GELU activation (immediate match) + if (actType_ == ACT_TYPE_GELU) { + return true; + } + + // Condition 2: Complex tuning configuration requiring: + // - K dimension outside normal vectorization range + // - Minimum tuning configuration threshold + // - Valid token/group size + const bool needs_double_vector = (maxK_ <= DOUBLE_VECTOT_THRESHOLD_K_LOWER) || + (maxK_ >= DOUBLE_VECTOT_THRESHOLD_K_UPPER); + const bool has_sufficient_tuning = (tuningConfig_ >= SMALL_TUNING_CONFIG_THRESHOLD); + const bool has_valid_workload = (perTokenOrPerGroupSize_ > 0U); + + return needs_double_vector && has_sufficient_tuning && has_valid_workload; +} + +bool GMMTiling::IsFixedAxisMoveCondition() { + bool isCorrectShape = (maxK_ == FIXAXISMOVE_K1 && maxN_ == FIXAXISMOVE_N1) || + (maxK_ == FIXAXISMOVE_K2 && maxN_ == FIXAXISMOVE_N2); + bool isGroupCorrect = (groupNum_ == FIXAXISMOVE_GROUP_NUM); + bool isTuningInRange = (tuningConfig_ >= FIXAXISMOVE_PERM_LOWER) && + (tuningConfig_ <= FIXAXISMOVE_PERM_UPPER); + bool isDataTypeCorrect = yDtype_ == ge::DT_FLOAT16 && scaleDtype_ == ge::DT_FLOAT && perTokenScaleDtype_ == ge::DT_FLOAT; + bool isConfigCorrect = !transposeX_ && (splitItem_ == FIXAXISMOVE_SPLIT_ITEM2 || splitItem_ == FIXAXISMOVE_SPLIT_ITEM3) + && (groupListType_ == FIXAXISMOVE_GROUP_LIST_TYPE) + && (groupType_ == FIXAXISMOVE_GROUP_TYPE) && (actType_ == 0) + && !transposeWeight_; + bool isWorkspaceValid = (FixedAxisMoveWorkspace_ <= tuningConfigWorkspace_) || + (tuningConfigWorkspace_ == -1); + bool isFormatValid = (wFormat_ == matmul_tiling::CubeFormat::NZ); + + return isCorrectShape && isTuningInRange && isGroupCorrect && isA8W8_ && + isDataTypeCorrect && isConfigCorrect && isWorkspaceValid && !hasBias_ && isFormatValid; +} + +bool GMMTiling::IsIntDataType() { + return yDtype_ == ge::DT_INT8 || yDtype_ == ge::DT_INT32; +} + +void GMMTiling::GMMSetTplTilingKey(gert::TilingContext* context) { + uint32_t isStaticTilingApi = 0; + uint32_t a8w4KernelTemplate = GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE; + uint32_t a16w8KernelTemplate = GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE; + uint32_t aivAicRatio = GROUPED_MATMUL_AIV_AIC_RATIO_1; + uint32_t isEnableFixedAxis = 0; + + if (isA8W4FakeA8W8_) { + a8w4KernelTemplate = static_cast(GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERCHANNEL_ANTIQUANT); + } + + if ((xDType_ == ge::DT_FLOAT16 || xDType_ == ge::DT_BF16) && weightDtype_ == ge::DT_INT8) { + a16w8KernelTemplate = static_cast(GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT); + if (isA16W8Msd_) { + a16w8KernelTemplate = static_cast(GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD); + } + } + + if (isA4W4_ || antiquantPerformance_) { + aivAicRatio = static_cast(GROUPED_MATMUL_AIV_AIC_RATIO_2); + } else if (isA8W8_) { + if (isFixedAxisMove_) { + aivAicRatio = static_cast(GROUPED_MATMUL_AIV_AIC_RATIO_1); + isEnableFixedAxis = 1; + } else if (IsAivAicRatioTwoRequired()) { + aivAicRatio = static_cast(GROUPED_MATMUL_AIV_AIC_RATIO_2); + } else if (IsIntDataType()) { + aivAicRatio = static_cast(GROUPED_MATMUL_CUBE_ONLY); + } + } else if (!transposeX_ && xDType_ == weightDtype_ && (xDType_ == ge::DT_FLOAT16 || xDType_ == ge::DT_BF16 || xDType_ == ge::DT_FLOAT)) { + aivAicRatio = static_cast(GROUPED_MATMUL_CUBE_ONLY); + } + + if (a8w4KernelTemplate == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE && + a16w8KernelTemplate == GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE && + aivAicRatio != GROUPED_MATMUL_AIV_AIC_RATIO_2 && + !isFixedAxisMove_ && + StaticTilingProcess(context)) { + isStaticTilingApi = 1U; + } + + const uint64_t tilingKey = GET_TPL_TILING_KEY(GetTplDataType(xDType_), + GetTplDataType(weightDtype_), + GetTplDataType(yDtype_), + transposeX_? 1UL : 0UL, + transposeWeight_? 1UL : 0UL, + groupListType_, + isStaticTilingApi, + a8w4KernelTemplate, + a16w8KernelTemplate, + aivAicRatio, + isEnableFixedAxis); + context->SetTilingKey(tilingKey); + + if (isA16W8Msd_ || antiquantPerformance_) { + context->SetScheduleMode(1); // set as batchmod for template using SyncAll + } +} + +ge::graphStatus GMMTiling::GMMGetAttrs(const gert::TilingContext* context) { + auto attr = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attr); // check attr is not null + const bool* transposeWeightPtr = attr->GetAttrPointer(ATTR_INDEX_TRANS_W); + const bool* transposeXPtr = attr->GetAttrPointer(ATTR_INDEX_TRANS_X); + const int32_t* groupTypePtr = attr->GetAttrPointer(ATTR_INDEX_GROUPTYPE); + const int64_t* splitItemPtr = attr->GetAttrPointer(ATTR_INDEX_SPLIT_ITEM); + const int64_t* actTypePtr = attr->GetAttrPointer(ATTR_INDEX_ACT_TYPE); + const uint32_t* groupListTypePtr = attr->GetAttrPointer(ATTR_INDEX_GROUP_LIST_TYPE); + const auto tuningConfigPtr = attr->GetAttrPointer(ATTR_INDEX_TUNING_CONFIG); + transposeWeight_ = transposeWeightPtr != nullptr ? *transposeWeightPtr : false; + transposeX_ = transposeXPtr != nullptr ? *transposeXPtr : false; + groupType_ = groupTypePtr != nullptr ? *groupTypePtr : NO_SPLIT; + splitItem_ = splitItemPtr != nullptr ? *splitItemPtr : 0U; // 0: 默认split_item + actType_ = actTypePtr != nullptr ? *actTypePtr : 0; + groupListType_ = groupListTypePtr != nullptr ? *groupListTypePtr : 0; + + auto xDesc = context->GetDynamicInputDesc(X_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); // check xDesc is not null + xDType_ = xDesc->GetDataType(); + auto w0Desc = context->GetDynamicInputDesc(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + weightDtype_ = w0Desc->GetDataType(); + if (xDType_ == ge::DT_INT8 && weightDtype_ == ge::DT_INT4) { + const uint64_t n = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDim(2); + const uint64_t k = context->GetDynamicInputTensor(X_INDEX, 0)->GetStorageShape().GetDim(1); + const uint64_t groupNum = context->GetDynamicInputTensor(WEIGHT_INDEX, 0)->GetStorageShape().GetDim(0); + const uint64_t quantGroupNum = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDim(1); + isA8W4FakeA8W8_ = true; + A8W4noMsdSpace_ = groupNum * k * n * sizeof(int8_t) + groupNum * n * sizeof(float); + tilingData.gmmBaseParams.set_groupNum(groupNum); + tilingData.gmmBaseParams.set_n(n); + tilingData.gmmBaseParams.set_k(k); + tilingData.gmmBaseParams.set_quantGroupNum(quantGroupNum); + } + isA8W8_ = (xDType_ == ge::DT_INT8 && weightDtype_ == ge::DT_INT8); + isA4W4_ = xDType_ == ge::DT_INT4 && weightDtype_ == ge::DT_INT4; + + auto compileInfoPtr = GetGMMCompileInfo(context); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null + if (groupListType_ == GROUP_LIST_SPARSE_M) { + OP_CHECK_IF((!(compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910B || + compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910_93)), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "This platform not support groupListType is 2"), + return ge::GRAPH_FAILED); + bool isNonQuantA16W16 = (xDType_ == weightDtype_) && + (xDType_ == ge::DT_BF16 || xDType_ == ge::DT_FLOAT16); + OP_CHECK_IF(!(isA8W8_ || isNonQuantA16W16), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "Only A8W8 or non-quant BF16/FP16 support groupListType is 2"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupType_ != SPLIT_M, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "When groupListType is 2 only support groupType 0, but get groupType %d", + groupType_), + return ge::GRAPH_FAILED); + } + + auto perTokenScalePtr = context->GetOptionalInputTensor(PER_TOKEN_SCALE_INDEX); + if (perTokenScalePtr != nullptr && perTokenScalePtr->GetStorageShape().GetShapeSize() != 0) { + perTokenOrPerGroupSize_ = 1U; + } + auto pertokenScaleDesc = context->GetOptionalInputDesc(PER_TOKEN_SCALE_INDEX); + if (pertokenScaleDesc != nullptr) { + perTokenScaleDtype_ = pertokenScaleDesc->GetDataType(); + } + tilingData.gmmBaseParams.set_quantParam(perTokenOrPerGroupSize_); + auto yDesc = context->GetOutputDesc(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, yDesc); + yDtype_ = yDesc->GetDataType(); + if ((weightDtype_ == ge::DT_INT8 && xDType_ == ge::DT_INT8 && yDtype_ != ge::DT_INT32) || isA8W4FakeA8W8_ || + (xDType_ == ge::DT_FLOAT8_E4M3FN) || (xDType_ == ge::DT_FLOAT8_E5M2)) { + auto scale0Desc = context->GetDynamicInputDesc(SCALE_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, scale0Desc); + scaleDtype_ = scale0Desc->GetDataType(); + } + auto wFormat0 = static_cast(ge::GetPrimaryFormat(w0Desc->GetStorageFormat())); + wFormat_ = wFormat0 == ge::FORMAT_FRACTAL_NZ ? matmul_tiling::CubeFormat::NZ : matmul_tiling::CubeFormat::ND; + tuningConfig_ = (tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > TUNING_CONFIG_TOKEN_PER_EXPECT_INDEX) ? + (reinterpret_cast(tuningConfigPtr->GetData()))[TUNING_CONFIG_TOKEN_PER_EXPECT_INDEX] : 0; + tuningConfigWorkspace_ = (tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > TUNING_CONFIG_ALLOW_WORKSPACE_INDEX) ? + (reinterpret_cast(tuningConfigPtr->GetData()))[TUNING_CONFIG_ALLOW_WORKSPACE_INDEX] : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::GMMSetUbDivideBlkAntiquant() { + if (isA16W8Msd_) { + ubDivideBlkNum_ = UB_A16W8_MSD_BLOCK_NUM; + ubIoBlkNum_ = UB_A16W8_MSD_IO_USED_BLOCK; + ubBlockAlign_ = UB_A16W8_MSD_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; + } + if (xDType_ == ge::DT_FLOAT16 && (weightDtype_ == ge::DT_INT8 || weightDtype_ == ge::DT_INT4)) { + if (weightDtype_ == ge::DT_INT8) { + ubDivideBlkNum_ = UB_A16W8_BLOCK_NUM_FP16; + ubIoBlkNum_ = UB_A16W8_IO_USED_BLOCK_FP16; + } else { // int4 + ubDivideBlkNum_ = UB_A16W4_BLOCK_NUM_FP16; + ubIoBlkNum_ = UB_A16W4_IO_USED_BLOCK_FP16; + } + ubBlockAlign_ = UB_ANTIQUANT_PER_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; + } + if (xDType_ == ge::DT_BF16 && (weightDtype_ == ge::DT_INT8 || weightDtype_ == ge::DT_INT4)) { + if (weightDtype_ == ge::DT_INT8) { + ubDivideBlkNum_ = UB_A16W8_BLOCK_NUM_BF16; + ubIoBlkNum_ = UB_A16W8_IO_USED_BLOCK_BF16; + } else { + ubDivideBlkNum_ = UB_A16W4_BLOCK_NUM_BF16; + ubIoBlkNum_ = UB_A16W4_IO_USED_BLOCK_BF16; + } + ubBlockAlign_ = UB_ANTIQUANT_PER_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_FAILED; +} + +ge::graphStatus GMMTiling::GMMSetUbDivideBlkQuant() { + if ((weightDtype_ == ge::DT_INT8 || isA8W4FakeA8W8_) && (perTokenOrPerGroupSize_ == 1 || actType_ != 0)) { + // include case per-token without activation, per-token with activation and per-tensor with activation + ubDivideBlkNum_ = UB_DYNAMIC_QUANT_BLOCK_NUM; + ubIoBlkNum_ = UB_DUNAMIC_QUANT_IO_USED_BLOCK; + ubBlockAlign_ = UB_QUANT_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; + } + if ((weightDtype_ == ge::DT_INT8 || isA8W4FakeA8W8_) && perTokenOrPerGroupSize_ != 1) { + // include case per-tensor without activation + if (yDtype_ == ge::DT_FLOAT16) { + ubDivideBlkNum_ = UB_STATIC_QUANT_BLOCK_NUM_FP16; + } else { + ubDivideBlkNum_ = UB_STATIC_QUANT_BLOCK_NUM_BF16; + } + ubIoBlkNum_ = UB_STATIC_QUANT_IO_USED_BLOCK; + ubBlockAlign_ = UB_QUANT_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_FAILED; +} + +ge::graphStatus GMMTiling::GMMSetUbDivideBlkA4W4() { + ubDivideBlkNum_ = UB_A4W4_BLOCK_NUM; + ubIoBlkNum_ = UB_A4W4_IO_USED_BLOCK_HALF; + ubBlockAlign_ = UB_A4W4_PER_BLOCK_ALIGN; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::GMMSetUbDivideBlk() { + ubDivideBlkNum_ = 0U; // init ubDivideBlkNum_ + ubIoBlkNum_ = 0U; // init ubIoBlkNum_ + ubBlockAlign_ = 0U; // init ubBlockAlign_ + if (xDType_ == ge::DT_INT8) { + return GMMSetUbDivideBlkQuant(); + } else if (isA4W4_) { + return GMMSetUbDivideBlkA4W4(); + } else { + return GMMSetUbDivideBlkAntiquant(); + } + return ge::GRAPH_FAILED; +} + +ge::graphStatus GMMTiling::SetBias(const gert::TilingContext* context, matmul_tiling::MultiCoreMatmulTiling& mm) const { + if (!hasBias_ || isA16W8Msd_ || isA4W4_) { + mm.SetBias(false); + } else { + mm.SetBias(true); + auto biasTensor = context->GetDynamicInputTensor(BIAS_INDEX, 0); + OP_CHECK_IF(biasTensor == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "Get bias tensor failed."), + return ge::GRAPH_FAILED); + mm.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, + static_cast(biasTensor->GetDataType())); + } + return ge::GRAPH_SUCCESS; +} + +static void InitPlatformInfo(const GMMCompileInfo* compileInfoPtr, matmul_tiling::PlatformInfo& platformInfo) { + platformInfo.socVersion = compileInfoPtr->socVersion; + platformInfo.l1Size = compileInfoPtr->l1Size; + platformInfo.l0CSize = compileInfoPtr->l0CSize; + platformInfo.ubSize = compileInfoPtr->ubSize; + platformInfo.l0ASize = compileInfoPtr->l0ASize; + platformInfo.l0BSize = compileInfoPtr->l0BSize; +} + +void GMMTiling::FullLoadK(const GMMCompileInfo* compileInfoPtr) { + if (wFormat_ == matmul_tiling::CubeFormat::ND && !transposeWeight_ && !transposeX_ && xDType_ == weightDtype_ && + (weightDtype_ == ge::DT_FLOAT16 || weightDtype_ == ge::DT_BF16) && + maxM_ >= FULL_K_M_E_THRESHOLD * groupNum_ && maxN_ == FULL_K_N_THRESHOLD && + maxK_ <= FULL_K_MAX_K_THRESHOLD && + maxK_ >= FULL_K_MIN_K_THRESHOLD) { + int64_t fullLoadStepKa = CeilDiv(maxK_ , baseK_); + int64_t fullLoadStepKb = fullLoadStepKa / static_cast(QUEUE_DOUBLE_BUFFER); + int64_t fullLoadDepthKa = fullLoadStepKa * static_cast(QUEUE_DOUBLE_BUFFER); + int64_t fullLoadDepthKb = fullLoadStepKb * static_cast(QUEUE_DOUBLE_BUFFER); + bool ifFullLoad = (maxM_ > FULL_K_M_THRESHOLD) && isAllSingleTensor_ && groupType_ == SPLIT_M && + (((baseM_ * baseK_ * static_cast(mmDataTypeSize_)) * fullLoadDepthKa + + (baseN_ * baseK_ * static_cast(mmDataTypeSize_)) * fullLoadDepthKb) <= + static_cast(compileInfoPtr->l1Size)); + if (ifFullLoad) { + tilingData.mmTilingData.set_stepKa(fullLoadStepKa); // set precomputed mmStepKa + tilingData.mmTilingData.set_depthA1(fullLoadDepthKa); // set precomputed mmDepthA1 + tilingData.mmTilingData.set_stepKb(fullLoadStepKb); // set precomputed mmStepKb + tilingData.mmTilingData.set_depthB1(fullLoadDepthKb); // set precomputed mmDepthB1 + tilingData.mmTilingData.set_iterateOrder(1); // set precomputed stepN + tilingData.gmmBaseParams.set_singleN(FULL_K_SINGLE_N); // 0 is the default value + } + } +} + +ge::graphStatus GMMTiling::CalcStepKaKb(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr, + int64_t mInMM, uint32_t& mmStepKa, uint32_t& mmStepKb) { + uint64_t availableL1Size = compileInfoPtr->l1Size; + if (isA8W8_ || isA8W4FakeA8W8_) { + availableL1Size -= static_cast(baseN_) * sizeof(uint64_t); + } + if (hasBias_) { + availableL1Size -= static_cast(baseN_) * static_cast(4); // 4: size of float32 or int32 + } + if (compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND310P) { + availableL1Size = BEST_L1_PARTA + BEST_L1_PARTB; + } + OP_CHECK_IF(availableL1Size < L1_PARTA_SIZE, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "availableL1Size is less than 256k"), + return ge::GRAPH_FAILED); + // according to double buffer, recompute the params used for data movement from GM to L1 + uint64_t l1ASize = baseM_ > baseN_ ? L1_PARTA_SIZE : availableL1Size - L1_PARTA_SIZE; + uint64_t l1BSize = availableL1Size - l1ASize; + if (isA4W4_) { + // 2: double buffer + mmStepKa = static_cast((l1ASize / 2UL) / static_cast(INT4_DATA_TYPE_SIZE * + static_cast(baseM_) * + static_cast(baseK_))); + // 2: double buffer + mmStepKb = static_cast((l1BSize / 2UL) / static_cast(INT4_DATA_TYPE_SIZE * + static_cast(baseN_) * + static_cast(baseK_))); + } else { + // 2: double buffer + mmStepKa = (l1ASize / 2UL) / (static_cast(baseM_) * baseK_ * mmDataTypeSize_); + // 2: double buffer + mmStepKb = (l1BSize / 2UL) / (static_cast(baseN_) * baseK_ * mmDataTypeSize_); + } + if (compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND310P && + wFormat_ == matmul_tiling::CubeFormat::NZ && mInMM <= baseM_) { + mmStepKa = std::min(mmStepKa, std::max(1, 128 / baseK_)); // 128: nz inner block size. In practice, baseK_*mmStepKa=128 makes performance better. + } + + OP_CHECK_IF(mmStepKa == 0 || mmStepKb == 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "stepka or stepkb cannot be 0"), + return ge::GRAPH_FAILED); + + if (mmStepKa > mmStepKb) { + mmStepKa = mmStepKa / mmStepKb * mmStepKb; + } else if (mmStepKa < mmStepKb) { + mmStepKb = mmStepKb / mmStepKa * mmStepKa; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GMMTiling::GMMSetMMTiling(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr) { + matmul_tiling::DataType matmulDtype = static_cast(mmDType_); + matmul_tiling::PlatformInfo platformInfo; + InitPlatformInfo(compileInfoPtr, platformInfo); + matmul_tiling::MultiCoreMatmulTiling mm(platformInfo); + int64_t mInMM = isA16W8Msd_ ? static_cast(A16W8_MSD_STEP) * maxM_ : maxM_; // if msd, m in matmul should mul steps + mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmulDtype, false); + mm.SetBType(matmul_tiling::TPosition::GM, wFormat_, matmulDtype, false); + mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND_ALIGN, matmul_tiling::DataType::DT_FLOAT16); + OP_CHECK_IF(SetBias(context, mm) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "SetBias failed."), return ge::GRAPH_FAILED); + mm.SetOrgShape(mInMM, maxN_, maxK_); + mm.SetShape(mInMM, baseN_, maxK_); + mm.SetFixSplit(baseM_, baseN_, baseK_); + mm.SetBufferSpace(compileInfoPtr->l1Size, compileInfoPtr->l0CSize, ubSize_); + OP_CHECK_IF(mm.GetTiling(tilingData.mmTilingData) == -1, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "matmul getTiling failed."), + return ge::GRAPH_FAILED); + + uint32_t mmStepKa = 1; + uint32_t mmStepKb = 1; + OP_CHECK_IF(CalcStepKaKb(context, compileInfoPtr, mInMM, mmStepKa, mmStepKb) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "matmul calc stepka or stepkb failed."), + return ge::GRAPH_FAILED); + + constexpr uint32_t stepM = 1; // 1: stepM set fixed value 1 + constexpr uint32_t stepN = 1; // 1: stepN set fixed value 1 + uint32_t mmDepthA1 = mmStepKa * DOUBLE_BUFFER_STEPKA_STEPKB * stepM; + uint32_t mmDepthB1 = mmStepKb * DOUBLE_BUFFER_STEPKA_STEPKB * stepN; + tilingData.mmTilingData.set_shareMode(0); + if (compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND310P) { + tilingData.mmTilingData.set_shareUbSize(0); + tilingData.mmTilingData.set_transLength(131072); // 131072: 128KB size + } + tilingData.mmTilingData.set_dbL0C(1); // disable double buffer for LOC + tilingData.mmTilingData.set_baseM(baseM_); // set precomputed baseM + tilingData.mmTilingData.set_baseN(baseN_); // set precomputed baseN + tilingData.mmTilingData.set_baseK(baseK_); // set precomputed baseK + tilingData.mmTilingData.set_stepKa(mmStepKa); // set precomputed mmStepKa + tilingData.mmTilingData.set_depthA1(mmDepthA1); // set precomputed mmDepthA1 + tilingData.mmTilingData.set_stepKb(mmStepKb); // set precomputed mmStepKb + tilingData.mmTilingData.set_depthB1(mmDepthB1); // set precomputed mmDepthB1 + tilingData.mmTilingData.set_stepM(stepM); // set precomputed stepM + tilingData.mmTilingData.set_stepN(stepN); // set precomputed stepN + SetMMPreTiling(); + OP_LOGI(context->GetNodeName(), "GMM_tiling: baseM is %d, baseK is %d, baseN is %d.", baseM_, baseK_, baseN_); + return ge::GRAPH_SUCCESS; +} + +void GMMTiling::SetMMPreTiling() { + uint64_t ispreTiling = 0; + int64_t isNz = wFormat_ == matmul_tiling::CubeFormat::NZ ? 1 : 0; + if (tuningConfig_ == 0L && (isA8W8_ || isA8W4FakeA8W8_) && groupNum_ == 1U && usedCoreNum_ == A3_AIC_NUM) { + ispreTiling = static_cast(1); // 1: pretiling key + } else { + tilingData.gmmBaseParams.set_isPreTiling(ispreTiling); + return; + } + std::array mKNList = {maxM_, maxK_, maxN_, isNz}; // 4: input shape info size + if (A8W8_PRETILING_WHITE_LIST.count(mKNList)) { + int64_t newBaseM = A8W8_PRETILING_WHITE_LIST.at(mKNList)[0]; + int64_t bestSingleN = A8W8_PRETILING_WHITE_LIST.at(mKNList)[1]; + tilingData.mmTilingData.set_baseM(newBaseM); // set pretiling baseM + tilingData.mmTilingData.set_singleCoreN(bestSingleN); // set pretiling singleN + ispreTiling = static_cast(2); // 2: white list pretiling key + } + tilingData.gmmBaseParams.set_isPreTiling(ispreTiling); + return; +} + +ge::graphStatus GMMTiling::CalMMTiling(const gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr) { + // if tuningConfig_ in (128, 256], recompute tiling. + // or y:int32 and tuningConfig_ in (0, 128], k,n:(7168, 2048) or k,n:(7680, 2048), which got by actual measurement) + constexpr int32_t tuningConfigBaseLowerLimit = 128; + constexpr int32_t tuningConfigBaseUpperLimit = 256; + constexpr int32_t maxKLimit = 7168; + constexpr int32_t altKLimit = 7680; + constexpr int32_t maxNLimit = 2048; + + bool tuningConfigFlag = (tuningConfig_ > tuningConfigBaseLowerLimit && tuningConfig_ <= tuningConfigBaseUpperLimit) + || (yDtype_ == ge::DT_INT32 && tuningConfig_ > 0 && tuningConfig_ <= tuningConfigBaseLowerLimit && (maxK_ == maxKLimit || maxK_ == altKLimit) && maxN_ == maxNLimit); + + baseN_ = BEST_BASEN; // init + // 2048: min n for a16w8 msd to set baseN 512 + if (isA16W8Msd_ && maxN_ >= 2048 && !transposeWeight_) { + baseN_ = BEST_BASEN_MSD; + } else if ((isA8W8_ || isA8W4FakeA8W8_) && tuningConfigFlag){ + baseN_ = BEST_BASEN_QUANT_ONE_GROUP; + baseM_ = BEST_BASEM_QUANT_ONE_GROUP; + baseK_ = BEST_BASEK_QUANT_ONE_GROUP; + baseM_ = baseM_ > maxM_ ? static_cast(SixteenAlign(maxM_, true)) : baseM_; + return ge::GRAPH_SUCCESS; + } else if (isA4W4_) { + baseN_ = tuningConfig_ > 64 ? BEST_BASEN : BEST_BASEN_A4W4; // 64 : when token in each group > 64, set baseN to 256 + } else { + baseN_ = BEST_BASEN; + } + // according to the double buffer enabled L0B, compute baseK + baseK_ = isA4W4_ ? static_cast((compileInfoPtr->l0BSize / DOUBLE_BUFFER_L0A_L0B) / + static_cast(baseN_ * INT4_DATA_TYPE_SIZE)) : + static_cast((compileInfoPtr->l0BSize / DOUBLE_BUFFER_L0A_L0B) / + (static_cast(baseN_) * mmDataTypeSize_)); + baseK_ = static_cast(SixteenAlign(static_cast(baseK_))); + OP_CHECK_IF(baseK_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "baseK_ cannot be 0."), + return ge::GRAPH_FAILED); + // according to the double buffer enabled L0A/L0C, compute baseM(cube) + uint32_t maxBaseM = static_cast(compileInfoPtr->l0CSize / + (static_cast(baseN_) * FP32_DATATYPE_SIZE)); + baseM_ = isA4W4_ ? + std::min((compileInfoPtr->l0ASize / DOUBLE_BUFFER_L0A_L0B) / + static_cast(baseK_ * INT4_DATA_TYPE_SIZE), maxBaseM) : + std::min((compileInfoPtr->l0ASize / DOUBLE_BUFFER_L0A_L0B) / + (static_cast(baseK_) * mmDataTypeSize_), maxBaseM); + + if (!isA16W8Msd_) { + baseM_ = baseM_ > maxM_ ? SixteenAlign(maxM_, true) : SixteenAlign(static_cast(baseM_)); + } else { + baseM_ = baseM_ > A16W8_MSD_STEP * maxM_ ? static_cast(SixteenAlign(static_cast(A16W8_MSD_STEP) * maxM_, true)) : + SixteenAlign(static_cast(baseM_)); + } + if (baseM_ > MAX_BASEM) { + baseM_ = MAX_BASEM; + } + OP_CHECK_IF(baseM_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "baseM_ cannot be 0."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +static void SetA8W4HPTiling(A8W4HPTiling *tiling_data, uint32_t aicNum) +{ + constexpr int SIZE_TWO = 2; + constexpr int SIZE_THREE = 3; + constexpr int IDX_ZERO = 0; + constexpr int IDX_ONE = 1; + constexpr int IDX_TWO = 2; + constexpr uint32_t SINGLE_CORE_TILING_0 = 128; + constexpr uint32_t SINGLE_CORE_TILING_1 = 256; + constexpr uint32_t SINGLE_CORE_BASE_TILING_0 = 128; + constexpr uint32_t SINGLE_CORE_BASE_TILING_1 = 256; + constexpr uint32_t TOTAL_K_THRESHOLD_6656 = 6656; + + uint32_t *ori_in0_shape = tiling_data->get_ori_in0_shape(); + uint32_t total_K = ori_in0_shape[IDX_ONE]; + uint32_t core_num = aicNum; + uint32_t splitRecord[SIZE_THREE] = {1, 1, 1}; + uint32_t single_core_tiling[SIZE_THREE] = {SINGLE_CORE_TILING_0, SINGLE_CORE_TILING_1, total_K}; + uint32_t single_core_base_tiling[SIZE_TWO] = {SINGLE_CORE_BASE_TILING_0, SINGLE_CORE_BASE_TILING_1}; + + tiling_data->set_required_core_num(core_num); + tiling_data->set_kernel_index(0); + tiling_data->set_splitTimes(0); + + if (total_K > TOTAL_K_THRESHOLD_6656) { + splitRecord[IDX_ZERO] = CeilDiv(total_K, TOTAL_K_THRESHOLD_6656); + single_core_tiling[IDX_TWO] = TOTAL_K_THRESHOLD_6656; + } + + tiling_data->set_splitRecord(splitRecord); + tiling_data->set_single_core_tiling(single_core_tiling); + tiling_data->set_single_core_base_tiling(single_core_base_tiling); +} + +static void PrintA8W4HPTiling(gert::TilingContext* context, A8W4HPTiling *data) +{ + constexpr int TWO = 2; + OP_LOGD(context->GetNodeName(), "Tiling data strategy: "); + OP_LOGD(context->GetNodeName(), " group_num=%u", data->get_group_num()); + OP_LOGD(context->GetNodeName(), " group_type=%hhd", data->get_group_type()); + OP_LOGD(context->GetNodeName(), " required_core_num=%u", data->get_required_core_num()); + OP_LOGD(context->GetNodeName(), " format_in=%f", data->get_format_in()); + OP_LOGD(context->GetNodeName(), " format_out=%f", data->get_format_out()); + OP_LOGD(context->GetNodeName(), " numAic=%u", data->get_numAic()); + OP_LOGD(context->GetNodeName(), " numAiv=%u", data->get_numAiv()); + OP_LOGD(context->GetNodeName(), " szUb=%llu", static_cast(data->get_szUb())); + OP_LOGD(context->GetNodeName(), " szL0A=%llu", static_cast(data->get_szL0A())); + OP_LOGD(context->GetNodeName(), " szL0C=%llu", static_cast(data->get_szL0C())); + OP_LOGD(context->GetNodeName(), " pattern=%hhu", data->get_pattern()); + OP_LOGD(context->GetNodeName(), " kernel_index=%hhu", data->get_kernel_index()); + OP_LOGD(context->GetNodeName(), " splitTimes=%u", data->get_splitTimes()); + OP_LOGD(context->GetNodeName(), " output_type=%hhd", data->get_output_type()); + OP_LOGD(context->GetNodeName(), " ori_in0_shape=[%u,%u]", data->get_ori_in0_shape()[0], + data->get_ori_in0_shape()[1]); + OP_LOGD(context->GetNodeName(), " ori_in1_shape=[%u,%u]", data->get_ori_in1_shape()[0], + data->get_ori_in1_shape()[1]); + OP_LOGD(context->GetNodeName(), " ori_out_shape=[%u,%u]", data->get_ori_out_shape()[0], + data->get_ori_out_shape()[1]); + OP_LOGD(context->GetNodeName(), " single_core_tiling=[%u,%u,%u]", data->get_single_core_tiling()[0], + data->get_single_core_tiling()[1], data->get_single_core_tiling()[TWO]); + OP_LOGD(context->GetNodeName(), " single_core_base_tiling=[%u,%u]", data->get_single_core_base_tiling()[0], + data->get_single_core_base_tiling()[1]); + OP_LOGD(context->GetNodeName(), " splitRecord=[%u,%u,%u]", data->get_splitRecord()[0], + data->get_splitRecord()[1], data->get_splitRecord()[TWO]); + OP_LOGD(context->GetNodeName(), " workspaceOffset=%llu", static_cast(data->get_workspaceOffset())); +} + +ge::graphStatus GMMTiling::A8W4Tiling(gert::TilingContext* context, const GMMCompileInfo* compileInfoPtr) { + auto yDesc = context->GetOutputDesc(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, yDesc); + uint32_t yDtype = GMM_TPL_FLOAT16; + if (yDesc->GetDataType() == ge::DT_BF16) { + yDtype = static_cast(GMM_TPL_BF16); + } + GMMTilingData tilingDataA8W4; + auto w0Desc = context->GetDynamicInputDesc(WEIGHT_INDEX, 0); + auto wFormat0 = static_cast(ge::GetPrimaryFormat(w0Desc->GetStorageFormat())); + bool wNZ = wFormat0 == ge::FORMAT_FRACTAL_NZ; + + constexpr uint32_t cvParallNum = 4; // for cv collaboration + constexpr uint32_t THIRTY_TWO = 32; + constexpr uint32_t UBCALSIZE = 32U * 256U; // for vector compute + constexpr uint32_t UBRESTBYTES = 9U * 32U * 256U; // for vector compute + constexpr uint32_t TWO = 2; + constexpr uint32_t EIGHT = 8; + constexpr uint32_t FIVE = 5; + uint32_t singleN = 256; + uint32_t singleM = 128; + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null + const uint32_t& aicNum = compileInfoPtr->aicNum; // get aicNum from compileInfo + if (aicNum == 0U) { // invaild value + return ge::GRAPH_FAILED; + } + + auto attr = context->GetAttrs(); + const auto tuningConfigPtr = attr != nullptr ? (attr->GetAttrPointer(ATTR_INDEX_TUNING_CONFIG)) : nullptr; + + bool useHighPerf = (tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > TUNING_CONFIG_A8W4_SPEC_SCENARIO_INDEX) ? + ((reinterpret_cast(tuningConfigPtr->GetData()))[TUNING_CONFIG_A8W4_SPEC_SCENARIO_INDEX] == 1) : false; + if (useHighPerf) { + OP_LOGD(context->GetNodeName(), "Enter GMM A8W4 MSD high performance path..."); + constexpr int CASE_ZERO = 0; + constexpr int CASE_ONE = 1; + constexpr int CASE_TWO = 2; + constexpr int CASE_THREE = 3; + constexpr float INT4_TYPE_COUNT = 0.5f; + constexpr float FP16_TYPE_COUNT = 2.0f; + uint32_t groupNum = context->GetDynamicInputTensor(WEIGHT_INDEX, 0)->GetStorageShape().GetDim(0); + uint32_t N = context->GetDynamicInputTensor(WEIGHT_INDEX, 0)->GetStorageShape().GetDim(TWO); + uint32_t K = context->GetDynamicInputTensor(X_INDEX, 0)->GetStorageShape().GetDim(1); + uint32_t M = TWO * context->GetDynamicInputTensor(X_INDEX, 0)->GetStorageShape().GetDim(0); + + const auto transposeWeightPtr = attr->GetAttrPointer(ATTR_INDEX_TRANS_W); + const auto transposeXPtr = attr->GetAttrPointer(ATTR_INDEX_TRANS_X); + auto transposeWeight = transposeWeightPtr != nullptr ? *transposeWeightPtr : false; + auto transposeX = transposeXPtr != nullptr ? *transposeXPtr : false; + // | transposeX | transposeWeight | pattern | + // | true | true | 2 | + // | true | false | 3 | + // | false | true | 0 | + // | false | false | 1 | + int pattern = transposeX ? (transposeWeight ? CASE_TWO: CASE_THREE) : (transposeWeight ? CASE_ZERO : CASE_ONE); + pattern = CASE_ZERO; + + tilingDataA8W4.hpTilingData.set_pattern(static_cast(pattern)); + tilingDataA8W4.hpTilingData.set_kernel_index(0); + tilingDataA8W4.hpTilingData.set_format_in(INT4_TYPE_COUNT); // int4 + tilingDataA8W4.hpTilingData.set_format_out(FP16_TYPE_COUNT); // fp16 + + std::vector ori_in0_shape; + std::vector ori_in1_shape; + std::vector ori_out_shape; + switch (pattern) { + case CASE_ZERO: + ori_in0_shape = {M, K}; + ori_in1_shape = {N, K}; + ori_out_shape = {M, N}; + break; + case CASE_ONE: + ori_in0_shape = {M, K}; + ori_in1_shape = {K, N}; + ori_out_shape = {M, N}; + break; + case CASE_TWO: + ori_in0_shape = {K, M}; + ori_in1_shape = {N, K}; + ori_out_shape = {M, N}; + break; + case CASE_THREE: + ori_in0_shape = {K, M}; + ori_in1_shape = {K, N}; + ori_out_shape = {M, N}; + break; + default: + // unreachable + return ge::GRAPH_FAILED; + } + tilingDataA8W4.hpTilingData.set_ori_in0_shape(ori_in0_shape.data()); + tilingDataA8W4.hpTilingData.set_ori_in1_shape(ori_in1_shape.data()); + tilingDataA8W4.hpTilingData.set_ori_out_shape(ori_out_shape.data()); + uint32_t aic = aicNum; + uint32_t aiv = compileInfoPtr->aivNum; + uint64_t szUB = compileInfoPtr->ubSize; + uint64_t szL0A = compileInfoPtr->l0ASize; + uint64_t szL0C = compileInfoPtr->l0CSize; + SetA8W4HPTiling(&tilingDataA8W4.hpTilingData, aic); + + // autotiling parameters + tilingDataA8W4.hpTilingData.set_group_num(groupNum); + tilingDataA8W4.hpTilingData.set_group_type(0); + aic = tilingDataA8W4.hpTilingData.get_required_core_num(); + if (aic == 0U) { // invaild value + return ge::GRAPH_FAILED; + } + context->SetBlockDim(aic); + tilingDataA8W4.hpTilingData.set_numAic(aic); + tilingDataA8W4.hpTilingData.set_numAiv(aiv); + tilingDataA8W4.hpTilingData.set_szUb(szUB); + tilingDataA8W4.hpTilingData.set_szL0A(szL0A); + tilingDataA8W4.hpTilingData.set_szL0C(szL0C); + auto yDtypeLocal = yDesc->GetDataType(); + if (yDtypeLocal == ge::DT_FLOAT16) { + tilingDataA8W4.hpTilingData.set_output_type(0); + } else { + tilingDataA8W4.hpTilingData.set_output_type(1); + } + + size_t workspaceSize = M * N * sizeof(int16_t) + + (static_cast(SixteenAlign(M, true)) * K / TWO * sizeof(uint8_t)); + context->SetScheduleMode(1); // set as batchmod for template using SyncAll + context->SetTilingKey(GET_TPL_TILING_KEY(GMM_TPL_INT8, GMM_TPL_INT4, yDtype, 0, 0, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, 0, + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_AUTOTILING, + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE, + GROUPED_MATMUL_AIV_AIC_RATIO_2, 0)); + + size_t *workspaces = context->GetWorkspaceSizes(1); // get second variable + workspaces[0] = SYS_WORKSPACE_SIZE; // default size + workspaces[0] += workspaceSize; + + tilingDataA8W4.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tilingDataA8W4.GetDataSize()); + PrintA8W4HPTiling(context, &tilingDataA8W4.hpTilingData); + return ge::GRAPH_SUCCESS; + } else { + const uint32_t n = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDim(TWO); + const uint32_t k = context->GetDynamicInputTensor(X_INDEX, 0)->GetStorageShape().GetDim(1); + const uint32_t m = context->GetDynamicInputTensor(X_INDEX, 0)->GetStorageShape().GetDim(0); + const uint32_t groupNum = context->GetDynamicInputTensor(WEIGHT_INDEX, 0)->GetStorageShape().GetDim(0); + const uint32_t quantGroupNum = context->GetDynamicInputTensor(SCALE_INDEX, 0)->GetStorageShape().GetDim(1); + std::array mKNList = {groupNum, m, k, n, wNZ}; // 5: input shape info size + + auto offset = context->GetDynamicInputTensor(OFFSET_INDEX, 0); + uint32_t withOffset = 0; + if (offset != nullptr) { + auto &offsetShape = offset->GetStorageShape(); + const size_t offsetDimNum = offsetShape.GetDimNum(); + auto offsetDim0 = offsetShape.GetDim(0); + auto offsetDim1 = offsetShape.GetDim(1); + auto offsetDim2 = offsetShape.GetDim(TWO); + if (offsetDimNum == OFFSET_DIM_A8W4 && offsetDim0 == groupNum && offsetDim1 == 1 && offsetDim2 == n) { + withOffset = 1U; + OP_LOGD(context->GetNodeName(), "GMM A8W4: offset is enable ."); + } else { + OP_LOGW(context->GetNodeName(), "GMM A8W4: offset's shape is invalid, If you want to enable offset, the expected shape is (%u,1,%u), and the current shape is (%ld,%ld,%ld).", + groupNum, n, offsetDim0, offsetDim1, offsetDim2); + } + } + const int is_in_a8w4_white_list = A8W4_PRETILING_WHITE_LIST.count(mKNList) + && quantGroupNum != 0 && k / quantGroupNum == 256 && k % quantGroupNum == 0 + && withOffset == 0; // 256: 新方案只支持256 pergroup + + tilingDataA8W4.gmmBaseParams.set_coreNum(aicNum); + tilingDataA8W4.gmmBaseParams.set_groupNum(groupNum); + tilingDataA8W4.gmmBaseParams.set_totalInGroup(m); + tilingDataA8W4.gmmBaseParams.set_k(k); + tilingDataA8W4.gmmBaseParams.set_n(n); + tilingDataA8W4.gmmBaseParams.set_vBaseM(THIRTY_TWO); + tilingDataA8W4.gmmBaseParams.set_ubCalSize(UBCALSIZE); + tilingDataA8W4.gmmBaseParams.set_ubRestBytes(UBRESTBYTES); + tilingDataA8W4.gmmBaseParams.set_parallNum(cvParallNum); + tilingDataA8W4.gmmBaseParams.set_quantGroupNum(quantGroupNum); + tilingDataA8W4.gmmBaseParams.set_m(m); + tilingDataA8W4.gmmBaseParams.set_withOffset(withOffset); + context->SetBlockDim(aicNum); + + if (quantGroupNum == 0U || k % quantGroupNum != 0U) { + OP_LOGE(context->GetNodeName(), "GMM_tiling: k should be divisible by quantGroupNum, but now k=%u and quantGroupNum=%u", + k, quantGroupNum); + return ge::GRAPH_FAILED; + } + const uint32_t K_UNIT = 64; // 64: int4 in 32B + if (k % K_UNIT != 0) { + OP_LOGE(context->GetNodeName(), "GMM_tiling: k should be divisible by 64, but now k=%u", k); + return ge::GRAPH_FAILED; + } + const uint32_t MAX_K_A8W4_MSD = 18432; // k is limited by pre process, a line of X should be able to put in UB + if (k > MAX_K_A8W4_MSD) { + OP_LOGE(context->GetNodeName(), "GMM_tiling: K should be less than 18432 on the A8W4 scenario, but now is %u", + k); + return ge::GRAPH_FAILED; + } + matmul_tiling::PlatformInfo platformInfo; + InitPlatformInfo(compileInfoPtr, platformInfo); + matmul_tiling::MultiCoreMatmulTiling mm(platformInfo); + //GEMM Tiling + int64_t tuningConfig = (tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > TUNING_CONFIG_TOKEN_PER_EXPECT_INDEX) ? + (reinterpret_cast(tuningConfigPtr->GetData()))[TUNING_CONFIG_TOKEN_PER_EXPECT_INDEX] : 0; + uint32_t calc_m = 1U; + if (groupNum != 0U) { + calc_m = m / groupNum; + } + const uint32_t avg_m = tuningConfig != 0L ? static_cast(tuningConfig) : calc_m; + const bool isPerchannel = quantGroupNum == 1U; + const bool isMSD = tuningConfig == 0L || avg_m == 0U || n / avg_m > 4U || withOffset == true; + if (!isMSD) { + constexpr uint32_t A8W4_BASE_M = 128; + constexpr uint32_t A8W4_BASE_K = 64; + constexpr uint32_t A8W4_BASE_N = 128; + mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT8, false); + if (wNZ) { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::NZ, matmul_tiling::DataType::DT_INT8, false); + } else { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT8, false); + } + mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + mm.SetBias(false); + mm.SetOrgShape(A8W4_BASE_M, n, k); + mm.SetShape(A8W4_BASE_M, A8W4_BASE_N, k); + mm.SetFixSplit(A8W4_BASE_M, A8W4_BASE_N, A8W4_BASE_K); + if (mm.GetTiling(tilingDataA8W4.mmTilingData) == -1) { + return ge::GRAPH_FAILED; + } + context->SetTilingKey(GET_TPL_TILING_KEY(GMM_TPL_INT8, GMM_TPL_INT4, yDtype, 0, 0, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, 0, + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERGROUP_ANTIQUANT, + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE, + GROUPED_MATMUL_AIV_AIC_RATIO_2, 0)); + tilingDataA8W4.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tilingDataA8W4.GetDataSize()); + + size_t* workspaces = context->GetWorkspaceSizes(1); // get second variable + OP_CHECK_NULL_WITH_CONTEXT(context, workspaces); // check workspaces is not null + + workspaces[0] = SYS_WORKSPACE_SIZE; // default size + workspaces[0] += static_cast(groupNum * k * n * static_cast(sizeof(int8_t)) + (cvParallNum * aicNum * singleN * singleM * static_cast(sizeof(int32_t)) * EIGHT)); + if (isPerchannel) { + return ge::GRAPH_PARAM_INVALID; // continue A8W8 + } else { + return ge::GRAPH_SUCCESS; + } + } else { + const bool isShortM = avg_m < 32U; + uint32_t A8W4_MSD_BASE_M = isShortM ? 64U : 128U; + constexpr uint32_t A8W4_MSD_BASE_M_NEW = 32; + constexpr uint32_t A8W4_MSD_BASE_K = 256; + uint32_t A8W4_MSD_BASE_N = isShortM ? 512U : 256U; + constexpr uint32_t A8W4_MSD_BASE_N_NEW = 512; + mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4, false); + if (wNZ) { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::NZ, matmul_tiling::DataType::DT_INT4, false); + } else { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4, false); + } + mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + mm.SetBias(false); + if (is_in_a8w4_white_list) { + mm.SetOrgShape(A8W4_MSD_BASE_M_NEW, n, k); + mm.SetShape(A8W4_MSD_BASE_M_NEW, n, k); + mm.SetFixSplit(A8W4_MSD_BASE_M_NEW, A8W4_MSD_BASE_N_NEW, A8W4_MSD_BASE_K); + OP_LOGI(context->GetNodeName(), "GMM A8W4 tiling: baseM is %u, baseN is %u, baseK is %u, tuningConfig is %ld.", A8W4_MSD_BASE_M_NEW, A8W4_MSD_BASE_N_NEW, A8W4_MSD_BASE_K, tuningConfig); + } else { + mm.SetOrgShape(A8W4_MSD_BASE_M, n, k); + mm.SetShape(A8W4_MSD_BASE_M, A8W4_MSD_BASE_N, k); + mm.SetFixSplit(A8W4_MSD_BASE_M, A8W4_MSD_BASE_N, A8W4_MSD_BASE_K); + OP_LOGI(context->GetNodeName(), "GMM A8W4 tiling: baseM is %u, baseN is %u, baseK is %u, tuningConfig is %ld.", A8W4_MSD_BASE_M, A8W4_MSD_BASE_N, A8W4_MSD_BASE_K, tuningConfig); + } + if (mm.GetTiling(tilingDataA8W4.mmTilingData) == -1){ + return ge::GRAPH_FAILED; + } + constexpr uint32_t FOUR = 4; + if (is_in_a8w4_white_list) { + tilingDataA8W4.mmTilingData.set_dbL0B(1); // disable double buffer for LOB + tilingDataA8W4.mmTilingData.set_dbL0C(1); // disable double buffer for LOC + tilingDataA8W4.mmTilingData.set_stepKa(FOUR); // set precomputed mmStepKa + tilingDataA8W4.mmTilingData.set_stepKb(TWO); // set precomputed mmStepKb + tilingDataA8W4.mmTilingData.set_depthA1(EIGHT); // set precomputed mmDepthA1 + tilingDataA8W4.mmTilingData.set_depthB1(FOUR); // set precomputed mmDepthB1 + tilingDataA8W4.mmTilingData.set_baseK(A8W4_MSD_BASE_K); + tilingDataA8W4.mmTilingData.set_stepM(1); // set precomputed stepM + tilingDataA8W4.mmTilingData.set_stepN(1); // set precomputed stepN + } else { + tilingDataA8W4.mmTilingData.set_dbL0C(1); // disable double buffer for LOC + tilingDataA8W4.mmTilingData.set_stepKa(FOUR); // set precomputed mmStepKa + tilingDataA8W4.mmTilingData.set_stepKb(FOUR); // set precomputed mmStepKb + tilingDataA8W4.mmTilingData.set_depthA1(EIGHT); // set precomputed mmDepthA1 + tilingDataA8W4.mmTilingData.set_depthB1(EIGHT); // set precomputed mmDepthB1 + tilingDataA8W4.mmTilingData.set_stepM(1); // set precomputed stepM + tilingDataA8W4.mmTilingData.set_stepN(1); // set precomputed stepN + } + OP_LOGI(context->GetNodeName(), "GMM_tiling: baseM is %u, baseK is %u, baseN is %u.", A8W4_MSD_BASE_M, A8W4_MSD_BASE_K, A8W4_MSD_BASE_N); + context->SetScheduleMode(1); // set as batchmod for template using SyncAll + uint32_t a8w4KernelTemplate = GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_API_DEQUANT; + if (is_in_a8w4_white_list) { + a8w4KernelTemplate = static_cast(GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_VECTOR_DEQUANT); + } + context->SetTilingKey(GET_TPL_TILING_KEY(GMM_TPL_INT8, GMM_TPL_INT4, yDtype, 0, 0, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, 0, + a8w4KernelTemplate, + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE, + GROUPED_MATMUL_AIV_AIC_RATIO_2, 0)); + tilingDataA8W4.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tilingDataA8W4.GetDataSize()); + + size_t* workspaces = context->GetWorkspaceSizes(1); // get second variable + OP_CHECK_NULL_WITH_CONTEXT(context, workspaces); // check workspaces is not null + workspaces[0] = SYS_WORKSPACE_SIZE; // default size + if (is_in_a8w4_white_list) { + workspaces[0] += static_cast((cvParallNum * aicNum * A8W4_MSD_BASE_M * A8W4_MSD_BASE_N * static_cast(sizeof(short))) * TWO); + } else { + workspaces[0] += static_cast((cvParallNum * aicNum * singleN * singleM * static_cast(sizeof(int32_t)) * EIGHT)); + } + return ge::GRAPH_SUCCESS; + } + } +} + +ASCENDC_EXTERN_C ge::graphStatus TilingGMM(gert::TilingContext* context) { + OP_CHECK_NULL_WITH_CONTEXT(context, context); + auto xDesc = context->GetDynamicInputDesc(X_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); // check xDesc is not null + ge::DataType xDType = xDesc->GetDataType(); + auto w0Desc = context->GetDynamicInputDesc(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, w0Desc); + ge::DataType weightDtype = w0Desc->GetDataType(); + auto compileInfoPtr = GetGMMCompileInfo(context); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); + if (compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910_95) { + // 全量化：双8bits或双4bits(不会有A4W2) + bool isQuant = xDType == ge::DT_FLOAT4_E1M2 || xDType == ge::DT_FLOAT4_E2M1 || xDType == ge::DT_INT4 || + (ge::GetSizeByDataType(xDType) == 1 && ge::GetSizeByDataType(weightDtype) == 1); + if (isQuant) { + return TilingRegistry::GetInstance().DoTilingImpl(context); + } else if (xDType != weightDtype) { + GroupedWeightQuantBatchMatmulTiling groupedWeightQuantTiling; + OP_CHECK_IF(!groupedWeightQuantTiling.SetTiling(context), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "SetTiling failed."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + bool isUnQuant = (xDType == ge::DT_FLOAT16 || xDType == ge::DT_BF16) && (xDType == weightDtype); + if (isUnQuant) { + GroupedNoQuantMatmulTiling groupedNoQuantMatmulTiling; + OP_CHECK_IF(!groupedNoQuantMatmulTiling.SetTiling(context), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "SetTiling failed."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + } + GMMTiling tiling; + if(xDType == ge::DT_INT8 && weightDtype == ge::DT_INT4) { // A8W4 Tiling + ge::graphStatus A8W4TilingResult = tiling.A8W4Tiling(context, compileInfoPtr); + if (A8W4TilingResult != ge::GRAPH_PARAM_INVALID) { + return A8W4TilingResult; + } + } + + OP_CHECK_IF(tiling.Init(context) != ge::GRAPH_SUCCESS, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "GMM tiling init failed"), + return ge::GRAPH_FAILED); + return tiling.RunFusionKernelTiling(context); +} + +ASCENDC_EXTERN_C ge::graphStatus TilingPrepareForGMM(gert::TilingParseContext* context) { + OP_CHECK_NULL_WITH_CONTEXT(context, context); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto compileInfoPtr = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + compileInfoPtr->aicNum = ascendcPlatform.GetCoreNumAic(); + compileInfoPtr->aivNum = ascendcPlatform.GetCoreNumAiv(); + compileInfoPtr->socVersion = ascendcPlatform.GetSocVersion(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, compileInfoPtr->l1Size); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_A, compileInfoPtr->l0ASize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, compileInfoPtr->l0BSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, compileInfoPtr->l0CSize); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L2, compileInfoPtr->l2Size); + + OP_CHECK_IF((compileInfoPtr->aicNum == 0 || compileInfoPtr->aivNum == 0 || compileInfoPtr->ubSize == 0 || \ + compileInfoPtr->l1Size == 0 || compileInfoPtr->l0CSize == 0 || compileInfoPtr->l0ASize == 0 || \ + compileInfoPtr->l0BSize == 0), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), + "platform info is invalid, aicNum=%u, aivNum=%u, ubSize=%lu, l1Size=%lu, l0CSize=%lu, l0ASize=%lu, l0BSize=%lu", + compileInfoPtr->aicNum, compileInfoPtr->aivNum, compileInfoPtr->ubSize, compileInfoPtr->l1Size, + compileInfoPtr->l0CSize, compileInfoPtr->l0ASize, compileInfoPtr->l0BSize), + return ge::GRAPH_FAILED); + + OP_LOGI(context->GetNodeName(), "Parse compile info success, soc: %d", + static_cast(compileInfoPtr->socVersion)); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(DlinferGroupedMatmulDirect) +.Tiling(TilingGMM) +.TilingParse(TilingPrepareForGMM); // regist into the framework +} // namespace optiling diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.h new file mode 100644 index 00000000..200169f1 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_host/op_tiling/grouped_matmul_tiling.h @@ -0,0 +1,224 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_tiling.h + * \brief + */ +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_H + +#include +#include + +#include "../grouped_matmul_host_util.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(GMMBaseParams) +TILING_DATA_FIELD_DEF(uint32_t, groupNum); +TILING_DATA_FIELD_DEF(uint32_t, coreNum); +TILING_DATA_FIELD_DEF(uint32_t, activeType); +TILING_DATA_FIELD_DEF(uint32_t, ubBaseK); +TILING_DATA_FIELD_DEF(uint32_t, ubBaseN); +TILING_DATA_FIELD_DEF(uint32_t, ubCalSize); +TILING_DATA_FIELD_DEF(uint32_t, ubRestBytes); +TILING_DATA_FIELD_DEF(uint32_t, singleWeight); +TILING_DATA_FIELD_DEF(uint32_t, singleX); +TILING_DATA_FIELD_DEF(uint32_t, singleY); +TILING_DATA_FIELD_DEF(int32_t, groupType); +TILING_DATA_FIELD_DEF(uint32_t, singleN); // If sequential write， the value should be zero! +TILING_DATA_FIELD_DEF(uint32_t, quantParam); // in quant case, PerToken: 1; in antiquant case, represents PerGroupSize +TILING_DATA_FIELD_DEF(uint32_t, groupListType); +TILING_DATA_FIELD_DEF(uint32_t, m); +TILING_DATA_FIELD_DEF(uint32_t, hasBias); +TILING_DATA_FIELD_DEF(uint64_t, workspaceSize); +TILING_DATA_FIELD_DEF(uint64_t, totalInGroup); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, k); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, n); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, vBaseM); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, parallNum); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, quantGroupNum); // for A8W4 MSD +TILING_DATA_FIELD_DEF(uint64_t, isPreTiling); +TILING_DATA_FIELD_DEF(uint32_t, withOffset); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMBaseParamsOp, GMMBaseParams) + +BEGIN_TILING_DATA_DEF(GMMArray) +TILING_DATA_FIELD_DEF_ARR(int32_t, 128, mList); // 128 ：MAX_TENSOR_CONT +TILING_DATA_FIELD_DEF_ARR(int32_t, 128, kList); +TILING_DATA_FIELD_DEF_ARR(int32_t, 128, nList); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMArrayOp, GMMArray) + +// for autotiling w4a8 +BEGIN_TILING_DATA_DEF(A8W4HPTiling) +TILING_DATA_FIELD_DEF(uint32_t, group_num); +TILING_DATA_FIELD_DEF(int8_t, group_type); +TILING_DATA_FIELD_DEF(uint32_t, required_core_num); +TILING_DATA_FIELD_DEF(float, format_in); +TILING_DATA_FIELD_DEF(float, format_out); +TILING_DATA_FIELD_DEF(uint32_t, numAic); +TILING_DATA_FIELD_DEF(uint32_t, numAiv); +TILING_DATA_FIELD_DEF(uint64_t, szUb); +TILING_DATA_FIELD_DEF(uint64_t, szL0A); +TILING_DATA_FIELD_DEF(uint64_t, szL0C); +TILING_DATA_FIELD_DEF(uint8_t, pattern); +TILING_DATA_FIELD_DEF(uint8_t, kernel_index); +TILING_DATA_FIELD_DEF(uint32_t, splitTimes); +TILING_DATA_FIELD_DEF(int8_t, output_type); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 2, ori_in0_shape); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 2, ori_in1_shape); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 2, ori_out_shape); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 3, single_core_tiling); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 2, single_core_base_tiling); +TILING_DATA_FIELD_DEF_ARR(uint32_t, 3, splitRecord); +TILING_DATA_FIELD_DEF(uint64_t, workspaceOffset); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(A8W4HPTilingOp, A8W4HPTiling) + +BEGIN_TILING_DATA_DEF(GMMTilingData) +TILING_DATA_FIELD_DEF_STRUCT(GMMBaseParams, gmmBaseParams); +TILING_DATA_FIELD_DEF_STRUCT(GMMArray, gmmArray); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData); +// for autotiling +TILING_DATA_FIELD_DEF_STRUCT(A8W4HPTiling, hpTilingData); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(DlinferGroupedMatmulDirect, GMMTilingData) + +struct GMMCompileInfo { + uint32_t aicNum; + uint32_t aivNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l2Size; + uint64_t l0CSize; + uint64_t l0ASize; + uint64_t l0BSize; + platform_ascendc::SocVersion socVersion; +}; + +class GMMTiling { +public: + GMMTilingData tilingData; + ge::graphStatus Init(const gert::TilingContext *context); + ge::graphStatus RunFusionKernelTiling(gert::TilingContext *context); + ge::graphStatus A8W4Tiling(gert::TilingContext *context, const GMMCompileInfo *compileInfoPtr); + +protected: + bool IsAivAicRatioTwoRequired(); + bool IsFixedAxisMoveCondition(); + bool IsIntDataType(); + ge::graphStatus CalMMTiling(const gert::TilingContext *context, const GMMCompileInfo *compileInfoPtr); + ge::graphStatus GMMSetMMTiling(const gert::TilingContext *context, const GMMCompileInfo *compileInfoPtr); + ge::graphStatus GMMGetAttrs(const gert::TilingContext *context); + ge::graphStatus GMMSetUbDivideBlk(); + ge::graphStatus GMMSetUbDivideBlkAntiquant(); + ge::graphStatus GMMSetUbDivideBlkQuant(); + ge::graphStatus GMMSetUbDivideBlkA4W4(); + ge::graphStatus GMMCalUbSize(const gert::TilingContext *context, uint32_t ubSize); + int64_t GMMGetBS(const gert::Shape &xShape) const; + ge::graphStatus PrepareTilingData(const gert::TilingContext *context); + ge::graphStatus CheckWeightNZShape(const gert::TilingContext *context, int64_t numInOneBlk) const; + ge::graphStatus GMMGetTensorShapeSplitM(const gert::TilingContext *context, const gert::Shape &xShape, + const gert::Shape &wShape); + ge::graphStatus GMMGetTensorShapeSplitK(const gert::TilingContext *context, const gert::Shape &xShape, + const gert::Shape &wShape); + ge::graphStatus SplitMSingleXSingleWeightSingleY(const gert::Shape &xShape, const gert::Shape &wShape); + ge::graphStatus SplitMSingleXSeparatedWeight(const gert::TilingContext *context, const gert::Shape &xShape); + ge::graphStatus SeparatedXSeparatedWeight(const gert::TilingContext *context); + ge::graphStatus SeparatedXSingleWeight(const gert::TilingContext *context, const gert::Shape &wShape); + ge::graphStatus SplitKSingleXSingleWeightSingleY(const gert::TilingContext *context, const gert::Shape &xShape, + const gert::Shape &wShape); + ge::graphStatus SplitKSingleXSeparatedWeight(const gert::TilingContext* context, const gert::Shape &xShape, + const gert::Shape &wShape); + ge::graphStatus DivideUbAndSetWorkspace(gert::TilingContext *context, const uint32_t &aicNum); + ge::graphStatus DynamicTilingSingleN(gert::TilingContext *context, const uint32_t &aicNum, const GMMCompileInfo *compileInfoPtr); + int32_t FindBestSingleN(const uint32_t &aicNum); + bool TryFullLoadA(int32_t baseM, const GMMCompileInfo *compileInfoPtr); + void DivideUbAndSetWorkspaceAntiquant(size_t *workspaces, const uint32_t &aicNum, uint32_t &ubSize); + ge::graphStatus CalcStepKaKb(const gert::TilingContext *context, const GMMCompileInfo *compileInfoPtr, + int64_t mInMM, uint32_t &mmStepKa, uint32_t &mmStepKb); + ge::graphStatus SetBias(const gert::TilingContext *context, matmul_tiling::MultiCoreMatmulTiling &mm) const; + int32_t FindBestSingleNPertoken(const uint32_t aicNum) const; + void FindBestUsedCoreNumOneGroup(const uint32_t aicNum); + ge::graphStatus SetWorkspscesPerTokenQuant(const uint32_t aicNum, size_t *workspaces); + void SetTilingDataIsSingleTensor(); + ge::graphStatus GetPerGroupNum(const gert::TilingContext *context); + ge::graphStatus CheckMKN(const gert::TilingContext *context); + void FullLoadK(const GMMCompileInfo *compileInfoPtr); + void SetMMPreTiling(); + bool StaticTilingProcess(gert::TilingContext *context); + bool CheckTilingMatchStaticValue(); + void PrintTilingInfo(gert::TilingContext *context); + void GMMSetTplTilingKey(gert::TilingContext *context); + uint32_t GetTplDataType(const ge::DataType &dtype); +private: + int32_t mList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t kList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int32_t nList_[DlinferGroupedMatmulDirect::MAX_TENSOR_CONT] = {0}; + int64_t maxM_ = 0L; + int64_t maxN_ = 0L; + int64_t maxK_ = 0L; + int32_t minK_ = INT32_MAX; + int32_t baseM_ = 0; + int32_t baseN_ = 0; + int32_t baseK_ = 0; + uint64_t ubSize_ = 0UL; + uint32_t mmDataTypeSize_ = 0; + uint32_t ubDivideBlkNum_ = 0; + uint32_t ubIoBlkNum_ = 0; + uint32_t ubBlockAlign_ = 0; + uint64_t workspacesSize_ = 0UL; // for antiquant + uint32_t groupNum_ = 0; + bool transposeWeight_ = false; + bool transposeX_ = false; + bool isSingleWeight_ = false; + bool isSingleX_ = false; + bool isSingleY_ = false; + bool isAllSingleTensor_ = false; + bool hasBias_ = false; + int32_t groupType_ = 0; + int64_t splitItem_ = 0L; + uint32_t groupListType_ = 0; + uint32_t xKDim_ = 0; + uint32_t weightNDim_ = 0; + uint32_t weightKDim_ = 0; + uint32_t xDimNum_ = 0; + bool antiquantPerformance_ = false; + uint32_t actType_ = 0; + uint32_t usedCoreNum_ = 0; + int64_t tuningConfig_ = 0L; + int64_t tuningConfigWorkspace_ = 0L; + uint64_t FixedAxisMoveWorkspace_ = 0L; + bool isA4W4_ = false; + bool isA8W4FakeA8W8_ = false; + bool isFixedAxisMove_ = false; + uint64_t A8W4noMsdSpace_ = 0; + + ge::DataType xDType_ = ge::DT_UNDEFINED; + ge::DataType mmDType_ = ge::DT_UNDEFINED; + ge::DataType weightDtype_ = ge::DT_UNDEFINED; + ge::DataType scaleDtype_ = ge::DT_UNDEFINED; + ge::DataType perTokenScaleDtype_ = ge::DT_UNDEFINED; + ge::DataType yDtype_ = ge::DT_UNDEFINED; + bool isA8W8_ = false; + // in quant case, it indicates pertoken flag; in antiquant case, it represents pergroup size + uint32_t perTokenOrPerGroupSize_ = 0; + bool isA16W8Msd_ = false; + uint32_t totalM_ = 0; + matmul_tiling::CubeFormat wFormat_; + int32_t nzFactor_; // for weight nz format +}; +} // namespace optiling + +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/CMakeLists.txt b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/CMakeLists.txt new file mode 100644 index 00000000..fc2b3681 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/CMakeLists.txt @@ -0,0 +1,12 @@ + +# set custom compile options +if ("${CMAKE_BUILD_TYPE}x" STREQUAL "Debugx") + add_ops_compile_options(ALL OPTIONS -g -O0 --cce-ignore-always-inline=true) +endif() + +# Multi-operator should add OpName and kernel file +add_kernel_compile(DlinferGroupedMatmulDirect ${CMAKE_CURRENT_SOURCE_DIR}/dlinfer_grouped_matmul_direct.cpp) + +if (ENABLE_TEST AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/testcases) + add_subdirectory(testcases) +endif() diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/grouped_matmul_tiling_data_apt.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/grouped_matmul_tiling_data_apt.h new file mode 100644 index 00000000..2cd36e35 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/grouped_matmul_tiling_data_apt.h @@ -0,0 +1,110 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_tiling_data_apt.h + * \brief + */ +#ifndef GROUPED_MATMUL_TILING_DATA_H +#define GROUPED_MATMUL_TILING_DATA_H +#include +#include "kernel_tiling/kernel_tiling.h" + +namespace DlinferGroupedMatmulDirectTilingData { +#pragma pack(push, 8) +struct GMMArray { + // DlinferGroupedMatmulDirect::MAX_TENSOR_CONT + int32_t mList[128] = {0}; + int32_t kList[128] = {0}; + int32_t nList[128] = {0}; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMNoQuantBaseParams { + uint32_t groupNum = 0; + uint32_t coreNum = 0; + uint32_t singleWeight = 0; + uint32_t singleX = 0; + uint32_t singleY = 0; + int32_t groupType = 0; + uint32_t groupListType = 0; + uint32_t hasBias = 0; + uint32_t mTailCnt = 0; + uint32_t nTailCnt = 0; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMQuantParams { + uint32_t groupNum = 0; + uint32_t activeType = 0; + uint32_t aQuantMode = 0; + uint32_t bQuantMode = 0; + uint8_t singleX = 0; + uint8_t singleW = 0; + uint8_t singleY = 0; + int8_t groupType = 0; + uint8_t groupListType = 0; + uint8_t hasBias = 0; + uint16_t reserved = 0; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMWeightQuantParam { + uint32_t groupNum = 0; + uint32_t coreNum = 0; + uint64_t kSize = 0; + uint64_t nSize = 0; + uint8_t singleX = 0; + uint8_t singleWeight = 0; + uint8_t singleY = 0; + int8_t groupType = 0; + uint8_t groupListType = 0; + uint8_t hasBias = 0; + uint8_t cubeBlockDimN = 0; + uint8_t reserved = 0; + uint32_t groupSize = 0; + uint32_t mainBlockSize = 0; + uint64_t mainBlockCount = 0; + uint16_t firstTailBlockSize = 0; + uint16_t secondTailBlockSize = 0; + uint16_t firstTailBlockCount = 0; + uint16_t secondTailBlockCount = 0; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMQuantTilingData { + GMMQuantParams gmmQuantParams; + GMMArray gmmArray; + TCubeTiling mmTilingData; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMNoQuantTilingData { + GMMNoQuantBaseParams gmmNoQuantParam; + GMMArray gmmArray; + TCubeTiling mmTilingData; +}; +#pragma pack(pop) + +#pragma pack(push, 8) +struct GMMWeightQuantTilingData { + GMMWeightQuantParam gmmWeightQuantParam; + GMMArray gmmArray; + TCubeTiling mmTilingData; +}; +#pragma pack(pop) + +} // DlinferGroupedMatmulDirectTilingData +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_basic_kernel.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_basic_kernel.h new file mode 100644 index 00000000..f309d17e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_basic_kernel.h @@ -0,0 +1,153 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file grouped_matmul_kernel.h +* \brief +*/ + +#ifndef NON_QUANT_GROUPED_MATMUL_BASIC_KERNEL_ACT +#define NON_QUANT_GROUPED_MATMUL_BASIC_KERNEL_ACT + +#include "../../../../common/groupedmatmul_act/kernel/kernel_grouped_matmul.h" +#include "../../../../common/groupedmatmul_act/block/block_scheduler_grouped_matmul_aswt.h" +#include "../grouped_matmul_tiling_data_apt.h" +#include "../../grouped_matmul_utils.h" + +using namespace Act::Gemm; +using namespace Act::Gemm::Kernel; +using GMMNoQuantTilingData = DlinferGroupedMatmulDirectTilingData::GMMNoQuantTilingData; + +namespace GROUPED_MATMUL { + +template +__aicore__ inline void GmmNoQuantAswt(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR groupList, GM_ADDR y, GM_ADDR tiling) +{ + GET_TILING_DATA_MEMBER(GMMNoQuantTilingData, gmmNoQuantParam, gmmBaseParams_, tiling); + GET_TILING_DATA_MEMBER(GMMNoQuantTilingData, mmTilingData, mmTilingData_, tiling); + GET_TILING_DATA_MEMBER_ADDR(GMMNoQuantTilingData, gmmArray, gmmArrayAddr_, tiling); + // 定义L1和L0的TileShape + using L1TileShape = AscendC::Shape<_0, _0, _0>; + using L0TileShape = AscendC::Shape<_0, _0, _0>; + // 定义矩阵的类型和布局 + using AType = DTYPE_X; + using BType = DTYPE_X; + using CType = DTYPE_Y; + using BiasType = DTYPE_BIAS; + using LayoutA = layoutA; + using LayoutB = layoutB; + using LayoutC = layout::RowMajor; + using LayoutBias = layout::RowMajor; + // 定义scheduler类型 + using BlockScheduler = DlinferGroupedMatmulDirectAswtScheduler; + // 定义MMAD类型 + using BlockMmad = Block::BlockDlinferGroupedMatmulDirectBuilder< + AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias, + L1TileShape, L0TileShape, BlockScheduler, MatmulMultiBlockBias<>>; + // 定义BlockEpilogue类型 + using BlockEpilogue = Block::BlockEpilogueEmpty; + // 定义shape的形状，tuple保存 m n k batch + using ProblemShape = MatmulShape; + // 定义Kernel类型 + using DlinferGroupedMatmulDirectKernel = Kernel::KernelDlinferGroupedMatmulDirect; + using Params = typename DlinferGroupedMatmulDirectKernel::Params; + using GMMTiling = typename DlinferGroupedMatmulDirectKernel::GMMTiling; + GMMTiling gmmParams {gmmBaseParams_.groupNum, gmmBaseParams_.groupType, gmmBaseParams_.groupListType, + mmTilingData_.baseM, mmTilingData_.baseN, mmTilingData_.baseK, + gmmBaseParams_.singleX, gmmBaseParams_.singleWeight, gmmBaseParams_.singleY, gmmBaseParams_.hasBias, + gmmBaseParams_.mTailCnt, gmmBaseParams_.nTailCnt}; + gmmParams.matmulTiling = &mmTilingData_; + gmmParams.gmmArrayAddrIn = gmmArrayAddr_; + Params params = { + // template shape, gmm shape can not get now + {1, 1, 1, 1}, + // mmad args + {x, weight, y, bias, groupList}, + // epilogue args + {}, + // gmm tiling data + gmmParams + }; + DlinferGroupedMatmulDirectKernel op; + op(params); +} + +__aicore__ inline int32_t GetSplitValue(uint32_t groupIdx, int32_t &preOffset, + const int32_t groupType, const uint32_t groupListType, + const GlobalTensor &groupListGm) { + int32_t splitValue = 0; + if (likely(groupType != -1)) { // -1: no need to split + if (groupListType == 0) { // 0: cumsum 1: count + int32_t offset = static_cast(groupListGm.GetValue(groupIdx)); + splitValue = offset - preOffset; + preOffset = offset; + } else { + splitValue = static_cast(groupListGm.GetValue(groupIdx)); + } + } + return splitValue; +} + +template +__aicore__ inline void EmptyTensor(GM_ADDR groupListPtr, GM_ADDR y, GM_ADDR tiling) { + GET_TILING_DATA_MEMBER(GMMNoQuantTilingData, gmmNoQuantParam, gmmBaseParams, tiling); + GET_TILING_DATA_MEMBER_ADDR(GMMNoQuantTilingData, gmmArray, gmmArrayAddr, tiling); + // In the V2 interface, grouptype is -1 after host is grouped. Thus, grouptype can be either -1 or 2. + if (groupListPtr == nullptr || gmmBaseParams.groupType == 0) { + return; + } + + GlobalTensor yGm; + GlobalTensor groupListGm; + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + if (groupListPtr != nullptr) { + groupListGm.SetGlobalBuffer((__gm__ int64_t*)groupListPtr); + } + uint64_t yBaseOffset = 0; + int32_t preOffset = 0; + uint32_t singleWeight = gmmBaseParams.singleWeight; + uint32_t singleX = gmmBaseParams.singleX; + uint32_t singleY = gmmBaseParams.singleY; + bool isAllSingleTensor = singleWeight == 1 && singleX == 1 && singleY == 1; + + TILING_TYPE *ubM = gmmArrayAddr; + TILING_TYPE *ubK = gmmArrayAddr + MKN_LIST_LEN; + TILING_TYPE *ubN = gmmArrayAddr + MKN_LIST_LEN * 2; + uint64_t coreIdx = GetBlockIdx(); + uint64_t coreRation = GetTaskRation(); + if (coreRation > 1) { + coreIdx /= coreRation; + } + + for (uint32_t groupIdx = 0; groupIdx < gmmBaseParams.groupNum; ++groupIdx) { + int32_t splitValue = GetSplitValue(groupIdx, preOffset, gmmBaseParams.groupType, + gmmBaseParams.groupListType, groupListGm); + uint32_t m = isAllSingleTensor && gmmBaseParams.groupType == 2 ? *ubM : *(ubM + groupIdx); // 2: split K + uint32_t k = *ubK < 0 && gmmBaseParams.groupType == 2 ? splitValue : *(ubK + groupIdx); + uint32_t n = isAllSingleTensor ? *ubN : *(ubN + groupIdx); + + if (k == 0) { + uint32_t singleM = Ceil(m, gmmBaseParams.coreNum); + singleM = AlignUp(singleM); + uint32_t cursingleM = singleM; + if (coreIdx * singleM >= m) { + yBaseOffset += static_cast(m) * n; + continue; + } else if (m - singleM * coreIdx < singleM) { + cursingleM = m - singleM * coreIdx; + } + InitOutput(yGm[yBaseOffset + coreIdx * singleM * n], static_cast(cursingleM) * n, 0); + } + yBaseOffset += static_cast(m) * n; + } +} + +} +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_tiling_key.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_tiling_key.h new file mode 100644 index 00000000..1e4b5ebe --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/non_quant/grouped_matmul_tiling_key.h @@ -0,0 +1,48 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file grouped_matmul_tiling_key.h +* \brief +*/ + +#ifndef __OP_KERNEL_NO_QUANT_GMM_TILING_KEY_H__ +#define __OP_KERNEL_NO_QUANT_GMM_TILING_KEY_H__ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define GMM_NO_TRANS 0 +#define GMM_TRANS 1 + +// 模板参数 +ASCENDC_TPL_ARGS_DECL( + DlinferGroupedMatmulDirect, // 算子OpType + ASCENDC_TPL_UINT_DECL(NO_QUANT_B_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS, GMM_TRANS), + ASCENDC_TPL_UINT_DECL(NO_QUANT_A_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS, GMM_TRANS) + ); + +// 模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIC_ONLY), + ASCENDC_TPL_UINT_SEL(NO_QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(NO_QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), + ASCENDC_TPL_UINT_SEL(NO_QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(NO_QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIC_ONLY), + ASCENDC_TPL_UINT_SEL(NO_QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(NO_QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS)) +); + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_act_pertile_kernel.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_act_pertile_kernel.h new file mode 100644 index 00000000..4c93072c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_act_pertile_kernel.h @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gqmm_act_pertile_kernel.h + * \brief + */ +#ifndef GQMM_ACT_PERTILE_KERNEL_H +#define GQMM_ACT_PERTILE_KERNEL_H + +#include "../../../../common/groupedmatmul_act/epilogue/block_epilogue_pertile.h" +#include "../../../../common/groupedmatmul_act/block/block_mmad_pertile.h" +#include "../../../../common/groupedmatmul_act/block/block_scheduler_gmm_aswt_with_tail_split.h" +#include "../../../../common/groupedmatmul_act/block/block_scheduler_policy.h" +#include "../../../../common/groupedmatmul_act/kernel/kernel_qgmm_pertile.h" +#include "../../../../common/groupedmatmul_act/policy/dispatch_policy.h" +#include "../../grouped_matmul_utils.h" +#include "../grouped_matmul_tiling_data_apt.h" +#include "quant_utils.h" +using GMMQuantParams = DlinferGroupedMatmulDirectTilingData::GMMQuantParams; + +template +__aicore__ inline void GmmActPerTileKernel(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMQuantParams *gmmBaseParamsIn, const TCubeTiling *mmTilingDataIn, + AscendC::TPipe *que) +{ + // 定义L1和L0的TileShape + using L1TileShape = AscendC::Shape; + using L0TileShape = AscendC::Shape; + + // 定义矩阵的类型和布局 + using AType = xType; + using BType = wType; + using CType = l0cType; + using BiasType = biasType; + using ScaleType = scaleType; + using PtScaleType = ptScaleType; + using YType = yType; + + using LayoutA = xLayout; + using LayoutB = wLayout; + using LayoutC = yLayout; + using LayoutY = yLayout; + using LayoutBias = yLayout; + + // 定义shape的形状，tuple保存 m n k batch + using ProblemShape = Act::Gemm::MatmulShape; + + // 定义scheduler类型 + using BlockScheduler = Act::Gemm::DlinferGroupedMatmulDirectAswtWithTailSplitScheduler; + + // 定义MMAD类型 + using BlockMmadPolicy = Act::Gemm::GMMPerTile<>; + using BlockMmad = Act::Gemm::Block::BlockMmadGmm; + + // 定义BlockEpilogue类型 + using BlockEpilogue = Act::Gemm::Block::BlockEpiloguePerTile; + + // 定义Kernel类型 + using GmmKernel = Act::Gemm::Kernel::QuantMmGroupedPerTile; + using Params = typename GmmKernel::Params; + using GMMTiling = typename GmmKernel::GMMTiling; + GMMTiling gmmParams{mmTilingDataIn->M, + mmTilingDataIn->N, + mmTilingDataIn->Ka, + mmTilingDataIn->baseM, + mmTilingDataIn->baseN, + mmTilingDataIn->baseK, + mmTilingDataIn->stepM, + mmTilingDataIn->stepN, + mmTilingDataIn->stepKa, + mmTilingDataIn->stepKb, + gmmBaseParamsIn->groupNum, + gmmBaseParamsIn->groupType, + gmmBaseParamsIn->groupListType}; + // G-B only support S-S-S + Params params = {{1, 1, 1, 1}, // shape + {x, weight, y, bias, groupList}, // gm addr + {(GM_ADDR)GROUPED_MATMUL::GetTensorAddr(0, y), + (GM_ADDR)GROUPED_MATMUL::GetTensorAddr(0, scale), perTokenScale, nullptr, + static_cast(mmTilingDataIn->baseM), static_cast(mmTilingDataIn->baseN), + static_cast(mmTilingDataIn->baseK), 1, QuantUtils::PER_BLOCK_SIZE, + QuantUtils::PER_BLOCK_SIZE}, // epilogue args + gmmParams}; + GmmKernel gmm; + gmm(params); +} +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_cube_on_the_fly.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_cube_on_the_fly.h new file mode 100644 index 00000000..e8bf6680 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_cube_on_the_fly.h @@ -0,0 +1,347 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gqmm_cube_on_the_fly.h + * \brief + */ +#ifndef GQMM_CUBE_ON_THE_FLY_H +#define GQMM_CUBE_ON_THE_FLY_H + +#include "quant_block_sch.h" +#include "quant_utils.h" +#include "../../grouped_matmul_utils.h" +#include "../grouped_matmul_tiling_data_apt.h" +using GMMQuantParams = DlinferGroupedMatmulDirectTilingData::GMMQuantParams; + +namespace AscendC { + +constexpr uint64_t DEQ_SCALE_MUL = 0xFFFFE000; + +LOCAL_TEMPLATE_CLASS_PARAMS +class GmmASWKernel { +public: + __aicore__ inline GmmASWKernel() {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, const GMMQuantParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn, TPipe *que); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void InitAddrAndParams(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, TILING_TYPE *gmmArrayAddrIn); + __aicore__ inline void UpdateMMGlobalAddr(uint32_t groupIdx); + __aicore__ inline void SetMNK(uint32_t groupIdx, int32_t &mSize, int32_t &nSize, int32_t &kSize); + __aicore__ inline void CalcTailTile(uint64_t mTail, uint64_t nTail); + __aicore__ inline void SetMMParaAndCompute(); + __aicore__ inline bool IsLastGroupAndNeedSplit(uint32_t groupIdx); + __aicore__ inline bool IsLastGroupAndRound(uint32_t groupIdx, uint64_t roundIdx); + + uint32_t blockIdx_; + uint32_t groupNum_; + int8_t groupType_; + uint8_t groupListType_; + int32_t preOffset_; + uint64_t scaleScalar_; + const TCubeTiling* __restrict mmTilingData_; + const GMMQuantParams* __restrict gmmQuantParams_; + + TILING_TYPE* mListGm_; + TILING_TYPE* kListGm_; + TILING_TYPE* nListGm_; + + GlobalTensor xGlobal_; + GlobalTensor wGlobal_; + GlobalTensor yGlobal_; + GlobalTensor biasGlobal_; + GlobalTensor groupListGlobal_; + GlobalTensor scaleAGlobal_; + GlobalTensor mxScaleBGlobal_; + GlobalTensor scaleBGlobal_; + + // dynamic输入输出需要地址，再按照group依次读取二级指针 + GM_ADDR xTensorPtr_; + GM_ADDR weightTensorPtr_; + GM_ADDR biasTensorPtr_; + GM_ADDR scaleTensorPtr_; + GM_ADDR groupListPtr_; + GM_ADDR perTokenScalePtr_; + GM_ADDR yTensorPtr_; + DlinferGroupedMatmulDirect::QuantASWBlockSch block_; + + // mxType定义 + using aType = typename AscendC::Conditional< + QuantUtils::IsMxType(), + matmul::MatmulTypeWithScale, + matmul::MatmulType>::type; + using bType = typename AscendC::Conditional< + QuantUtils::IsMxType(), + matmul::MatmulTypeWithScale, + matmul::MatmulType>::type; + using cType = matmul::MatmulType; + using biasMatmulType = matmul::MatmulType; + using MmType = typename AscendC::Conditional< + QuantUtils::IsMxType(), + matmul::MatmulImpl, AscendC::Impl::Detail::MatmulWithScalePolicy>, + matmul::MatmulImpl>::type; + MmType mm_; +}; + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, + const GMMQuantParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, + TILING_TYPE* gmmArrayAddrIn, TPipe *que) +{ + if ASCEND_IS_AIV { + return; + } + mmTilingData_ = mmTilingDataIn; + gmmQuantParams_ = gmmBaseParamsIn; + blockIdx_ = GetBlockIdx(); + mm_.SetSubBlockIdx(0); + mm_.Init(mmTilingData_, que); + InitAddrAndParams(x, weight, bias, scale, groupList, perTokenScale, y, gmmArrayAddrIn); +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::InitAddrAndParams(GM_ADDR x, GM_ADDR weight, + GM_ADDR bias, GM_ADDR scale, + GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, + TILING_TYPE *gmmArrayAddrIn) +{ + groupNum_ = gmmQuantParams_->groupNum; + groupType_ = gmmQuantParams_->groupType; + groupListType_ = gmmQuantParams_->groupListType; + block_.template Init(mmTilingData_, blockIdx_); + xTensorPtr_ = x; + weightTensorPtr_ = weight; + biasTensorPtr_ = bias; + scaleTensorPtr_ = scale; + groupListPtr_ = groupList; + perTokenScalePtr_ = perTokenScale; + yTensorPtr_ = y; + if constexpr (QuantUtils::IsMxType()) { + scaleAGlobal_.SetGlobalBuffer((__gm__ fp8_e8m0_t*)perTokenScale); + } + if (groupList != nullptr) { + groupListGlobal_.SetGlobalBuffer((__gm__ int64_t*)groupList); + } + mListGm_ = gmmArrayAddrIn; + kListGm_ = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN; + nListGm_ = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN * 2; // 2: mListGm_ + kListGm_ +} + +// 更新每个group的global地址 +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::UpdateMMGlobalAddr(uint32_t groupIdx) +{ + if constexpr (QuantUtils::IsMxType()) { + mxScaleBGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, scaleTensorPtr_) + + block_.params_.wScaleGroupAddrOffset); + } else { + __gm__ scaleType* scaleB = GROUPED_MATMUL::GetTensorAddr(0, scaleTensorPtr_) + groupIdx; + if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE) && + gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { // doubleScale, M_SPLIT + float scaleBValue = *((__gm__ float *)scaleB); + float scaleAValue = *((__gm__ float *)perTokenScalePtr_ + groupIdx); + float deqScale = scaleBValue * scaleAValue; + uint32_t uint32Scale = *(reinterpret_cast(&deqScale)); + scaleScalar_ = uint32Scale & DEQ_SCALE_MUL; // fixpipe只能取高19位 + } else if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::DEFAULT) && + gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { // pertensor, M_SPLIT + if constexpr (!IsSameType::value && !IsSameType::value) { + uint32_t uint32Scale = 0; + if constexpr (IsSameType::value) { + uint16_t uint16Scale = *((__gm__ uint16_t *)scaleB); + // 16为将uint16的数据前移16位，依据硬件逻辑，将scale变为uint32类型 + uint32Scale = uint16Scale << 16; + } else { + uint32Scale = *((__gm__ uint32_t *)scaleB); + } + scaleScalar_ = uint32Scale & DEQ_SCALE_MUL; + } else { + scaleScalar_ = *((__gm__ uint64_t*)scaleB); + } + } else if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERCHANNEL_MODE)) { // perChannel, M_SPLIT + scaleBGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, scaleTensorPtr_) + + block_.params_.wScaleGroupAddrOffset); + } + } + xGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, xTensorPtr_) + block_.params_.aGroupAddrOffset); + wGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, weightTensorPtr_) + + block_.params_.bGroupAddrOffset); + yGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, yTensorPtr_) + block_.params_.cGroupAddrOffset); + if (gmmQuantParams_->hasBias) { + biasGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, biasTensorPtr_) + + block_.params_.biasGroupAddrOffset); + } +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::SetMNK(uint32_t groupIdx, int32_t &mSize, + int32_t &nSize, int32_t &kSize) +{ + int32_t splitValue = + QuantUtils::GetSplitValueFromGroupList(groupIdx, preOffset_, groupType_, groupListType_, groupListGlobal_); + switch (groupType_) { + case (QuantUtils::SPLIT_M): { + mSize = splitValue; + uint32_t valueIdx = gmmQuantParams_->singleW == 1 ? 0 : groupIdx; + kSize = kListGm_[valueIdx]; + nSize = nListGm_[valueIdx]; + break; + } + case (QuantUtils::SPLIT_K): { + mSize = gmmQuantParams_->singleX == 1 ? mListGm_[0] : mListGm_[groupIdx]; + kSize = splitValue; + nSize = gmmQuantParams_->singleW == 1 ? nListGm_[0] : nListGm_[groupIdx]; + break; + } + default: { + mSize = mListGm_[groupIdx]; + kSize = kListGm_[groupIdx]; + nSize = nListGm_[groupIdx]; + } + } + return; +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::CalcTailTile(uint64_t mTail, uint64_t nTail) +{ + // 计算实际 base 块数 + uint64_t tailCnt = block_.GetEndBlockIdx() + 1; + // 计算可切分数 + uint64_t remainTile = mmTilingData_->usedCoreNum / tailCnt; + if (remainTile <= 1) { + return; + } + + // 初始化最小 tile 大小 + uint64_t mMin = QuantUtils::CUBE_BLOCK; + uint64_t nMin = QuantUtils::CUBE_BLOCK; + + // 根据矩阵是否转置调整最小 tile 大小 + if constexpr (aTrans) { + mMin = QuantUtils::INNER_AXIS_MIN_SPLIT_VAL; + } + if constexpr (!bTrans) { + nMin = QuantUtils::INNER_AXIS_MIN_SPLIT_VAL; + } + + // 计算 mTile 和 nTile，尽可能让m,n方向切分数一致 + uint64_t mTile = GROUPED_MATMUL::Min(QuantUtils::CeilDiv(mTail, mMin), remainTile); + uint64_t nTile = GROUPED_MATMUL::Min(QuantUtils::CeilDiv(nTail, nMin), remainTile); + while (mTile * nTile > remainTile) { + if (mTile >= nTile) { + mTile -= 1; + } else { + nTile -= 1; + } + } + block_.params_.mTailTile = mTile; + block_.params_.nTailTile = nTile; +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline bool GmmASWKernel::IsLastGroupAndNeedSplit(uint32_t groupIdx) +{ + // 2: 剩一半及以上核数时才考虑尾块切分 + return groupIdx == groupNum_ - 1 && (block_.GetEndBlockIdx() + 1) <= mmTilingData_->usedCoreNum / 2; +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline bool GmmASWKernel::IsLastGroupAndRound(uint32_t groupIdx, + uint64_t roundIdx) +{ + return groupIdx == groupNum_ - 1 && roundIdx == block_.params_.round - 1 && blockIdx_ <= block_.GetEndBlockIdx(); +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::Process() +{ + if ASCEND_IS_AIV { + return; + } + if (groupType_ != -1) { // -1: no split + if (unlikely(groupListPtr_ == nullptr)) { + return; + } + preOffset_ = 0; + } + + for (uint32_t groupIdx = 0; groupIdx < groupNum_; ++groupIdx) { + int32_t mSize; + int32_t nSize; + int32_t kSize; + // 更新group内的输入参数M,N,K + SetMNK(groupIdx, mSize, nSize, kSize); + block_.template UpdateGroupOffset(mSize, nSize, kSize, groupIdx); + if (mSize <= 0 || kSize <= 0 || nSize <= 0) { + continue; + } + block_.template UpdateGroupParams(); + // 最后一个group最后一轮是否进一步切分以使用更多的核数 + if (IsLastGroupAndNeedSplit(groupIdx)) { + CalcTailTile(block_.params_.mBaseTail, block_.params_.nBaseTail); + block_.UpdateTailTile(); + } + + UpdateMMGlobalAddr(groupIdx); + for (uint64_t roundIdx = 0; roundIdx < block_.params_.round; ++roundIdx) { + bool isLastGroupRound = IsLastGroupAndRound(groupIdx, roundIdx); + block_.template UpdateBasicIndex(roundIdx, isLastGroupRound); + // 1. Set single core param + block_.template UpdateBlockParams(roundIdx, isLastGroupRound); + if (block_.params_.singleCoreM <= 0 || block_.params_.singleCoreN <= 0) { + continue; + } + // 2. compute offset + block_.template CalcGMOffset(); + // 3. set offset and compute + SetMMParaAndCompute(); + } + } +} + +LOCAL_TEMPLATE_CLASS_PARAMS +__aicore__ inline void GmmASWKernel::SetMMParaAndCompute() +{ + if ASCEND_IS_AIV { + return; + } + mm_.SetSingleShape(block_.params_.singleCoreM, block_.params_.singleCoreN, block_.params_.k); + if constexpr (QuantUtils::IsMxType()) { + mm_.SetTensorScaleA(scaleAGlobal_[block_.offset_.offsetPerTokenScale], aTrans); + mm_.SetTensorScaleB(mxScaleBGlobal_[block_.offset_.offsetScale], bTrans); + } else { + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { // perTensor && doubleScale + mm_.SetQuantScalar(scaleScalar_); + } else { + mm_.SetQuantVector(scaleBGlobal_[block_.offset_.offsetScale]); + } + } + if (gmmQuantParams_->hasBias) { + mm_.SetBias(biasGlobal_[block_.offset_.offsetBias]); + } + mm_.SetTensorA(xGlobal_[block_.offset_.offsetA], aTrans); + mm_.SetTensorB(wGlobal_[block_.offset_.offsetB], bTrans); + mm_.Iterate(); + mm_.GetTensorC(yGlobal_[block_.offset_.offsetC]); +} +} +#endif // GQMM_CUBE_ON_THE_FLY_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_init_output.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_init_output.h new file mode 100644 index 00000000..4a0f8339 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_init_output.h @@ -0,0 +1,74 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gqmm_init_output.h + * \brief + */ +#ifndef GQMM_INIT_OUTPUT_H +#define GQMM_INIT_OUTPUT_H + +#include "quant_utils.h" +#include "../grouped_matmul_tiling_data_apt.h" +using GMMQuantParams = DlinferGroupedMatmulDirectTilingData::GMMQuantParams; + +namespace AscendC { +template +__aicore__ inline void GQmmEmptyTensor(GM_ADDR groupListPtr, GM_ADDR y, const GMMQuantParams *__restrict gmmQuantParams, + TILING_TYPE *gmmArrayAddrIn, int32_t usedCoreNum, TPipe *pipe) +{ + if (GetSubBlockIdx() > 1) { + return; + } + // 只有K轴分组k=0时才需要创建空tensor + if (groupListPtr == nullptr || gmmQuantParams->groupType != QuantUtils::SPLIT_K) { + return; + } + + GlobalTensor yGm; + GlobalTensor groupListGm; + TBuf initBuff; + LocalTensor initLocal; + yGm.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, y)); + groupListGm.SetGlobalBuffer((__gm__ int64_t*)groupListPtr); + if (GetSubBlockIdx() == 0) { + pipe->InitBuffer(initBuff, QuantUtils::MAX_REPEAT_TIMES * QuantUtils::UB_ALIGN_SIZE); + initLocal = initBuff.Get(); + } + uint64_t yBaseOffset = 0; + int32_t preOffset = 0; + bool isSingleX = static_cast(gmmQuantParams->singleX); + bool isSingleW = static_cast(gmmQuantParams->singleW); + + TILING_TYPE *mList = gmmArrayAddrIn; + TILING_TYPE *kList = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN; + TILING_TYPE *nList = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN * 2; // 2: mListGm_ + kListGm_ + + bool isKZeroInit = false; + for (uint32_t groupIdx = 0; groupIdx < gmmQuantParams->groupNum; ++groupIdx) { + int32_t splitValue = QuantUtils::GetSplitValueFromGroupList(groupIdx, preOffset, gmmQuantParams->groupType, + gmmQuantParams->groupListType, groupListGm); + int32_t mSize = isSingleX ? mList[0] : mList[groupIdx]; + int32_t kSize = splitValue; + int32_t nSize = isSingleW ? nList[0] : nList[groupIdx]; + if (mSize <= 0 || nSize <= 0) { + continue; + } + uint64_t ySize = static_cast(mSize) * nSize; + if (kSize == 0) { + QuantUtils::InitOutputWithZero(yGm[yBaseOffset], initLocal, ySize, usedCoreNum, isKZeroInit); + } + yBaseOffset += ySize; + } +} + +} // namespace GROUPED_MATMUL + +#endif // GQMM_INIT_OUTPUT_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_mix_online_dynamic.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_mix_online_dynamic.h new file mode 100644 index 00000000..7dc7dd80 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_mix_online_dynamic.h @@ -0,0 +1,739 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file gqmm_mix_online_dynamic.h + * \brief + */ +#ifndef GQMM_MIX_ONLINE_DYNAMIC_H +#define GQMM_MIX_ONLINE_DYNAMIC_H + +#include "mm_extension_interface/gqmm_custom_mm_policy.h" +#include "quant_block_sch.h" +#include "quant_utils.h" +#include "../../grouped_matmul_utils.h" +#include "../grouped_matmul_tiling_data_apt.h" +using GMMQuantParams = DlinferGroupedMatmulDirectTilingData::GMMQuantParams; + +#define LOCAL_TEMPLATE_CLASS_MIX_PARAMS \ + template +#define LOCAL_TEMPLATE_FUNC_MIX_PARAMS \ + xType, wType, biasType, scaleType, ptScaleType, yType, wFormat, aTrans, bTrans, l0cType + +namespace AscendC { + +constexpr MicroAPI::CastTrait ctInt322Fp32 = {MicroAPI::RegLayout::UNKNOWN, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + +constexpr MicroAPI::CastTrait ctFp322Half = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + +constexpr MicroAPI::CastTrait ctHalf2Fp32Zero = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +constexpr MicroAPI::CastTrait ctHalf2Fp32One = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +class GQmmMixRegbaseKernel { +public: + __aicore__ inline GQmmMixRegbaseKernel() {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMQuantParams *__restrict gmmBaseParamsIn, + const TCubeTiling *__restrict mmTilingDataIn, TILING_TYPE *gmmArrayAddrIn, TPipe *pipe); + __aicore__ inline void Process(); + +public: + using aT = MatmulType; + using bT = MatmulType; + using biasT = MatmulType; + using cT = MatmulType; + MatmulImpl, GQmmCustomMatmulPolicy> + mm; + +protected: + __aicore__ inline void InitAddrAndParams(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, TILING_TYPE *gmmArrayAddrIn); + __aicore__ inline void UpdateMMGlobalAddr(uint32_t groupIdx); + __aicore__ inline void SetMNK(uint32_t groupIdx, int32_t &mSize, int32_t &nSize, int32_t &kSize); + __aicore__ inline void CalcTailTile(uint64_t mTail, uint64_t nTail); + __aicore__ inline bool IsLastGroupAndNeedSplit(uint32_t groupIdx); + __aicore__ inline bool IsLastGroupAndRound(uint32_t groupIdx, uint64_t roundIdx); + __aicore__ inline void ProcessSingleGroup(uint32_t groupIdx); + __aicore__ inline void MMCompute(); + __aicore__ inline void DequantCompute(); + __aicore__ inline void End(); + __aicore__ inline void VFDoDequantWithAPertoken(__ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, + uint64_t offsetPtScale, uint16_t mSize); + __aicore__ inline void VFDoDequantWithAPertensor(__ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, + uint16_t mSize); + __aicore__ inline void VFDoDequantWithoutAScale(__ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, + uint16_t mSize); + template + __aicore__ inline void VFDoDequant(__ubuf__ yType *dst, __ubuf__ l0cType *l0cOut, __ubuf__ scaleType *scale, + __ubuf__ ptScaleType *perTokenScale, __ubuf__ biasType *bias, uint16_t mSize, + uint16_t nSize); + __aicore__ inline void NotifyCube() + { + CrossCoreSetFlag(QuantUtils::AIV_SYNC_AIC_FLAG); + } + __aicore__ inline void WaitForVector() + { + CrossCoreWaitFlag(QuantUtils::AIV_SYNC_AIC_FLAG); + CrossCoreWaitFlag(QuantUtils::AIV_SYNC_AIC_FLAG + + QuantUtils::FLAG_ID_MAX); + } + __aicore__ inline void NotifyVector() + { + CrossCoreSetFlag(QuantUtils::AIC_SYNC_AIV_FLAG); + CrossCoreSetFlag(QuantUtils::AIC_SYNC_AIV_FLAG + + QuantUtils::FLAG_ID_MAX); + } + __aicore__ inline void WaitForCube() + { + CrossCoreWaitFlag(QuantUtils::AIC_SYNC_AIV_FLAG); + } + __aicore__ inline void CopyDataFromGm2Ub(); + __aicore__ inline void CopyX1ScaleFromGm2Ub(LocalTensor &dst, uint64_t blockLen, uint64_t offset); + __aicore__ inline void CopyX2ScaleFromGm2Ub(LocalTensor &dst); + __aicore__ inline void CopyBiasFromGm2Ub(LocalTensor &dst); + __aicore__ inline void CopyDequantResFromUb2Gm(uint64_t blockCount, uint64_t offset, LocalTensor &src); + __aicore__ inline void FreeUbTensor(); + +protected: + uint32_t blockIdx_; + uint32_t subBlockIdx_; + int32_t preOffset_; + float scaleScalar_; + float perTokenScaleScalar_; + uint32_t groupNum_; + int8_t groupType_; + uint8_t groupListType_; + bool isVecSetSyncCom_ = false; + bool isBiasEpilogue_; + DlinferGroupedMatmulDirect::QuantASWBlockSch block_; + + const TCubeTiling *__restrict mmTilingData_; + const GMMQuantParams *__restrict gmmQuantParams_; + + TILING_TYPE *mListGm_; + TILING_TYPE *kListGm_; + TILING_TYPE *nListGm_; + + GlobalTensor xGlobal_; + GlobalTensor wGlobal_; + GlobalTensor yGlobal_; + GlobalTensor biasGlobal_; + GlobalTensor scaleBGlobal_; + GlobalTensor groupListGlobal_; + GlobalTensor scaleAGlobal_; + + // dynamic输入输出需要地址，再按照group依次读取二级指针 + GM_ADDR xTensorPtr_; + GM_ADDR weightTensorPtr_; + GM_ADDR biasTensorPtr_; + GM_ADDR scaleTensorPtr_; + GM_ADDR groupListPtr_; + GM_ADDR perTokenScalePtr_; + GM_ADDR yTensorPtr_; + + LocalTensor l0cOutUb_; + LocalTensor scaleUb_; + LocalTensor ptScaleUb_; + LocalTensor biasUb_; + LocalTensor initLocal_; + + TQue vecQueMMRes_; + TQue vecQueScale_; + TQue vecQuePertokenScale_; + TQue vecQueBias_; + TQue vecQueOut_; + TBuf initBuff_; +}; + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::Init( + GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, const GMMQuantParams *__restrict gmmBaseParamsIn, const TCubeTiling *__restrict mmTilingDataIn, + TILING_TYPE *gmmArrayAddrIn, TPipe *pipe) +{ + blockIdx_ = GetBlockIdx(); + if ASCEND_IS_AIV { + blockIdx_ = blockIdx_ / GetTaskRation(); + subBlockIdx_ = GetSubBlockIdx(); + } + mmTilingData_ = mmTilingDataIn; + gmmQuantParams_ = gmmBaseParamsIn; + mm.SetSubBlockIdx(0); + mm.Init(mmTilingData_, pipe); + InitAddrAndParams(x, weight, bias, scale, groupList, perTokenScale, y, gmmArrayAddrIn); + + uint32_t mForSingleVec = QuantUtils::CeilDiv(mmTilingData_->baseM, GetTaskRation()); + pipe->InitBuffer(vecQueMMRes_, 1, mForSingleVec * mmTilingData_->baseN * sizeof(float)); // int32/fp32 + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERCHANNEL_MODE)) { + pipe->InitBuffer(vecQueScale_, 1, mmTilingData_->baseN * sizeof(scaleType)); + } + if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::PERTOKEN_MODE)) { + pipe->InitBuffer(vecQuePertokenScale_, 1, + QuantUtils::Align(mForSingleVec * sizeof(ptScaleType), QuantUtils::UB_ALIGN_SIZE)); + } + isBiasEpilogue_ = QuantUtils::IsBiasEpilogue() && gmmQuantParams_->hasBias == 1; + if (isBiasEpilogue_) { + pipe->InitBuffer(vecQueBias_, 1, mmTilingData_->baseN * sizeof(biasType)); + } + + // fp16/bf16分两次输出，fp32分四次输出 + pipe->InitBuffer(vecQueOut_, QuantUtils::BUFFER_SWITCH, + QuantUtils::CeilDiv(mForSingleVec, QuantUtils::FP32_OUTPUT_TIMES) * mmTilingData_->baseN * + sizeof(yType)); + l0cOutUb_ = vecQueMMRes_.AllocTensor(); + + // k = 0, init out + if ASCEND_IS_AIV { + if (subBlockIdx_ == 0) { + pipe->InitBuffer(initBuff_, QuantUtils::MAX_REPEAT_TIMES * QuantUtils::UB_ALIGN_SIZE); + initLocal_ = initBuff_.Get(); + } + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::InitAddrAndParams( + GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + TILING_TYPE *gmmArrayAddrIn) +{ + groupNum_ = gmmQuantParams_->groupNum; + groupType_ = gmmQuantParams_->groupType; + groupListType_ = gmmQuantParams_->groupListType; + block_.template Init(mmTilingData_, blockIdx_); + xTensorPtr_ = x; + weightTensorPtr_ = weight; + biasTensorPtr_ = bias; + scaleTensorPtr_ = scale; + groupListPtr_ = groupList; + perTokenScalePtr_ = perTokenScale; + yTensorPtr_ = y; + scaleAGlobal_.SetGlobalBuffer((__gm__ ptScaleType *)perTokenScale); + if (groupList != nullptr) { + groupListGlobal_.SetGlobalBuffer((__gm__ int64_t *)groupList); + } + mListGm_ = gmmArrayAddrIn; + kListGm_ = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN; + nListGm_ = gmmArrayAddrIn + GROUPED_MATMUL::MKN_LIST_LEN * 2; // 2: mListGm_ + kListGm_ +} + +// 更新每个group的dynamic global地址 +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::UpdateMMGlobalAddr(uint32_t groupIdx) +{ + xGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, xTensorPtr_) + block_.params_.aGroupAddrOffset); + wGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, weightTensorPtr_) + + block_.params_.bGroupAddrOffset); + if (static_cast(gmmQuantParams_->hasBias)) { + biasGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, biasTensorPtr_) + + block_.params_.biasGroupAddrOffset); + } + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { + scaleType scaleBValue = + *((__gm__ scaleType *)(GROUPED_MATMUL::GetTensorAddr(0, scaleTensorPtr_) + groupIdx)); + if constexpr (IsSameType::value) { + scaleScalar_ = ToFloat(scaleBValue); + } else { + scaleScalar_ = scaleBValue; + } + } else { + scaleBGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, scaleTensorPtr_) + + block_.params_.wScaleGroupAddrOffset); + } + if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { + perTokenScaleScalar_ = *((__gm__ ptScaleType *)perTokenScalePtr_ + groupIdx); + } + yGlobal_.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, yTensorPtr_) + block_.params_.cGroupAddrOffset); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::SetMNK(uint32_t groupIdx, int32_t &mSize, + int32_t &nSize, int32_t &kSize) +{ + int32_t splitValue = + QuantUtils::GetSplitValueFromGroupList(groupIdx, preOffset_, groupType_, groupListType_, groupListGlobal_); + switch (groupType_) { + case (QuantUtils::SPLIT_M): + { + mSize = splitValue; + uint32_t valueIdx = gmmQuantParams_->singleW == 1 ? 0 : groupIdx; + kSize = kListGm_[valueIdx]; + nSize = nListGm_[valueIdx]; + break; + } + case (QuantUtils::SPLIT_K): + { + mSize = gmmQuantParams_->singleX == 1 ? mListGm_[0] : mListGm_[groupIdx]; + kSize = splitValue; + nSize = gmmQuantParams_->singleW == 1 ? nListGm_[0] : nListGm_[groupIdx]; + break; + } + default: + { + mSize = mListGm_[groupIdx]; + kSize = kListGm_[groupIdx]; + nSize = nListGm_[groupIdx]; + } + } + return; +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::CalcTailTile(uint64_t mTail, + uint64_t nTail) +{ + // 计算实际 base 块数 + uint64_t tailCnt = block_.GetEndBlockIdx() + 1; + // 计算可切分数 + uint64_t remainTile = mmTilingData_->usedCoreNum / tailCnt; + if (remainTile <= 1) { + return; + } + + // 初始化最小 tile 大小 + uint64_t mMin = QuantUtils::CUBE_BLOCK; + uint64_t nMin = QuantUtils::CUBE_BLOCK; + + // 根据矩阵是否转置调整最小 tile 大小 + if constexpr (aTrans) { + mMin = QuantUtils::INNER_AXIS_MIN_SPLIT_VAL; + } + if constexpr (!bTrans) { + nMin = QuantUtils::INNER_AXIS_MIN_SPLIT_VAL; + } + + // 计算 mTile 和 nTile，尽可能让m,n方向切分数一致 + uint64_t mTile = GROUPED_MATMUL::Min(QuantUtils::CeilDiv(mTail, mMin), remainTile); + uint64_t nTile = GROUPED_MATMUL::Min(QuantUtils::CeilDiv(nTail, nMin), remainTile); + while (mTile * nTile > remainTile) { + if (mTile >= nTile) { + mTile -= 1; + } else { + nTile -= 1; + } + } + block_.params_.mTailTile = mTile; + block_.params_.nTailTile = nTile; +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline bool GQmmMixRegbaseKernel::IsLastGroupAndNeedSplit(uint32_t groupIdx) +{ + // 2: 剩一半及以上核数时才考虑尾块切分 + return groupIdx == groupNum_ - 1 && (block_.GetEndBlockIdx() + 1) <= mmTilingData_->usedCoreNum / 2; +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline bool GQmmMixRegbaseKernel::IsLastGroupAndRound(uint32_t groupIdx, + uint64_t roundIdx) +{ + return groupIdx == groupNum_ - 1 && roundIdx == block_.params_.round - 1 && blockIdx_ <= block_.GetEndBlockIdx(); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::Process() +{ + if (groupType_ != -1) { // -1: no split + if (unlikely(groupListPtr_ == nullptr)) { + return; + } + } + preOffset_ = 0; + bool isKZeroInit = false; + for (uint32_t groupIdx = 0; groupIdx < groupNum_; ++groupIdx) { + int32_t mSize; + int32_t nSize; + int32_t kSize; + // 更新group内的输入参数M,N,K + SetMNK(groupIdx, mSize, nSize, kSize); + block_.template UpdateGroupOffset(mSize, nSize, kSize, groupIdx); + if (mSize <= 0 || nSize <= 0) { + continue; + } + if (kSize <= 0) { + if (groupType_ == QuantUtils::SPLIT_K) { // k轴分组时，有(m,n)需要输出。int8输入无K轴分组。K轴分组无bias全0 + GlobalTensor yInitGlobal; + yInitGlobal.SetGlobalBuffer(GROUPED_MATMUL::GetTensorAddr(0, yTensorPtr_) + + block_.params_.cGroupAddrOffset); + QuantUtils::InitOutputWithZero(yInitGlobal, initLocal_, static_cast(mSize) * nSize, + mmTilingData_->usedCoreNum, isKZeroInit); + } + continue; + } + block_.template UpdateGroupParams(); + // 最后一个group最后一轮是否进一步切分以使用更多的核数 + if (IsLastGroupAndNeedSplit(groupIdx)) { + CalcTailTile(block_.params_.mBaseTail, block_.params_.nBaseTail); + block_.UpdateTailTile(); + } + + UpdateMMGlobalAddr(groupIdx); + ProcessSingleGroup(groupIdx); + } + End(); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::ProcessSingleGroup(uint32_t groupIdx) +{ + for (uint64_t roundIdx = 0; roundIdx < block_.params_.round; ++roundIdx) { + bool isLastGroupRound = IsLastGroupAndRound(groupIdx, roundIdx); + block_.template UpdateBasicIndex(roundIdx, isLastGroupRound); + // 1. Set single core param + block_.template UpdateBlockParams(roundIdx, isLastGroupRound); + if (block_.params_.singleCoreM <= 0 || block_.params_.singleCoreN <= 0) { + return; + } + mm.SetSingleShape(block_.params_.singleCoreM, block_.params_.singleCoreN, block_.params_.k); + block_.template CalcGMOffset(); + if ASCEND_IS_AIC { + if (isVecSetSyncCom_) { + WaitForVector(); + } + MMCompute(); + NotifyVector(); + } + isVecSetSyncCom_ = true; + if ASCEND_IS_AIV { + WaitForCube(); + DequantCompute(); + NotifyCube(); + } + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::MMCompute() +{ + mm.SetTensorA(xGlobal_[block_.offset_.offsetA], aTrans); + mm.SetTensorB(wGlobal_[block_.offset_.offsetB], bTrans); + if (static_cast(gmmQuantParams_->hasBias) && !isBiasEpilogue_) { + mm.SetBias(biasGlobal_[block_.offset_.offsetBias]); + } + mm.Iterate(); + mm.GetTensorC(l0cOutUb_, 0, true); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::DequantCompute() +{ + auto halfSingleM = QuantUtils::CeilDiv(block_.params_.singleCoreM, static_cast(2)); // 分配给2个AIV计算 + auto singleMInVec = subBlockIdx_ == 1 ? block_.params_.singleCoreM - halfSingleM : halfSingleM; + if (singleMInVec == 0) { + return; + } + uint64_t mOffset = subBlockIdx_ * halfSingleM; + CopyDataFromGm2Ub(); + // UB空间受限，分四次输出 + uint16_t splitNumOfOut = singleMInVec >= 4 ? 4 : singleMInVec; + auto mSizeForOnce = QuantUtils::CeilDiv(singleMInVec, static_cast(splitNumOfOut)); + for (uint16_t i = 0; i < splitNumOfOut; i++) { + // do dequant in vector + uint64_t offsetL0c = i * mSizeForOnce * + QuantUtils::Align(block_.params_.singleCoreN, + static_cast(QuantUtils::UB_ALIGN_SIZE / sizeof(l0cType))); + if (i * mSizeForOnce >= singleMInVec) { + break; + } + auto mSize = singleMInVec - i * mSizeForOnce >= mSizeForOnce ? mSizeForOnce : singleMInVec - i * mSizeForOnce; + LocalTensor dequantOutInUB = vecQueOut_.AllocTensor(); + + __ubuf__ yType *dequantOutInUbAddr = (__ubuf__ yType *)dequantOutInUB.GetPhyAddr(); + __ubuf__ l0cType *l0cOutUbAddr = (__ubuf__ l0cType *)l0cOutUb_.GetPhyAddr(); + l0cOutUbAddr = l0cOutUbAddr + offsetL0c; + + switch (gmmQuantParams_->aQuantMode) { + case (static_cast(QuantUtils::QuantMode::PERTOKEN_MODE)): + { + uint64_t offsetPtScale = i * mSizeForOnce; + VFDoDequantWithAPertoken(dequantOutInUbAddr, l0cOutUbAddr, offsetPtScale, mSize); + break; + } + case (static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)): + { + VFDoDequantWithAPertensor(dequantOutInUbAddr, l0cOutUbAddr, mSize); + break; + } + default: + { + VFDoDequantWithoutAScale(dequantOutInUbAddr, l0cOutUbAddr, mSize); + } + } + vecQueOut_.EnQue(dequantOutInUB); + // mmDequant result: UB -> GM + dequantOutInUB = vecQueOut_.DeQue(); + CopyDequantResFromUb2Gm(mSize, (mOffset + i * mSizeForOnce) * mmTilingData_->N, dequantOutInUB); + vecQueOut_.FreeTensor(dequantOutInUB); + } + FreeUbTensor(); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::CopyDataFromGm2Ub() +{ + auto halfSingleM = QuantUtils::CeilDiv(block_.params_.singleCoreM, static_cast(2)); // 分配给2个AIV计算 + auto singleMInVec = subBlockIdx_ == 1 ? block_.params_.singleCoreM - halfSingleM : halfSingleM; + // scale: GM -> UB + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERCHANNEL_MODE)) { + scaleUb_ = vecQueScale_.AllocTensor(); + CopyX2ScaleFromGm2Ub(scaleUb_); + vecQueScale_.EnQue(scaleUb_); + scaleUb_ = vecQueScale_.DeQue(); + } + + uint64_t mOffset = subBlockIdx_ * halfSingleM; + // perTokenScale: GM -> UB + if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::PERTOKEN_MODE)) { + ptScaleUb_ = vecQuePertokenScale_.AllocTensor(); + CopyX1ScaleFromGm2Ub(ptScaleUb_, singleMInVec * sizeof(ptScaleType), + block_.offset_.offsetPerTokenScale + mOffset); + vecQuePertokenScale_.EnQue(ptScaleUb_); + ptScaleUb_ = vecQuePertokenScale_.DeQue(); + } + if (isBiasEpilogue_) { + biasUb_ = vecQueBias_.AllocTensor(); + CopyBiasFromGm2Ub(biasUb_); + vecQueBias_.EnQue(biasUb_); + biasUb_ = vecQueBias_.DeQue(); + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void +GQmmMixRegbaseKernel::CopyX1ScaleFromGm2Ub(LocalTensor &dst, + uint64_t blockLen, uint64_t offset) +{ + DataCopyParams ptScale2UbParams{1, 0, 0, 0}; + DataCopyPadParams padParams; + ptScale2UbParams.blockLen = blockLen; + DataCopyPad(dst, scaleAGlobal_[offset], ptScale2UbParams, padParams); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void +GQmmMixRegbaseKernel::CopyX2ScaleFromGm2Ub(LocalTensor &dst) +{ + DataCopyParams scale2UbParams{1, 0, 0, 0}; + DataCopyPadParams padParams; + scale2UbParams.blockLen = block_.params_.singleCoreN * sizeof(scaleType); + DataCopyPad(dst, scaleBGlobal_[block_.offset_.offsetScale], scale2UbParams, padParams); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void +GQmmMixRegbaseKernel::CopyBiasFromGm2Ub(LocalTensor &dst) +{ + DataCopyParams bias2UbParams{1, 0, 0, 0}; + DataCopyPadParams padParams; + bias2UbParams.blockLen = block_.params_.singleCoreN * sizeof(biasType); + DataCopyPad(dst, biasGlobal_[block_.offset_.offsetBias], bias2UbParams, padParams); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void +GQmmMixRegbaseKernel::CopyDequantResFromUb2Gm(uint64_t blockCount, uint64_t offset, + LocalTensor &src) +{ + DataCopyExtParams ub2GmParams{1, 0, 0, 0, 0}; + ub2GmParams.blockLen = block_.params_.singleCoreN * sizeof(yType); + ub2GmParams.blockCount = blockCount; + ub2GmParams.dstStride = (mmTilingData_->N - block_.params_.singleCoreN) * sizeof(yType); + DataCopyPad(yGlobal_[block_.offset_.offsetC + offset], src, ub2GmParams); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::FreeUbTensor() +{ + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERCHANNEL_MODE)) { + vecQueScale_.FreeTensor(scaleUb_); + } + + if (gmmQuantParams_->aQuantMode == static_cast(QuantUtils::QuantMode::PERTOKEN_MODE)) { + vecQuePertokenScale_.FreeTensor(ptScaleUb_); + } + + if (isBiasEpilogue_) { + vecQueBias_.FreeTensor(biasUb_); + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::VFDoDequantWithAPertoken( + __ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, uint64_t offsetPtScale, uint16_t mSize) +{ + __ubuf__ ptScaleType *ptScaleUbAddr = (__ubuf__ ptScaleType *)ptScaleUb_.GetPhyAddr(); + ptScaleUbAddr = ptScaleUbAddr + offsetPtScale; + if (!isBiasEpilogue_) { + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, nullptr, ptScaleUbAddr, nullptr, mSize, block_.params_.singleCoreN); + } else { + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, (__ubuf__ scaleType *)scaleUb_.GetPhyAddr(), ptScaleUbAddr, nullptr, + mSize, block_.params_.singleCoreN); + } + } else { + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, nullptr, ptScaleUbAddr, (__ubuf__ biasType *)biasUb_.GetPhyAddr(), + mSize, block_.params_.singleCoreN); + } else { + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, (__ubuf__ scaleType *)scaleUb_.GetPhyAddr(), ptScaleUbAddr, + (__ubuf__ biasType *)biasUb_.GetPhyAddr(), mSize, block_.params_.singleCoreN); + } + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::VFDoDequantWithAPertensor( + __ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, uint16_t mSize) +{ + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, (__ubuf__ scaleType *)scaleUb_.GetPhyAddr(), nullptr, nullptr, mSize, + block_.params_.singleCoreN); +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::VFDoDequantWithoutAScale( + __ubuf__ yType *dequantOutInUbAddr, __ubuf__ l0cType *l0cOutUbAddr, uint16_t mSize) +{ + if (!isBiasEpilogue_) { + VFDoDequant(dequantOutInUbAddr, l0cOutUbAddr, + (__ubuf__ scaleType *)scaleUb_.GetPhyAddr(), nullptr, + nullptr, mSize, block_.params_.singleCoreN); + } else { + if (gmmQuantParams_->bQuantMode == static_cast(QuantUtils::QuantMode::PERTENSOR_MODE)) { + VFDoDequant(dequantOutInUbAddr, l0cOutUbAddr, nullptr, nullptr, + (__ubuf__ biasType *)biasUb_.GetPhyAddr(), mSize, + block_.params_.singleCoreN); + } else { + VFDoDequant( + dequantOutInUbAddr, l0cOutUbAddr, (__ubuf__ scaleType *)scaleUb_.GetPhyAddr(), nullptr, + (__ubuf__ biasType *)biasUb_.GetPhyAddr(), mSize, block_.params_.singleCoreN); + } + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +template +__aicore__ inline void GQmmMixRegbaseKernel::VFDoDequant( + __ubuf__ yType *dst, __ubuf__ l0cType *l0cOut, __ubuf__ scaleType *scale, __ubuf__ ptScaleType *perTokenScale, + __ubuf__ biasType *bias, uint16_t mSize, uint16_t nSize) +{ + uint32_t eleNumPerVf = QuantUtils::GetVRegSize() / sizeof(l0cType); + uint32_t nSrcUbAligned = + QuantUtils::Align(nSize, static_cast(QuantUtils::UB_ALIGN_SIZE / sizeof(l0cType))); + uint32_t nDstUbAligned = QuantUtils::Align(nSize, static_cast(QuantUtils::UB_ALIGN_SIZE / sizeof(yType))); + uint16_t nLoopCnt = (nSize + eleNumPerVf - 1) / eleNumPerVf; + __VEC_SCOPE__ + { + MicroAPI::MaskReg maskN4B16 = MicroAPI::CreateMask(); + for (uint16_t mIdx = 0; mIdx < mSize; mIdx++) { + uint32_t elementNum = nSize; + for (uint16_t vfBlockIdx = 0; vfBlockIdx < nLoopCnt; vfBlockIdx++) { + MicroAPI::RegTensor l0cOutReg; + MicroAPI::RegTensor scaleReg; + MicroAPI::RegTensor perTokenScaleReg; + MicroAPI::RegTensor biasReg; + MicroAPI::RegTensor castSrcOutReg, castScaleReg, castScaleOneReg, mulScaleOutReg, + mulPtScaleOutReg, castBiasReg, castBiasOneReg, addBiasOutReg; + MicroAPI::RegTensor castResultOutReg; + MicroAPI::MaskReg maskN = MicroAPI::UpdateMask(elementNum); + // copy input from ub to register, addr of ub should align to 32B + uint32_t l0cOutOffset = mIdx * nSrcUbAligned + vfBlockIdx * eleNumPerVf; + MicroAPI::DataCopy(l0cOutReg, l0cOut + l0cOutOffset); + // cast l0cOut from int32 to float + if constexpr (IsSameType::value) { + MicroAPI::Cast(castSrcOutReg, l0cOutReg, maskN); + } else { + castSrcOutReg = l0cOutReg; + } + // l0c_out * scale + if constexpr (isPertensor) { + MicroAPI::Muls(mulScaleOutReg, castSrcOutReg, scaleScalar_, maskN); + } else { + MicroAPI::DataCopy(scaleReg, scale + vfBlockIdx * eleNumPerVf); + if constexpr (!IsSameType::value) { // cast scale from bf16 to float + MicroAPI::Cast(castScaleReg, scaleReg, maskN); + MicroAPI::Cast(castScaleOneReg, scaleReg, maskN4B16); + MicroAPI::Interleave(castScaleReg, castScaleOneReg, castScaleReg, castScaleOneReg); + } else { + castScaleReg = scaleReg; + } + MicroAPI::Mul(mulScaleOutReg, castSrcOutReg, castScaleReg, maskN); + } + // out * perTokenScale + if constexpr (aQuantMode == QuantUtils::QuantMode::PERTENSOR_MODE) { + AscendC::MicroAPI::Muls(mulPtScaleOutReg, mulScaleOutReg, perTokenScaleScalar_, maskN); + } else if constexpr (aQuantMode == QuantUtils::QuantMode::PERTOKEN_MODE) { + MicroAPI::DataCopy(perTokenScaleReg, + perTokenScale + mIdx); + MicroAPI::Mul(mulPtScaleOutReg, mulScaleOutReg, perTokenScaleReg, maskN); + } else { + mulPtScaleOutReg = mulScaleOutReg; + } + // out + bias + if constexpr (isBiasEpilogue) { + MicroAPI::DataCopy(biasReg, bias + vfBlockIdx * eleNumPerVf); + // cast bias from bf16/fp16 to float + if constexpr (IsSameType::value || IsSameType::value) { + MicroAPI::Cast(castBiasReg, biasReg, maskN); + // bf16/fp16共用maskN4B16 + MicroAPI::Cast(castBiasOneReg, biasReg, maskN4B16); + MicroAPI::Interleave(castBiasReg, castBiasOneReg, castBiasReg, castBiasOneReg); + } else if constexpr (IsSameType::value) { + castBiasReg = biasReg; + } + MicroAPI::Add(addBiasOutReg, mulPtScaleOutReg, castBiasReg, maskN); + } else { + addBiasOutReg = mulPtScaleOutReg; + } + // cast dequant result from float to fp16/bf16 + if constexpr (!IsSameType::value) { + MicroAPI::Cast(castResultOutReg, addBiasOutReg, maskN); + } else { + castResultOutReg = addBiasOutReg; + } + // copy out from register to ub + uint32_t dstUbOffset = mIdx * nDstUbAligned + vfBlockIdx * eleNumPerVf; + if constexpr (IsSameType::value) { + MicroAPI::DataCopy(dst + dstUbOffset, castResultOutReg, + maskN); + } else { + MicroAPI::DataCopy(dst + dstUbOffset, castResultOutReg, + maskN); + } + } + } + } +} + +LOCAL_TEMPLATE_CLASS_MIX_PARAMS +__aicore__ inline void GQmmMixRegbaseKernel::End() +{ + // 由于vec最后一次会通过NotifyCube多发一次硬同步，所以cube侧需要额外加一次硬同步 + if ASCEND_IS_AIC { + if (isVecSetSyncCom_) { + WaitForVector(); + } + } +} +} // namespace AscendC + +#endif // GQMM_MIX_ONLINE_DYNAMIC_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_tiling_key.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_tiling_key.h new file mode 100644 index 00000000..7e7a58c6 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/gqmm_tiling_key.h @@ -0,0 +1,87 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file gqmm_tiling_key.h +* \brief +*/ + +#ifndef __OP_KERNEL_GQMM_TILING_KEY_H__ +#define __OP_KERNEL_GQMM_TILING_KEY_H__ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define GMM_NO_TRANS 0 +#define GMM_TRANS 1 + +#define GMM_DEQUANT_FIXP 0 +#define GMM_DEQUANT_VECTOR 1 +#define GMM_PERGROUP_PERBLOCK 2 + +// 模板参数 +ASCENDC_TPL_ARGS_DECL( + DlinferGroupedMatmulDirect, // 算子OpType + ASCENDC_TPL_UINT_DECL(QUANT_B_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS, GMM_TRANS), + ASCENDC_TPL_UINT_DECL(QUANT_A_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS, GMM_TRANS), + ASCENDC_TPL_UINT_DECL(KERNEL_TYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_FIXP, GMM_DEQUANT_VECTOR, + GMM_PERGROUP_PERBLOCK) + ); + +// 模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIC_ONLY), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_FIXP)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIC_ONLY), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_FIXP)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_FIXP)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_VECTOR)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_VECTOR)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_DEQUANT_VECTOR)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_PERGROUP_PERBLOCK)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_PERGROUP_PERBLOCK)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_TRANS), + ASCENDC_TPL_UINT_SEL(KERNEL_TYPE, ASCENDC_TPL_UI_LIST, GMM_PERGROUP_PERBLOCK)) +); + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_copy_cube_out.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_copy_cube_out.h new file mode 100644 index 00000000..2878487a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_copy_cube_out.h @@ -0,0 +1,222 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gqmm_copy_cube_out.h + * \brief + */ +#ifndef GQMM_COPY_CUBE_OUT_H +#define GQMM_COPY_CUBE_OUT_H +#include "lib/matmul_intf.h" + +namespace AscendC { + + template + class GQmmCustomCopyCubeOut { + using DstT = typename C_TYPE::T; + using SrcT = typename GetMmDstType::Type; + using FixpipeAdaptor = + AscendC::Impl::Detail::FixpipeParamsUtil::GetFixpipeParamsType()>; + + MATMUL_USE_MODULE(Context); + MATMUL_USE_MODULE(MatmulQuantProcessor); + MATMUL_USE_MODULE(MatmulShapeInfo); + MATMUL_USE_MODULE(MatmulShapeTiling); + MATMUL_USE_MODULE(MatmulUserDefineInfo); + MATMUL_USE_MODULE(MatmulSubBlockInfo); + + public: + __aicore__ inline GQmmCustomCopyCubeOut() = default; + + template + __aicore__ inline void Copy(const GlobalTensor& gm, const LocalTensor& co1Local, int32_t curRow, + int32_t curCol, int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth, const ScheduleContext& context = 0) + { + if constexpr (ToMatmulConfig(MM_CFG).intraBlockPartSum) { + if (!MATMUL_MODULE(MatmulSubBlockInfo)->GetFakeMsg()) { + CopyOutImpl, true>(gm, co1Local, curRow, curCol, baseHeight, + baseWidth, baseBlockHeight, baseBlockWidth); + return; + } + } + CopyOutImpl, false>(gm, co1Local, curRow, curCol, baseHeight, + baseWidth, baseBlockHeight, baseBlockWidth); + } + + template + __aicore__ inline void Copy(const LocalTensor& co2Local, const LocalTensor& co1Local, int32_t curRow, + int32_t curCol, int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth, const ScheduleContext& context = 0) + { + if constexpr (FIXPIPE_MODE == McgShfMode::DUAL_DST_SPLIT_M) { + baseHeight = (baseHeight + 1) / 2 * 2; // 指令要求m必须是偶数，向上对齐到2的倍数 + } + CopyOutImpl(co2Local, co1Local, curRow, curCol, baseHeight, baseWidth, baseBlockHeight, + baseBlockWidth); + } + + private: + template + __aicore__ inline void CopyOutImpl(const T& dst, const LocalTensor& co1Local, int32_t curRow, + int32_t curCol, int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth) + { + if constexpr (FIXPIPE_MODE == McgShfMode::DUAL_DST_SPLIT_M && PhyPosIsUB(C_TYPE::pos)) { + ASCENDC_ASSERT((baseHeight % 2 == 0), // 校验tileHeight是否2对齐 + {KERNEL_LOG(KERNEL_ERROR, "If split M when copy cube out, baseHeight must be even.");}); + } + if constexpr (C_TYPE::format == CubeFormat::ND || C_TYPE::format == CubeFormat::ND_ALIGN) { + CopyOutNZ2ND(dst, co1Local, curRow, curCol, baseHeight, + baseWidth, baseBlockHeight, baseBlockWidth); + } else { + ASCENDC_ASSERT(false, {KERNEL_LOG(KERNEL_ERROR, "Copy: unsupport Matmul format type.");}); + } + } + + template + __aicore__ inline void CopyOutNZ2ND(const T& dst, const LocalTensor& co1Local, int32_t curRow, int32_t curCol, + int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth) + { + auto stride = baseWidth; + int64_t dstOffset = 0; + if constexpr (!enSequentialWrite) { + stride = GetOrgWidth(); + if constexpr (!IsBasic(MM_CFG)) { + dstOffset = GetDstOffset(curRow, curCol, baseHeight, stride); + } + } + if constexpr (FIXPIPE_MODE == McgShfMode::DUAL_DST_SPLIT_N && PhyPosIsUB(C_TYPE::pos)) { + stride = stride >> 1; + } + + FixpipeAdaptor fixpipe(baseWidth, baseHeight, baseBlockWidth, baseBlockHeight, + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseM(), stride); + SetFixpipeParams(fixpipe); + CopyTensor(dst[dstOffset], co1Local, fixpipe, curCol, baseWidth); + } + + __aicore__ inline int64_t GetDstOffset(int32_t curRow, int32_t curCol, int32_t baseHeight, int32_t stride) + { + int64_t dstOffset = 0; + if constexpr (AscendC::Impl::Detail::MatmulFeatureTrait::IsSupportL0CToUB() && PhyPosIsUB(C_TYPE::pos) + && ((A_TYPE::ibShare && B_TYPE::ibShare) || FIXPIPE_MODE == McgShfMode::DUAL_DST_SPLIT_M)) { + dstOffset = (static_cast(static_cast( + curRow * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseM() * stride)) >> 1) + + static_cast(curCol * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseN()); + } else if constexpr (AscendC::Impl::Detail::MatmulFeatureTrait::IsSupportL0CToUB() && + PhyPosIsUB(C_TYPE::pos) && FIXPIPE_MODE == McgShfMode::DUAL_DST_SPLIT_N) { + dstOffset = + static_cast(curRow * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseM() * stride) + + static_cast(curCol * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseN() * baseHeight); + dstOffset = dstOffset >> 1; + } else { + dstOffset = + static_cast(curRow * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseM() * stride) + + static_cast(curCol * MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseN()); + } + return dstOffset; + } + + __aicore__ inline void SetFixpipeParams(FixpipeAdaptor& fixpipe) { + if constexpr (PhyPosIsUB(C_TYPE::pos) && + AscendC::Impl::Detail::MatmulFeatureTrait::IsSupportL0CToUB()) { + fixpipe.SetSubBlockId(MATMUL_MODULE(MatmulSubBlockInfo)->GetSubBlockIdx()); + if constexpr (A_TYPE::ibShare && B_TYPE::ibShare) { + fixpipe.SetMcgShfMode(McgShfMode::DUAL_DST_SPLIT_M); + } else { + fixpipe.SetMcgShfMode(FIXPIPE_MODE); + } + } + } + + template + __aicore__ inline void CopyTensor(const T& dst, const LocalTensor& co1Local, + FixpipeAdaptor& fixpipe, const int32_t curN = 0, const int32_t baseUseN = 0) + { + if (MATMUL_MODULE(MatmulQuantProcessor)->IsQuantSenario()) { + fixpipe.SetQuantMode(MATMUL_MODULE(MatmulQuantProcessor)->GetMatmulQuantMode()); + if (MATMUL_MODULE(MatmulQuantProcessor)->IsPerChannelSenario()) { + LocalTensor quantTensor; + MATMUL_MODULE(MatmulQuantProcessor)->CopyQuantTensor(quantTensor, curN, baseUseN); + fixpipe.template FixpipeOut(dst, co1Local, quantTensor); + MATMUL_MODULE(MatmulQuantProcessor)->FreeQuantTensor(quantTensor); + } else { + fixpipe.SetQuantScalar(MATMUL_MODULE(MatmulQuantProcessor)->GetQuantScalarValue()); + fixpipe.template FixpipeOut(dst, co1Local); + } + } else { + fixpipe.SetCastMode(); + fixpipe.template FixpipeOut(dst, co1Local); + } + } + + template + __aicore__ inline uint32_t GetOrgWidth() + { + uint32_t dimN = GetOrgN(); + if (GetOrgKc() != 0) { + dimN = GetOrgKc(); + } + constexpr uint32_t blockCount = ONE_BLK_SIZE / sizeof(DstT); + if constexpr (C_TYPE::format == CubeFormat::ND_ALIGN) { + dimN = Ceil(dimN, blockCount) * blockCount; + } + return dimN; + } + + template + __aicore__ inline uint32_t GetOrgKc() + { + if constexpr ((C_TYPE::layout == LayoutMode::SBNGD) || (C_TYPE::layout == LayoutMode::BSNGD)) { + return 0; + } else { + return MATMUL_MODULE(MatmulShapeInfo)->template GetOrgKc(); + } + } + + template + __aicore__ inline uint32_t GetOrgM() + { + if constexpr (C_TYPE::layout == LayoutMode::SBNGD) { + return MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoB() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoS1(); + } else if constexpr (C_TYPE::layout == LayoutMode::BSNGD) { + return MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoS1(); + } else if constexpr (ToMatmulConfig(MM_CFG).isEnableChannelSplit && A_TYPE::format == CubeFormat::ND && + C_TYPE::format == CubeFormat::NZ) { + return Ceil(MATMUL_MODULE(MatmulShapeInfo)->template GetOrgM(), BLOCK_CUBE) * BLOCK_CUBE; + } else { + return MATMUL_MODULE(MatmulShapeInfo)->template GetOrgM(); + } + } + + template + __aicore__ inline uint32_t GetOrgN() + { + if constexpr (C_TYPE::layout == LayoutMode::SBNGD) { + return MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoG() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoS2() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoN() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoB(); + } else if constexpr (C_TYPE::layout == LayoutMode::BSNGD) { + return MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoG() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoS2() * + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetCLayoutInfoN(); + } else { + return MATMUL_MODULE(MatmulShapeInfo)->template GetOrgN(); + } + } + }; + + } // namespace AscendC +#endif // GQMM_COPY_CUBE_OUT_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_custom_mm_policy.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_custom_mm_policy.h new file mode 100644 index 00000000..e8df99ef --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/mm_extension_interface/gqmm_custom_mm_policy.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gqmm_custom_mm_policy.h + * \brief + */ +#ifndef GQMM_CUSTOM_MM_POLICY_H +#define GQMM_CUSTOM_MM_POLICY_H + +#include "lib/matmul_intf.h" +#include "gqmm_copy_cube_out.h" + +namespace AscendC { +template +class GQmmCustomMatmulPolicy : public AscendC::Impl::Detail::MatmulPolicy { +public: + using CopyCubeOut = GQmmCustomCopyCubeOut; +}; +} // namespace AscendC +#endif // GQMM_CUSTOM_MM_POLICY_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_block_sch.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_block_sch.h new file mode 100644 index 00000000..10f603c1 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_block_sch.h @@ -0,0 +1,335 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file quant_block_sch.h + * \brief + */ +#ifndef GROUPED_MATMUL_QUANT_BLOCK_SCH_H +#define GROUPED_MATMUL_QUANT_BLOCK_SCH_H + +#include "quant_utils.h" + +namespace DlinferGroupedMatmulDirect { +struct ASWTilingParam { + uint64_t m; + uint64_t n; + uint64_t k; + uint64_t singleCoreM; + uint64_t singleCoreN; + uint64_t mCnt; + uint64_t nCnt; + uint64_t totalCnt; + uint64_t mCoreNum; + uint64_t mTailCoreNum; + uint64_t mBaseTail; + uint64_t nBaseTail; + uint64_t mIndex; + uint64_t nIndex; + uint64_t mSplitAddrOffset; + uint64_t nSplitAddrOffset; + uint64_t aGroupAddrOffset; + uint64_t bGroupAddrOffset; + uint64_t cGroupAddrOffset; + uint64_t xScaleGroupAddrOffset; + uint64_t wScaleGroupAddrOffset; + uint64_t biasGroupAddrOffset; + uint64_t mTailTile; + uint64_t nTailTile; + uint64_t mainRow; + uint64_t round; + uint64_t index; +}; + +// 对于每一个GlobalTensor的偏移 +struct ASWOffsetParam { + uint64_t offsetA; + uint64_t offsetB; + uint64_t offsetC; + uint64_t offsetScale; + uint64_t offsetBias; + uint64_t offsetPerTokenScale; +}; + +class QuantASWBlockSch { +public: + __aicore__ inline QuantASWBlockSch() {} + template + __aicore__ inline void Init(const TCubeTiling* __restrict &tilingData, uint32_t blockIdx); + // 每一个group需要更新mm的group偏移和MNK + template + __aicore__ inline void UpdateGroupOffset(int32_t m, int32_t n, int32_t k, uint32_t groupIdx); + template + __aicore__ inline void UpdateGroupParams(); // 每一个group需要更新mm的参数 + __aicore__ inline void UpdateTailTile(); + template + __aicore__ inline void UpdateBasicIndex(uint64_t roundIdx, bool isLastGroupRound); + template + __aicore__ inline void UpdateBlockParams(uint64_t roundIdx, bool isLastGroupRound = true); + template + __aicore__ inline void CalcGMOffset(); + __aicore__ inline void ResetAddressOffsets(); + __aicore__ inline uint32_t GetStartBlockIdx() const; + __aicore__ inline uint32_t GetEndBlockIdx() const; + +public: + ASWTilingParam params_; + ASWOffsetParam offset_; + const TCubeTiling* __restrict tilingData_; + +private: + const uint64_t WINDOW_LEN = 4; + uint32_t blockIdx_; + uint32_t startBlockIdx_; + uint32_t endBlockIdx_; +}; + +template +__aicore__ inline void QuantASWBlockSch::Init(const TCubeTiling* __restrict &tilingData, uint32_t blockIdx) +{ + blockIdx_ = blockIdx; + tilingData_ = tilingData; + params_.mSplitAddrOffset = 0; + params_.nSplitAddrOffset = 0; + params_.xScaleGroupAddrOffset = 0; // xScale is optional, mm不需要，与GMM代码归一 + params_.mTailTile = 1; // 1 说明不切分, mm可直接从mm tiling里获取，但目前仓不同，无法拿到 + params_.nTailTile = 1; // 1 说明不切分 + startBlockIdx_ = 0; // 每个group核使用开始的索引, 单mm默认从0开始 + + if constexpr (isGmm) { // GMM + params_.m = 0; + params_.n = 0; + params_.k = 0; + // 不同group的偏移量, 在dynamic输入输出的取二级指针偏移所需，对于globalTensor偏移不需要 + params_.aGroupAddrOffset = 0; + params_.bGroupAddrOffset = 0; + params_.cGroupAddrOffset = 0; + params_.wScaleGroupAddrOffset = 0; + params_.biasGroupAddrOffset = 0; + // 标记每个group的结束核 + endBlockIdx_ = tilingData_->usedCoreNum - 1; // 上个group核使用结束的索引 + } else { // MM + params_.m = tilingData_->M; // mm可以直接从tilingData里取值 + params_.n = tilingData_->N; + params_.k = tilingData_->Ka; + UpdateGroupParams(); + } +} + +template +__aicore__ inline void QuantASWBlockSch::UpdateGroupOffset(int32_t m, int32_t n, int32_t k, uint32_t groupIdx) +{ + // 用初始化或上个group的mm的m,k,n值更新group矩阵的偏移量。group内2维mm。 + if (groupIdx > 0) { // groupIdx==0时，起始点均为0，无需计算，减少scalar + if constexpr (QuantUtils::IsFp4()) { // 2: fp4为半个字节 + params_.aGroupAddrOffset += params_.m * params_.k / 2; + params_.bGroupAddrOffset += params_.n * params_.k / 2; + } else { + params_.aGroupAddrOffset += params_.m * params_.k; + params_.bGroupAddrOffset += params_.n * params_.k; + } + params_.cGroupAddrOffset += params_.m * params_.n; + if constexpr (QuantUtils::IsMxType()) { + uint64_t scaleK = QuantUtils::MXFP_MULTI_BASE_SIZE; + if constexpr (!aTrans) { // mx (m, ceil(k / 64), 2) + scaleK *= QuantUtils::CeilDiv(params_.k, QuantUtils::MXFP_DIVISOR_SIZE); + params_.xScaleGroupAddrOffset += params_.m * scaleK; + params_.wScaleGroupAddrOffset += params_.n * scaleK; + } else if constexpr (aTrans && !bTrans) { // mx (k / 64 + G, m, 2) + // scaleK from (k0 + k1 + k2 + ... + k_{i - 1}) / 64 + Gi, cumsum + // n在host侧已保证不会为0 + scaleK *= (params_.bGroupAddrOffset / params_.n / QuantUtils::MXFP_DIVISOR_SIZE + groupIdx); + params_.xScaleGroupAddrOffset = params_.m * scaleK; + params_.wScaleGroupAddrOffset = params_.n * scaleK; + } + } else { // 当perChannel/perToken（重点场景）计算offset，kernel侧在perTensor场景下直接使用groupIdx偏移，减少分支判断 + params_.xScaleGroupAddrOffset += params_.m; + params_.wScaleGroupAddrOffset += params_.n; + } + params_.biasGroupAddrOffset += params_.n; + } + + // 需要kernel传参m,n,k, 兼容group_type=0,2和多tensor情况 + params_.m = m; + params_.n = n; + params_.k = k; +} + +// 兼容GMM和MM的更新 +template +__aicore__ inline void QuantASWBlockSch::UpdateGroupParams() +{ + params_.mCnt = QuantUtils::CeilDiv(params_.m, tilingData_->baseM); + params_.nCnt = QuantUtils::CeilDiv(params_.n, tilingData_->baseN); + params_.totalCnt = params_.mCnt * params_.nCnt; + params_.mBaseTail = params_.m - (params_.mCnt - 1) * tilingData_->baseM; + params_.nBaseTail = params_.n - (params_.nCnt - 1) * tilingData_->baseN; + params_.mCoreNum = QuantUtils::Min(WINDOW_LEN, params_.mCnt); + // 计算round数还是按照实际，使用核数按照startBlockIdx开始 + params_.round = QuantUtils::CeilDiv(params_.totalCnt, tilingData_->usedCoreNum); + params_.mainRow = params_.mCnt / params_.mCoreNum - 1; + params_.mTailCoreNum = params_.mCnt - params_.mCoreNum * params_.mainRow; + if constexpr (isGmm) { + // 计算当前group的mm最后一轮计算空闲的核的索引 + // 新group开始的空闲的核索引 + startBlockIdx_ = endBlockIdx_ == tilingData_->usedCoreNum - 1 ? 0 : (endBlockIdx_ + 1); + // 当前group结束的空闲的核索引 + endBlockIdx_ = (params_.totalCnt + startBlockIdx_ - 1) % tilingData_->usedCoreNum; + // 如果当前group不是最后一个group，则空闲核不参与最后一轮计算，需要留给下一个group使用 + if (startBlockIdx_ > endBlockIdx_ && (blockIdx_ > endBlockIdx_ && blockIdx_ < startBlockIdx_)) { + params_.round -= 1; + } else if (startBlockIdx_ <= endBlockIdx_ && (blockIdx_ > endBlockIdx_ || blockIdx_ < startBlockIdx_)) { + params_.round -= 1; + } + } +} + +// 尾块切分后，更新结束的核索引和round数 +__aicore__ inline void QuantASWBlockSch::UpdateTailTile() +{ + uint64_t newEndBlockIdx = params_.mTailTile * params_.nTailTile * (endBlockIdx_ + 1) - 1; + if (blockIdx_ > endBlockIdx_ && blockIdx_ <= newEndBlockIdx) + { + params_.round += 1; + } + endBlockIdx_ = newEndBlockIdx; +} + +template +__aicore__ inline void QuantASWBlockSch::UpdateBasicIndex(uint64_t roundIdx, bool isLastGroupRound) +{ + uint64_t newBlockIdx = isLastGroupRound ? (blockIdx_ / (params_.mTailTile * params_.nTailTile)) : blockIdx_; + params_.index = newBlockIdx + roundIdx * tilingData_->usedCoreNum; + // GMM当前group的startBlockIdx_不一定从0开始，要进行计算过的base块数修正 + if constexpr (isGmm) { + if (blockIdx_ < startBlockIdx_) { + params_.index += tilingData_->usedCoreNum - startBlockIdx_; // 加上最开始从startBlockIdx_开始的base数 + } else { + params_.index -= startBlockIdx_; // 减去roundIdx * CoreNum的未包含的最开始0-(startBlockIdx_-1)的base数 + } + } + uint64_t rowIdx = params_.index / params_.nCnt / params_.mCoreNum; + if (rowIdx < params_.mainRow) { + params_.mIndex = rowIdx * params_.mCoreNum + params_.index % params_.mCoreNum; + params_.nIndex = (params_.index / params_.mCoreNum) % params_.nCnt; + } else { + rowIdx = params_.mainRow; + uint64_t tailIndex = params_.index - params_.mainRow * params_.mCoreNum * params_.nCnt; + params_.mIndex = params_.mainRow * params_.mCoreNum + tailIndex % params_.mTailCoreNum; + params_.nIndex = (tailIndex / params_.mTailCoreNum) % params_.nCnt; + } + + if (rowIdx & 1) { + params_.nIndex = params_.nCnt - 1 - params_.nIndex; + } +} + +// 返回给kernel的开始空闲核索引，帮助最后一个group的尾块细分 +__aicore__ inline uint32_t QuantASWBlockSch::GetStartBlockIdx() const +{ + return startBlockIdx_; +} + +// 返回给kernel的开始空闲核索引，帮助最后一个group的尾块细分 +__aicore__ inline uint32_t QuantASWBlockSch::GetEndBlockIdx() const +{ + return endBlockIdx_; +} + +template +__aicore__ inline void QuantASWBlockSch::UpdateBlockParams(uint64_t roundIdx, bool isLastGroupRound) +{ + params_.singleCoreM = params_.mIndex != (params_.mCnt - 1) ? tilingData_->baseM : params_.mBaseTail; + params_.singleCoreN = params_.nIndex != (params_.nCnt - 1) ? tilingData_->baseN : params_.nBaseTail; + if (!isLastGroupRound || (params_.mTailTile == 1 && params_.nTailTile == 1)) { + return; + } + + if (roundIdx == params_.round - 1) { + uint64_t singleCoreMSplit = (params_.singleCoreM + params_.mTailTile - 1) / + params_.mTailTile; + uint64_t singleCoreNSplit = (params_.singleCoreN + params_.nTailTile - 1) / + params_.nTailTile; + if constexpr (aTrans) { // (k, m) + singleCoreMSplit = QuantUtils::Align(singleCoreMSplit, QuantUtils::INNER_AXIS_MIN_SPLIT_VAL); + } + if constexpr (!bTrans) { // (k, n) + singleCoreNSplit = QuantUtils::Align(singleCoreNSplit, QuantUtils::INNER_AXIS_MIN_SPLIT_VAL); + } + uint64_t totalTailTile = params_.mTailTile * params_.nTailTile; + uint64_t mSplitIdx = (blockIdx_ % totalTailTile) % params_.mTailTile; + uint64_t nSplitIdx = (blockIdx_ % totalTailTile) / params_.mTailTile; + params_.mSplitAddrOffset = mSplitIdx * singleCoreMSplit; + params_.nSplitAddrOffset = nSplitIdx * singleCoreNSplit; + if (params_.mSplitAddrOffset >= params_.singleCoreM || params_.nSplitAddrOffset >= params_.singleCoreN) { + params_.singleCoreM = 0; + params_.singleCoreN = 0; + return; + } + if (params_.mSplitAddrOffset + singleCoreMSplit > params_.singleCoreM) { + params_.singleCoreM = params_.singleCoreM - singleCoreMSplit * mSplitIdx; + } else { + params_.singleCoreM = singleCoreMSplit; + } + if (params_.nSplitAddrOffset + singleCoreNSplit > params_.singleCoreN) { + params_.singleCoreN = params_.singleCoreN - singleCoreNSplit * nSplitIdx; + } else { + params_.singleCoreN = singleCoreNSplit; + } + } +} + +__aicore__ inline void QuantASWBlockSch::ResetAddressOffsets() +{ + // 尾块细分的m方向偏移量和n方向偏移量需要重置 + params_.mSplitAddrOffset = 0; + params_.nSplitAddrOffset = 0; +} + +template +__aicore__ inline void QuantASWBlockSch::CalcGMOffset() +{ + uint64_t mOffset = params_.mIndex * tilingData_->baseM + params_.mSplitAddrOffset; + uint64_t nOffset = params_.nIndex * tilingData_->baseN + params_.nSplitAddrOffset; + if constexpr (aTrans) { + offset_.offsetA = mOffset; + } else { + offset_.offsetA = mOffset * params_.k; + } + + if constexpr (bTrans) { + offset_.offsetB = nOffset * params_.k; + } else { + offset_.offsetB = nOffset; + } + + offset_.offsetC = mOffset * params_.n + nOffset; + + if constexpr (QuantUtils::IsMxType()) { + uint64_t pertokenScaleK = QuantUtils::MXFP_MULTI_BASE_SIZE; + uint64_t scaleK = QuantUtils::MXFP_MULTI_BASE_SIZE; + if constexpr (!aTrans) { // mx (m, ceil(k / 64), 2) + pertokenScaleK *= QuantUtils::CeilDiv(params_.k, QuantUtils::MXFP_DIVISOR_SIZE); + } + if constexpr (bTrans) { // mx (n, ceil(k / 64), 2) + scaleK *= QuantUtils::CeilDiv(params_.k, QuantUtils::MXFP_DIVISOR_SIZE); + } + offset_.offsetPerTokenScale = mOffset * pertokenScaleK; + offset_.offsetScale = nOffset * scaleK; + } else { + offset_.offsetPerTokenScale = mOffset; + offset_.offsetScale = nOffset; + } + // pertoken是optional，不是dynamic，要累加分组的偏移量 + offset_.offsetPerTokenScale += params_.xScaleGroupAddrOffset; + offset_.offsetBias = nOffset; +} +} // namespace DlinferGroupedMatmulDirect +#endif // GROUPED_MATMUL_QUANT_BLOCK_SCH_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_utils.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_utils.h new file mode 100644 index 00000000..2c087b52 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/quant_adaptive_sliding_window_templates/quant_utils.h @@ -0,0 +1,184 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_utils.h + * \brief + */ +#ifndef ASCENDC_QUANT_UTILS_H +#define ASCENDC_QUANT_UTILS_H + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +#define LOCAL_TEMPLATE_CLASS_PARAMS \ + template +#define LOCAL_TEMPLATE_FUNC_PARAMS xType, wType, biasType, scaleType, yType, wFormat, aTrans, bTrans + +namespace QuantUtils { + +constexpr uint32_t PER_BLOCK_SIZE = 128; +constexpr int32_t MXFP_DIVISOR_SIZE = 64; +constexpr int32_t MXFP_MULTI_BASE_SIZE = 2; +constexpr uint32_t MAX_REPEAT_TIMES = 255; +constexpr uint32_t UB_ALIGN_SIZE = 32; +const uint32_t FP32_OUTPUT_TIMES = 4; +constexpr uint8_t BUFFER_SWITCH = 2; +constexpr uint8_t SPLIT_M = 0; +constexpr uint8_t SPLIT_K = 2; +constexpr uint64_t CUBE_BLOCK = 16; +constexpr uint64_t INNER_AXIS_MIN_SPLIT_VAL = 128; // ND2NZ cacheline 128 + +constexpr uint8_t SYNC_AIC_AIV_MODE = 4; +constexpr uint16_t FLAG_ID_MAX = 16; +constexpr uint16_t AIC_SYNC_AIV_FLAG = 4; +constexpr uint16_t AIV_SYNC_AIC_FLAG = 6; +constexpr MatmulConfig MM_CFG_NO_PRELOAD_OPEN_UNIT_FLAG = GetMDLConfig(false, false, 0, false, false, false, true); + +enum class QuantMode : uint32_t { + DEFAULT = 0x0U, + PERTENSOR_MODE = 0x1U, + PERCHANNEL_MODE = 0x1U << 1, + PERTOKEN_MODE = 0x1U << 2, + MX_PERGROUP_MODE = 0x1U << 3, + PERGROUP_MODE = 0x1U << 4, + PERBLOCK_MODE = 0x1U << 5, +}; + +template +__aicore__ inline T Max(T a, T b) { + return a > b ? a : b; +} + +template +__aicore__ inline T Min(T a, T b) { + return a > b ? b : a; +} + +__aicore__ inline uint64_t CeilDiv(uint64_t a, uint64_t b) { + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline uint64_t Align(uint64_t a, uint64_t b) { + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline constexpr bool IsMxType() +{ + return AscendC::IsSameType::value; +} + +template +__aicore__ inline constexpr bool IsFp4() +{ + return (AscendC::IsSameType::value || AscendC::IsSameType::value); +} + +template +__aicore__ inline constexpr bool IsBiasEpilogue() +{ + return AscendC::IsSameType::value && + (AscendC::IsSameType::value || AscendC::IsSameType::value || + AscendC::IsSameType::value); +} + +/** + * Get the size of vector registers in bytes + */ +__aicore__ inline constexpr uint32_t GetVRegSize() +{ +#if __CCE_AICORE__ == 310 + return AscendC::VECTOR_REG_WIDTH; +#else + return 256U; +#endif +} + +/** + * Get the aiv corenum in different platforms + */ +__aicore__ inline constexpr uint32_t GetTaskRation() +{ +#if __CCE_AICORE__ == 310 + return 2U; // aiv corenum is 2 in C310 platform +#else + return 1U; +#endif +} + +__aicore__ inline int32_t GetSplitValueFromGroupList(uint32_t groupIdx, int32_t &preOffset, + int32_t groupType, uint32_t groupListType, + const AscendC::GlobalTensor &groupListGm) { + int32_t splitValue = 0; + if (likely(groupType != -1)) { // -1: no need to split + if (groupListType == 0) { + int32_t offset = static_cast(groupListGm.GetValue(groupIdx)); + splitValue = offset - preOffset; + preOffset = offset; + } else { + splitValue = static_cast(groupListGm.GetValue(groupIdx)); + } + } + return splitValue; +} + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ >= 310 +template +__aicore__ inline void InitOutputWithZero(AscendC::GlobalTensor yInitGlobal, AscendC::LocalTensor &initLocal, + uint64_t ySize, int32_t usedCoreNum, bool &isKZeroInit) +{ + if ASCEND_IS_AIC { + return; + } + if (AscendC::GetSubBlockIdx() >= 1) { // 搬运共享带宽，只需要把aiv0的0搬到GM即可，也兼容aic:aiv=1:1的场景 + return; + } + uint32_t blockIdx = AscendC::GetBlockIdx() / AscendC::GetTaskRation(); + // 仿照InitOutput接口取值 + uint64_t initSize = (QuantUtils::MAX_REPEAT_TIMES * AscendC::ONE_BLK_SIZE) / sizeof(T); // 能存放输出dtype的多少个元素 + uint64_t perCoreSize = QuantUtils::CeilDiv(ySize, usedCoreNum); + perCoreSize = GROUPED_MATMUL::AlignUp(perCoreSize * sizeof(T)) / sizeof(T); + initSize = QuantUtils::Min(initSize, perCoreSize); + uint64_t realCoreNum = + QuantUtils::Min(QuantUtils::CeilDiv(ySize, initSize), static_cast(usedCoreNum)); + if (blockIdx >= realCoreNum) { // 多余核数返回，每个核上最少32B + return; + } + if (!isKZeroInit) { // 第一次k==0时，需要初始化ub中buffer全0 + AscendC::Duplicate(initLocal, 0, initSize); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_MTE3)); + AscendC::SetFlag(eventIdVToMte3); + AscendC::WaitFlag(eventIdVToMte3); + isKZeroInit = true; + } + uint64_t yOffset = perCoreSize * blockIdx; + uint64_t outCurSize = (blockIdx == realCoreNum - 1) ? (ySize - yOffset) : perCoreSize; + uint64_t movRound = outCurSize / initSize; + uint64_t movTail = outCurSize - movRound * initSize; + + AscendC::DataCopyExtParams ub2GmParams{1, static_cast(initSize * sizeof(T)), 0, 0, 0}; + for (uint64_t i = 0; i < movRound; ++i) { + AscendC::DataCopyPad(yInitGlobal[yOffset], initLocal, ub2GmParams); + yOffset += initSize; + } + if (movTail != 0) { // mov tail zero data + ub2GmParams.blockLen = static_cast(movTail * sizeof(T)); + AscendC::DataCopyPad(yInitGlobal[yOffset], initLocal, ub2GmParams); + } +} +#endif +} // namespace QUANT_UTILS + +#endif // ASCENDC_QUANT_UTILS_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/anti_quant_y_vf.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/anti_quant_y_vf.h new file mode 100644 index 00000000..0f969d72 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/anti_quant_y_vf.h @@ -0,0 +1,92 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file anti_quant_y_vf.h + * \brief + */ +#ifndef GROUPED_MATMUL_ANTI_QUANT_Y_VF_H +#define GROUPED_MATMUL_ANTI_QUANT_Y_VF_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "tool.h" + +namespace MicroAPI = AscendC::MicroAPI; +using AscendC::MicroAPI::RegTensor; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +template +struct LocalAddressYParam { + __local_mem__ int32_t *yOriginPhyAddr; + __local_mem__ float *cScalePhyAddr; + __local_mem__ float *kScalePhyAddr; + __local_mem__ float *biasPhyAddr; + __local_mem__ yType *yPhyAddr; +}; + +static constexpr MicroAPI::CastTrait C32_TO_FP32_TRAIT = {MicroAPI::RegLayout::UNKNOWN, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT}; + +static constexpr MicroAPI::CastTrait FP32_TO_F16 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; + +template +__aicore__ inline void AntiQuantYB32(LocalAddressYParam &localAddressParam, uint64_t nRealL0Size, + uint64_t nRealVRegSize, uint16_t ubLoopN, uint16_t ubLoopM) +{ + __VEC_SCOPE__ + { + RegTensor yOriginVreg; + RegTensor yVreg; + RegTensor antiQuantCScaleVreg; + RegTensor antiQuantKScaleVreg; + RegTensor biasVreg; + RegTensor yF16Vreg; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + + uint32_t nRealL0Temp = nRealL0Size; + for (uint16_t nId = 0; nId < ubLoopN; nId++) { + if constexpr (hasBias) { + MicroAPI::DataCopy( + biasVreg, localAddressParam.biasPhyAddr + nId * nRealVRegSize); + } + MicroAPI::DataCopy( + antiQuantCScaleVreg, localAddressParam.cScalePhyAddr + nId * nRealVRegSize); + + MicroAPI::MaskReg yResultMask = MicroAPI::UpdateMask(nRealL0Temp); + for (uint16_t mId = 0; mId < ubLoopM; mId++) { + uint64_t yOffset = nId * nRealVRegSize + mId * nRealL0Size; + uint64_t ubYOffset = (nId * nRealVRegSize >> 1) + mId * nRealL0Size; + MicroAPI::DataCopy(yOriginVreg, + localAddressParam.yOriginPhyAddr + yOffset); + MicroAPI::DataCopy(antiQuantKScaleVreg, + localAddressParam.kScalePhyAddr + mId); + + MicroAPI::Cast(yVreg, yOriginVreg, maskAll); + MicroAPI::Mul(yVreg, yVreg, antiQuantCScaleVreg, maskAll); + MicroAPI::Mul(yVreg, yVreg, antiQuantKScaleVreg, maskAll); + if constexpr (hasBias) { + MicroAPI::Add(yVreg, yVreg, biasVreg, maskAll); + } + + MicroAPI::Cast(yF16Vreg, yVreg, maskAll); + + MicroAPI::DataCopy( + localAddressParam.yPhyAddr + ubYOffset * 2, yF16Vreg, yResultMask); + } + } + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // GROUPED_MATMUL_ANTI_QUANT_Y_VF_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_config.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_config.h new file mode 100644 index 00000000..b75d0f14 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_config.h @@ -0,0 +1,260 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file basic_block_config.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_CONFIG_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_CONFIG_H + +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "lib/matmul_intf.h" +#include "tool.h" + +namespace WeightQuantBatchMatmulV2::Arch35 { + +constexpr static uint16_t WEIGHT_F16_UB_NZ_STRIDE = 65; +constexpr int16_t SHIFT_FOR_BF16 = 1; + +struct WqmmConfig { + bool aTrans; + bool bTrans; + QuantType antiQuantType; + bool hasAntiQuantOffset; + QuantType quantType; + CubeFormat weightFormat; +}; + +static constexpr WqmmConfig S8S4_NZKN_G = {false, false, QuantType::PER_GROUP, false, QuantType::NONE, CubeFormat::NZ}; +static constexpr WqmmConfig A16MXF4_NZKN = {false, false, QuantType::MX, false, QuantType::NONE, CubeFormat::NZ}; +static constexpr WqmmConfig MXA8W4_NZNK = {false, true, QuantType::MX, false, QuantType::NONE, CubeFormat::NZ}; + +struct BasicBlockControlParam { + uint64_t processId; + uint64_t mSize; + uint64_t mL1Size; + uint64_t curBasicBlockId; + uint64_t basicBlockLimit; + uint64_t mOffset; + uint64_t nOffset; +}; + +struct BasicBlockOffsetParam { + uint64_t mL1Size; + uint64_t kaL1Size; + uint64_t kbL1Size; + uint64_t nL1Size; + + uint64_t mOffset; + uint64_t nOffset; + + uint64_t mSize; + uint64_t kSize; + uint64_t nSize; + uint64_t kAlign; + uint64_t nAlign; + + int8_t scaleAFactor; + int8_t scaleBFactor; + + GM_ADDR yGmAddr; +}; + +struct VecAntiQuantConfig { + uint64_t ubMte2BufferNum = 2; + uint64_t ubMte2InnerSize = 512; +}; + +struct UbConsumeConfig { + uint64_t ubVfBufferNum; + uint64_t l1RequireVfComputeRealK; + uint64_t l1RequireVfComputeRealN; + uint64_t kWeightLowBitUbOffset; + uint64_t nWeightLowBitUbOffset; +}; + +struct L1ConsumeConfig { + uint64_t l1SplitTwoVecExternalOffset; + uint64_t l1RealExternalLen; +}; + +struct UbBufferInfo { + uint64_t ubWeightOutputHighBitBufferNum; + uint64_t weightInputLowbitUbTotalSize; + uint64_t highBitDataUbTotalSize; + uint64_t antiQuantScaleUbTotalSize; + uint64_t antiQuantScaleAfterCastUbTotalSize; + uint64_t antiQuantOffsetUbTotalSize; + uint64_t weightInputLowBitUbSingleBufferSize; + uint64_t antiQuantScaleUbSingleBufferSize; + uint64_t antiQuantScaleAfterCastUbSingleBufferSize; + uint64_t antiQuantOffsetUbSingleBufferSize; + uint64_t highBitDataUbSingleBufferSize; + uint32_t antiQuantScaleMaskBufferSize; +}; + +template +__aicore__ constexpr UbBufferInfo GetNzBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = QUADRUPLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 112 * GetKBUnit(), // 112KB + .highBitDataUbTotalSize = 128 * GetKBUnit(), // 128KB + .antiQuantScaleUbTotalSize = 4 * GetKBUnit(), // 4KB + .antiQuantScaleAfterCastUbTotalSize = 0, + .antiQuantOffsetUbTotalSize = 4 * GetKBUnit(), // 4KB + .weightInputLowBitUbSingleBufferSize = 112 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleUbSingleBufferSize = 4 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleAfterCastUbSingleBufferSize = 0, + .antiQuantOffsetUbSingleBufferSize = 4 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .highBitDataUbSingleBufferSize = 128 * GetKBUnit() / QUADRUPLE_BUFFER_NUM, + .antiQuantScaleMaskBufferSize = 0}; +} + +template +__aicore__ constexpr UbBufferInfo GetS8S4NzBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = QUADRUPLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 96 * GetKBUnit(), // 96KB + .highBitDataUbTotalSize = 128 * GetKBUnit(), // 128KB + .antiQuantScaleUbTotalSize = 12 * GetKBUnit(), // 12KB + .antiQuantScaleAfterCastUbTotalSize = 0, + .antiQuantOffsetUbTotalSize = 0, + .weightInputLowBitUbSingleBufferSize = 96 * GetKBUnit() / vecConfig.ubMte2BufferNum, // 32KB + .antiQuantScaleUbSingleBufferSize = 12 * GetKBUnit() / vecConfig.ubMte2BufferNum, // 4KB + .antiQuantScaleAfterCastUbSingleBufferSize = 0, + .antiQuantOffsetUbSingleBufferSize = 0, + .highBitDataUbSingleBufferSize = 128 * GetKBUnit() / QUADRUPLE_BUFFER_NUM, // 32KB + .antiQuantScaleMaskBufferSize = 32 / sizeof(uint64_t)}; // 32B (4个uint64) +} + +template +__aicore__ constexpr UbBufferInfo GetNdBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = DOUBLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 174 * GetKBUnit(), // 174KB + .highBitDataUbTotalSize = 66 * GetKBUnit(), // 66KB + .antiQuantScaleUbTotalSize = 4 * GetKBUnit(), // 4KB + .antiQuantScaleAfterCastUbTotalSize = 0, + .antiQuantOffsetUbTotalSize = 4 * GetKBUnit(), // 4KB + .weightInputLowBitUbSingleBufferSize = 174 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleUbSingleBufferSize = 4 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleAfterCastUbSingleBufferSize = 0, + .antiQuantOffsetUbSingleBufferSize = 4 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .highBitDataUbSingleBufferSize = 66 * GetKBUnit() / DOUBLE_BUFFER_NUM, + .antiQuantScaleMaskBufferSize = 0}; +} + +template +__aicore__ constexpr UbBufferInfo GetMxFp4NdBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = DOUBLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 128 * GetKBUnit(), // 128KB + .highBitDataUbTotalSize = 66 * GetKBUnit(), // 66KB + .antiQuantScaleUbTotalSize = 8 * GetKBUnit(), // 8KB + .antiQuantScaleAfterCastUbTotalSize = 32 * GetKBUnit(), // 32KB + .antiQuantOffsetUbTotalSize = 0, + .weightInputLowBitUbSingleBufferSize = 128 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleUbSingleBufferSize = 8 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleAfterCastUbSingleBufferSize = 32 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantOffsetUbSingleBufferSize = 0, + .highBitDataUbSingleBufferSize = 66 * GetKBUnit() / DOUBLE_BUFFER_NUM, + .antiQuantScaleMaskBufferSize = 0}; +} + +template +__aicore__ constexpr UbBufferInfo GetMxFp4NzBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = QUADRUPLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 64 * GetKBUnit(), // 64KB + .highBitDataUbTotalSize = 128 * GetKBUnit(), // 128KB + .antiQuantScaleUbTotalSize = 8 * GetKBUnit(), // 8KB + .antiQuantScaleAfterCastUbTotalSize = 16 * GetKBUnit(), // 16KB + .antiQuantOffsetUbTotalSize = 0, + .weightInputLowBitUbSingleBufferSize = 64 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleUbSingleBufferSize = 8 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleAfterCastUbSingleBufferSize = 16 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantOffsetUbSingleBufferSize = 0, + .highBitDataUbSingleBufferSize = 128 * GetKBUnit() / QUADRUPLE_BUFFER_NUM, + .antiQuantScaleMaskBufferSize = 0}; +} + +template +__aicore__ constexpr UbBufferInfo GetMxA8W4NzBufferInfo() +{ + return {.ubWeightOutputHighBitBufferNum = QUADRUPLE_BUFFER_NUM, + .weightInputLowbitUbTotalSize = 64 * GetKBUnit(), // 64KB + .highBitDataUbTotalSize = 64 * GetKBUnit(), // 64KB + .antiQuantScaleUbTotalSize = 0, + .antiQuantScaleAfterCastUbTotalSize = 0, + .antiQuantOffsetUbTotalSize = 0, + .weightInputLowBitUbSingleBufferSize = 64 * GetKBUnit() / vecConfig.ubMte2BufferNum, + .antiQuantScaleUbSingleBufferSize = 0, + .antiQuantScaleAfterCastUbSingleBufferSize = 0, + .antiQuantOffsetUbSingleBufferSize = 0, + .highBitDataUbSingleBufferSize = 64 * GetKBUnit() / QUADRUPLE_BUFFER_NUM, + .antiQuantScaleMaskBufferSize = 0}; +} + +template +__aicore__ constexpr UbBufferInfo GetBufferConfig() +{ + if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + return GetMxFp4NdBufferInfo(); + } else if constexpr (IsSameType::value) { + return GetMxA8W4NzBufferInfo(); + } else { + return GetMxFp4NzBufferInfo(); + } + } + + if constexpr (wqmmConfig.weightFormat == CubeFormat::ND) { + return GetNdBufferInfo(); + } else { + if constexpr (IsSameType::value) { + return GetS8S4NzBufferInfo(); + } else { + return GetNzBufferInfo(); + } + } +} + +struct VfConfig { + uint64_t vfNStandardLen; + uint64_t vfKStandardLen; +}; + +template +__aicore__ constexpr VfConfig GetVfConfig() +{ + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ && wqmmConfig.bTrans) { + return {.vfNStandardLen = 64, .vfKStandardLen = 256}; + } else if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ && !wqmmConfig.bTrans) { + return {.vfNStandardLen = 256, .vfKStandardLen = 64}; + } else if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + if constexpr (IsSameType::value) { + return {.vfNStandardLen = 256, .vfKStandardLen = 64}; + } else { + return {.vfNStandardLen = 32 * GetKBUnit() / vecConfig.ubMte2InnerSize, + .vfKStandardLen = vecConfig.ubMte2InnerSize}; + } + } else { + // NZ transB=False + if constexpr (IsSameType::value) { + return {.vfNStandardLen = 64, .vfKStandardLen = vecConfig.ubMte2InnerSize}; + } else { + return {.vfNStandardLen = 64, .vfKStandardLen = 256}; + } + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_CONFIG_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_mx.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_mx.h new file mode 100644 index 00000000..bc50e81c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_mx.h @@ -0,0 +1,366 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file basic_block_vf_mx.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_MX_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_MX_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" + +namespace MicroAPI = AscendC::MicroAPI; +using AscendC::BLOCK_CUBE; +using AscendC::VECTOR_REG_WIDTH; +using AscendC::MicroAPI::AddrReg; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +template +struct MxFp4NdScaleParams { + uint64_t ubLoopExternalAxis; + __local_mem__ uint8_t *antiQuantScaleBasePhyAddr; + __local_mem__ xType *antiQuantScaleF16PhyAddr0; + __local_mem__ xType *antiQuantScaleF16PhyAddr1; +}; + +template +struct MxFp4NdWeightParams { + uint64_t ubLoopExternalAxis; + __local_mem__ wType *weightLowBitPhyAddr0; + __local_mem__ wType *weightLowBitPhyAddr1; + __local_mem__ xType *weightF16PhyAddr0; + __local_mem__ xType *weightF16PhyAddr1; + __local_mem__ xType *antiQuantScaleF16PhyAddr0; + __local_mem__ xType *antiQuantScaleF16PhyAddr1; +}; + +template +struct Fp4NzParams { + uint64_t loopN1; + uint64_t loopGroupNum; + uint64_t loopInnerNum; + uint64_t innerDstStride; + uint64_t groupDstStride; + uint64_t loopN1DstStride; + __local_mem__ xType *antiQuantScaleBasePhyAddr; + __local_mem__ wType *weightLowBitPhyAddr; + __local_mem__ xType *weightHighBitPhyAddr; +}; + +template +struct MxA8W4NzParams { + uint64_t loopKNum; + uint64_t innerLoopNum; + uint64_t loopKDstStride; + uint64_t innerDstStride; + __local_mem__ wType *weightLowBitPhyAddr; + __local_mem__ xType *weightHighBitPhyAddr; +}; + +static constexpr MicroAPI::CastTrait CAST_BF16_TO_FP16_TRAIT = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT}; + +static constexpr MicroAPI::CastTrait CAST_FP4_TO_BF16_TRAIT = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; +static constexpr uint32_t E2M1_SHIFT_RIGHT_SIZE = 0x2; +static constexpr uint32_t SHIFT_LEFT_SIZE = 0x4; +static constexpr uint32_t E2M1_AND_MASK = 0x9C; + +template +__aicore__ inline void MxScaleVf(RegTensor &antiQuantScaleE8m0Vreg0, + RegTensor &antiQuantScaleE8m0Vreg1, RegTensor &antiQuantScaleF16Vreg0, + RegTensor &antiQuantScaleF16Vreg1, MicroAPI::MaskReg &maskAll) +{ + RegTensor zeroVreg; + MicroAPI::Duplicate(zeroVreg, 0); + + // 通过数据重排指令， 交织 antiQuantScaleE8m0Vreg0 和 zeroVreg , Interleave后变为 + // antiQuantScaleE8m0Vreg0 + // Vn s1 0 s2 0 s3 0 s4 0 s5 0 s6 0 s7 0 s8 0....... s127 0 s128 0 + // antiQuantScaleE8m0Vreg1 + // Vd s128 0 s129 0 s130 0 s131 0 s132 0 s133 0....... s255 0 s256 0 + MicroAPI::Interleave(antiQuantScaleE8m0Vreg0, antiQuantScaleE8m0Vreg1, zeroVreg, antiQuantScaleE8m0Vreg0); + + if constexpr (!IsSameType::T, vector_f16>::value) { + // 逻辑右移一位,得到BF16{S1E8M7}的排列,比如s1 0 =[8bit, 8bit]= [10101011 00000000] + // 变为 S1 = [16bit] = [0101010110000000] + // antiQuantScaleF16Vreg0： + // S1 S2 S3 S4 S5 S6 S7 S8 ..... S127 S128 + // antiQuantScaleF16Vreg1: + // S128 S129 S130 S131 S132 S133 S134 S135 ..... S255 S256 + MicroAPI::ShiftRights((MicroAPI::RegTensor &)antiQuantScaleF16Vreg0, + (MicroAPI::RegTensor &)antiQuantScaleE8m0Vreg0, SHIFT_FOR_BF16, maskAll); + MicroAPI::ShiftRights((MicroAPI::RegTensor &)antiQuantScaleF16Vreg1, + (MicroAPI::RegTensor &)antiQuantScaleE8m0Vreg1, SHIFT_FOR_BF16, maskAll); + } else { + // 需要转换为FP16格式时，先转换为BF16再转换为FP16 + RegTensor antiQuantScaleBF16Vreg0; + RegTensor antiQuantScaleBF16Vreg1; + MicroAPI::ShiftRights((MicroAPI::RegTensor &)antiQuantScaleBF16Vreg0, + (MicroAPI::RegTensor &)antiQuantScaleE8m0Vreg0, SHIFT_FOR_BF16, maskAll); + MicroAPI::ShiftRights((MicroAPI::RegTensor &)antiQuantScaleBF16Vreg1, + (MicroAPI::RegTensor &)antiQuantScaleE8m0Vreg1, SHIFT_FOR_BF16, maskAll); + MicroAPI::Cast( + antiQuantScaleF16Vreg0, (MicroAPI::RegTensor &)antiQuantScaleBF16Vreg0, maskAll); + MicroAPI::Cast( + antiQuantScaleF16Vreg1, (MicroAPI::RegTensor &)antiQuantScaleBF16Vreg1, maskAll); + } +} + +template +__aicore__ inline void CastFp4ToF16(RegTensor &f16Vreg, RegTensor &fp4Vreg, MicroAPI::MaskReg &maskAll) +{ + if constexpr (!IsSameType::T, vector_f16>::value) { + // CAST_FP4_TO_BF16_TRAIT中设置RegLayout::ZERO 表示按照如下形式处理做cast： + // Vn 1 2 0 0 0 0 0 0 3 4 0 0 0 0 0 0 + // Vd 1 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4 + MicroAPI::Cast(f16Vreg, fp4Vreg, maskAll); + } else { + // FP4-->FP16指令不支持，需要转换为BF16再转换为FP16 + RegTensor bF16Vreg; + MicroAPI::Cast(bF16Vreg, fp4Vreg, maskAll); + MicroAPI::Cast(f16Vreg, bF16Vreg, maskAll); + } +} + +template +__aicore__ inline void MxNkScaleVf(MxFp4NdScaleParams &mxFp4NdNkScaleParams) +{ + // NK mte2搬运的antiquantscale的标准大小为(128, 32), 按照4行的粒度处理，每次处理128个E8M0, 得到256个F16的数字 + RegTensor antiQuantScaleE8m0Vreg0; + RegTensor antiQuantScaleE8m0Vreg1; + RegTensor antiQuantScaleF16Vreg0; + RegTensor antiQuantScaleF16Vreg1; + MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopNIdx = 0; ubLoopNIdx < mxFp4NdNkScaleParams.ubLoopExternalAxis; ubLoopNIdx++) { + // 搬运128个E8M0的antiquantscale, 通过两倍上采样变成256个E8M0， DIST_US_B8表示搬运模式如下： + // Vn s1 s2 s3 s4 s5 s6 s7 s8 s9 ...... s125 s126 s127 s128 + // Vd s1 s1 s2 s2 s3 s3 s4 s4 s5 ...... s125 s125 s126 s126 s127 s127 s128 s128 + MicroAPI::DataCopy( + antiQuantScaleE8m0Vreg0, mxFp4NdNkScaleParams.antiQuantScaleBasePhyAddr + ubLoopNIdx * 128); + + MxScaleVf(antiQuantScaleE8m0Vreg0, antiQuantScaleE8m0Vreg1, antiQuantScaleF16Vreg0, antiQuantScaleF16Vreg1, + maskAll); + + MicroAPI::DataCopy( + mxFp4NdNkScaleParams.antiQuantScaleF16PhyAddr0 + ubLoopNIdx * VECTOR_REG_WIDTH, antiQuantScaleF16Vreg0, + maskAll); + MicroAPI::DataCopy( + mxFp4NdNkScaleParams.antiQuantScaleF16PhyAddr1 + ubLoopNIdx * VECTOR_REG_WIDTH, antiQuantScaleF16Vreg1, + maskAll); + } +} + +template +__aicore__ inline void MxKnScaleVf(MxFp4NdScaleParams &mxFp4NdScaleParams) +{ + // KN mte2搬运的antiquantscale的标准大小为(4，256), 按照一行的粒度处理 + RegTensor antiQuantScaleE8m0Vreg0; + RegTensor antiQuantScaleE8m0Vreg1; + RegTensor antiQuantScaleF16Vreg0; + RegTensor antiQuantScaleF16Vreg1; + MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopKIdx = 0; ubLoopKIdx < mxFp4NdScaleParams.ubLoopExternalAxis; ubLoopKIdx++) { + // 搬运256个E8M0的antiquantscale, DIST_NORM表示搬运模式如下： + // Vn s1 s2 s3 s4 s5 s6 s7 s8 s9 ...... s254 s255 s256 + // Vd s1 s2 s3 s4 s5 s6 s7 s8 s9 ...... s254 s255 s256 + MicroAPI::DataCopy( + antiQuantScaleE8m0Vreg0, mxFp4NdScaleParams.antiQuantScaleBasePhyAddr + ubLoopKIdx * VECTOR_REG_WIDTH); + + MxScaleVf(antiQuantScaleE8m0Vreg0, antiQuantScaleE8m0Vreg1, antiQuantScaleF16Vreg0, antiQuantScaleF16Vreg1, + maskAll); + + MicroAPI::DataCopy( + mxFp4NdScaleParams.antiQuantScaleF16PhyAddr0 + ubLoopKIdx * VECTOR_REG_WIDTH, antiQuantScaleF16Vreg0, + maskAll); + MicroAPI::DataCopy( + mxFp4NdScaleParams.antiQuantScaleF16PhyAddr1 + ubLoopKIdx * VECTOR_REG_WIDTH, antiQuantScaleF16Vreg1, + maskAll); + } +} + +template +__aicore__ inline void MxNkWeightVf(MxFp4NdWeightParams &mxFp4NdNkWeightParams) +{ + // 每次处理一行， 一行为256个FP4的数， 分两条指令处理，每条指令处理128个FP4的数 + RegTensor antiQuantScaleF16Vreg0; + RegTensor antiQuantScaleF16Vreg1; + RegTensor weightFp4Vreg0; + RegTensor weightFp4Vreg1; + RegTensor weightF16Vreg0; + RegTensor weightF16Vreg1; + + MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopNIdx = 0; ubLoopNIdx < mxFp4NdNkWeightParams.ubLoopExternalAxis; ubLoopNIdx++) { + // DIST_E2B_B16 表示搬运模式如下, 将一个f16的数扩展成16个 + // Vn 1 2 3 4 5 6 7 8 + // Vd + // 11111111111111112222222222222233333333333333334444444444444444455555555555555666666666666666677777777777777 + MicroAPI::DataCopy( + antiQuantScaleF16Vreg0, + mxFp4NdNkWeightParams.antiQuantScaleF16PhyAddr0 + ubLoopNIdx * (VECTOR_REG_WIDTH >> 2)); + MicroAPI::DataCopy( + antiQuantScaleF16Vreg1, + mxFp4NdNkWeightParams.antiQuantScaleF16PhyAddr1 + ubLoopNIdx * (VECTOR_REG_WIDTH >> 2)); + + // DIST_UNPK_B8 表示按照如下形式载入, 其中Vn中一个数字为4bit: + // Vn 1 2 3 4 5 6 7 8 9 a b c d e f g ..... + // Vd 1 2 x x x x x x 3 4 x x x x x x ..... + MicroAPI::DataCopy( + weightFp4Vreg0, (__local_mem__ wType *)(mxFp4NdNkWeightParams.weightLowBitPhyAddr0 + + ubLoopNIdx * (vecConfig.ubMte2InnerSize >> 1))); + + MicroAPI::DataCopy( + weightFp4Vreg1, (__local_mem__ wType *)(mxFp4NdNkWeightParams.weightLowBitPhyAddr1 + + ubLoopNIdx * (vecConfig.ubMte2InnerSize >> 1))); + + CastFp4ToF16(weightF16Vreg0, weightFp4Vreg0, maskAll); + CastFp4ToF16(weightF16Vreg1, weightFp4Vreg1, maskAll); + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleF16Vreg0, maskAll); + MicroAPI::Mul(weightF16Vreg1, weightF16Vreg1, antiQuantScaleF16Vreg1, maskAll); + MicroAPI::DataCopy( + mxFp4NdNkWeightParams.weightF16PhyAddr0, weightF16Vreg0, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + MicroAPI::DataCopy( + mxFp4NdNkWeightParams.weightF16PhyAddr1, weightF16Vreg1, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + } +} + +template +__aicore__ inline void MxKnWeightVf(MxFp4NdWeightParams &mxFp4NdKnWeightParams) +{ + // 每次处理一行， 一行为256个FP4的数， 分两条指令处理，每条指令处理128个FP4的数 + RegTensor antiQuantScaleF16Vreg0; + RegTensor antiQuantScaleF16Vreg1; + RegTensor weightFp4Vreg0; + RegTensor weightFp4Vreg1; + RegTensor weightF16Vreg0; + RegTensor weightF16Vreg1; + + MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopKIdx = 0; ubLoopKIdx < mxFp4NdKnWeightParams.ubLoopExternalAxis; ubLoopKIdx++) { + MicroAPI::DataCopy( + antiQuantScaleF16Vreg0, mxFp4NdKnWeightParams.antiQuantScaleF16PhyAddr0 + ubLoopKIdx * VECTOR_REG_WIDTH); + MicroAPI::DataCopy( + antiQuantScaleF16Vreg1, mxFp4NdKnWeightParams.antiQuantScaleF16PhyAddr1 + ubLoopKIdx * VECTOR_REG_WIDTH); + + for (uint16_t groupIdx = 0; groupIdx < MX_GROUPSIZE; groupIdx++) { + MicroAPI::DataCopy( + weightFp4Vreg0, (__local_mem__ wType *)(mxFp4NdKnWeightParams.weightLowBitPhyAddr0 + + ubLoopKIdx * (vecConfig.ubMte2InnerSize >> 1) * MX_GROUPSIZE + + groupIdx * (vecConfig.ubMte2InnerSize >> 1))); + + MicroAPI::DataCopy( + weightFp4Vreg1, (__local_mem__ wType *)(mxFp4NdKnWeightParams.weightLowBitPhyAddr1 + + ubLoopKIdx * (vecConfig.ubMte2InnerSize >> 1) * MX_GROUPSIZE + + groupIdx * (vecConfig.ubMte2InnerSize >> 1))); + + CastFp4ToF16(weightF16Vreg0, weightFp4Vreg0, maskAll); + CastFp4ToF16(weightF16Vreg1, weightFp4Vreg1, maskAll); + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleF16Vreg0, maskAll); + MicroAPI::Mul(weightF16Vreg1, weightF16Vreg1, antiQuantScaleF16Vreg1, maskAll); + MicroAPI::DataCopy( + mxFp4NdKnWeightParams.weightF16PhyAddr0, weightF16Vreg0, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + MicroAPI::DataCopy( + mxFp4NdKnWeightParams.weightF16PhyAddr1, weightF16Vreg1, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + } + } +} + +template +__aicore__ inline void AntiQuantFp4NzKnVf(Fp4NzParams &fp4NzParams) +{ + RegTensor antiQuantScaleVreg; + RegTensor weightFp4Vreg; + RegTensor weightF16Vreg; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + __local_mem__ xType *antiQuantScaleBasePhyAddr; + + for (uint16_t loopN1Idx = 0; loopN1Idx < fp4NzParams.loopN1; loopN1Idx++) { + for (uint16_t loopGroupIdx = 0; loopGroupIdx < fp4NzParams.loopGroupNum; loopGroupIdx++) { + // DIST_BLK 的含义为读取一个32B(即16个数)的数据，广播到256B(即128个数) + antiQuantScaleBasePhyAddr = fp4NzParams.antiQuantScaleBasePhyAddr + loopN1Idx * BLOCK_CUBE + + loopGroupIdx * 128; // 当前方案一次对齐到128 + MicroAPI::DataCopy(antiQuantScaleVreg, antiQuantScaleBasePhyAddr); + + for (uint16_t loopGroupInnerIdx = 0; loopGroupInnerIdx < fp4NzParams.loopInnerNum; loopGroupInnerIdx++) { + // DIST_UNPACK4_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 8 9 a b c d e f + // Vd 0 1 x x x x x x 2 3 x x x x x x + // 4bit物理地址位移 = 逻辑索引 >> 1 + AddrReg weightFp4AddrReg = MicroAPI::CreateAddrReg(loopN1Idx, BLOCK_CUBE * ubMte2KSize >> 1, + loopGroupIdx, MX_GROUPSIZE * BLOCK_CUBE >> 1, + loopGroupInnerIdx, VEC_MAX_ELEM_B16 >> 1); + MicroAPI::DataCopy( + weightFp4Vreg, (__local_mem__ wType *)fp4NzParams.weightLowBitPhyAddr, weightFp4AddrReg); + + CastFp4ToF16(weightF16Vreg, weightFp4Vreg, maskAll); + MicroAPI::Mul(weightF16Vreg, weightF16Vreg, antiQuantScaleVreg, maskAll); + + AddrReg weightHighBitPhyAddrReg = MicroAPI::CreateAddrReg( + loopN1Idx, fp4NzParams.loopN1DstStride, loopGroupIdx, fp4NzParams.groupDstStride, loopGroupInnerIdx, + fp4NzParams.innerDstStride); + MicroAPI::DataCopy( + fp4NzParams.weightHighBitPhyAddr, weightF16Vreg, weightHighBitPhyAddrReg, maskAll); + } + } + } +} + +template +__aicore__ inline void AntiQuantMxA8W4NzNkVf(MxA8W4NzParams &mxA8W4NzParams) +{ + MicroAPI::RegTensor wShrReg, wShlReg, wAndReg, wLoad, wShl, wShr0, wShr1, wSel, wAnd; + MicroAPI::MaskReg preg = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregVsel = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(wShrReg, E2M1_SHIFT_RIGHT_SIZE, preg); + MicroAPI::Duplicate(wShlReg, SHIFT_LEFT_SIZE, preg); + MicroAPI::Duplicate(wAndReg, E2M1_AND_MASK, preg); + + for (uint16_t loopKIdx = 0; loopKIdx < mxA8W4NzParams.loopKNum; ++loopKIdx) { + for (uint16_t innerLoopIdx = 0; innerLoopIdx < mxA8W4NzParams.innerLoopNum; ++innerLoopIdx) { + // DIST_US_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 + // Vd 0 1 0 1 2 3 2 3 4 5 4 5 6 7 6 7 + // 4bit物理地址位移 = 逻辑索引 >> 1 + MicroAPI::AddrReg aregWeightB8In = MicroAPI::CreateAddrReg( + loopKIdx, (C0_SIZE_B8 * ubMte2InnerSize) >> 1, innerLoopIdx, VECTOR_REG_WIDTH >> 1); + MicroAPI::DataCopy( + (MicroAPI::RegTensor &)wLoad, (__local_mem__ uint8_t *&)mxA8W4NzParams.weightLowBitPhyAddr, + aregWeightB8In); + + MicroAPI::ShiftRight(wShr0, wLoad, wShrReg, preg); + MicroAPI::ShiftLeft(wShl, wLoad, wShlReg, preg); + MicroAPI::ShiftRight(wShr1, wShl, wShrReg, preg); + MicroAPI::Select(wSel, wShr1, wShr0, pregVsel); + MicroAPI::And(wAnd, wSel, wAndReg, preg); + + MicroAPI::AddrReg aregWeightB8Out = MicroAPI::CreateAddrReg( + loopKIdx, mxA8W4NzParams.loopKDstStride, innerLoopIdx, mxA8W4NzParams.innerDstStride); + MicroAPI::DataCopy( + (__local_mem__ uint8_t *&)mxA8W4NzParams.weightHighBitPhyAddr, (MicroAPI::RegTensor &)wAnd, + aregWeightB8Out, preg); + } + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_MX_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nd.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nd.h new file mode 100644 index 00000000..ad610622 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nd.h @@ -0,0 +1,566 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file basic_block_vf_nd.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_ND_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_ND_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "tool.h" + +namespace MicroAPI = AscendC::MicroAPI; +using AscendC::VECTOR_REG_WIDTH; +using AscendC::MicroAPI::AddrReg; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +template +struct LocalAddressParam { + __local_mem__ xType *antiQuantScaleBasePhyAddr; + __local_mem__ xType *antiQuantScaleBasePhyAddr1; + __local_mem__ xType *antiQuantOffsetBasePhyAddr; + __local_mem__ xType *antiQuantOffsetBasePhyAddr1; + __local_mem__ wType *weightLowBitPhyAddr0; + __local_mem__ wType *weightLowBitPhyAddr1; + __local_mem__ xType *weightF16PhyAddr0; + __local_mem__ xType *weightF16PhyAddr1; +}; + +template +struct CalculateParam { + xType offsetValue; + xType scaleValue; + uint64_t ubLoop; +}; + +static constexpr MicroAPI::CastTrait FP8_TO_FP32_TRAIT_0 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; +static constexpr MicroAPI::CastTrait FP8_TO_FP32_TRAIT_2 = {MicroAPI::RegLayout::TWO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; + +static constexpr MicroAPI::CastTrait FP32_TO_F16_ODD = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND}; +static constexpr MicroAPI::CastTrait FP32_TO_F16_EVEN = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND}; + +static constexpr MicroAPI::CastTrait S8_TO_FP16_TRAIT_ODD = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; +static constexpr MicroAPI::CastTrait FP16_TO_BF16_TRAIT = {MicroAPI::RegLayout::UNKNOWN, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT}; + +template +__aicore__ inline void AntiQuantFP8NdNkVfLoadScaleOffset(RegTensor &antiQuantScaleVreg, + RegTensor &antiQuantOffsetVreg, + LocalAddressParam &localAddressParam, + const CalculateParam &calculateParam, + uint16_t nIdx) +{ + if constexpr (hasAntiQuantOffset) { + MicroAPI::DataCopy( + antiQuantOffsetVreg, localAddressParam.antiQuantOffsetBasePhyAddr + nIdx); + } + if constexpr (antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Duplicate(antiQuantScaleVreg, calculateParam.scaleValue); + } else { + MicroAPI::DataCopy(antiQuantScaleVreg, + localAddressParam.antiQuantScaleBasePhyAddr + nIdx); + } +} + +template +__aicore__ inline void AntiQuantFP8NdNkVfLoadWeight(RegTensor &weightF8Vreg0, RegTensor &weightF8Vreg1, + LocalAddressParam &localAddressParam, uint16_t nIdx) +{ + // UNPK_B8 表示按照如下形式载入: + // Vn 1 2 3 4 5 6 7 8 9 a b c d e f g ..... + // Vd 1 0 2 0 3 0 4 0 5 0 6 0 7 0 8 0 ..... + MicroAPI::DataCopy( + weightF8Vreg0, localAddressParam.weightLowBitPhyAddr0 + nIdx * ubMte2InnerSize); + MicroAPI::DataCopy( + weightF8Vreg1, localAddressParam.weightLowBitPhyAddr1 + nIdx * ubMte2InnerSize); +} + +template +__aicore__ inline void AntiQuantFP8NdNkVf(LocalAddressParam &localAddressParam, + const CalculateParam &calculateParam) +{ + RegTensor antiQuantScaleVreg, antiQuantOffsetVreg; + RegTensor weightF8Vreg0, weightF8Vreg1; + RegTensor weightF16Vreg0, weightF16Vreg1, weightF16Vreg2, weightF16Vreg3; + RegTensor weightF32Vreg0, weightF32Vreg1, weightF32Vreg2, weightF32Vreg3; + MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopNIdx = 0; ubLoopNIdx < calculateParam.ubLoop; ubLoopNIdx++) { + AntiQuantFP8NdNkVfLoadScaleOffset( + antiQuantScaleVreg, antiQuantOffsetVreg, localAddressParam, calculateParam, ubLoopNIdx); + AntiQuantFP8NdNkVfLoadWeight(weightF8Vreg0, weightF8Vreg1, + localAddressParam, ubLoopNIdx); + // 奇数、偶数位置分散到2个fp32寄存器存储 + MicroAPI::Cast(weightF32Vreg0, weightF8Vreg0, maskAll); + MicroAPI::Cast(weightF32Vreg2, weightF8Vreg1, maskAll); + MicroAPI::Cast(weightF32Vreg1, weightF8Vreg0, maskAll); + MicroAPI::Cast(weightF32Vreg3, weightF8Vreg1, maskAll); + + MicroAPI::Cast(weightF16Vreg0, weightF32Vreg0, maskAll); + MicroAPI::Cast(weightF16Vreg2, weightF32Vreg2, maskAll); + MicroAPI::Cast(weightF16Vreg1, weightF32Vreg1, maskAll); + MicroAPI::Cast(weightF16Vreg3, weightF32Vreg3, maskAll); + + MicroAPI::Or((RegTensor &)weightF16Vreg2, + (RegTensor &)weightF16Vreg2, + (RegTensor &)weightF16Vreg3, maskAll); + MicroAPI::Or((RegTensor &)weightF16Vreg0, + (RegTensor &)weightF16Vreg0, + (RegTensor &)weightF16Vreg1, maskAll); + + if constexpr (hasAntiQuantOffset) { + if constexpr (antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Adds(weightF16Vreg0, weightF16Vreg0, calculateParam.offsetValue, maskAll); + MicroAPI::Adds(weightF16Vreg2, weightF16Vreg2, calculateParam.offsetValue, maskAll); + } else { + MicroAPI::Add(weightF16Vreg0, weightF16Vreg0, antiQuantOffsetVreg, maskAll); + MicroAPI::Add(weightF16Vreg2, weightF16Vreg2, antiQuantOffsetVreg, maskAll); + } + } + + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleVreg, maskAll); + MicroAPI::Mul(weightF16Vreg2, weightF16Vreg2, antiQuantScaleVreg, maskAll); + + MicroAPI::DataCopy( + localAddressParam.weightF16PhyAddr0, weightF16Vreg0, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + MicroAPI::DataCopy( + localAddressParam.weightF16PhyAddr1, weightF16Vreg2, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + } +} + +template +__aicore__ inline void AntiQuantFP8NdKnVfLoadOffset(RegTensor &antiQuantOffsetVreg0, + RegTensor &antiQuantOffsetVreg1, + LocalAddressParam &localAddressParam) +{ + if constexpr (hasAntiQuantOffset) { + MicroAPI::DataCopy(antiQuantOffsetVreg0, + localAddressParam.antiQuantOffsetBasePhyAddr); + MicroAPI::DataCopy(antiQuantOffsetVreg1, + localAddressParam.antiQuantOffsetBasePhyAddr1); + } +} + +template +__aicore__ inline void AntiQuantFP8NdKnVfLoadScale(RegTensor &antiQuantScaleVreg0, + RegTensor &antiQuantScaleVreg1, + LocalAddressParam &localAddressParam, + const CalculateParam &calculateParam) +{ + if constexpr (antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Duplicate(antiQuantScaleVreg0, calculateParam.scaleValue); + MicroAPI::Duplicate(antiQuantScaleVreg1, calculateParam.scaleValue); + } else { + MicroAPI::DataCopy(antiQuantScaleVreg0, + localAddressParam.antiQuantScaleBasePhyAddr); + MicroAPI::DataCopy(antiQuantScaleVreg1, + localAddressParam.antiQuantScaleBasePhyAddr1); + } +} + +template +__aicore__ inline void AntiQuantFP8NdKnVfLoadWeight(RegTensor &weightF8Vreg0, RegTensor &weightF8Vreg1, + LocalAddressParam &localAddressParam, uint16_t kIdx) +{ + // UNPK_B8 表示按照如下形式载入: + // Vn 1 2 3 4 5 6 7 8 9 a b c d e f g ..... + // Vd 1 0 2 0 3 0 4 0 5 0 6 0 7 0 8 0 ..... + MicroAPI::DataCopy( + weightF8Vreg0, localAddressParam.weightLowBitPhyAddr0 + kIdx * ubMte2InnerSize); + MicroAPI::DataCopy( + weightF8Vreg1, localAddressParam.weightLowBitPhyAddr1 + kIdx * ubMte2InnerSize); +} + +template +__aicore__ inline void AntiQuantFP8NdKnVfStoreWeight(RegTensor &weightF16Vreg0, RegTensor &weightF16Vreg2, + LocalAddressParam &localAddressParam, MaskReg &mask) +{ + MicroAPI::DataCopy( + localAddressParam.weightF16PhyAddr0, weightF16Vreg0, WEIGHT_F16_UB_NZ_STRIDE, 1, mask); + MicroAPI::DataCopy( + localAddressParam.weightF16PhyAddr1, weightF16Vreg2, WEIGHT_F16_UB_NZ_STRIDE, 1, mask); +} + +template +__aicore__ inline void AntiQuantFP8NdKnVf(LocalAddressParam &localAddressParam, + const CalculateParam &calculateParam) +{ + RegTensor antiQuantScaleVreg0, antiQuantScaleVreg1, antiQuantOffsetVreg0, antiQuantOffsetVreg1; + RegTensor weightF8Vreg0, weightF8Vreg1; + RegTensor weightF16Vreg0, weightF16Vreg1, weightF16Vreg2, weightF16Vreg3; + RegTensor weightF32Vreg0, weightF32Vreg1, weightF32Vreg2, weightF32Vreg3; + MaskReg maskAll = MicroAPI::CreateMask(); + AntiQuantFP8NdKnVfLoadOffset(antiQuantOffsetVreg0, antiQuantOffsetVreg1, + localAddressParam); + AntiQuantFP8NdKnVfLoadScale( + antiQuantScaleVreg0, antiQuantScaleVreg1, localAddressParam, calculateParam); + for (uint16_t ubLoopKIdx = 0; ubLoopKIdx < calculateParam.ubLoop; ubLoopKIdx++) { + AntiQuantFP8NdKnVfLoadWeight(weightF8Vreg0, weightF8Vreg1, + localAddressParam, ubLoopKIdx); + MicroAPI::Cast(weightF32Vreg0, weightF8Vreg0, maskAll); + MicroAPI::Cast(weightF32Vreg1, weightF8Vreg0, maskAll); + MicroAPI::Cast(weightF32Vreg2, weightF8Vreg1, maskAll); + MicroAPI::Cast(weightF32Vreg3, weightF8Vreg1, maskAll); + MicroAPI::Cast(weightF16Vreg0, weightF32Vreg0, maskAll); + MicroAPI::Cast(weightF16Vreg1, weightF32Vreg1, maskAll); + MicroAPI::Cast(weightF16Vreg2, weightF32Vreg2, maskAll); + MicroAPI::Cast(weightF16Vreg3, weightF32Vreg3, maskAll); + MicroAPI::Or((RegTensor &)weightF16Vreg0, + (RegTensor &)weightF16Vreg0, + (RegTensor &)weightF16Vreg1, maskAll); + MicroAPI::Or((RegTensor &)weightF16Vreg2, + (RegTensor &)weightF16Vreg2, + (RegTensor &)weightF16Vreg3, maskAll); + if constexpr (hasAntiQuantOffset) { + if constexpr (antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Adds(weightF16Vreg0, weightF16Vreg0, calculateParam.offsetValue, maskAll); + MicroAPI::Adds(weightF16Vreg2, weightF16Vreg2, calculateParam.offsetValue, maskAll); + } else { + MicroAPI::Add(weightF16Vreg0, weightF16Vreg0, antiQuantOffsetVreg0, maskAll); + MicroAPI::Add(weightF16Vreg2, weightF16Vreg2, antiQuantOffsetVreg1, maskAll); + } + } + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleVreg0, maskAll); + MicroAPI::Mul(weightF16Vreg2, weightF16Vreg2, antiQuantScaleVreg1, maskAll); + AntiQuantFP8NdKnVfStoreWeight(weightF16Vreg0, weightF16Vreg2, localAddressParam, maskAll); + } +} + +template +__aicore__ inline void CastS8RegTensorToF16RegTensor(RegTensor &weightF16Vreg, RegTensor &weightS8Vreg, + RegTensor &weightFp16Vreg, MicroAPI::MaskReg &maskAll) +{ + // PART_EVEN 表示按照如下形式处理做cast： + // Vn 1 0 2 0 3 0 4 0 5 0 6 0 7 0 8 0 ..... + // Vd 1 1 2 2 3 3 4 4 5 5 6 6 7 7 8 8 ..... + if constexpr (!IsSameType::T, vector_f16>::value) { + MicroAPI::Cast(weightFp16Vreg, weightS8Vreg, maskAll); + MicroAPI::Cast(weightF16Vreg, weightFp16Vreg, maskAll); + } else { + MicroAPI::Cast(weightF16Vreg, weightS8Vreg, maskAll); + } +} + +template +__aicore__ inline void AddMulWeightF16RegTensorNdKn(RegTensor &weightF16Vreg0, RegTensor &weightF16Vreg1, + RegTensor &antiQuantScaleVreg, + RegTensor &antiQuantOffsetVreg, + RegTensor &antiQuantScaleVreg1, + RegTensor &antiQuantOffsetVreg1, MicroAPI::MaskReg &maskAll, + const xType &offsetValue) +{ + if constexpr (wqmmConfig.hasAntiQuantOffset) { + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Adds(weightF16Vreg0, weightF16Vreg0, offsetValue, maskAll); + MicroAPI::Adds(weightF16Vreg1, weightF16Vreg1, offsetValue, maskAll); + } else { + MicroAPI::Add(weightF16Vreg0, weightF16Vreg0, antiQuantOffsetVreg, maskAll); + MicroAPI::Add(weightF16Vreg1, weightF16Vreg1, antiQuantOffsetVreg1, maskAll); + } + } + + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleVreg, maskAll); + MicroAPI::Mul(weightF16Vreg1, weightF16Vreg1, antiQuantScaleVreg1, maskAll); +} + +template +__aicore__ inline void WeightF16NdRegTensorToNzUb(__local_mem__ xType *&weightF16PhyAddr0, + __local_mem__ xType *&weightF16PhyAddr1, + RegTensor &weightF16Vreg0, RegTensor &weightF16Vreg1, + MicroAPI::MaskReg &maskAll) +{ + MicroAPI::DataCopy( + weightF16PhyAddr0, weightF16Vreg0, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); + MicroAPI::DataCopy( + weightF16PhyAddr1, weightF16Vreg1, WEIGHT_F16_UB_NZ_STRIDE, 1, maskAll); +} + +template +__aicore__ inline void AddMulWeightF16RegTensorNdNk(RegTensor &weightF16Vreg0, RegTensor &weightF16Vreg1, + RegTensor &antiQuantScaleVreg, + RegTensor &antiQuantOffsetVreg, MicroAPI::MaskReg &maskAll, + const xType &scaleValue, const xType &offsetValue) +{ + if constexpr (wqmmConfig.hasAntiQuantOffset) { + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Adds(weightF16Vreg0, weightF16Vreg0, offsetValue, maskAll); + MicroAPI::Adds(weightF16Vreg1, weightF16Vreg1, offsetValue, maskAll); + } else { + MicroAPI::Add(weightF16Vreg0, weightF16Vreg0, antiQuantOffsetVreg, maskAll); + MicroAPI::Add(weightF16Vreg1, weightF16Vreg1, antiQuantOffsetVreg, maskAll); + } + } + + MicroAPI::Mul(weightF16Vreg0, weightF16Vreg0, antiQuantScaleVreg, maskAll); + MicroAPI::Mul(weightF16Vreg1, weightF16Vreg1, antiQuantScaleVreg, maskAll); +} + +template +__aicore__ inline void NdNkLoadScaleOffset(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + RegTensor &antiQuantScaleVreg, + RegTensor &antiQuantOffsetVreg, + xType scaleValue, uint64_t ubLoopNIdx) +{ + if constexpr (wqmmConfig.hasAntiQuantOffset) { + MicroAPI::DataCopy(antiQuantOffsetVreg, + antiQuantOffsetBasePhyAddr + ubLoopNIdx); + } + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Duplicate(antiQuantScaleVreg, scaleValue); + } else { + MicroAPI::DataCopy(antiQuantScaleVreg, + antiQuantScaleBasePhyAddr + ubLoopNIdx); + } +} + +template +__aicore__ inline void NdKnLoadScaleOffset(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantScaleBasePhyAddr1, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr1, + RegTensor &antiQuantScaleVreg, + RegTensor &antiQuantScaleVreg1, + RegTensor &antiQuantOffsetVreg, + RegTensor &antiQuantOffsetVreg1, + xType scaleValue) +{ + if constexpr (wqmmConfig.hasAntiQuantOffset) { + MicroAPI::DataCopy(antiQuantOffsetVreg, antiQuantOffsetBasePhyAddr); + MicroAPI::DataCopy(antiQuantOffsetVreg1, antiQuantOffsetBasePhyAddr1); + } + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_TENSOR) { + MicroAPI::Duplicate(antiQuantScaleVreg, scaleValue); + MicroAPI::Duplicate(antiQuantScaleVreg1, scaleValue); + } else { + MicroAPI::DataCopy(antiQuantScaleVreg, antiQuantScaleBasePhyAddr); + MicroAPI::DataCopy(antiQuantScaleVreg1, antiQuantScaleBasePhyAddr1); + } +} + +template +__aicore__ inline void AntiQuantB8CommonNdKn(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + __local_mem__ wType *weightLowBitPhyAddr0, + __local_mem__ xType *weightF16PhyAddr0, xType scaleValue, + xType offsetValue, uint64_t ubLoopK) +{ + __local_mem__ xType *antiQuantScaleBasePhyAddr1 = antiQuantScaleBasePhyAddr + VEC_MAX_ELEM_B16; + __local_mem__ xType *antiQuantOffsetBasePhyAddr1 = antiQuantOffsetBasePhyAddr + VEC_MAX_ELEM_B16; + __local_mem__ wType *weightLowBitPhyAddr1 = weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 1); + __local_mem__ xType *weightF16PhyAddr1 = weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); + + __VEC_SCOPE__ + { + RegTensor antiQuantScaleVreg; + RegTensor antiQuantOffsetVreg; + RegTensor antiQuantScaleVreg1; + RegTensor antiQuantOffsetVreg1; + RegTensor weightFp16Vreg0; + RegTensor weightFp16Vreg1; + RegTensor weightS8Vreg0; + RegTensor weightS8Vreg1; + RegTensor weightF16Vreg0; + RegTensor weightF16Vreg1; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + + NdKnLoadScaleOffset(antiQuantScaleBasePhyAddr, antiQuantScaleBasePhyAddr1, + antiQuantOffsetBasePhyAddr, antiQuantOffsetBasePhyAddr1, + antiQuantScaleVreg, antiQuantScaleVreg1, + antiQuantOffsetVreg, antiQuantOffsetVreg1, + scaleValue); + + for (uint16_t ubLoopKIdx = 0; ubLoopKIdx < static_cast(ubLoopK); ubLoopKIdx++) { + // UNPK_B8 表示按照如下形式载入: + // Vn 1 2 3 4 5 6 7 8 9 a b c d e f g ..... + // Vd 1 0 2 0 3 0 4 0 5 0 6 0 7 0 8 0 ..... + MicroAPI::DataCopy( + weightS8Vreg0, weightLowBitPhyAddr0 + ubLoopKIdx * vecConfig.ubMte2InnerSize); + MicroAPI::DataCopy( + weightS8Vreg1, weightLowBitPhyAddr1 + ubLoopKIdx * vecConfig.ubMte2InnerSize); + CastS8RegTensorToF16RegTensor(weightF16Vreg0, weightS8Vreg0, weightFp16Vreg0, maskAll); + CastS8RegTensorToF16RegTensor(weightF16Vreg1, weightS8Vreg1, weightFp16Vreg1, maskAll); + AddMulWeightF16RegTensorNdKn( + weightF16Vreg0, weightF16Vreg1, antiQuantScaleVreg, antiQuantOffsetVreg, antiQuantScaleVreg1, + antiQuantOffsetVreg1, maskAll, offsetValue); + WeightF16NdRegTensorToNzUb(weightF16PhyAddr0, weightF16PhyAddr1, weightF16Vreg0, + weightF16Vreg1, maskAll); + } + } +} + +template +__aicore__ inline void AntiQuantB8CommonNdNk(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + __local_mem__ wType *weightLowBitPhyAddr0, + __local_mem__ xType *weightF16PhyAddr0, xType scaleValue, + xType offsetValue, uint64_t ubLoopN) +{ + __local_mem__ wType *weightLowBitPhyAddr1 = weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 1); + __local_mem__ xType *weightF16PhyAddr1 = weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); + + __VEC_SCOPE__ + { + RegTensor antiQuantScaleVreg; + RegTensor antiQuantOffsetVreg; + RegTensor weightS8Vreg0; + RegTensor weightS8Vreg1; + RegTensor weightFp16Vreg0; + RegTensor weightFp16Vreg1; + RegTensor weightF16Vreg0; + RegTensor weightF16Vreg1; + + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t ubLoopNIdx = 0; ubLoopNIdx < static_cast(ubLoopN); ubLoopNIdx++) { + NdNkLoadScaleOffset(antiQuantScaleBasePhyAddr, antiQuantOffsetBasePhyAddr, + antiQuantScaleVreg, antiQuantOffsetVreg, scaleValue, ubLoopNIdx); + // UNPK_B8 表示按照如下形式载入: + // Vn 1 2 3 4 5 6 7 8 9 a b c d e f g ..... + // Vd 1 0 2 0 3 0 4 0 5 0 6 0 7 0 8 0 ..... + MicroAPI::DataCopy( + weightS8Vreg0, weightLowBitPhyAddr0 + ubLoopNIdx * vecConfig.ubMte2InnerSize); + MicroAPI::DataCopy( + weightS8Vreg1, weightLowBitPhyAddr1 + ubLoopNIdx * vecConfig.ubMte2InnerSize); + CastS8RegTensorToF16RegTensor(weightF16Vreg0, weightS8Vreg0, weightFp16Vreg0, maskAll); + CastS8RegTensorToF16RegTensor(weightF16Vreg1, weightS8Vreg1, weightFp16Vreg1, maskAll); + AddMulWeightF16RegTensorNdNk(weightF16Vreg0, weightF16Vreg1, antiQuantScaleVreg, + antiQuantOffsetVreg, maskAll, scaleValue, + offsetValue); + WeightF16NdRegTensorToNzUb(weightF16PhyAddr0, weightF16PhyAddr1, weightF16Vreg0, + weightF16Vreg1, maskAll); + } + } +} + +template +__aicore__ inline void AntiQuantInt4NdNk(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + __local_mem__ wType *weightLowBitPhyAddr0, + __local_mem__ xType *weightF16PhyAddr0, xType scaleValue, xType offsetValue, + uint64_t ubLoopN) +{ + // int4每次处理128个数即为64B, 256>>2=64 + __local_mem__ wType *weightLowBitPhyAddr1 = weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 2); + __local_mem__ xType *weightF16PhyAddr1 = weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); + __VEC_SCOPE__ + { + RegTensor antiQuantScaleVreg; + RegTensor antiQuantOffsetVreg; + RegTensor weightS4Vreg0; + RegTensor weightS4Vreg1; + RegTensor weightF16Vreg0; + RegTensor weightF16Vreg1; + + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + static constexpr MicroAPI::CastTrait castS4ToF16Trait = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; + for (uint16_t ubLoopNIdx = 0; ubLoopNIdx < static_cast(ubLoopN); ubLoopNIdx++) { + NdNkLoadScaleOffset(antiQuantScaleBasePhyAddr, antiQuantOffsetBasePhyAddr, + antiQuantScaleVreg, antiQuantOffsetVreg, scaleValue, ubLoopNIdx); + // DIST_UNPACK4_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 8 9 a b c d e f + // Vd 0 1 x x x x x x 2 3 x x x x x x + MicroAPI::DataCopy( + weightS4Vreg0, + (__local_mem__ int4x2_t *)(weightLowBitPhyAddr0 + ubLoopNIdx * (vecConfig.ubMte2InnerSize >> 1))); + + MicroAPI::DataCopy( + weightS4Vreg1, + (__local_mem__ int4x2_t *)(weightLowBitPhyAddr1 + ubLoopNIdx * (vecConfig.ubMte2InnerSize >> 1))); + + // PART_P0 表示按照如下形式处理做cast： + // Vn 1 2 0 0 0 0 0 0 3 4 0 0 0 0 0 0 + // Vd 1 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4 + MicroAPI::Cast(weightF16Vreg0, weightS4Vreg0, maskAll); + MicroAPI::Cast(weightF16Vreg1, weightS4Vreg1, maskAll); + AddMulWeightF16RegTensorNdNk(weightF16Vreg0, weightF16Vreg1, antiQuantScaleVreg, + antiQuantOffsetVreg, maskAll, scaleValue, + offsetValue); + WeightF16NdRegTensorToNzUb(weightF16PhyAddr0, weightF16PhyAddr1, weightF16Vreg0, + weightF16Vreg1, maskAll); + } + } +} + +template +__aicore__ inline void AntiQuantInt4NdKn(__local_mem__ xType *antiQuantScaleBasePhyAddr, + __local_mem__ xType *antiQuantOffsetBasePhyAddr, + __local_mem__ wType *weightLowBitPhyAddr0, + __local_mem__ xType *weightF16PhyAddr0, xType scaleValue, xType offsetValue, + uint64_t ubLoopK) +{ + __local_mem__ xType *antiQuantScaleBasePhyAddr1 = antiQuantScaleBasePhyAddr + VEC_MAX_ELEM_B16; + __local_mem__ xType *antiQuantOffsetBasePhyAddr1 = antiQuantOffsetBasePhyAddr + VEC_MAX_ELEM_B16; + __local_mem__ wType *weightLowBitPhyAddr1 = weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 2); + __local_mem__ xType *weightF16PhyAddr1 = weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); + __VEC_SCOPE__ + { + RegTensor antiQuantScaleVreg, antiQuantOffsetVreg, antiQuantScaleVreg1, antiQuantOffsetVreg1; + RegTensor weightS4Vreg0, weightS4Vreg1; + RegTensor weightF16Vreg0, weightF16Vreg1; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + static constexpr MicroAPI::CastTrait castS4ToF16Trait = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; + NdKnLoadScaleOffset(antiQuantScaleBasePhyAddr, antiQuantScaleBasePhyAddr1, + antiQuantOffsetBasePhyAddr, antiQuantOffsetBasePhyAddr1, + antiQuantScaleVreg, antiQuantScaleVreg1, + antiQuantOffsetVreg, antiQuantOffsetVreg1, + scaleValue); + + for (uint16_t ubLoopKIdx = 0; ubLoopKIdx < static_cast(ubLoopK); ubLoopKIdx++) { + // DIST_UNPACK4_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 8 9 a b c d e f + // Vd 0 1 x x x x x x 2 3 x x x x x x + MicroAPI::DataCopy( + weightS4Vreg0, + (__local_mem__ int4x2_t *)(weightLowBitPhyAddr0 + ubLoopKIdx * (vecConfig.ubMte2InnerSize >> 1))); + + MicroAPI::DataCopy( + weightS4Vreg1, + (__local_mem__ int4x2_t *)(weightLowBitPhyAddr1 + ubLoopKIdx * (vecConfig.ubMte2InnerSize >> 1))); + + // PART_P0 表示按照如下形式处理做cast： + // Vn 1 2 0 0 0 0 0 0 3 4 0 0 0 0 0 0 + // Vd 1 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4 + MicroAPI::Cast(weightF16Vreg0, weightS4Vreg0, maskAll); + MicroAPI::Cast(weightF16Vreg1, weightS4Vreg1, maskAll); + AddMulWeightF16RegTensorNdKn( + weightF16Vreg0, weightF16Vreg1, antiQuantScaleVreg, antiQuantOffsetVreg, antiQuantScaleVreg1, + antiQuantOffsetVreg1, maskAll, offsetValue); + + WeightF16NdRegTensorToNzUb(weightF16PhyAddr0, weightF16PhyAddr1, weightF16Vreg0, + weightF16Vreg1, maskAll); + } + } +} + +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_ND_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nz.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nz.h new file mode 100644 index 00000000..c014021c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/basic_block_vf_nz.h @@ -0,0 +1,165 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file basic_block_vf_nz.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_NZ_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_NZ_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "tool.h" + +namespace MicroAPI = AscendC::MicroAPI; +using AscendC::BLOCK_CUBE; +using AscendC::VECTOR_REG_WIDTH; +using AscendC::MicroAPI::AddrReg; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +template +struct Int4NzParams { + uint64_t loopN1; + uint64_t loopGroupNum; + uint64_t loopInnerNum; + uint64_t innerDstStride; + uint64_t groupDstStride; + uint64_t loopN1DstStride; + uint64_t antiQuantGroupSize; + __local_mem__ antiQuantScaleType *antiQuantScaleBasePhyAddr; + __local_mem__ xType *antiQuantOffsetBasePhyAddr; + __local_mem__ wType *weightLowBitPhyAddr; + __local_mem__ xType *weightHighBitPhyAddr; + __local_mem__ uint8_t *antiQuantScaleMaskPhyAddr; +}; + +static constexpr MicroAPI::CastTrait CAST_S4_TO_F16_TRAIT = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN}; + +static constexpr MicroAPI::CastTrait CAST_F16_TO_S8_TRAIT = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT}; + +static constexpr int64_t ONE_BLK_ELEM_B16 = ONE_BLK_SIZE / sizeof(half); + +template +__aicore__ inline void AntiQuantInt4NzKnVf(Int4NzParams &int4NzParams) +{ + RegTensor antiQuantScaleVreg; + RegTensor antiQuantOffsetVreg; + RegTensor weightS4Vreg; + RegTensor weightF16Vreg; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + + for (uint16_t LoopN1Idx = 0; LoopN1Idx < int4NzParams.loopN1; LoopN1Idx++) { + // DIST_BLK 的含义为读取一个32B(即16个数)的数据，广播到256B(即128个数) + MicroAPI::DataCopy( + antiQuantScaleVreg, int4NzParams.antiQuantScaleBasePhyAddr + LoopN1Idx * BLOCK_CUBE); + if constexpr (hasAntiQuantOffset) { + MicroAPI::DataCopy( + antiQuantOffsetVreg, int4NzParams.antiQuantOffsetBasePhyAddr + LoopN1Idx * BLOCK_CUBE); + } + + for (uint16_t LoopInnerNumIdx = 0; LoopInnerNumIdx < int4NzParams.loopInnerNum; LoopInnerNumIdx++) { + // DIST_UNPACK4_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 8 9 a b c d e f + // Vd 0 1 x x x x x x 2 3 x x x x x x + MicroAPI::DataCopy( + weightS4Vreg, (__local_mem__ int4x2_t *)(int4NzParams.weightLowBitPhyAddr + + LoopN1Idx * (BLOCK_CUBE >> 1) * ubMte2InnerSize + + LoopInnerNumIdx * (VECTOR_REG_WIDTH >> 2))); + // PART_P0 表示按照如下形式处理做cast： + // Vn 0 1 x x x x x x 2 3 x x x x x x + // Vd 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 + MicroAPI::Cast(weightF16Vreg, weightS4Vreg, maskAll); + if constexpr (hasAntiQuantOffset) { + MicroAPI::Add(weightF16Vreg, weightF16Vreg, antiQuantOffsetVreg, maskAll); + } + MicroAPI::Mul(weightF16Vreg, weightF16Vreg, antiQuantScaleVreg, maskAll); + + if constexpr (useVag) { + AddrReg weightF16PhyAddrReg = MicroAPI::CreateAddrReg( + LoopN1Idx, int4NzParams.loopN1DstStride, LoopInnerNumIdx, int4NzParams.innerDstStride); + MicroAPI::DataCopy( + int4NzParams.weightHighBitPhyAddr, weightF16Vreg, weightF16PhyAddrReg, maskAll); + } else { + MicroAPI::DataCopy( + int4NzParams.weightHighBitPhyAddr + LoopN1Idx * int4NzParams.loopN1DstStride + + LoopInnerNumIdx * int4NzParams.innerDstStride, + weightF16Vreg, maskAll); + } + } + } +} + +template +__aicore__ inline void AntiQuantS8S4NzKnGroupVf(Int4NzParams &int4NzParams) +{ + RegTensor antiQuantScaleVreg; + RegTensor antiQuantScaleVreg1; + RegTensor weightS4Vreg; + RegTensor weightF16Vreg; + RegTensor weightS8Vreg; + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskSelect = MicroAPI::CreateMask(); + MicroAPI::DataCopy(maskSelect, int4NzParams.antiQuantScaleMaskPhyAddr); + __local_mem__ antiQuantScaleType *antiQuantScaleBasePhyAddr; + + for (uint16_t loopN1Idx = 0; loopN1Idx < int4NzParams.loopN1; loopN1Idx++) { + for (uint16_t loopGroupIdx = 0; loopGroupIdx < int4NzParams.loopGroupNum; loopGroupIdx++) { + // DIST_BLK 的含义为读取一个32B(即16个数)的数据，广播到256B(即128个数) + antiQuantScaleBasePhyAddr = + int4NzParams.antiQuantScaleBasePhyAddr + loopN1Idx * C0_SIZE_B8 + loopGroupIdx * VEC_MAX_ELEM_B16; + MicroAPI::DataCopy(antiQuantScaleVreg, + antiQuantScaleBasePhyAddr); + MicroAPI::DataCopy( + antiQuantScaleVreg1, antiQuantScaleBasePhyAddr + ONE_BLK_ELEM_B16); + MicroAPI::Select(antiQuantScaleVreg, antiQuantScaleVreg, antiQuantScaleVreg1, maskSelect); + for (uint16_t loopGroupInnerIdx = 0; loopGroupInnerIdx < int4NzParams.loopInnerNum; loopGroupInnerIdx++) { + // DIST_UNPACK4_B8 表示搬运模式如下，Vn中一个数字4bit(0.5Byte)： + // Vn 0 1 2 3 4 5 6 7 8 9 a b c d e f + // Vd 0 1 x x x x x x 2 3 x x x x x x + // 地址偏移以B记，对C0_SIZE_B8和VEC_MAX_ELEM_B16除以2实现正确偏移 + MicroAPI::DataCopy( + weightS4Vreg, + (__local_mem__ int4x2_t *)(int4NzParams.weightLowBitPhyAddr + + loopN1Idx * (C0_SIZE_B8 >> 1) * ubMte2InnerSize + + loopGroupIdx * int4NzParams.antiQuantGroupSize * (C0_SIZE_B8 >> 1) + + loopGroupInnerIdx * (VEC_MAX_ELEM_B16 >> 1))); + // S4_TO_F16按如下模式cast + // Vn 00000012 00000034 + // Vd 3c004000 42004400 + MicroAPI::Cast(weightF16Vreg, weightS4Vreg, maskAll); + MicroAPI::Mul(weightF16Vreg, weightF16Vreg, antiQuantScaleVreg, maskAll); + // F16_TO_S8按如下模式cast + // Vn 3c004000 42004400 + // Vd 00010002 00030004 + MicroAPI::Cast(weightS8Vreg, weightF16Vreg, maskAll); + + AddrReg weightHighBitPhyAddrReg = MicroAPI::CreateAddrReg( + loopN1Idx, int4NzParams.loopN1DstStride, + loopGroupIdx, int4NzParams.groupDstStride, + loopGroupInnerIdx, int4NzParams.innerDstStride); + MicroAPI::DataCopy( + int4NzParams.weightHighBitPhyAddr, weightS8Vreg, weightHighBitPhyAddrReg, maskAll); + } + } + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_VF_NZ_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_copy_cube_out.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_copy_cube_out.h new file mode 100644 index 00000000..c1848b59 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_copy_cube_out.h @@ -0,0 +1,98 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file wqbmm_copy_cube_out.h + * \brief + */ +#ifndef WQBMM_COPY_CUBE_OUT_H +#define WQBMM_COPY_CUBE_OUT_H +#include "lib/matmul_intf.h" + +namespace WeightQuantBatchMatmulV2::Arch35 { +template ::IsNeedUB()>> +class WQBmmCustomCopyCubeOut { + using DstT = typename C_TYPE::T; + using SrcT = typename AscendC::GetMmDstType::Type; + using FixpipeAdaptor = AscendC::Impl::Detail::FixpipeParamsUtil< + A_TYPE, C_TYPE, MM_CFG, AscendC::Impl::Detail::MatmulFeatureTrait::GetFixpipeParamsType()>; + + MATMUL_USE_MODULE(MatmulShapeTiling); + MATMUL_USE_MODULE(MatmulSubBlockInfo); + +public: + __aicore__ inline WQBmmCustomCopyCubeOut() = default; + + template + __aicore__ inline void Copy(const GlobalTensor &gm, const LocalTensor &co1Local, int32_t curRow, + int32_t curCol, int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth, const ScheduleContext &context = 0) + { + if constexpr (ToMatmulConfig(MM_CFG).intraBlockPartSum) { + if (!MATMUL_MODULE(MatmulSubBlockInfo)->GetFakeMsg()) { + CopyOutImpl, true>( + gm, co1Local, curRow, curCol, baseHeight, baseWidth, baseBlockHeight, baseBlockWidth); + return; + } + } + CopyOutImpl, false>(gm, co1Local, curRow, curCol, baseHeight, + baseWidth, baseBlockHeight, baseBlockWidth); + } + + template + __aicore__ inline void Copy(const LocalTensor &co2Local, const LocalTensor &co1Local, int32_t curRow, + int32_t curCol, int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth, const ScheduleContext &context = 0) + { + baseHeight = (baseHeight + 1) / 2 * 2; // 2含义：指令要求m必须是偶数，向上对齐到偶数 + CopyOutImpl>(co2Local, co1Local, curRow, curCol, baseHeight, + baseWidth, baseBlockHeight, baseBlockWidth); + } + +private: + template + __aicore__ inline void CopyOutImpl(const T &dst, const LocalTensor &co1Local, int32_t curRow, int32_t curCol, + int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth) + { + // 2含义：指令要求m必须是偶数，向上对齐到偶数 + ASCENDC_ASSERT((baseHeight % 2 == 0), + { KERNEL_LOG(KERNEL_ERROR, "If split M when copy cube out, baseHeight must be even"); }); + if constexpr (C_TYPE::format == CubeFormat::ND || C_TYPE::format == CubeFormat::ND_ALIGN) { + CopyOutNZ2ND(dst, co1Local, curRow, curCol, baseHeight, baseWidth, + baseBlockHeight, baseBlockWidth); + } else { + ASCENDC_ASSERT(false, { KERNEL_LOG(KERNEL_ERROR, "CopyOut: unsupport Matmul format type."); }); + } + } + + template + __aicore__ inline void CopyOutNZ2ND(const T &dst, const LocalTensor &co1Local, int32_t curRow, int32_t curCol, + int32_t baseHeight, int32_t baseWidth, int32_t baseBlockHeight, + int32_t baseBlockWidth) + { + int64_t stride = baseWidth; + int64_t dstOffset = 0; + if constexpr (C_TYPE::format == CubeFormat::ND_ALIGN) { + constexpr int64_t blockCount = AscendC::VECTOR_REG_WIDTH / sizeof(SrcT); + stride = (stride + blockCount - 1) / blockCount * blockCount; + } + + FixpipeAdaptor fixpipe(baseWidth, baseHeight, baseBlockWidth, baseBlockHeight, + MATMUL_MODULE(MatmulShapeTiling)->GetTiling().GetBaseM(), stride); + fixpipe.SetMcgShfMode(AscendC::McgShfMode::DUAL_DST_SPLIT_M); + fixpipe.SetCastMode(); + fixpipe.template FixpipeOut(dst, co1Local); + } +}; + +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // WQBMM_COPY_CUBE_OUT_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_custom_policy.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_custom_policy.h new file mode 100644 index 00000000..d3987ac3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/custom_policy/wqbmm_custom_policy.h @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file wqbmm_custom_policy.h + * \brief + */ +#ifndef WQBMM_CUSTOM_POLICY_H +#define WQBMM_CUSTOM_POLICY_H + +#include "lib/matmul_intf.h" +#include "wqbmm_copy_cube_out.h" + +namespace WeightQuantBatchMatmulV2::Arch35 { +template +struct WQBmmCustomPolicy : public AscendC::Impl::Detail::MatmulPolicy { +public: + using CopyCubeOut = WQBmmCustomCopyCubeOut; +}; +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif // WQBMM_CUSTOM_POLICY_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_basic_controller.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_basic_controller.h new file mode 100644 index 00000000..d1f518c4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_basic_controller.h @@ -0,0 +1,261 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_weight_quant_basic_controller.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_CONTROLLER_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_CONTROLLER_H + +#include "weight_quant_basic_block.h" +#include "../grouped_matmul_tiling_data_apt.h" + + +using WeightQuantBatchMatmulV2::Arch35::BasicBlockOffsetParam; +using WeightQuantBatchMatmulV2::Arch35::CeilDivide; +using WeightQuantBatchMatmulV2::Arch35::DOUBLE_BUFFER_NUM; +using WeightQuantBatchMatmulV2::Arch35::QUADRUPLE_BUFFER_NUM; +using WeightQuantBatchMatmulV2::Arch35::VecAntiQuantConfig; +using WeightQuantBatchMatmulV2::Arch35::WeightQuantMatmulBasicBlock; +using WeightQuantBatchMatmulV2::Arch35::WqmmConfig; +using GMMWeightQuantParam = DlinferGroupedMatmulDirectTilingData::GMMWeightQuantParam; + +namespace GROUPED_MATMUL { +template +class GMMWeightQuantBasicController { +public: + __aicore__ inline GMMWeightQuantBasicController(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, + GM_ADDR bias, GM_ADDR groupList, GM_ADDR y, + const GMMWeightQuantParam *__restrict baseTiling, + const TCubeTiling *__restrict mmTiling, GM_ADDR tiling, TILING_TYPE *gmmArrayAddrIn, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitOffsetParam(BasicBlockOffsetParam &offsetParam); + __aicore__ inline void SetMKN(uint64_t groupIdx, uint64_t &preOffset, BasicBlockOffsetParam &offsetParam); + __aicore__ inline void SetGmAddr(uint64_t groupIdx, uint64_t &xBaseOffset, uint64_t &weightBaseOffset, + uint64_t &yBaseOffset, uint64_t &antiquantParamsBaseOffset, + const BasicBlockOffsetParam &offsetParam); + __aicore__ inline uint64_t GetSplitValueFromGroupList(uint64_t groupIdx, uint64_t &preOffset); + + const GMMWeightQuantParam *gmmBaseTiling_; + const TCubeTiling *mmTiling_; + + GM_ADDR xGm_; + GM_ADDR weightGm_; + GM_ADDR antiquantScaleGm_; + GM_ADDR antiquantOffsetGm_; + GM_ADDR biasGm_; + GM_ADDR yGm_; + GlobalTensor groupListGm_; + TILING_TYPE *mListGm_; + TILING_TYPE *nListGm_; + TILING_TYPE *kListGm_; + + WeightQuantMatmulBasicBlock + wqmmBasicBlock_; +}; + +template +__aicore__ inline void GMMWeightQuantBasicController::Init( + GM_ADDR x, GM_ADDR weight, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR bias, GM_ADDR groupList, + GM_ADDR y, const GMMWeightQuantParam *__restrict baseTiling, const TCubeTiling *__restrict mmTiling, GM_ADDR tiling, + TILING_TYPE *gmmArrayAddrIn, TPipe *tPipe) +{ + gmmBaseTiling_ = baseTiling; + mmTiling_ = mmTiling; + + mListGm_ = gmmArrayAddrIn; + kListGm_ = gmmArrayAddrIn + MKN_LIST_LEN; + nListGm_ = gmmArrayAddrIn + MKN_LIST_LEN * 2; + + xGm_ = x; + weightGm_ = weight; + antiquantScaleGm_ = antiquantScale; + antiquantOffsetGm_ = antiquantOffset; + biasGm_ = bias; + yGm_ = y; + if (groupList != nullptr) { + groupListGm_.SetGlobalBuffer((__gm__ int64_t *)groupList); + } + wqmmBasicBlock_.Init(gmmBaseTiling_->hasBias, gmmBaseTiling_->groupSize, 0, mmTiling_, + tPipe); // gmm场景不确定group是否激活，prefetch size设定为0 +} + +template +__aicore__ inline void GMMWeightQuantBasicController::Process() +{ + uint32_t cubeBlockIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + cubeBlockIdx = cubeBlockIdx >> 1; + } + + uint64_t preOffset = 0; + uint64_t xBaseOffset = 0; + uint64_t weightBaseOffset = 0; + uint64_t yBaseOffset = 0; + uint64_t antiquantParamsBaseOffset = 0; + + BasicBlockOffsetParam offsetParam; + InitOffsetParam(offsetParam); + for (uint32_t groupIdx = 0, count = 0; groupIdx < gmmBaseTiling_->groupNum; ++groupIdx) { + SetMKN(groupIdx, preOffset, offsetParam); // 每个核都必须知道之前group的信息，在单场景下作为baseOffset + SetGmAddr(groupIdx, xBaseOffset, weightBaseOffset, yBaseOffset, antiquantParamsBaseOffset, offsetParam); + if (offsetParam.mSize == 0) { + continue; + } + /* + * 1.在mSize小于mmTiling_.baseM或者mSize大于mmTiling_.baseM * 4时，baseM设置为mmTiling_.baseM + * 2.在mSize大于mmTiling_.baseM且小于等于mmTiling_.baseM * 2时，baseM使用mSize / 2向上取整 + * 3.在mSize大于mmTiling_.baseM * 2且小于等于mmTiling_.baseM * 4时，baseM使用mSize / 4向上取整 + */ + uint64_t baseM = + offsetParam.mSize <= mmTiling_->baseM || offsetParam.mSize >= mmTiling_->baseM * QUADRUPLE_BUFFER_NUM + ? mmTiling_->baseM + : (offsetParam.mSize <= DOUBLE_BUFFER_NUM * mmTiling_->baseM + ? CeilDivide(offsetParam.mSize, (uint64_t)DOUBLE_BUFFER_NUM) + : CeilDivide(offsetParam.mSize, (uint64_t)QUADRUPLE_BUFFER_NUM)); + uint64_t baseN = offsetParam.nSize >= mmTiling_->baseN * gmmBaseTiling_->coreNum ? mmTiling_->baseN + : (mmTiling_->baseN >> 1); + + uint32_t mBlockNum = CeilDivide(offsetParam.mSize, baseM); + uint32_t nBlockNum = CeilDivide(offsetParam.nSize, baseN); + + uint32_t curCount = count + mBlockNum * nBlockNum; + uint32_t curBlock = cubeBlockIdx >= count ? cubeBlockIdx : cubeBlockIdx + gmmBaseTiling_->coreNum; + + while (curBlock < curCount) { + offsetParam.mOffset = ((curBlock - count) / nBlockNum) * baseM; + offsetParam.nOffset = ((curBlock - count) % nBlockNum) * baseN; + offsetParam.mL1Size = + offsetParam.mOffset + baseM > offsetParam.mSize ? offsetParam.mSize - offsetParam.mOffset : baseM; + offsetParam.nL1Size = + offsetParam.nOffset + baseN > offsetParam.nSize ? offsetParam.nSize - offsetParam.nOffset : baseN; + + wqmmBasicBlock_.ComputeBasicBlock(offsetParam, offsetParam); + curBlock += gmmBaseTiling_->coreNum; + } + count = curCount % gmmBaseTiling_->coreNum; + } + wqmmBasicBlock_.End(offsetParam); +} + +template +__aicore__ inline void GMMWeightQuantBasicController::InitOffsetParam(BasicBlockOffsetParam &offsetParam) +{ + // 单的场景不需要频繁取值，n/k轴在开始的时候获取一次即可 + if (gmmBaseTiling_->singleWeight == 1) { + offsetParam.kSize = kListGm_[0]; + offsetParam.nSize = nListGm_[0]; + } + + offsetParam.kbL1Size = mmTiling_->baseK * mmTiling_->stepKb; + offsetParam.kaL1Size = offsetParam.kbL1Size; // 当前实现a矩阵切分保持b矩阵一致 +} + +template +__aicore__ inline void GMMWeightQuantBasicController::SetMKN( + uint64_t groupIdx, uint64_t &preOffset, BasicBlockOffsetParam &offsetParam) +{ + uint64_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset); + if (gmmBaseTiling_->groupType == 0) { + offsetParam.mSize = splitValue; + offsetParam.kSize = gmmBaseTiling_->singleWeight == 1 ? offsetParam.kSize : kListGm_[groupIdx]; + offsetParam.kAlign = + WeightQuantBatchMatmulV2::Arch35::CeilAlign(offsetParam.kSize, static_cast(BLOCK_CUBE)); + offsetParam.nSize = gmmBaseTiling_->singleWeight == 1 ? offsetParam.nSize : nListGm_[groupIdx]; + return; + } + + offsetParam.mSize = mListGm_[groupIdx]; + offsetParam.kSize = kListGm_[groupIdx]; + offsetParam.kAlign = + WeightQuantBatchMatmulV2::Arch35::CeilAlign(offsetParam.kSize, static_cast(BLOCK_CUBE)); + offsetParam.nSize = nListGm_[groupIdx]; +} + +template +__aicore__ inline void GMMWeightQuantBasicController::SetGmAddr( + uint64_t groupIdx, uint64_t &xBaseOffset, uint64_t &weightBaseOffset, uint64_t &yBaseOffset, + uint64_t &antiquantParamsBaseOffset, const BasicBlockOffsetParam &offsetParam) +{ + __gm__ xType *xGm; + __gm__ wType *weightGm; + __gm__ xType *antiquantScaleGm; + __gm__ xType *antiquantOffsetGm; + __gm__ biasType *biasGm; + __gm__ yType *yGm; + if (gmmBaseTiling_->singleX == 0) { + xGm = GetTensorAddr(groupIdx, xGm_); + } else { + xGm = GetTensorAddr(0, xGm_) + xBaseOffset; + } + + if (gmmBaseTiling_->singleY == 0) { + yGm = GetTensorAddr(groupIdx, yGm_); + } else { + yGm = GetTensorAddr(0, yGm_) + yBaseOffset; + } + + if (gmmBaseTiling_->singleWeight == 0) { + weightGm = GetTensorAddr(groupIdx, weightGm_); + antiquantScaleGm = GetTensorAddr(groupIdx, antiquantScaleGm_); + antiquantOffsetGm = GetTensorAddr(groupIdx, antiquantOffsetGm_); + biasGm = GetTensorAddr(groupIdx, biasGm_); + } else { + weightGm = GetTensorAddr(0, weightGm_) + weightBaseOffset; + antiquantScaleGm = GetTensorAddr(0, antiquantScaleGm_) + antiquantParamsBaseOffset; + antiquantOffsetGm = GetTensorAddr(0, antiquantOffsetGm_) + antiquantParamsBaseOffset; + biasGm = GetTensorAddr(0, biasGm_) + antiquantParamsBaseOffset; + } + wqmmBasicBlock_.UpdateGlobalAddr(xGm, weightGm, antiquantScaleGm, antiquantOffsetGm, nullptr, nullptr, biasGm, + yGm, mmTiling_->isBias, true); + xBaseOffset += offsetParam.mSize * offsetParam.kSize; + if constexpr (IsSameType::value) { + weightBaseOffset += (offsetParam.nSize * offsetParam.kSize) >> 1; + } else { + weightBaseOffset += offsetParam.nSize * offsetParam.kSize; + } + antiquantParamsBaseOffset += offsetParam.nSize; + yBaseOffset += offsetParam.mSize * offsetParam.nSize; +} + +template +__aicore__ inline uint64_t GMMWeightQuantBasicController::GetSplitValueFromGroupList(uint64_t groupIdx, + uint64_t &preOffset) +{ + uint64_t splitValue = 0; + if (likely(gmmBaseTiling_->groupType != -1)) { + if (gmmBaseTiling_->groupListType == 0) { + uint64_t offset = static_cast(groupListGm_.GetValue(groupIdx)); + splitValue = offset - preOffset; + preOffset = offset; + } else { + splitValue = static_cast(groupListGm_.GetValue(groupIdx)); + } + } + return splitValue; +} + +} // namespace GROUPED_MATMUL + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_CONTROLLER_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_resplit_controller.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_resplit_controller.h new file mode 100644 index 00000000..2ef09b41 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/grouped_matmul_weight_quant_resplit_controller.h @@ -0,0 +1,319 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_weight_quant_resplit_controller.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_RESPLIT_CONTROLLER_H +#define GROUPED_MATMUL_WEIGHT_QUANT_RESPLIT_CONTROLLER_H + +#include "weight_quant_vcv_basic_block_base.h" +#include "../grouped_matmul_tiling_data_apt.h" + +using WeightQuantBatchMatmulV2::Arch35::A_L1_MAX_SIZE_WITH_BIAS_QUANT; +using WeightQuantBatchMatmulV2::Arch35::BASIC_BLOCK_PROCESS_NUM; +using WeightQuantBatchMatmulV2::Arch35::BasicBlockControlParam; +using WeightQuantBatchMatmulV2::Arch35::BasicBlockOffsetParam; +using WeightQuantBatchMatmulV2::Arch35::CeilDivide; +using WeightQuantBatchMatmulV2::Arch35::DOUBLE_BUFFER_NUM; +using WeightQuantBatchMatmulV2::Arch35::IsMxA8W4; +using WeightQuantBatchMatmulV2::Arch35::QUADRUPLE_BUFFER_NUM; +using WeightQuantBatchMatmulV2::Arch35::QuantType; +using WeightQuantBatchMatmulV2::Arch35::SCALE_FACTOR_B_BIT; +using WeightQuantBatchMatmulV2::Arch35::VecAntiQuantConfig; +using WeightQuantBatchMatmulV2::Arch35::WeightQuantVcvMatmulBasicBlockBaseClass; +using WeightQuantBatchMatmulV2::Arch35::WqmmConfig; +using GMMWeightQuantParam = DlinferGroupedMatmulDirectTilingData::GMMWeightQuantParam; + +namespace GROUPED_MATMUL { +#define GMM_WQ_BASIC_BLOCK_TEMPLATE_CLASS \ + template \ + class +#define GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM \ + template + +#define GMM_WQ_RESPLIT_CONTROLLER_CLASS \ + GMMWeightQuantResplitController + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +class GMMWeightQuantResplitController { +public: + __aicore__ inline GMMWeightQuantResplitController(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR antiquantScale, + GM_ADDR antiquantOffset, GM_ADDR bias, GM_ADDR groupList, GM_ADDR perTokenScale, + GM_ADDR y, const GMMWeightQuantParam *__restrict baseTiling, + const TCubeTiling *__restrict mmTiling, GM_ADDR tiling, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitOffsetParam(BasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM]); + __aicore__ inline void SplitNByMultiCore(BasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM], + BasicBlockControlParam &ctrlParam, uint64_t basicBlockCount, + uint64_t basicBlockSize); + __aicore__ inline uint64_t GetSplitValueFromGroupList(uint64_t groupIdx); + __aicore__ inline void UpdateGmAddr(uint64_t mSize, uint64_t kSize, uint64_t nSize); + __aicore__ inline void PrefetchA(uint64_t mSize, uint64_t kSize); + __aicore__ inline uint64_t GetSwitchedProcessId(const BasicBlockControlParam &ctrlParam); + + const GMMWeightQuantParam *gmmBaseTiling_; + const TCubeTiling *mmTiling_; + + __gm__ xType *xGm_; + __gm__ wType *weightGm_; + __gm__ antiQuantScaleType *antiquantScaleGm_; + __gm__ xType *antiquantOffsetGm_; + __gm__ biasType *biasGm_; + __gm__ yType *yGm_; + __gm__ perTokenScaleType *perTokenScaleGm_; + __gm__ scaleType *scaleGm_; + GlobalTensor groupListGm_; + BasicBlock + basicBlock_; + + uint64_t preOffset_ = 0; +}; + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::Init( + GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR bias, + GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, const GMMWeightQuantParam *__restrict baseTiling, + const TCubeTiling *__restrict mmTiling, GM_ADDR tiling, TPipe *tPipe) +{ + gmmBaseTiling_ = baseTiling; + + mmTiling_ = mmTiling; + + xGm_ = GetTensorAddr(0, x); + weightGm_ = GetTensorAddr(0, weight); + antiquantScaleGm_ = GetTensorAddr(0, antiquantScale); + antiquantOffsetGm_ = GetTensorAddr(0, antiquantOffset); + biasGm_ = GetTensorAddr(0, bias); + scaleGm_ = GetTensorAddr(0, scale); + perTokenScaleGm_ = reinterpret_cast<__gm__ perTokenScaleType *>(perTokenScale); + yGm_ = GetTensorAddr(0, y); + if (groupList != nullptr) { + groupListGm_.SetGlobalBuffer((__gm__ int64_t *)groupList); + } + basicBlock_.Init(gmmBaseTiling_->hasBias, gmmBaseTiling_->groupSize, 0, mmTiling_, + tPipe); // gmm场景不确定group是否激活，Init中的prefetch size设定为0，在Process中做prefetch +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::Process() +{ + uint32_t cubeBlockIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + cubeBlockIdx = cubeBlockIdx >> 1; + } + + BasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM]; + InitOffsetParam(offsetParam); + + bool isCacheLineUnaligned = offsetParam[0].kSize % 128 != 0; // 缓存大小128B，对应8bit为128个元素 + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value) { + isCacheLineUnaligned = offsetParam[0].kSize % 256 != 0; // 缓存大小128B，对应4bit为256个元素 + } + + BasicBlockControlParam ctrlParam; + ctrlParam.processId = 0; + for (uint32_t groupIdx = 0, startBasicBlockId = 0; groupIdx < gmmBaseTiling_->groupNum; ++groupIdx) { + ctrlParam.mSize = GetSplitValueFromGroupList(groupIdx); + if (ctrlParam.mSize > 0) { + /* + * 1.在mSize小于mmTiling_.baseM或者mSize大于mmTiling_.baseM * 4时，mL1Size设置为mmTiling_.baseM + * 2.在mSize大于mmTiling_.baseM且小于等于mmTiling_.baseM * 2时，mL1Size使用mSize / 2向上取整 + * 3.在mSize大于mmTiling_.baseM * 2且小于等于mmTiling_.baseM * 4时，mL1Size使用mSize / 4向上取整 + */ + ctrlParam.mL1Size = + ctrlParam.mSize <= mmTiling_->baseM || ctrlParam.mSize >= mmTiling_->baseM * QUADRUPLE_BUFFER_NUM + ? mmTiling_->baseM + : (ctrlParam.mSize <= DOUBLE_BUFFER_NUM * mmTiling_->baseM + ? CeilDivide(ctrlParam.mSize, (uint64_t)DOUBLE_BUFFER_NUM) + : CeilDivide(ctrlParam.mSize, (uint64_t)QUADRUPLE_BUFFER_NUM)); + basicBlock_.UpdateGlobalAddr(xGm_, weightGm_, antiquantScaleGm_, antiquantOffsetGm_, scaleGm_, + perTokenScaleGm_, biasGm_, yGm_, mmTiling_->isBias, + ctrlParam.mL1Size < ctrlParam.mSize || isCacheLineUnaligned); + PrefetchA(ctrlParam.mSize, offsetParam[0].kSize); + ctrlParam.curBasicBlockId = + cubeBlockIdx >= startBasicBlockId ? cubeBlockIdx : cubeBlockIdx + gmmBaseTiling_->coreNum; + ctrlParam.basicBlockLimit = startBasicBlockId; + for (ctrlParam.mOffset = 0; ctrlParam.mOffset < ctrlParam.mSize; ctrlParam.mOffset += ctrlParam.mL1Size) { + ctrlParam.nOffset = 0; + + // 主块 + SplitNByMultiCore(offsetParam, ctrlParam, gmmBaseTiling_->mainBlockCount, + gmmBaseTiling_->mainBlockSize); + ctrlParam.basicBlockLimit += gmmBaseTiling_->mainBlockCount; + + // 第一段尾块 + SplitNByMultiCore(offsetParam, ctrlParam, gmmBaseTiling_->firstTailBlockCount, + gmmBaseTiling_->firstTailBlockSize); + ctrlParam.basicBlockLimit += gmmBaseTiling_->firstTailBlockCount; + + // 第二段尾块 + SplitNByMultiCore(offsetParam, ctrlParam, gmmBaseTiling_->secondTailBlockCount, + gmmBaseTiling_->secondTailBlockSize); + ctrlParam.basicBlockLimit += gmmBaseTiling_->secondTailBlockCount; + } + startBasicBlockId = ctrlParam.basicBlockLimit % gmmBaseTiling_->coreNum; + } + UpdateGmAddr(ctrlParam.mSize, offsetParam[0].kSize, offsetParam[0].nSize); + } + + basicBlock_.End(offsetParam[GetSwitchedProcessId(ctrlParam)]); +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::InitOffsetParam( + BasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM]) +{ + offsetParam[0].kbL1Size = mmTiling_->baseK * mmTiling_->stepKb; + offsetParam[0].kaL1Size = offsetParam[0].kbL1Size; // 当前实现a矩阵切分保持b矩阵一致 + offsetParam[0].kSize = gmmBaseTiling_->kSize; + offsetParam[0].nSize = gmmBaseTiling_->nSize; + offsetParam[0].kAlign = CeilAlign(gmmBaseTiling_->kSize, static_cast(BLOCK_CUBE)); + offsetParam[1].kbL1Size = mmTiling_->baseK * mmTiling_->stepKb; + offsetParam[1].kaL1Size = offsetParam[0].kbL1Size; // 当前实现a矩阵切分保持b矩阵一致 + offsetParam[1].kSize = offsetParam[0].kSize; + offsetParam[1].nSize = offsetParam[0].nSize; + offsetParam[1].kAlign = offsetParam[0].kAlign; + + if constexpr (IsMxA8W4()) { + offsetParam[0].nAlign = CeilAlign(gmmBaseTiling_->nSize, static_cast(BLOCK_CUBE)); + offsetParam[0].scaleAFactor = mmTiling_->mxTypePara & 0xff; + offsetParam[0].scaleBFactor = (mmTiling_->mxTypePara >> SCALE_FACTOR_B_BIT) & 0xff; + offsetParam[1].nAlign = offsetParam[0].nAlign; + offsetParam[1].scaleAFactor = offsetParam[0].scaleAFactor; + offsetParam[1].scaleBFactor = offsetParam[0].scaleBFactor; + } +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::SplitNByMultiCore( + BasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM], BasicBlockControlParam &ctrlParam, + uint64_t basicBlockCount, uint64_t basicBlockSize) +{ + for (; ctrlParam.curBasicBlockId < ctrlParam.basicBlockLimit + basicBlockCount; + ctrlParam.curBasicBlockId += gmmBaseTiling_->coreNum) { + offsetParam[ctrlParam.processId].mSize = ctrlParam.mSize; + offsetParam[ctrlParam.processId].mOffset = ctrlParam.mOffset; + offsetParam[ctrlParam.processId].mL1Size = ctrlParam.mOffset + ctrlParam.mL1Size > ctrlParam.mSize + ? ctrlParam.mSize - ctrlParam.mOffset + : ctrlParam.mL1Size; + offsetParam[ctrlParam.processId].nOffset = + ctrlParam.nOffset + + ((ctrlParam.curBasicBlockId - ctrlParam.basicBlockLimit) % basicBlockCount) * basicBlockSize; + offsetParam[ctrlParam.processId].nL1Size = + offsetParam[ctrlParam.processId].nOffset + basicBlockSize > gmmBaseTiling_->nSize + ? gmmBaseTiling_->nSize - offsetParam[ctrlParam.processId].nOffset + : basicBlockSize; + offsetParam[ctrlParam.processId].yGmAddr = reinterpret_cast(yGm_); + basicBlock_.ComputeBasicBlock(offsetParam[ctrlParam.processId], offsetParam[GetSwitchedProcessId(ctrlParam)]); + ctrlParam.processId = GetSwitchedProcessId(ctrlParam); + } + ctrlParam.nOffset += basicBlockSize * basicBlockCount; +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::UpdateGmAddr(uint64_t mSize, uint64_t kSize, uint64_t nSize) +{ + xGm_ += mSize * kSize; + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value) { + weightGm_ += (nSize * kSize) >> 1; + } else { + weightGm_ += nSize * kSize; + } + + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_GROUP || wqmmConfig.antiQuantType == QuantType::MX) { + antiquantScaleGm_ += nSize * CeilDivide(kSize, static_cast(gmmBaseTiling_->groupSize)); + antiquantOffsetGm_ += nSize * CeilDivide(kSize, static_cast(gmmBaseTiling_->groupSize)); + } else { + antiquantScaleGm_ += nSize; + antiquantOffsetGm_ += nSize; + } + + scaleGm_ += nSize; + + if constexpr (IsMxA8W4()) { + perTokenScaleGm_ += mSize * CeilDivide(kSize, static_cast(gmmBaseTiling_->groupSize)); + } else { + perTokenScaleGm_ += mSize; + } + + biasGm_ += nSize; + yGm_ += mSize * nSize; +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_RESPLIT_CONTROLLER_CLASS::PrefetchA(uint64_t mSize, uint64_t kSize) +{ + if ASCEND_IS_AIV { + return; + } + + uint64_t aSize = mSize * kSize * sizeof(xType); + + /* + * 准入条件： + * 1. m <= 512 + * 2. A的大小在cubeBlockDimN上可被一条mte2指令均分载入 + * 3. 核数是N分核数的倍数 + */ + if (mSize <= 512 && aSize <= static_cast(gmmBaseTiling_->cubeBlockDimN) * A_L1_MAX_SIZE_WITH_BIAS_QUANT && + (gmmBaseTiling_->coreNum % gmmBaseTiling_->cubeBlockDimN == 0)) { + uint64_t aPrefetchSize = + CeilAlign(CeilDivide(mSize * kSize, static_cast(gmmBaseTiling_->cubeBlockDimN)), + 64UL); // 64 表示128B的cacheline对齐 + basicBlock_.PrefetchA(aPrefetchSize, mSize * kSize); + } +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline uint64_t GMM_WQ_RESPLIT_CONTROLLER_CLASS::GetSplitValueFromGroupList(uint64_t groupIdx) +{ + uint64_t splitValue = 0; + if (likely(gmmBaseTiling_->groupType != -1)) { + if (gmmBaseTiling_->groupListType == 0) { + uint64_t offset = static_cast(groupListGm_.GetValue(groupIdx)); + splitValue = offset - preOffset_; + preOffset_ = offset; + } else { + splitValue = static_cast(groupListGm_.GetValue(groupIdx)); + } + } + return splitValue; +} + +GMM_WQ_RESPLIT_CONTROLLER_TEMPLATE_PARAM +__aicore__ inline uint64_t GMM_WQ_RESPLIT_CONTROLLER_CLASS::GetSwitchedProcessId( + const BasicBlockControlParam &ctrlParam) +{ + // vcv流水0/1倒换，vc流水ctrlParam.processId始终取0 + if constexpr (std::is_base_of_v>) { + return 1 - ctrlParam.processId; + } else { + return ctrlParam.processId; + } +} + +} // namespace GROUPED_MATMUL + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_RESPLIT_CONTROLLER_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/tool.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/tool.h new file mode 100644 index 00000000..73c8ef1b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/tool.h @@ -0,0 +1,182 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file tool.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_TOOL_H +#define GROUPED_MATMUL_WEIGHT_QUANT_TOOL_H + +#include "kernel_log.h" +#include "kernel_operator.h" +#include "kernel_utils.h" + +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; +using AscendC::DataCopyExtParams; +using AscendC::DataCopyPadExtParams; +using AscendC::fp8_e8m0_t; +using AscendC::GetUserWorkspace; +using AscendC::GlobalTensor; +using AscendC::int4b_t; +using AscendC::IsSameType; +using AscendC::LocalTensor; +using AscendC::ONE_BLK_SIZE; +using AscendC::TPosition; +using AscendC::VECTOR_REG_WIDTH; +using matmul::MatmulCallBackFunc; +using matmul::MatmulImpl; +using matmul::MatmulType; +using matmul::MatmulTypeWithScale; + +#define SHORT_MIX_LOG(format, ...) + +namespace WeightQuantBatchMatmulV2::Arch35 { +enum class QuantType { + NONE = 0, + PER_TENSOR = 1, + PER_CHANNEL = 2, + PER_GROUP = 3, + MX = 4, +}; + +// buffer相关定义 +static constexpr int32_t QUADRUPLE_BUFFER_NUM = 4; +static constexpr int32_t DOUBLE_BUFFER_NUM = 2; +static constexpr int32_t SINGLE_BUFFER_NUM = 1; +static constexpr int64_t L1_SIZE = 512; +static constexpr int64_t L1_SIZE_BYTE = L1_SIZE * 1024; +static constexpr int64_t L1_HALF_SIZE = L1_SIZE / 2; +static constexpr int64_t L1_SIZE_WITH_QUANTSCALE = 504; +static constexpr int64_t L1_SIZE_WITH_QUANTSCALE_BYTE = L1_SIZE_WITH_QUANTSCALE * 1024; +static constexpr int64_t BIAS_L1_SIZE = 4; +static constexpr int64_t MX_SCALE_L1_SIZE = 20; +static constexpr uint64_t A_L1_MAX_SIZE_WITH_BIAS_QUANT = 240UL * 1024UL; + +// 控制参数定义 +static constexpr int32_t BASIC_BLOCK_PROCESS_NUM = 2; +static constexpr uint64_t SCALE_COPY_GROUP_SIZE = 2; +static constexpr int32_t SCALE_COPY_DEFAULT_STRIDE = 0; +static constexpr int32_t SCALE_COPY_DEFAULT_N_STRIDE = 1; + +// 参数约束定义 +static constexpr uint64_t MX_GROUPSIZE = 32; +static constexpr uint64_t VEC_MAX_ELEM_B16 = VECTOR_REG_WIDTH / sizeof(half); +static constexpr uint32_t FP32_BLOCK_SIZE = 8; +static constexpr int32_t C0_SIZE_B8 = 32; +static constexpr uint32_t SCALE_FACTOR_B_BIT = 8; + +// 同步定义 +static constexpr uint64_t SYNC_AIV_AIC_FLAG = 8; +static constexpr uint64_t SYNC_AIC_AIV_FLAG = 9; +static constexpr uint64_t SYNC_AIC_FIX_AIV_VF_FLAG = 3; +static constexpr uint64_t SYNC_AIV_MTE3_AIC_FIX_FLAG = 4; +static constexpr uint64_t SYNC_MODE4 = 4; +static constexpr uint64_t FLAG_ID_MAX = 16; + +// 函数定义 +template +__aicore__ inline T CeilAlign(T a, T b) +{ + ASCENDC_ASSERT(b != 0, { KERNEL_LOG(KERNEL_ERROR, "Division by zero error!"); }); + return (a + b - 1) / b * b; +} + +__aicore__ inline uint32_t CeilAlign(uint32_t a, uint32_t b) +{ + ASCENDC_ASSERT(a <= (std::numeric_limits::max() - b), + { KERNEL_LOG(KERNEL_ERROR, "CeilAlign uint32 over limit."); }); + ASCENDC_ASSERT(b != 0, { KERNEL_LOG(KERNEL_ERROR, "Division by zero error!"); }); + return (a + b - 1) / b * b; +} + +template +__aicore__ inline T CeilDivide(T a, T b) +{ + ASCENDC_ASSERT(b != 0, { KERNEL_LOG(KERNEL_ERROR, "Division by zero error!"); }); + return (a + b - 1) / b; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a < b ? a : b; +} + +template +__aicore__ inline void DataCopyPad2D(const LocalTensor &dst, const GlobalTensor &src, uint32_t blockCount, + uint32_t blockLen, uint32_t dstInnerLength, uint32_t srcInnerLength) +{ + DataCopyExtParams params; + params.blockCount = blockCount; + params.blockLen = blockLen * sizeof(T); + params.srcStride = (srcInnerLength - blockLen) * sizeof(T); + params.dstStride = (dstInnerLength - blockLen) * sizeof(T) / ONE_BLK_SIZE; + DataCopyPadExtParams padParams; + if (blockLen % (32 / sizeof(T)) != 0) { + padParams.isPad = true; + padParams.rightPadding = CeilAlign(blockLen, static_cast(32 / sizeof(T))) - blockLen; + padParams.paddingValue = 0; + } + + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value) { + // 4bit场景下， 跳转的步长、数据长度等需要除2 + params.blockLen = params.blockLen >> 1; + params.srcStride = params.srcStride >> 1; + params.dstStride = params.dstStride >> 1; + padParams.rightPadding = padParams.rightPadding >> 1; + } + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2D(const GlobalTensor &dst, const LocalTensor &src, uint32_t dim1, + uint32_t dim0, uint32_t srcFullDim0, uint32_t dstFullDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = CeilDivide((srcFullDim0 - dim0) * sizeof(T), static_cast(ONE_BLK_SIZE)); + params.dstStride = (dstFullDim0 - dim0) * sizeof(T); + SHORT_MIX_LOG("dim1 %d dim0 %d dstFullDim0 %d blockCount %d blockLen %d srcStride %d dstStride %d", dim1, dim0, + dstFullDim0, params.blockCount, params.blockLen, params.srcStride, params.dstStride); + DataCopyPad(dst, src, params); +} + +template +__aicore__ constexpr uint32_t GetKBUnit() +{ + if constexpr (IsSameType::value) { + return 2048; // 2048个int4是1kb + } + if constexpr (IsSameType::value || IsSameType::value) { + return 1024; // 1024个B8是1kb + } + if constexpr (IsSameType::value) { + return 256; // 256个float是1kb + } + return 512; // 512个half是1kb +} + +template +__aicore__ constexpr bool IsMxA8W4() +{ + return antiQuantType == QuantType::MX && IsSameType::value; +} + +template +struct MatmulL1GmType : MatmulType { + constexpr static TPosition srcPos = TPosition::GM; +}; +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block.h new file mode 100644 index 00000000..31a330fa --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block.h @@ -0,0 +1,346 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file weight_quant_basic_block.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "lib/matmul_intf.h" +#include "tool.h" +#include "weight_quant_basic_block_base.h" +#include "weight_quant_cube_compute.h" +#include "weight_quant_vec_compute.h" + +using AscendC::Conditional; +using AscendC::GetSubBlockIdx; +using AscendC::IsSameType; +using AscendC::LocalTensor; +using AscendC::TBuf; +using AscendC::TPipe; +using AscendC::TPosition; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +#define GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM \ + template + +#define GMM_WQ_BASIC_BLOCK_CLASS \ + WeightQuantMatmulBasicBlock + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +class WeightQuantMatmulBasicBlock : public WeightQuantMatmulBasicBlockBaseClass { +public: + __aicore__ inline WeightQuantMatmulBasicBlock(){}; + __aicore__ inline void Init(bool hasBias, uint64_t antiQuantGroupSize, uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, TPipe *tPipe); + __aicore__ inline void UpdateGlobalAddr(__gm__ xType *x, __gm__ wType *weight, + __gm__ antiQuantScaleType *antiquantScale, __gm__ xType *antiquantOffset, + __gm__ scaleType *scale, __gm__ perTokenScaleType *perTokenScale, + __gm__ biasType *bias, __gm__ yType *y, const bool hasBias, + const bool weightL2Cacheable); + __aicore__ inline void ComputeBasicBlock(const BasicBlockOffsetParam &offsetParam, + const BasicBlockOffsetParam &lastOffsetParam); + __aicore__ inline void PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit); + __aicore__ inline void End(const BasicBlockOffsetParam &offsetParam); + +protected: + __aicore__ inline void SetAivToAic(); + __aicore__ inline void WaitAivToAic(); + __aicore__ inline void SetAicToAiv(); + __aicore__ inline void WaitAicToAiv(); + __aicore__ inline void ComputeBasicBlockAivNdNkNzKn(const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void ComputeBasicBlockAivNdKnNzNk(const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void ComputeBasicBlockAic(const BasicBlockOffsetParam &offsetParam); + + BasicBlockLibVectorAntiQuantCompute vectorCompute_; + + using aType = typename Conditional< + IsMxA8W4(), + MatmulTypeWithScale, + MatmulL1GmType>::type; + using bType = typename Conditional< + IsMxA8W4(), + MatmulTypeWithScale, + MatmulL1GmType>::type; + using cType = MatmulType; + using biasMatmulType = MatmulType; + + using MMImpl = typename Conditional< + IsMxA8W4(), + MatmulImpl, + AscendC::Impl::Detail::MatmulWithScalePolicy>, + MatmulImpl>::type; + WeightQuantBatchMatmulV2CubeCompute + cubeCompute_; + + uint64_t cvLoopIdx_ = 0; + + LocalTensor weightL1_; + uint64_t weightL1DbOffset_; +}; + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::Init(bool hasBias, uint64_t antiQuantGroupSize, uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, TPipe *tPipe) +{ + TBuf l1Tbuf; + uint64_t weightL1Space = matmulTiling->baseN * matmulTiling->stepKb * matmulTiling->baseK; // weight单块大小 + if constexpr (IsSameType::value) { + tPipe->InitBuffer(l1Tbuf, L1_SIZE_WITH_QUANTSCALE_BYTE); // 除去quantScale, 共使用504KB + weightL1DbOffset_ = L1_SIZE_WITH_QUANTSCALE * GetKBUnit() - weightL1Space; + } else if constexpr (IsMxA8W4()) { + tPipe->InitBuffer(l1Tbuf, L1_SIZE_BYTE); + weightL1DbOffset_ = L1_SIZE * GetKBUnit() - weightL1Space; + } else { + tPipe->InitBuffer(l1Tbuf, L1_SIZE_BYTE); + weightL1DbOffset_ = L1_SIZE * GetKBUnit() - weightL1Space; + } + weightL1_ = l1Tbuf.Get(); + if ASCEND_IS_AIC { + cubeCompute_.Init(l1Tbuf, weightL1Space, aPrefetchSize, matmulTiling, tPipe); + } else { + vectorCompute_.Init(tPipe); + } +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::UpdateGlobalAddr( + __gm__ xType *x, __gm__ wType *weight, __gm__ antiQuantScaleType *antiquantScale, __gm__ xType *antiquantOffset, + __gm__ scaleType *scale, __gm__ perTokenScaleType *perTokenScale, __gm__ biasType *bias, __gm__ yType *y, + const bool hasBias, const bool weightL2Cacheable) +{ + if ASCEND_IS_AIC { + cubeCompute_.UpdateGlobalAddr(x, y, bias, antiquantScale, scale, perTokenScale, hasBias); + } else { + vectorCompute_.UpdateGlobalAddr(weight, antiquantScale, antiquantOffset, nullptr, nullptr, nullptr, + weightL2Cacheable); + } +} + +/* + * 该函数作用为根据转置属性，L1上shape大小以及vecconfig确定vec核的实际搬运量 + * ND TransB = True 两个vec核的mte2搬运量为 (curVecCoreMte2RealN, curVecCoreMte2RealK) + * NZ TransB = False 两个vec核的mte2搬运量为 (CeilDiv(curVecCoreMte2RealN,16), CeilDiv(curVecCoreMte2RealK,16), 16 ,16) + */ +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::ComputeBasicBlockAivNdNkNzKn(const BasicBlockOffsetParam &offsetParam) +{ + UbConsumeConfig ubConsumeConfig; + L1ConsumeConfig l1ConsumeConfig; + ubConsumeConfig.nWeightLowBitUbOffset = 0; + l1ConsumeConfig.l1RealExternalLen = offsetParam.nL1Size; + +#if defined(__DAV_310R6__) + uint64_t curVecCoreMte2RealN = offsetParam.nL1Size; +#else + // 初值为一半的nl1 + uint64_t curVecCoreMte2RealN = offsetParam.nL1Size >> 1; + if constexpr (wqmmConfig.weightFormat == CubeFormat::NZ) { + curVecCoreMte2RealN = offsetParam.nL1Size > BLOCK_CUBE + ? CeilAlign(curVecCoreMte2RealN, static_cast(BLOCK_CUBE)) + : offsetParam.nL1Size; + } +#endif + // 实际值需要根据vec核来确定 + ubConsumeConfig.l1RequireVfComputeRealN = + GetSubBlockIdx() == 0 ? curVecCoreMte2RealN : offsetParam.nL1Size - curVecCoreMte2RealN; + l1ConsumeConfig.l1SplitTwoVecExternalOffset = GetSubBlockIdx() * curVecCoreMte2RealN; + + for (uint64_t kMte2Offset = 0; kMte2Offset < offsetParam.kSize; kMte2Offset += vecConfig.ubMte2InnerSize) { + uint64_t mte2RealK = (kMte2Offset + vecConfig.ubMte2InnerSize) > offsetParam.kSize + ? offsetParam.kSize - kMte2Offset + : vecConfig.ubMte2InnerSize; // vec总共需要搬运的K方向的实际量（考虑尾块） + vectorCompute_.WaitVToMTE2(); + vectorCompute_.CopyGmToUb(ubConsumeConfig.l1RequireVfComputeRealN, mte2RealK, + offsetParam.nOffset + l1ConsumeConfig.l1SplitTwoVecExternalOffset, kMte2Offset, + offsetParam); + + // 当前方案下，不会出现N方向计算量小于载入量的情况，所以没有N的循环 + for (ubConsumeConfig.kWeightLowBitUbOffset = 0; ubConsumeConfig.kWeightLowBitUbOffset < mte2RealK; + ubConsumeConfig.kWeightLowBitUbOffset += offsetParam.kbL1Size, cvLoopIdx_++) { + ubConsumeConfig.l1RequireVfComputeRealK = + (ubConsumeConfig.kWeightLowBitUbOffset + offsetParam.kbL1Size) >= mte2RealK + ? mte2RealK - ubConsumeConfig.kWeightLowBitUbOffset + : offsetParam.kbL1Size; + if (cvLoopIdx_ > 1) { + WaitAicToAiv(); + } + vectorCompute_.WeightAntiQuantCompute(ubConsumeConfig, weightL1_[(cvLoopIdx_ & 1) * weightL1DbOffset_], + l1ConsumeConfig); + SetAivToAic(); + } + vectorCompute_.SetVToMTE2(); + } +} + +/* + * 该函数作用为根据转置属性，L1上shape大小以及vecconfig确定vec核的实际搬运量 + * ND TransB = False 两个vec核的mte2搬运量为 (curVecCoreMte2RealK, curVecCoreMte2RealN) + * NZ TransB = True 两个vec核的mte2搬运量为 (CeilDiv(curVecCoreMte2RealK, C0), CeilDiv(curVecCoreMte2RealN, 16), 16, C0) + */ +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::ComputeBasicBlockAivNdKnNzNk(const BasicBlockOffsetParam &offsetParam) +{ +// nd-kn场景下，搬运消费比为1:1 +#if defined(__DAV_310R6__) + // 此时c:v为1：1, cube所需数据由一个v提供 + uint64_t kMte2BaseSize = offsetParam.kbL1Size; +#else + // 此时c:v为1：2, cube所需数据由两个v提供 + uint64_t kMte2BaseSize = offsetParam.kbL1Size >> 1; +#endif + + UbConsumeConfig ubConsumeConfig; + L1ConsumeConfig l1ConsumeConfig; + ubConsumeConfig.l1RequireVfComputeRealN = offsetParam.nL1Size; + ubConsumeConfig.kWeightLowBitUbOffset = 0; + ubConsumeConfig.nWeightLowBitUbOffset = 0; + l1ConsumeConfig.l1SplitTwoVecExternalOffset = GetSubBlockIdx() * kMte2BaseSize; + + for (uint64_t kMte2Offset = 0; kMte2Offset < offsetParam.kSize; kMte2Offset += offsetParam.kbL1Size, cvLoopIdx_++) { + l1ConsumeConfig.l1RealExternalLen = (kMte2Offset + offsetParam.kbL1Size) > offsetParam.kSize + ? offsetParam.kSize - kMte2Offset + : offsetParam.kbL1Size; + /* + * 场景1：当前core为v0，mte2实际搬运的k值为mte2方向搬运标准值(kMte2BaseSize)和l1上k实际值的最小值 + * 场景2: 当前core为v1, 且l1实际的k值比core v1搬运值大，则mte2实际搬运的k值为两者之差 + * 场景3：当前core为v1, 且l1实际的k值比core v1搬运值小，则mte2无需搬运,设置为0即可 + */ + uint64_t mte2RealK = GetSubBlockIdx() == 0 ? min(kMte2BaseSize, l1ConsumeConfig.l1RealExternalLen) + : l1ConsumeConfig.l1RealExternalLen > kMte2BaseSize + ? l1ConsumeConfig.l1RealExternalLen - kMte2BaseSize + : 0; + vectorCompute_.WaitVToMTE2(); + vectorCompute_.CopyGmToUb(offsetParam.nL1Size, mte2RealK, offsetParam.nOffset, + kMte2Offset + GetSubBlockIdx() * kMte2BaseSize, offsetParam); + + if (cvLoopIdx_ > 1) { + WaitAicToAiv(); + } + ubConsumeConfig.l1RequireVfComputeRealK = mte2RealK; + vectorCompute_.WeightAntiQuantCompute(ubConsumeConfig, weightL1_[(cvLoopIdx_ & 1) * weightL1DbOffset_], + l1ConsumeConfig); + SetAivToAic(); + vectorCompute_.SetVToMTE2(); + } +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::ComputeBasicBlockAic(const BasicBlockOffsetParam &offsetParam) +{ + // 当前方案下，不会出现N方向计算量小于载入量的情况，所以没有N的循环 + for (uint64_t kbL1Offset = 0; kbL1Offset < offsetParam.kSize; kbL1Offset += offsetParam.kbL1Size, cvLoopIdx_++) { + uint64_t kbL1RealSize = (kbL1Offset + offsetParam.kbL1Size) >= offsetParam.kSize + ? offsetParam.kSize - kbL1Offset + : offsetParam.kbL1Size; + cubeCompute_.WaitMTE1ToMTE2(cvLoopIdx_); + if constexpr (IsMxA8W4()) { + // 当前仅支持scale单倍载入 + cubeCompute_.CopyMxScaleGmToL1(offsetParam, kbL1Offset, cvLoopIdx_); + } + cubeCompute_.CopyAAndBiasGmToL1(offsetParam, kbL1Offset, kbL1RealSize, offsetParam.nL1Size, cvLoopIdx_); + WaitAivToAic(); + cubeCompute_.LaunchMatmul(weightL1_[(cvLoopIdx_ & 1) * weightL1DbOffset_], kbL1Offset, kbL1RealSize, + offsetParam, cvLoopIdx_); // mte1 mmad fixp流水 + cubeCompute_.SetMTE1ToMTE2(cvLoopIdx_); + SetAicToAiv(); + } + cubeCompute_.GetTensorC(offsetParam); + cubeCompute_.ClearAFullLoadFlag(); // 清除A全载时之前循环的set同步标记 +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::ComputeBasicBlock(const BasicBlockOffsetParam &offsetParam, + const BasicBlockOffsetParam &lastOffsetParam) +{ + if ASCEND_IS_AIV { + if constexpr ((wqmmConfig.weightFormat != CubeFormat::NZ && !wqmmConfig.bTrans) || + (wqmmConfig.weightFormat == CubeFormat::NZ && wqmmConfig.bTrans)) { + ComputeBasicBlockAivNdKnNzNk(offsetParam); + } else { + ComputeBasicBlockAivNdNkNzKn(offsetParam); + } + } else { + ComputeBasicBlockAic(offsetParam); + } +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit) +{ + cubeCompute_.PrefetchA(aPrefetchSize, xSizeLimit); +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::End(const BasicBlockOffsetParam &offsetParam) +{ + if ASCEND_IS_AIC { + cubeCompute_.EndSync(cvLoopIdx_); + } else { + if (cvLoopIdx_ > 0) { + WaitAicToAiv(); + } + if (cvLoopIdx_ > 1) { + WaitAicToAiv(); + } + vectorCompute_.End(); + } +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::SetAivToAic() +{ +#ifndef __CCE_KT_TEST__ + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); +#endif +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::WaitAivToAic() +{ +#ifndef __CCE_KT_TEST__ +#ifndef __DAV_310R6__ + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); +#endif + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG); +#endif +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::SetAicToAiv() +{ +#ifndef __CCE_KT_TEST__ +#if !defined(__DAV_310R6__) + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG + FLAG_ID_MAX); +#endif + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG); +#endif +} + +GMM_WQ_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_BASIC_BLOCK_CLASS::WaitAicToAiv() +{ +#ifndef __CCE_KT_TEST__ + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); +#endif +} +} // namespace WeightQuantBatchMatmulV2::Arch35 + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block_base.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block_base.h new file mode 100644 index 00000000..775b735b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_basic_block_base.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file weight_quant_basic_block_base.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_BASE_H +#define GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_BASE_H + +namespace WeightQuantBatchMatmulV2::Arch35 { + +class WeightQuantMatmulBasicBlockBaseClass { +public: + __aicore__ inline WeightQuantMatmulBasicBlockBaseClass(){}; +}; + +} // namespace WeightQuantBatchMatmulV2::Arch35 + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_BASIC_BLOCK_BASE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_cube_compute.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_cube_compute.h new file mode 100644 index 00000000..cb283166 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_cube_compute.h @@ -0,0 +1,548 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file weight_quant_cube_compute.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_CUBE_COMPUTE_H +#define GROUPED_MATMUL_WEIGHT_QUANT_CUBE_COMPUTE_H + +#include "basic_block_config.h" +#include "custom_policy/wqbmm_custom_policy.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "lib/matmul_intf.h" +#include "tool.h" + +using AscendC::BLOCK_CUBE; +using AscendC::Dn2NzParams; +using AscendC::GetBlockIdx; +using AscendC::GlobalTensor; +using AscendC::HardEvent; +using AscendC::IsSameType; +using AscendC::LocalTensor; +using AscendC::PipeBarrier; +using AscendC::SetFlag; +using AscendC::TBuf; +using AscendC::TPosition; +using AscendC::WaitFlag; + +namespace WeightQuantBatchMatmulV2::Arch35 { + +#define WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM \ + template + +#define WQBMM_CUBE_COMPUTE_CLASS \ + WeightQuantBatchMatmulV2CubeCompute + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +class WeightQuantBatchMatmulV2CubeCompute { +public: + __aicore__ inline WeightQuantBatchMatmulV2CubeCompute(){}; + __aicore__ inline void UpdateGlobalAddr(__gm__ xType *x, __gm__ yType *y, __gm__ biasType *bias, + __gm__ antiQuantScaleType *antiquantScale, __gm__ uint64_t *quantScale, + __gm__ perTokenScaleType *perTokenScale, const bool isBias); + __aicore__ inline void Init(TBuf &l1Tbuf, uint64_t weightL1Space, uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, AscendC::TPipe *tPipe); + __aicore__ inline void LaunchMatmul(const LocalTensor &weightL1, int64_t kbOffset, uint64_t kbL1RealSize, + const BasicBlockOffsetParam ¶m, uint64_t cvLoopIdx); + __aicore__ inline void WaitMTE1ToMTE2(uint64_t cvLoopIdx); + __aicore__ inline void SetMTE1ToMTE2(uint64_t cvLoopIdx); + __aicore__ inline void CopyAAndBiasGmToL1(const BasicBlockOffsetParam ¶m, int64_t kaGmOffset, + int64_t kbL1RealSize, int64_t biasRealN, uint64_t cvLoopIdx); + __aicore__ inline void CopyMxScaleGmToL1(const BasicBlockOffsetParam ¶m, uint64_t kbL1Offset, + uint64_t cvLoopIdx); + __aicore__ inline void GetTensorC(const BasicBlockOffsetParam ¶m); + __aicore__ inline void GetTensorC(LocalTensor &yUb); + __aicore__ inline void EndSync(uint64_t cvLoopIdx); + __aicore__ inline void ClearAFullLoadFlag(); + __aicore__ inline void PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit); + +private: + __aicore__ inline void PrefetchA(uint64_t aPrefetchSize, const LocalTensor &perloadBuffer, + const TCubeTiling *__restrict matmulTiling); + __aicore__ inline void InitSync(); + __aicore__ inline uint64_t CheckMaxSpace(const BasicBlockOffsetParam ¶m); + __aicore__ inline void CopyAGmToL1SingleBuffer(const BasicBlockOffsetParam ¶m, int64_t kaGmOffset, + int64_t kbL1RealSize, int64_t biasRealN, uint64_t cvLoopIdx, + int64_t aGmOffset); + __aicore__ inline void ConfigScaleDn2NzParams(uint64_t rowNum, uint64_t scaleKGmSize, uint64_t scaleKL1Stride, + uint64_t scaleKL1RealSize, Dn2NzParams &dn2NzParams); + + int8_t aL1DbNum_; + bool isBias_; + uint64_t quantScaleValue_; + static constexpr uint32_t KB_UNIT = GetKBUnit(); + + MatmulImplType mmObj_; + + uint64_t aL1Count_; + uint64_t aL1MaxHalfCount_; + + AscendC::TEventID cubeEventIdsMte1ToMte2_[DOUBLE_BUFFER_NUM]; + AscendC::TEventID cubeEventIdMte2ToMte1_; + GlobalTensor xGlobal_; + GlobalTensor biasGlobal_; + GlobalTensor mxScaleAGlobal_; + GlobalTensor mxScaleBGlobal_; + GlobalTensor quantScaleGlobal_; + GlobalTensor yGlobal_; + + LocalTensor aL1_; + uint64_t aL1DbOffset_; + + LocalTensor biasL1_; + uint64_t biasL1DbOffset_; + + LocalTensor mxScaleAL1_; + uint64_t mxScaleAL1DbOffset_; + + LocalTensor mxScaleBL1_; + uint64_t mxScaleBL1DbOffset_; +}; + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline uint64_t WQBMM_CUBE_COMPUTE_CLASS::CheckMaxSpace(const BasicBlockOffsetParam ¶m) +{ + uint64_t maxSpace = aL1MaxHalfCount_ * param.kbL1Size * CeilAlign(param.mL1Size, static_cast(BLOCK_CUBE)); + if (param.kbL1Size > 0 && param.kSize % param.kbL1Size == 0 && !wqmmConfig.aTrans && maxSpace <= aL1DbOffset_) { + return maxSpace; + } + return 0; +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::LaunchMatmul(const LocalTensor &weightL1, int64_t kbOffset, + uint64_t kbL1RealSize, const BasicBlockOffsetParam ¶m, + uint64_t cvLoopIdx) +{ + mmObj_.SetOrgShape(CeilAlign(param.mL1Size, static_cast(BLOCK_CUBE)), + CeilAlign(param.nL1Size, static_cast(BLOCK_CUBE)), + CeilAlign(kbL1RealSize, static_cast(BLOCK_CUBE)), + CeilAlign(kbL1RealSize, static_cast(BLOCK_CUBE)), param.nSize); + if (aL1DbNum_ == SINGLE_BUFFER_NUM) { + uint64_t maxSpace = CheckMaxSpace(param); + if (maxSpace > 0) { + // block = kbOffset / param.kbL1Size 计算是在第几块 + // blockOffset = block / 2 确定从A0还是A1读取数据后，在块内的偏移，单位是块 + // k = blockOffset * param.kbL1Size 当前块内的偏移量kOffset，单位是元素 + // 将block和blockOffset带入，计算k + // k = (kbOffset / param.kbL1Size) / 2 * param.kbL1Size + // 块内偏移量 = m * k + // 举例： + // L1A: |0|2|4| |1|3| + // |A0:0~128KB |A1:128KB~256KB| + // 第5块（block = 5），在A1中偏移为3（blockOffset = 3），块内偏移量为m * (3 * k) + mmObj_.SetTensorA(aL1_[(cvLoopIdx & 1) * aL1DbOffset_ + + CeilAlign(param.mL1Size, static_cast(BLOCK_CUBE)) * + (static_cast(kbOffset) / (param.kbL1Size * 2) * param.kbL1Size)], + wqmmConfig.aTrans); + } else { + mmObj_.SetTensorA(aL1_[CeilAlign(param.mL1Size, static_cast(BLOCK_CUBE)) * kbOffset], + wqmmConfig.aTrans); + } + } else { + mmObj_.SetTensorA(aL1_[(cvLoopIdx & 1) * aL1DbOffset_], wqmmConfig.aTrans); + } + + mmObj_.SetTensorB(weightL1, wqmmConfig.bTrans); + + if constexpr (IsMxA8W4()) { + mmObj_.SetTensorScaleA(mxScaleAL1_[(cvLoopIdx & 1) * mxScaleAL1DbOffset_], wqmmConfig.aTrans); + mmObj_.SetTensorScaleB(mxScaleBL1_[(cvLoopIdx & 1) * mxScaleBL1DbOffset_], wqmmConfig.bTrans); + } + + if (isBias_) { + mmObj_.SetBias(biasL1_[(cvLoopIdx & 1) * biasL1DbOffset_]); + } + + mmObj_.SetTail(param.mL1Size, param.nL1Size, kbL1RealSize); + + if constexpr (IsSameType::value) { + if constexpr (wqmmConfig.quantType == QuantType::PER_TENSOR) { + mmObj_.SetQuantScalar(quantScaleValue_); + } else { + mmObj_.SetQuantVector(quantScaleGlobal_[param.nOffset]); + } + } + + mmObj_.Iterate(kbOffset != 0); +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::WaitMTE1ToMTE2(uint64_t cvLoopIdx) +{ + // 编译器对成员变量数组访问优化能力较弱，会引入大量scalar，此处抽取局部变量，规避编译器优化问题 + AscendC::TEventID tempEventIdsMte1ToMte2[DOUBLE_BUFFER_NUM] = {cubeEventIdsMte1ToMte2_[0], + cubeEventIdsMte1ToMte2_[1]}; + // 单buffer时保证了A一次全载不需要Wait，Double buffer时首次使用不需要Wait + if (aL1DbNum_ > SINGLE_BUFFER_NUM && cvLoopIdx >= DOUBLE_BUFFER_NUM) { + WaitFlag(tempEventIdsMte1ToMte2[cvLoopIdx & 1]); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::SetMTE1ToMTE2(uint64_t cvLoopIdx) +{ + // 编译器对成员变量数组访问优化能力较弱，会引入大量scalar，此处抽取局部变量，规避编译器优化问题 + AscendC::TEventID tempEventIdsMte1ToMte2[DOUBLE_BUFFER_NUM] = {cubeEventIdsMte1ToMte2_[0], + cubeEventIdsMte1ToMte2_[1]}; + if (aL1DbNum_ > SINGLE_BUFFER_NUM) { + SetFlag(tempEventIdsMte1ToMte2[cvLoopIdx & 1]); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::CopyAGmToL1SingleBuffer(const BasicBlockOffsetParam ¶m, + int64_t kaGmOffset, int64_t kbL1RealSize, + int64_t biasRealN, uint64_t cvLoopIdx, + int64_t aGmOffset) +{ + AscendC::Nd2NzParams nd2nzParams; + uint64_t maxSpace = CheckMaxSpace(param); + if (maxSpace > 0) { + nd2nzParams.ndNum = aL1MaxHalfCount_; + nd2nzParams.nValue = param.mL1Size; + nd2nzParams.dValue = param.kbL1Size; + nd2nzParams.srcDValue = param.kSize; + nd2nzParams.srcNdMatrixStride = 2 * nd2nzParams.dValue; + nd2nzParams.dstNzC0Stride = CeilAlign(nd2nzParams.nValue, static_cast(BLOCK_CUBE)); + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = + nd2nzParams.dstNzC0Stride * CeilAlign(nd2nzParams.dValue, static_cast(BLOCK_CUBE)); + DataCopy(aL1_[(cvLoopIdx & 1) * aL1DbOffset_], xGlobal_[aGmOffset], nd2nzParams); + + nd2nzParams.ndNum = aL1Count_ - aL1MaxHalfCount_; + DataCopy(aL1_[((cvLoopIdx + 1) & 1) * aL1DbOffset_], xGlobal_[aGmOffset + nd2nzParams.dValue], nd2nzParams); + } else { + nd2nzParams.ndNum = 1; + if constexpr (wqmmConfig.aTrans) { + nd2nzParams.nValue = param.kSize; + nd2nzParams.dValue = param.mL1Size; + nd2nzParams.srcDValue = param.mSize; + } else { + nd2nzParams.nValue = param.mL1Size; + nd2nzParams.dValue = param.kSize; + nd2nzParams.srcDValue = param.kSize; + } + nd2nzParams.srcNdMatrixStride = 0; + nd2nzParams.dstNzC0Stride = CeilAlign(nd2nzParams.nValue, static_cast(BLOCK_CUBE)); + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 0; + + DataCopy(aL1_, xGlobal_[aGmOffset], nd2nzParams); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::CopyAAndBiasGmToL1(const BasicBlockOffsetParam ¶m, + int64_t kaGmOffset, int64_t kbL1RealSize, + int64_t biasRealN, uint64_t cvLoopIdx) +{ + int64_t aGmOffset; + if constexpr (!wqmmConfig.aTrans) { + aGmOffset = param.mOffset * param.kSize + kaGmOffset; + } else { + aGmOffset = kaGmOffset * param.mSize + param.mOffset; + } + + if (aL1DbNum_ > SINGLE_BUFFER_NUM) { + AscendC::Nd2NzParams nd2nzParams; + nd2nzParams.ndNum = 1; + if constexpr (wqmmConfig.aTrans) { + nd2nzParams.nValue = kbL1RealSize; + nd2nzParams.dValue = param.mL1Size; + nd2nzParams.srcDValue = param.mSize; + } else { + nd2nzParams.nValue = param.mL1Size; + nd2nzParams.dValue = kbL1RealSize; + nd2nzParams.srcDValue = param.kSize; + } + nd2nzParams.srcNdMatrixStride = 0; + nd2nzParams.dstNzC0Stride = CeilAlign(nd2nzParams.nValue, static_cast(BLOCK_CUBE)); + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 0; + + DataCopy(aL1_[(cvLoopIdx & 1) * aL1DbOffset_], xGlobal_[aGmOffset], nd2nzParams); + } else if (aL1DbNum_ == SINGLE_BUFFER_NUM && kaGmOffset == 0) { + CopyAGmToL1SingleBuffer(param, kaGmOffset, kbL1RealSize, biasRealN, cvLoopIdx, aGmOffset); + } + + // bias仅与n有关，与k无关，所以只需要拷贝一次 + if (isBias_ && kaGmOffset == 0) { + DataCopyPad2D(biasL1_[(cvLoopIdx & 1) * biasL1DbOffset_], biasGlobal_[param.nOffset], 1, biasRealN, biasRealN, + biasRealN); + } + + SetFlag(cubeEventIdMte2ToMte1_); + WaitFlag(cubeEventIdMte2ToMte1_); +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::CopyMxScaleGmToL1(const BasicBlockOffsetParam ¶m, + uint64_t kbL1Offset, uint64_t cvLoopIdx) +{ + uint64_t scaleKGmSize = param.kSize / MX_GROUPSIZE; + // 当前scaleFactor为1，暂不考虑scaleFactor相关计算 + uint64_t scaleKL1StandardLen = param.kbL1Size / MX_GROUPSIZE; + uint64_t scaleKL1RealSize = + (kbL1Offset + param.kbL1Size) > param.kSize ? (param.kSize - kbL1Offset) / MX_GROUPSIZE : scaleKL1StandardLen; + + // copy mxScaleA + Dn2NzParams scaleAdn2NzParams; + ConfigScaleDn2NzParams(param.mL1Size, scaleKGmSize, scaleKL1RealSize, scaleKL1RealSize, scaleAdn2NzParams); + + int64_t scaleAGmOffset = param.mOffset * scaleKGmSize + kbL1Offset / MX_GROUPSIZE; + GlobalTensor f16ScaleAGlobal; + f16ScaleAGlobal.SetGlobalBuffer((__gm__ half *)mxScaleAGlobal_[scaleAGmOffset].GetPhyAddr(), + (param.mL1Size * scaleKL1RealSize) >> 1); + auto f16ScaleALocal = mxScaleAL1_[(cvLoopIdx & 1) * mxScaleAL1DbOffset_].template ReinterpretCast(); + + DataCopy(f16ScaleALocal, f16ScaleAGlobal, scaleAdn2NzParams); + + // copy mxScaleB + Dn2NzParams scaleBdn2NzParams; + ConfigScaleDn2NzParams(param.nL1Size, scaleKGmSize, scaleKL1RealSize, scaleKL1RealSize, scaleBdn2NzParams); + + int64_t scaleBGmOffset = param.nOffset * scaleKGmSize + kbL1Offset / MX_GROUPSIZE; + GlobalTensor f16ScaleBGlobal; + f16ScaleBGlobal.SetGlobalBuffer((__gm__ half *)mxScaleBGlobal_[scaleBGmOffset].GetPhyAddr(), + (param.nL1Size * scaleKL1RealSize) >> 1); + auto f16ScaleBLocal = mxScaleBL1_[(cvLoopIdx & 1) * mxScaleBL1DbOffset_].template ReinterpretCast(); + + DataCopy(f16ScaleBLocal, f16ScaleBGlobal, scaleBdn2NzParams); + + // scale和搬入时A的生命周期相同，该函数在CopyAAndBiasGmToL1之前调用，共用CopyAAndBiasGmToL1的set/wait +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::ConfigScaleDn2NzParams(uint64_t rowNum, uint64_t scaleKGmSize, + uint64_t scaleKL1Stride, + uint64_t scaleKL1RealSize, + Dn2NzParams &dn2NzParams) +{ + dn2NzParams.dnNum = 1; + dn2NzParams.dValue = rowNum; // 矩阵的行数，即待搬运的mxScaleA的m或mxScaleB的n + dn2NzParams.nValue = CeilDivide(scaleKL1RealSize, SCALE_COPY_GROUP_SIZE); // 矩阵的列数，使用B16搬B8需要除以2向上取整 + dn2NzParams.srcDnMatrixStride = SCALE_COPY_DEFAULT_STRIDE; + dn2NzParams.srcDValue = CeilDivide(scaleKGmSize, SCALE_COPY_GROUP_SIZE); // 源矩阵一行所含B16元素个数 + // 目标矩阵行方向两个相邻分形起始地址之间的间隔，单位32B + dn2NzParams.dstNzC0Stride = CeilDivide(scaleKL1Stride, SCALE_COPY_GROUP_SIZE); + // 目标矩阵列方向两个相邻分形起始地址之间的间隔，单位32B + dn2NzParams.dstNzNStride = SCALE_COPY_DEFAULT_N_STRIDE; + dn2NzParams.dstNzMatrixStride = SCALE_COPY_DEFAULT_STRIDE; +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::EndSync(uint64_t cvLoopIdx) +{ + AscendC::TEventID tempEventIdsMte1ToMte2[DOUBLE_BUFFER_NUM] = {cubeEventIdsMte1ToMte2_[0], + cubeEventIdsMte1ToMte2_[1]}; + + // 考虑到只循环一次时， 只需要同步wait第0块缓存。 不止1次时， 2个同步块都需要wait + if (cvLoopIdx > 1 && aL1DbNum_ > SINGLE_BUFFER_NUM) { + WaitFlag(tempEventIdsMte1ToMte2[cvLoopIdx & 1]); + } + + if (cvLoopIdx > 0 && aL1DbNum_ > SINGLE_BUFFER_NUM) { + WaitFlag(tempEventIdsMte1ToMte2[(cvLoopIdx + 1) & 1]); + } + + GetTPipePtr()->ReleaseEventID(tempEventIdsMte1ToMte2[0]); + GetTPipePtr()->ReleaseEventID(tempEventIdsMte1ToMte2[1]); + GetTPipePtr()->ReleaseEventID(cubeEventIdMte2ToMte1_); +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::ClearAFullLoadFlag() +{ + if (aL1DbNum_ == SINGLE_BUFFER_NUM) { + SetFlag(cubeEventIdsMte1ToMte2_[0]); + WaitFlag(cubeEventIdsMte1ToMte2_[0]); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::InitSync() +{ + cubeEventIdsMte1ToMte2_[0] = GetTPipePtr()->AllocEventID(); + cubeEventIdsMte1ToMte2_[1] = GetTPipePtr()->AllocEventID(); + cubeEventIdMte2ToMte1_ = GetTPipePtr()->AllocEventID(); +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::UpdateGlobalAddr( + __gm__ xType *x, __gm__ yType *y, __gm__ biasType *bias, __gm__ antiQuantScaleType *antiquantScale, + __gm__ uint64_t *quantScale, __gm__ perTokenScaleType *perTokenScale, const bool isBias) +{ + isBias_ = isBias; + xGlobal_.SetGlobalBuffer(x); + yGlobal_.SetGlobalBuffer(y); + if (isBias_) { + biasGlobal_.SetGlobalBuffer(bias); + } + if constexpr (IsSameType::value) { + quantScaleGlobal_.SetGlobalBuffer(quantScale); + } + if constexpr (IsMxA8W4()) { + mxScaleAGlobal_.SetGlobalBuffer(perTokenScale); + mxScaleBGlobal_.SetGlobalBuffer(antiquantScale); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::PrefetchA(uint64_t aPrefetchSize, + const LocalTensor &perloadBuffer, + const TCubeTiling *__restrict matmulTiling) +{ + uint64_t xOffset = GetBlockIdx() * aPrefetchSize; + uint64_t xSizeLimit = matmulTiling->M * matmulTiling->Ka; + if (aPrefetchSize == 0 || xOffset >= xSizeLimit) { + return; + } + DataCopyPadExtParams extParams; + DataCopyExtParams param; + param.blockCount = 1; + param.blockLen = (xOffset + aPrefetchSize > xSizeLimit ? xSizeLimit - xOffset : aPrefetchSize) * sizeof(xType); + param.srcStride = 0; + param.dstStride = 0; + DataCopyPad(perloadBuffer, xGlobal_[xOffset], param, extParams); + PipeBarrier(); +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit) +{ + uint64_t xOffset = GetBlockIdx() * aPrefetchSize; + if (aPrefetchSize == 0 || xOffset >= xSizeLimit) { + return; + } + DataCopyPadExtParams extParams; + DataCopyExtParams param; + param.blockCount = 1; + param.blockLen = (xOffset + aPrefetchSize > xSizeLimit ? xSizeLimit - xOffset : aPrefetchSize) * sizeof(xType); + param.srcStride = 0; + param.dstStride = 0; + event_t eventIdMTE1ToMTE2 = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdMTE1ToMTE2); + WaitFlag(eventIdMTE1ToMTE2); + + if constexpr (IsMxA8W4()) { + // 不支持直接搬运fp8，转成uint8搬 + DataCopyPadExtParams extParams; + GlobalTensor uint8XGlobal; + uint8XGlobal.SetGlobalBuffer((__gm__ uint8_t *)xGlobal_[xOffset].GetPhyAddr(), aPrefetchSize); + DataCopyPad(aL1_.template ReinterpretCast(), uint8XGlobal, param, extParams); + } else { + DataCopyPadExtParams extParams; + DataCopyPad(aL1_, xGlobal_[xOffset], param, extParams); + } + PipeBarrier(); +} + +// 场景1： 使能a prefetch。必须先更新地址再init +// 场景2： gm地址变化需要实时获取场景，必须先init再更新地址 +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::Init(TBuf &l1Tbuf, uint64_t weightL1Space, + uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, AscendC::TPipe *tPipe) +{ + // (1) y 数据类型为 int8 时，quantScale需要预留一份空间 + // ① 有bias + // L1 (0~512KB): WeightL1_P0(128KB) | Bias_P0(4KB) | AL1_P0(120KB) | AL1_P1(120KB) | Bias_P1(4KB) | WeightL1_P1(128KB) | quantScale(8KB) + // ② 无bias + // L1 (0~512KB): WeightL1_P0(128KB) | AL1_P0(124KB) | AL1_P1(124KB) | WeightL1_P1(128KB) | quantScale(8KB) + // (2) MxA8W4场景: + // L1 (0~512KB): WeightL1_P0(64KB) | Bias_P0(4KB) | ScaleAL1_P0(20KB) | ScaleBL1_P0(20KB) | AL1_P0(148KB) | + // | AL1_P1(148KB) | ScaleAL1_P1(20KB) | ScaleBL1_P1(20KB) | Bias_P1(4KB) | WeightL1_P1(64KB) + // (3) L1 上有bias时: + // L1 (0~512KB): WeightL1_P0(128KB) | Bias_P0(4KB) | AL1_P0(124KB) | AL1_P1(124KB) | Bias_P1(4KB) | WeightL1_P1(128KB) + // (4) 其他场景时 + // L1 (0~512KB): WeightL1_P0(128KB) | AL1_P0(128KB) | AL1_P1(128KB) | WeightL1_P1(128KB) + + uint64_t biasL1Space = matmulTiling->isBias ? BIAS_L1_SIZE * KB_UNIT : 0; // bias单块分配4K空间 + uint64_t aL1Offset = weightL1Space + biasL1Space; // A要跳过WeightL1_P0 + Bias_P0 + if constexpr (IsSameType::value) { + uint64_t aL1Space = L1_SIZE_WITH_QUANTSCALE * KB_UNIT - DOUBLE_BUFFER_NUM * aL1Offset; // L1上A可占据剩余空间 + aL1DbOffset_ = aL1Space >> 1; + } else if constexpr (IsMxA8W4()) { + biasL1Space = BIAS_L1_SIZE * KB_UNIT; + uint64_t mxScaleL1Space = MX_SCALE_L1_SIZE * KB_UNIT; // scaleA/B单块分配空间 + + aL1Offset = weightL1Space + biasL1Space + (mxScaleL1Space << 1); + uint64_t aL1Space = L1_SIZE * KB_UNIT - DOUBLE_BUFFER_NUM * aL1Offset; // L1上A可占据剩余空间 + aL1DbOffset_ = aL1Space >> 1; + + // MxA8W4场景bias类型为B16，各项l1Space均以B8元素个数计，计算B16偏移需除以2 + biasL1_ = l1Tbuf.Get()[weightL1Space >> 1]; + biasL1DbOffset_ = ((aL1Space + biasL1Space) >> 1) + (mxScaleL1Space << 1); + + mxScaleAL1_ = l1Tbuf.Get()[weightL1Space + biasL1Space]; + mxScaleAL1DbOffset_ = (mxScaleL1Space << 1) + aL1Space; + + mxScaleBL1_ = l1Tbuf.Get()[weightL1Space + biasL1Space + mxScaleL1Space]; + mxScaleBL1DbOffset_ = (mxScaleL1Space << 1) + aL1Space; + } else if (matmulTiling->isBias) { + uint64_t aL1Space = L1_SIZE * KB_UNIT - DOUBLE_BUFFER_NUM * aL1Offset; // L1上A可占据剩余空间 + aL1DbOffset_ = aL1Space >> 1; + if constexpr (IsSameType::value) { + biasL1_ = l1Tbuf.Get()[weightL1Space >> 1]; + biasL1DbOffset_ = (aL1Space + biasL1Space) >> 1; + } else { + biasL1_ = l1Tbuf.Get()[weightL1Space]; + biasL1DbOffset_ = aL1Space + biasL1Space; + } + } else { + aL1DbOffset_ = L1_HALF_SIZE * KB_UNIT - weightL1Space; + } + aL1_ = l1Tbuf.Get()[aL1Offset]; + aL1Count_ = matmulTiling->Ka / (matmulTiling->baseK * matmulTiling->stepKb); + aL1MaxHalfCount_ = CeilDivide(aL1Count_, static_cast(DOUBLE_BUFFER_NUM)); + + PrefetchA(aPrefetchSize, aL1_, matmulTiling); + // 当前tiling策略的细分场景： + // 1. stepKa <= stepKb 当前限制baseM的最大值，因此该场景下在L1上A矩阵大小<=128k。可以固定走db分支，保证a矩阵的db载入 + // 2. stepKa > stepKb 当前在m小k大的情况下才会出现该场景，走全载分支 + // 3. gmm场景，不知道真实的m值，tiling采取保守策略，恒定走db分支 + if (matmulTiling->stepKa > matmulTiling->stepKb) { + aL1DbNum_ = SINGLE_BUFFER_NUM; + } else { + aL1DbNum_ = DOUBLE_BUFFER_NUM; + } + mmObj_.SetSubBlockIdx(0); + mmObj_.Init(matmulTiling, tPipe); + InitSync(); + + if constexpr (IsSameType::value && wqmmConfig.quantType == QuantType::PER_TENSOR) { + quantScaleValue_ = this->quantScaleGlobal_.GetValue(0); + } +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::GetTensorC(const BasicBlockOffsetParam ¶m) +{ + uint64_t outOffset = param.mOffset * param.nSize + param.nOffset; +#ifndef __CCE_KT_TEST__ + mmObj_.GetTensorC(yGlobal_[outOffset]); +#endif +} + +WQBMM_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void WQBMM_CUBE_COMPUTE_CLASS::GetTensorC(LocalTensor &yUb) +{ +#ifndef __CCE_KT_TEST__ + mmObj_.GetTensorC(yUb, 0, true); +#endif +} +} // namespace WeightQuantBatchMatmulV2::Arch35 +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_tiling_key.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_tiling_key.h new file mode 100644 index 00000000..9c01213f --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_tiling_key.h @@ -0,0 +1,193 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file weight_quant_tiling_key.h +* \brief +*/ + +#ifndef __OP_KERNEL_WQGMM_TILING_KEY_H__ +#define __OP_KERNEL_WQGMM_TILING_KEY_H__ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define WQGMM_VECTOR_ANTIQUANT 0 +#define WQGMM_MULTI_SCALE_DEQUANT 1 +#define WQGMM_CUBE_ANTIQUANT 2 + +#define WQGMM_VDEFAULT 0 +#define WQGMM_SPLIT_K 1 +#define WQGMM_N_FIRST_TAIL_RESPLIT 2 +#define WQGMM_N_FIRST_BASIC_BLOCK 3 + +#define WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_2 0 +#define WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_4 1 +#define WQGMM_MTE2_INNER_SIZE_1024_BUF_NUM_2 2 +#define WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 3 +#define WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT 4 +#define WQGMM_MTE2_INNER_SIZE_384_BUF_NUM_3 5 + +#define WQGMM_NO_TRANS 0 +#define WQGMM_TRANS 1 + +#define WQGMM_NONE 0 +#define WQGMM_PER_TENSOR 1 +#define WQGMM_PER_CHANNEL 2 +#define WQGMM_PER_GROUP 3 +#define WQGMM_MX 4 + +#define WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST 0 +#define WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_EXIST 1 +#define WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST 2 +#define WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_EXIST 3 +#define WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_FP32_EXIST 4 +#define WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_FP32_EXIST 5 + +#define WQGMM_ND 0 +#define WQGMM_FRACTAL_NZ 1 + +// 模板参数 +ASCENDC_TPL_ARGS_DECL( + DlinferGroupedMatmulDirect, // 算子OpType + ASCENDC_TPL_UINT_DECL(W_TYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, WQGMM_ND, WQGMM_FRACTAL_NZ), + ASCENDC_TPL_UINT_DECL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_EXIST, + WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST, WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_EXIST, + WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_FP32_EXIST, WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_FP32_EXIST), + ASCENDC_TPL_UINT_DECL(C_QUANT_TYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, WQGMM_NONE, WQGMM_PER_TENSOR, + WQGMM_PER_CHANNEL, WQGMM_PER_GROUP, WQGMM_MX), + ASCENDC_TPL_UINT_DECL(W_QUANT_TYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, WQGMM_NONE, WQGMM_PER_TENSOR, + WQGMM_PER_CHANNEL, WQGMM_PER_GROUP, WQGMM_MX), + ASCENDC_TPL_UINT_DECL(WQ_B_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS, WQGMM_TRANS), + ASCENDC_TPL_UINT_DECL(WQ_A_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS, WQGMM_TRANS), + ASCENDC_TPL_UINT_DECL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_2, WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_4, WQGMM_MTE2_INNER_SIZE_1024_BUF_NUM_2, + WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4, WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT, + WQGMM_MTE2_INNER_SIZE_384_BUF_NUM_3), + ASCENDC_TPL_UINT_DECL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, WQGMM_VDEFAULT, WQGMM_SPLIT_K, + WQGMM_N_FIRST_TAIL_RESPLIT, WQGMM_N_FIRST_BASIC_BLOCK), + ASCENDC_TPL_UINT_DECL(ALGORITHM_CATEGORY, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT, + WQGMM_MULTI_SCALE_DEQUANT, WQGMM_CUBE_ANTIQUANT) + ); +// 模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_FRACTAL_NZ), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_FRACTAL_NZ), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_384_BUF_NUM_3), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_FRACTAL_NZ), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_MX), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_FRACTAL_NZ), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_MX), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_TAIL_RESPLIT), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_BASIC_BLOCK), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_BASIC_BLOCK), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_BASIC_BLOCK), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(W_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_ND), + ASCENDC_TPL_UINT_SEL(OFFSET_OR_BIAS_EXIT, ASCENDC_TPL_UI_LIST, WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST), + ASCENDC_TPL_UINT_SEL(C_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_NONE), + ASCENDC_TPL_UINT_SEL(W_QUANT_TYPE, ASCENDC_TPL_UI_LIST, WQGMM_PER_CHANNEL), + ASCENDC_TPL_UINT_SEL(WQ_B_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(WQ_A_TRANS, ASCENDC_TPL_UI_LIST, WQGMM_NO_TRANS), + ASCENDC_TPL_UINT_SEL(TEMPLATE_CUSTOM_SC, ASCENDC_TPL_UI_LIST, WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4), + ASCENDC_TPL_UINT_SEL(ALGORITHM_SUB_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_N_FIRST_BASIC_BLOCK), + ASCENDC_TPL_UINT_SEL(ALGORITHM_CATEGORY, ASCENDC_TPL_UI_LIST, WQGMM_VECTOR_ANTIQUANT)) +); +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block.h new file mode 100644 index 00000000..a8e751ee --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block.h @@ -0,0 +1,305 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file weight_quant_vcv_basic_block.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_H +#define GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_H + +#include "basic_block_config.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "lib/matmul_intf.h" +#include "tool.h" +#include "weight_quant_cube_compute.h" +#include "weight_quant_vcv_basic_block_base.h" +#include "weight_quant_vec_compute.h" + +using AscendC::GetSubBlockIdx; +using AscendC::LocalTensor; +using AscendC::TBuf; +using AscendC::TPipe; +using AscendC::TPosition; + +namespace WeightQuantBatchMatmulV2::Arch35 { +#define GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM \ + template + +#define GMM_WQ_VCV_BASIC_BLOCK_CLASS \ + WeightQuantVcvMatmulBasicBlock + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +class WeightQuantVcvMatmulBasicBlock : public WeightQuantVcvMatmulBasicBlockBaseClass { +public: + __aicore__ inline WeightQuantVcvMatmulBasicBlock(){}; + __aicore__ inline void Init(bool hasBias, uint64_t antiQuantGroupSize, uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, TPipe *tPipe); + __aicore__ inline void UpdateGlobalAddr(__gm__ xType *x, __gm__ wType *weight, + __gm__ antiQuantScaleType *antiquantScale, __gm__ xType *antiquantOffset, + __gm__ scaleType *scale, __gm__ perTokenScaleType *perTokenScale, + __gm__ biasType *bias, __gm__ yType *y, const bool hasBias, + const bool weightL2Cacheable); + __aicore__ inline void ComputeBasicBlock(const BasicBlockOffsetParam &curOffsetParam, + const BasicBlockOffsetParam &lastOffsetParam); + __aicore__ inline void PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit); + __aicore__ inline void End(const BasicBlockOffsetParam &curOffsetParam); + +protected: + __aicore__ inline void ComputeBasicBlockAivNzKn(const BasicBlockOffsetParam &curOffsetParam, + const BasicBlockOffsetParam &lastOffsetParam); + __aicore__ inline void IterateNzKnWithKAiv(uint64_t &kMte2Offset, uint64_t kMte2Limit, uint64_t mte2RealN, + uint64_t nL1Offset, const BasicBlockOffsetParam &curOffsetParam); + __aicore__ inline void IterateNzKnWithKAic(const BasicBlockOffsetParam &curOffsetParam); + + template + __aicore__ inline void SetAivToAic(uint64_t syncFlag) + { +#ifndef __CCE_KT_TEST__ + CrossCoreSetFlag(syncFlag); +#endif + }; + + template + __aicore__ inline void WaitAivToAic(uint64_t syncFlag) + { +#ifndef __CCE_KT_TEST__ + CrossCoreWaitFlag(syncFlag + FLAG_ID_MAX); + CrossCoreWaitFlag(syncFlag); +#endif + }; + + template + __aicore__ inline void SetAicToAiv(uint64_t syncFlag) + { +#ifndef __CCE_KT_TEST__ + CrossCoreSetFlag(syncFlag + FLAG_ID_MAX); + CrossCoreSetFlag(syncFlag); +#endif + }; + + template + __aicore__ inline void WaitAicToAiv(uint64_t syncFlag) + { +#ifndef __CCE_KT_TEST__ + CrossCoreWaitFlag(syncFlag); +#endif + }; + + BasicBlockLibVectorAntiQuantCompute vecCompute_; + using MMImpl = MatmulImpl, + MatmulL1GmType, + MatmulType, + MatmulType, CFG_MDL, + matmul::MatmulCallBackFunc, WQBmmCustomPolicy>; + WeightQuantBatchMatmulV2CubeCompute + cubeCompute_; + + uint64_t cvLoopIdx_ = 0; + uint64_t weightS8L1DbOffset_ = 0; + + LocalTensor weightS8L1_; + LocalTensor ubOutputS32Buffer_; +}; + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::Init(bool hasBias, uint64_t antiQuantGroupSize, + uint64_t aPrefetchSize, + const TCubeTiling *__restrict matmulTiling, TPipe *tPipe) +{ + uint64_t weightL1Space = matmulTiling->baseN * matmulTiling->stepKb * matmulTiling->baseK; // weight单块大小 + + TBuf l1Tbuf; + tPipe->InitBuffer(l1Tbuf, 512 * 1024); + weightS8L1DbOffset_ = 512 * GetKBUnit() - weightL1Space; + weightS8L1_ = l1Tbuf.Get(); + + TBuf<> ubBuffer; + tPipe->InitBuffer(ubBuffer, 248 * 1024); + ubOutputS32Buffer_ = ubBuffer.Get(); + + if ASCEND_IS_AIC { + cubeCompute_.Init(l1Tbuf, weightL1Space, aPrefetchSize, matmulTiling, tPipe); + } else { + LocalTensor ubWeightS8Buffer = ubOutputS32Buffer_.template ReinterpretCast(); + vecCompute_.InitKCG(antiQuantGroupSize, hasBias, ubBuffer, ubWeightS8Buffer, 128 * 1024); + } + cvLoopIdx_ = 0; +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::UpdateGlobalAddr( + __gm__ xType *x, __gm__ wType *weight, __gm__ antiQuantScaleType *antiquantScale, __gm__ xType *antiquantOffset, + __gm__ scaleType *scale, __gm__ perTokenScaleType *perTokenScale, __gm__ biasType *bias, __gm__ yType *y, + const bool hasBias, const bool weightL2Cacheable) +{ + if ASCEND_IS_AIC { + cubeCompute_.UpdateGlobalAddr(x, nullptr, nullptr, nullptr, nullptr, nullptr, hasBias); + } else { + vecCompute_.UpdateGlobalAddr(weight, antiquantScale, nullptr, perTokenScale, scale, bias, weightL2Cacheable); + } +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::ComputeBasicBlock(const BasicBlockOffsetParam &curOffsetParam, + const BasicBlockOffsetParam &lastOffsetParam) +{ + if ASCEND_IS_AIV { + ComputeBasicBlockAivNzKn(curOffsetParam, lastOffsetParam); + } else { + if (cvLoopIdx_ > 0) { + WaitAivToAic(SYNC_AIV_MTE3_AIC_FIX_FLAG); + cubeCompute_.GetTensorC(ubOutputS32Buffer_); + SetAicToAiv(SYNC_AIC_FIX_AIV_VF_FLAG); + cubeCompute_.ClearAFullLoadFlag(); // 清除A全载时之前循环的set同步标记 + } + IterateNzKnWithKAic(curOffsetParam); + } +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::ComputeBasicBlockAivNzKn( + const BasicBlockOffsetParam &curOffsetParam, const BasicBlockOffsetParam &lastOffsetParam) +{ + uint64_t kMte2Offset = 0; + uint64_t curCvLoopIdx = cvLoopIdx_; + // vec上两core切n + uint64_t mte2RealN = curOffsetParam.nL1Size > C0_SIZE_B8 + ? CeilAlign(curOffsetParam.nL1Size >> 1, static_cast(C0_SIZE_B8)) + : curOffsetParam.nL1Size; + uint64_t nL1Offset = GetSubBlockIdx() * mte2RealN; + mte2RealN = GetSubBlockIdx() == 0 ? mte2RealN : curOffsetParam.nL1Size - mte2RealN; + + IterateNzKnWithKAiv(kMte2Offset, Min(curOffsetParam.kSize, DOUBLE_BUFFER_NUM * curOffsetParam.kbL1Size), mte2RealN, + nL1Offset, curOffsetParam); + + if (curCvLoopIdx > 0) { + // y反量化在vec上两core切m + uint64_t lastBasicBlockMSize = lastOffsetParam.mL1Size - (lastOffsetParam.mL1Size >> 1); + uint64_t lastBasicBlockMOffset = GetSubBlockIdx() == 0 ? 0 : lastBasicBlockMSize; + lastBasicBlockMSize = GetSubBlockIdx() == 0 ? lastBasicBlockMSize : (lastOffsetParam.mL1Size >> 1); + + SetAivToAic(SYNC_AIV_MTE3_AIC_FIX_FLAG); + WaitAicToAiv(SYNC_AIC_FIX_AIV_VF_FLAG); + vecCompute_.AntiQuantYWithKc(lastOffsetParam.nL1Size, lastBasicBlockMSize); + vecCompute_.CopyYUbToGm(lastOffsetParam.nL1Size, lastBasicBlockMSize, + reinterpret_cast<__gm__ half *>(lastOffsetParam.yGmAddr), lastOffsetParam, + lastBasicBlockMOffset); + } + uint64_t antiquantYMSize = curOffsetParam.mL1Size - (curOffsetParam.mL1Size >> 1); + uint64_t antiquantYMOffset = GetSubBlockIdx() == 0 ? 0 : antiquantYMSize; + antiquantYMSize = GetSubBlockIdx() == 0 ? antiquantYMSize : (curOffsetParam.mL1Size >> 1); + IterateNzKnWithKAiv(kMte2Offset, curOffsetParam.kSize, mte2RealN, nL1Offset, curOffsetParam); + vecCompute_.CopyKcScaleBiasGmToUb(curOffsetParam.nL1Size, antiquantYMSize, curOffsetParam.nOffset, + curOffsetParam.mOffset + antiquantYMOffset); +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::IterateNzKnWithKAiv(uint64_t &kMte2Offset, uint64_t kMte2Limit, + uint64_t mte2RealN, uint64_t nL1Offset, + const BasicBlockOffsetParam &curOffsetParam) +{ + UbConsumeConfig ubConsumeConfig; + ubConsumeConfig.l1RequireVfComputeRealN = mte2RealN; + ubConsumeConfig.nWeightLowBitUbOffset = 0; + L1ConsumeConfig l1ConsumeConfig; + l1ConsumeConfig.l1SplitTwoVecExternalOffset = nL1Offset; + l1ConsumeConfig.l1RealExternalLen = curOffsetParam.nL1Size; + + for (; kMte2Offset < kMte2Limit; kMte2Offset += vecConfig.ubMte2InnerSize) { + uint64_t mte2RealK = (kMte2Offset + vecConfig.ubMte2InnerSize) > kMte2Limit + ? kMte2Limit - kMte2Offset + : vecConfig.ubMte2InnerSize; // vec总共需要搬运的K方向的实际量（考虑尾块） + vecCompute_.WaitVToMTE2(); + vecCompute_.CopyGmToUb(mte2RealN, mte2RealK, curOffsetParam.nOffset + nL1Offset, kMte2Offset, curOffsetParam); + for (uint64_t antiquantKOffset = 0; antiquantKOffset < mte2RealK; + antiquantKOffset += curOffsetParam.kbL1Size, cvLoopIdx_++) { + uint64_t antiquantRealK = (antiquantKOffset + curOffsetParam.kbL1Size) >= mte2RealK + ? mte2RealK - antiquantKOffset + : curOffsetParam.kbL1Size; + if (cvLoopIdx_ > 1) { + WaitAicToAiv(SYNC_AIC_AIV_FLAG); + } + + ubConsumeConfig.l1RequireVfComputeRealK = antiquantRealK; + ubConsumeConfig.kWeightLowBitUbOffset = antiquantKOffset; + vecCompute_.WeightAntiQuantCompute(ubConsumeConfig, weightS8L1_[(cvLoopIdx_ & 1) * weightS8L1DbOffset_], + l1ConsumeConfig); + + SetAivToAic(SYNC_AIV_AIC_FLAG); + } + vecCompute_.SetVToMTE2(); + } +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::IterateNzKnWithKAic(const BasicBlockOffsetParam &curOffsetParam) +{ + for (uint64_t kbL1Offset = 0; kbL1Offset < curOffsetParam.kSize; + kbL1Offset += curOffsetParam.kbL1Size, cvLoopIdx_++) { + uint64_t kbL1RealSize = (kbL1Offset + curOffsetParam.kbL1Size) >= curOffsetParam.kSize + ? curOffsetParam.kSize - kbL1Offset + : curOffsetParam.kbL1Size; + cubeCompute_.WaitMTE1ToMTE2(cvLoopIdx_); + cubeCompute_.CopyAAndBiasGmToL1(curOffsetParam, kbL1Offset, kbL1RealSize, curOffsetParam.nL1Size, cvLoopIdx_); + WaitAivToAic(SYNC_AIV_AIC_FLAG); + + cubeCompute_.LaunchMatmul(weightS8L1_[(cvLoopIdx_ & 1) * weightS8L1DbOffset_], kbL1Offset, kbL1RealSize, + curOffsetParam, + cvLoopIdx_); // mte1 mmad fixp流水 + + cubeCompute_.SetMTE1ToMTE2(cvLoopIdx_); + SetAicToAiv(SYNC_AIC_AIV_FLAG); + } +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::PrefetchA(uint64_t aPrefetchSize, uint64_t xSizeLimit) +{ + cubeCompute_.PrefetchA(aPrefetchSize, xSizeLimit); +} + +GMM_WQ_VCV_BASIC_BLOCK_TEMPLATE_PARAM +__aicore__ inline void GMM_WQ_VCV_BASIC_BLOCK_CLASS::End(const BasicBlockOffsetParam &curOffsetParam) +{ + if ASCEND_IS_AIC { + if (cvLoopIdx_ > 0) { + WaitAivToAic(SYNC_AIV_MTE3_AIC_FIX_FLAG); + cubeCompute_.GetTensorC(ubOutputS32Buffer_); + SetAicToAiv(SYNC_AIC_FIX_AIV_VF_FLAG); + } + cubeCompute_.EndSync(cvLoopIdx_); + } else { + if (cvLoopIdx_ > 0) { + SetAivToAic(SYNC_AIV_MTE3_AIC_FIX_FLAG); + uint64_t lastBasicBlockMSize = curOffsetParam.mL1Size - (curOffsetParam.mL1Size >> 1); + uint64_t lastBasicBlockMOffset = GetSubBlockIdx() == 0 ? 0 : lastBasicBlockMSize; + lastBasicBlockMSize = GetSubBlockIdx() == 0 ? lastBasicBlockMSize : (curOffsetParam.mL1Size >> 1); + WaitAicToAiv(SYNC_AIC_FIX_AIV_VF_FLAG); + vecCompute_.AntiQuantYWithKc(curOffsetParam.nL1Size, lastBasicBlockMSize); + vecCompute_.CopyYUbToGm(curOffsetParam.nL1Size, lastBasicBlockMSize, + reinterpret_cast<__gm__ half *>(curOffsetParam.yGmAddr), curOffsetParam, + lastBasicBlockMOffset); + WaitAicToAiv(SYNC_AIC_AIV_FLAG); + } + if (cvLoopIdx_ > 1) { + WaitAicToAiv(SYNC_AIC_AIV_FLAG); + } + vecCompute_.End(); + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block_base.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block_base.h new file mode 100644 index 00000000..c348ae5b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vcv_basic_block_base.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file weight_quant_vcv_basic_block_base.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_BASE_H +#define GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_BASE_H + +namespace WeightQuantBatchMatmulV2::Arch35 { + +class WeightQuantVcvMatmulBasicBlockBaseClass { +public: + __aicore__ inline WeightQuantVcvMatmulBasicBlockBaseClass(){}; +}; + +} // namespace WeightQuantBatchMatmulV2::Arch35 + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_VCV_BASIC_BLOCK_BASE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vec_compute.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vec_compute.h new file mode 100644 index 00000000..be4a690e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/arch35/weight_quant_basic_block/weight_quant_vec_compute.h @@ -0,0 +1,1281 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file weight_quant_vec_compute.h + * \brief + */ +#ifndef GROUPED_MATMUL_WEIGHT_QUANT_VEC_COMPUTE_H +#define GROUPED_MATMUL_WEIGHT_QUANT_VEC_COMPUTE_H + +#include "anti_quant_y_vf.h" +#include "basic_block_config.h" +#include "basic_block_vf_mx.h" +#include "basic_block_vf_nd.h" +#include "basic_block_vf_nz.h" +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "tool.h" + +using AscendC::BLOCK_CUBE; +using AscendC::CacheMode; +using AscendC::DataCopyParams; +using AscendC::GlobalTensor; +using AscendC::HardEvent; +using AscendC::IsSameType; +using AscendC::LocalTensor; +using AscendC::ONE_BLK_SIZE; +using AscendC::SetFlag; +using AscendC::TBuf; +using AscendC::TEventID; +using AscendC::TPipe; +using AscendC::VECTOR_REG_WIDTH; +using AscendC::WaitFlag; +namespace MicroAPI = AscendC::MicroAPI; +using AscendC::MicroAPI::AddrReg; +using AscendC::MicroAPI::GetRound; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; +using AscendC::MicroAPI::TypeGet; + +namespace WeightQuantBatchMatmulV2::Arch35 { +template +class BasicBlockLibVectorAntiQuantCompute { +public: + __aicore__ inline BasicBlockLibVectorAntiQuantCompute(){}; + + __aicore__ inline void UpdateGlobalAddr(__gm__ wType *weight, __gm__ antiQuantScaleType *antiQuantScale, + __gm__ xType *antiQuantOffset, __gm__ float *perTokenScale, + __gm__ float *perChannelScale, __gm__ float *bias, + const bool weightL2Cacheable); + __aicore__ inline void Init(TPipe *tPipe); + __aicore__ inline void InitKCG(uint32_t antiQuantGroupSize, bool hasBias, TBuf<> ubBuffer, + const LocalTensor &ubHighBitTotalBuffer, uint64_t highBitUbOffset); + __aicore__ inline void WaitVToMTE2(); + __aicore__ inline void SetVToMTE2(); + __aicore__ inline void CopyGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, uint64_t ubMte2NOffset, + uint64_t ubMte2KOffset, const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void WeightAntiQuantCompute(const UbConsumeConfig &ubConsumeConfig, + const LocalTensor &weightHighBitL1, + const L1ConsumeConfig &l1ConsumeConfig); + __aicore__ inline void CopyKcScaleBiasGmToUb(uint64_t nRealL0Size, uint64_t mRealL0Size, uint64_t nOffset, + uint64_t mOffset); + __aicore__ inline void AntiQuantYWithKc(uint64_t nRealL0Size, uint64_t mRealL0Size); + __aicore__ inline void CopyYUbToGm(uint64_t nRealL0Size, uint64_t mRealL0Size, __gm__ half *yGm, + const BasicBlockOffsetParam &offsetParam, uint64_t aivMOffset); + __aicore__ inline void End(); + +private: + __aicore__ inline void InitMx(TBuf<> &ubBuffer); + __aicore__ inline void CopyWeightGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, uint64_t ubMte2NOffset, + uint64_t ubMte2KOffset, const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void CopyAntiQuantParamsGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, uint64_t ubMte2NOffset, + uint64_t ubMte2KOffset, const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void WeightAntiQuantProcess(uint64_t nRealLen, uint64_t kRealLen, uint64_t antiQuantNOffset, + uint64_t antiQuantKOffset, const UbConsumeConfig &ubConsumeConfig); + __aicore__ inline void AntiQuantProcess(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void AntiQuantProcessNd(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void AntiQuantProcessNz(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void CalLocalAddrForVf(uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset, + LocalAddressParam &localAddressParam); + __aicore__ inline void CalInt4NzKnLocalAddrForVf(uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset, + Int4NzParams &int4NzParams); + __aicore__ inline void WeightHighBitUbToL1(uint64_t weightHighBitL1Offset, uint64_t antiQuantRealN, + uint64_t antiQuantRealK, const LocalTensor &weightHighBitL1, + uint64_t l1RealExternalLen); + __aicore__ inline uint64_t ComputeWeightHighBitL1Offset(uint64_t antiQuantNOffset, uint64_t antiQuantKOffset, + uint64_t nRealLen, uint64_t kRealLen, + const L1ConsumeConfig &l1ConsumeConfig); + __aicore__ inline void CopyMxAntiQuantParamsGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, + uint64_t ubMte2NOffset, uint64_t ubMte2KOffset, + const BasicBlockOffsetParam &offsetParam); + __aicore__ inline void MxScaleProcess(uint64_t ubMte2NSize, uint64_t ubMte2KSize); + __aicore__ inline void AntiQuantProcessNdMx(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void AntiQuantProcessNzMx(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void AntiQuantProcessNzMxA8W4(uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset); + __aicore__ inline void CalLocalAddrForYVf(LocalAddressYParam &localAddressParam); + __aicore__ inline void CopyWeightHighBitForAligned(uint64_t weightHighBitL1Offset, uint64_t antiQuantRealN, + const uint64_t antiQuantRealK, + const LocalTensor &weightHighBitL1); + __aicore__ inline void CopyWeightHighBitForUnaligned(uint64_t weightHighBitL1Offset, uint64_t antiQuantRealN, + uint64_t antiQuantRealK, + const LocalTensor &weightHighBitL1); + + // mte2搬运计数，用于控制weight输入的buffer和 mte2&&V间同步控制 + uint64_t ubMte2LoopIdx_ = 0; + // mte2搬运计数，用于控制antiquantY输入的buffer和 mte2&&V间同步控制 + uint64_t ubMte2AntiquantYLoopIdx_ = 0; + // vf中标准计算单元(vfNStandardLen, vfKStandardLen)的计数，用于控制weight反量化后输出的buffer和V&&mte3间同步控制 + uint64_t ubComputeLoopIdx_ = 0; + uint64_t ubAntiquantYLoopIdx_ = 0; + + TEventID vecEventIdVToMte2_[QUADRUPLE_BUFFER_NUM]; + TEventID vecEventIdMte3ToV_[QUADRUPLE_BUFFER_NUM]; + + xType scaleValue_; + xType offsetValue_; + + GlobalTensor wGlobal_; + GlobalTensor antiQuantOffsetGlobal_; + GlobalTensor antiQuantScaleGlobal_; + GlobalTensor antiQuantYPerTokenScaleGlobal_; + GlobalTensor antiQuantYPerChannelScaleGlobal_; + GlobalTensor antiQuantYBiasGlobal_; + GlobalTensor antiQuantYF16Global_; + + LocalTensor ubWeightInputLowBitTotalBuffer_; + LocalTensor ubHighBitTotalBuffer_; + LocalTensor ubAntiQuantScaleTotalBuffer_; + LocalTensor ubAntiQuantScaleAfterCastTotalBuffer_; + LocalTensor ubAntiQuantOffsetTotalBuffer_; + LocalTensor ubAntiQuantYPerTokenScaleTotalBuffer_; + LocalTensor ubAntiQuantYPerChannelScaleTotalBuffer_; + LocalTensor ubAntiQuantYBiasTotalBuffer_; + + LocalTensor ubAntiQuantScaleMaskBuffer_; + + uint64_t antiQuantGroupSize_; + bool hasBias_; + + constexpr static uint32_t C0_SIZE = + (IsSameType::value || IsSameType::value) ? C0_SIZE_B8 : BLOCK_CUBE; + constexpr static uint64_t VEC_REG_ELEM = + IsMxA8W4() ? VECTOR_REG_WIDTH : VEC_MAX_ELEM_B16; + + constexpr static uint64_t UB_AVAILABLE_SIZE = 248 * GetKBUnit(); + + constexpr static uint64_t ANTI_QUANT_Y_PER_TOKEN_SCALE_TOTAL_BUFFER_SIZE = 2 * GetKBUnit(); + constexpr static uint64_t ANTI_QUANT_Y_PER_CHANNEL_SCALE_TOTAL_BUFFER_SIZE = 2 * GetKBUnit(); + constexpr static uint64_t ANTI_QUANT_Y_BIAS_TOTAL_BUFFER_SIZE = 2 * GetKBUnit(); + + constexpr static uint64_t UB_ANTI_QUANT_Y_BUFFER_NUM = DOUBLE_BUFFER_NUM; + + constexpr static uint64_t ANTI_QUANT_Y_PER_TOKEN_SCALE_SINGLE_BUFFER_SIZE = + ANTI_QUANT_Y_PER_TOKEN_SCALE_TOTAL_BUFFER_SIZE / UB_ANTI_QUANT_Y_BUFFER_NUM; + constexpr static uint64_t ANTI_QUANT_Y_PER_CHANNEL_SCALE_SINGLE_BUFFER_SIZE = + ANTI_QUANT_Y_PER_CHANNEL_SCALE_TOTAL_BUFFER_SIZE / UB_ANTI_QUANT_Y_BUFFER_NUM; + constexpr static uint64_t ANTI_QUANT_Y_BIAS_SINGLE_BUFFER_SIZE = + ANTI_QUANT_Y_BIAS_TOTAL_BUFFER_SIZE / UB_ANTI_QUANT_Y_BUFFER_NUM; + + TEventID vecEventIdAntiQuantYVToMte2_[UB_ANTI_QUANT_Y_BUFFER_NUM]; + + constexpr static UbBufferInfo UB_BUFFER_INFO = GetBufferConfig(); + constexpr static VfConfig VF_CONFIG = GetVfConfig(); + + constexpr static uint64_t ANTIQUANT_Y_STANDARD_N_SIZE = VECTOR_REG_WIDTH / sizeof(int32_t); +}; + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::UpdateGlobalAddr( + __gm__ wType *weight, __gm__ antiQuantScaleType *antiQuantScale, __gm__ xType *antiQuantOffset, + __gm__ float *perTokenScale, __gm__ float *perChannelScale, __gm__ float *bias, const bool weightL2Cacheable) +{ + wGlobal_.SetGlobalBuffer(weight); + antiQuantScaleGlobal_.SetGlobalBuffer(antiQuantScale); + + if constexpr (IsSameType::value) { + antiQuantYPerTokenScaleGlobal_.SetGlobalBuffer(perTokenScale); + antiQuantYPerChannelScaleGlobal_.SetGlobalBuffer(perChannelScale); + antiQuantYBiasGlobal_.SetGlobalBuffer(bias); + } else { + antiQuantOffsetGlobal_.SetGlobalBuffer(antiQuantOffset); + } + + if (!weightL2Cacheable) { + wGlobal_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } +} + +/* + * 初始化buffer和同步所需的EventID + */ +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::Init(TPipe *tPipe) +{ + TBuf<> ubBuffer; + tPipe->InitBuffer(ubBuffer, UB_AVAILABLE_SIZE); + + if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + InitMx(ubBuffer); + } else { + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, 0); // 174KB + ubHighBitTotalBuffer_ = ubBuffer.template GetWithOffset(UB_BUFFER_INFO.highBitDataUbTotalSize, + 174 * GetKBUnit()); // 33KB*2 + ubAntiQuantScaleTotalBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantScaleUbTotalSize, 240 * GetKBUnit()); // 4KB + ubAntiQuantOffsetTotalBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantOffsetUbTotalSize, 244 * GetKBUnit()); // 4KB + } else { + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, 0); // 112KB + ubHighBitTotalBuffer_ = ubBuffer.template GetWithOffset(UB_BUFFER_INFO.highBitDataUbTotalSize, + 112 * GetKBUnit()); // 32KB*4 + ubAntiQuantScaleTotalBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantScaleUbTotalSize, 240 * GetKBUnit()); // 4KB + ubAntiQuantOffsetTotalBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantOffsetUbTotalSize, 244 * GetKBUnit()); // 4KB + } + } + + for (uint16_t i = 0; i < vecConfig.ubMte2BufferNum; ++i) { + vecEventIdVToMte2_[i] = GetTPipePtr()->AllocEventID(); + } + + for (uint16_t i = 0; i < UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; ++i) { + vecEventIdMte3ToV_[i] = GetTPipePtr()->AllocEventID(); + } +} + +/* + * 初始化mx场景buffer 封装防止Init超长 + */ +template +__aicore__ inline void BasicBlockLibVectorAntiQuantCompute::InitMx(TBuf<> &ubBuffer) +{ + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, 0); // 64KB*2 = 128KB + ubHighBitTotalBuffer_ = ubBuffer.template GetWithOffset(UB_BUFFER_INFO.highBitDataUbTotalSize, + 128 * GetKBUnit()); // 33KB*2 = 66KB + ubAntiQuantScaleTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.antiQuantScaleUbTotalSize, + 194 * GetKBUnit()); // 8KB + + ubAntiQuantScaleAfterCastTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.antiQuantScaleAfterCastUbTotalSize, + 202 * GetKBUnit()); // 32KB + } else if constexpr (IsSameType::value) { + // MxA8W4 + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, 0); // 16KB * 4 = 64KB + ubHighBitTotalBuffer_ = ubBuffer.template GetWithOffset(UB_BUFFER_INFO.highBitDataUbTotalSize, + 64 * GetKBUnit()); // 16KB * 4 = 64KB + } else { + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, 0); // 16KB * 4 = 64KB + ubHighBitTotalBuffer_ = ubBuffer.template GetWithOffset(UB_BUFFER_INFO.highBitDataUbTotalSize, + 64 * GetKBUnit()); // 128KB + ubAntiQuantScaleTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.antiQuantScaleUbTotalSize, + 192 * GetKBUnit()); // 8KB + + ubAntiQuantScaleAfterCastTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.antiQuantScaleAfterCastUbTotalSize, + 200 * GetKBUnit()); // 16KB + } +} + +/** + * @brief KCG 初始化buffer和同步所需的EventID + * @param antiQuantgroupSize per_group伪量化的groupSize + * @param hasBias Y反量化是否存在bias + * @param ubBuffer 用于分配buffer的TBuf对象 + * @param ubHighBitTotalBuffer weightS8/cS32/cF16 复用的UB + * @param highBitUbOffset weightS8/cS32/cF16 复用UB的偏移长度 + */ +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::InitKCG( + uint32_t antiQuantGroupSize, bool hasBias, TBuf<> ubBuffer, const LocalTensor &ubHighBitTotalBuffer, + uint64_t highBitUbOffset) +{ + antiQuantGroupSize_ = antiQuantGroupSize; + hasBias_ = hasBias; + + ubHighBitTotalBuffer_ = ubHighBitTotalBuffer; + ubWeightInputLowBitTotalBuffer_ = + ubBuffer.template GetWithOffset(UB_BUFFER_INFO.weightInputLowbitUbTotalSize, + highBitUbOffset); // 32 * 3KB = 96KB + ubAntiQuantScaleTotalBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantScaleUbTotalSize, + highBitUbOffset + 96 * GetKBUnit()); // 4 * 3 = 12KB + ubAntiQuantYPerTokenScaleTotalBuffer_ = + ubBuffer.template GetWithOffset(ANTI_QUANT_Y_PER_TOKEN_SCALE_TOTAL_BUFFER_SIZE, + highBitUbOffset + 108 * GetKBUnit()); // 1 * 2 = 2KB + ubAntiQuantYPerChannelScaleTotalBuffer_ = + ubBuffer.template GetWithOffset(ANTI_QUANT_Y_PER_CHANNEL_SCALE_TOTAL_BUFFER_SIZE, + highBitUbOffset + 110 * GetKBUnit()); // 1 * 2 = 2KB + ubAntiQuantYBiasTotalBuffer_ = + ubBuffer.template GetWithOffset(ANTI_QUANT_Y_BIAS_TOTAL_BUFFER_SIZE, + highBitUbOffset + 112 * GetKBUnit()); // 1 * 2 = 2KB + ubAntiQuantScaleMaskBuffer_ = ubBuffer.template GetWithOffset( + UB_BUFFER_INFO.antiQuantScaleMaskBufferSize, highBitUbOffset + 114 * GetKBUnit()); + ubAntiQuantScaleMaskBuffer_.SetValue(0, 0x00000000ffffffff); + ubAntiQuantScaleMaskBuffer_.SetValue(1, 0x00000000ffffffff); + ubAntiQuantScaleMaskBuffer_.SetValue(2, 0x00000000ffffffff); + ubAntiQuantScaleMaskBuffer_.SetValue(3, 0x00000000ffffffff); + + for (uint16_t i = 0; i < vecConfig.ubMte2BufferNum; ++i) { + vecEventIdVToMte2_[i] = GetTPipePtr()->AllocEventID(); + } + + for (uint16_t i = 0; i < UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; ++i) { + vecEventIdMte3ToV_[i] = GetTPipePtr()->AllocEventID(); + } + + for (uint16_t i = 0; i < UB_ANTI_QUANT_Y_BUFFER_NUM; ++i) { + vecEventIdAntiQuantYVToMte2_[i] = GetTPipePtr()->AllocEventID(); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::WaitVToMTE2() +{ + // 用临时变量接一下，优化编译的作用 + TEventID vecEventIdVToMte2[QUADRUPLE_BUFFER_NUM] = {vecEventIdVToMte2_[0], vecEventIdVToMte2_[1], + vecEventIdVToMte2_[2], vecEventIdVToMte2_[3]}; + if (likely(ubMte2LoopIdx_ > vecConfig.ubMte2BufferNum - 1)) { + if constexpr (vecConfig.ubMte2BufferNum == 2 || vecConfig.ubMte2BufferNum == 4) { + WaitFlag(vecEventIdVToMte2[ubMte2LoopIdx_ & (vecConfig.ubMte2BufferNum - 1)]); + } else { + WaitFlag(vecEventIdVToMte2[ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum]); + } + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::SetVToMTE2() +{ + // 用临时变量接一下，优化编译的作用 + TEventID vecEventIdVToMte2[QUADRUPLE_BUFFER_NUM] = {vecEventIdVToMte2_[0], vecEventIdVToMte2_[1], + vecEventIdVToMte2_[2], vecEventIdVToMte2_[3]}; + if constexpr (vecConfig.ubMte2BufferNum == 2 || vecConfig.ubMte2BufferNum == 4) { + SetFlag(vecEventIdVToMte2[(ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1)]); + } else { + SetFlag(vecEventIdVToMte2[(ubMte2LoopIdx_ - 1) % vecConfig.ubMte2BufferNum]); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyGmToUb( + uint64_t ubMte2NSize, uint64_t ubMte2KSize, uint64_t ubMte2NOffset, uint64_t ubMte2KOffset, + const BasicBlockOffsetParam &offsetParam) +{ + // ubMte2NSize和ubMte2KSize为实际MTE2搬运到UB的有效数据， + // 其按照ubMte2InnerSize进行跳写，垃圾数据无需操作，搬出的时搬运有效数据即可。 + if (ubMte2NSize == 0 || ubMte2KSize == 0) { + ubMte2LoopIdx_++; // 避免当前核无任务时，SetVToMTE2()对同一个flagID重复SetFlag的问题 + return; + } + + if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + if constexpr (IsSameType::value) { + CopyWeightGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + } else { + CopyAntiQuantParamsGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + CopyWeightGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + } + } else { + CopyWeightGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + CopyAntiQuantParamsGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + } + + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); + ubMte2LoopIdx_++; +} + +/** + * @brief 该函数搬运weight数据从GM进入UB上ubWeightInputLowBitTotalBuffer_中 + * 其分为vecConfig.ubMte2BufferNum块 每块大小为weightInputLowBitUbSingleBufferSize_ + * @param ubMte2NSize 从GM上搬运到UB的N方向大小 + * @param ubMte2KSize 从GM上搬运到UB的K方向大小 + * @param ubMte2NOffset 从GM上搬运到UB时, GM上N方向的偏移 + * @param ubMte2KOffset 从GM上搬运到UB时, GM上N方向的偏移 + * @param offsetParam 存储Weight矩阵的原始N,K,kAlign信息,用于搬运时GM上地址偏移的计算 + */ +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyWeightGmToUb( + uint64_t ubMte2NSize, uint64_t ubMte2KSize, uint64_t ubMte2NOffset, uint64_t ubMte2KOffset, + const BasicBlockOffsetParam &offsetParam) +{ + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + if constexpr (!wqmmConfig.bTrans) { + DataCopyPad2D(ubWeightInputLowBitTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .template ReinterpretCast(), + wGlobal_[ubMte2KOffset * offsetParam.nSize + ubMte2NOffset], ubMte2KSize, ubMte2NSize, + vecConfig.ubMte2InnerSize, offsetParam.nSize); + } else { + DataCopyPad2D(ubWeightInputLowBitTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .template ReinterpretCast(), + wGlobal_[ubMte2NOffset * offsetParam.kSize + ubMte2KOffset], ubMte2NSize, ubMte2KSize, + vecConfig.ubMte2InnerSize, offsetParam.kSize); + } + } else { + if constexpr (!wqmmConfig.bTrans) { + DataCopyPad2D(ubWeightInputLowBitTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .template ReinterpretCast(), + wGlobal_[ubMte2NOffset * offsetParam.kAlign + ubMte2KOffset * static_cast(C0_SIZE)], + CeilDivide(ubMte2NSize, static_cast(C0_SIZE)), + CeilAlign(ubMte2KSize, static_cast(BLOCK_CUBE)) * C0_SIZE, + vecConfig.ubMte2InnerSize * C0_SIZE, offsetParam.kAlign * C0_SIZE); + } else { + DataCopyPad2D(ubWeightInputLowBitTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .template ReinterpretCast(), + wGlobal_[ubMte2KOffset * offsetParam.nAlign + ubMte2NOffset * static_cast(C0_SIZE)], + CeilDivide(ubMte2KSize, static_cast(C0_SIZE)), + CeilAlign(ubMte2NSize, static_cast(BLOCK_CUBE)) * C0_SIZE, + vecConfig.ubMte2InnerSize * C0_SIZE, offsetParam.nAlign * C0_SIZE); + } + } +} + +/** + * @brief 该函数搬运scale和offset数据从GM进入UB上ubAntiQuantScaleTotalBuffer_中 + * pertensor场景读取单个数据即可，perchannel场景下搬运[1, ubMte2NSize]大小, 按照VECTOR_REG_WIDTH(256)对齐写入 + * A8W4场景将perGroupScale搬入antiQuantPerGroupScaleGlobal_中， + * 搬运[CeilDiv(ubMte2KOffset, antiQuantGroupSize), ubMte2NSize]大小，N方向按照128对齐写入 + * @param ubMte2NSize 从GM上搬运到UB的N方向大小 + * @param ubMte2KSize 从GM上搬运到UB的K方向大小 + * @param ubMte2NOffset 从GM上搬运到UB时, GM上N方向的偏移 + * @param ubMte2KOffset 从GM上搬运到UB时, GM上N方向的偏移 + * @param offsetParam 存储Weight矩阵的原始N,K,kAlign信息,用于搬运时GM上地址偏移的计算 + */ +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyAntiQuantParamsGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, + uint64_t ubMte2NOffset, + uint64_t ubMte2KOffset, + const BasicBlockOffsetParam &offsetParam) +{ + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_CHANNEL) { + DataCopyPad2D(ubAntiQuantScaleTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize], + antiQuantScaleGlobal_[ubMte2NOffset], 1, ubMte2NSize, + CeilAlign(ubMte2NSize, static_cast(VECTOR_REG_WIDTH)), offsetParam.nSize); + if constexpr (wqmmConfig.hasAntiQuantOffset) { + DataCopyPad2D(ubAntiQuantOffsetTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.antiQuantOffsetUbSingleBufferSize], + antiQuantOffsetGlobal_[ubMte2NOffset], 1, ubMte2NSize, + CeilAlign(ubMte2NSize, static_cast(VECTOR_REG_WIDTH)), offsetParam.nSize); + } + } else if constexpr (wqmmConfig.antiQuantType == QuantType::PER_TENSOR) { + scaleValue_ = antiQuantScaleGlobal_.GetValue(0); + if constexpr (wqmmConfig.hasAntiQuantOffset) { + offsetValue_ = antiQuantOffsetGlobal_.GetValue(0); + } + } else if constexpr (wqmmConfig.antiQuantType == QuantType::PER_GROUP) { + DataCopyPad2D(ubAntiQuantScaleTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize], + antiQuantScaleGlobal_[ubMte2KOffset / antiQuantGroupSize_ * offsetParam.nSize + ubMte2NOffset], + CeilDivide(ubMte2KSize, antiQuantGroupSize_), ubMte2NSize, VEC_MAX_ELEM_B16, offsetParam.nSize); + } else if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + CopyMxAntiQuantParamsGmToUb(ubMte2NSize, ubMte2KSize, ubMte2NOffset, ubMte2KOffset, offsetParam); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyMxAntiQuantParamsGmToUb(uint64_t ubMte2NSize, uint64_t ubMte2KSize, + uint64_t ubMte2NOffset, + uint64_t ubMte2KOffset, + const BasicBlockOffsetParam &offsetParam) +{ + uint64_t mxGroupNum = CeilDivide(offsetParam.kSize, MX_GROUPSIZE); + LocalTensor ubAntiQuantScaleBuffer = + ubAntiQuantScaleTotalBuffer_[(ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum) * + UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize] + .template ReinterpretCast(); + if constexpr (wqmmConfig.bTrans) { + DataCopyPad2D(ubAntiQuantScaleBuffer, + antiQuantScaleGlobal_[ubMte2NOffset * mxGroupNum + CeilDivide(ubMte2KOffset, MX_GROUPSIZE)], + ubMte2NSize, CeilDivide(ubMte2KSize, MX_GROUPSIZE), + CeilDivide(vecConfig.ubMte2InnerSize, MX_GROUPSIZE), mxGroupNum); + } else { + uint64_t scaleInnerSize = wqmmConfig.weightFormat != CubeFormat::NZ + ? vecConfig.ubMte2InnerSize + : 128UL; // nz场景当前不会合并，固定对齐到128即可 + DataCopyPad2D( + ubAntiQuantScaleBuffer, + antiQuantScaleGlobal_[CeilDivide(ubMte2KOffset, MX_GROUPSIZE) * offsetParam.nSize + ubMte2NOffset], + CeilDivide(ubMte2KSize, MX_GROUPSIZE), ubMte2NSize, scaleInnerSize, offsetParam.nSize); + } + + event_t eventIdScaleMTE2ToV = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdScaleMTE2ToV); + WaitFlag(eventIdScaleMTE2ToV); + MxScaleProcess(ubMte2NSize, ubMte2KSize); +} + +template +__aicore__ inline void BasicBlockLibVectorAntiQuantCompute::MxScaleProcess(uint64_t ubMte2NSize, + uint64_t ubMte2KSize) +{ + uint64_t ubMte2BufferIdx = ubMte2LoopIdx_ % vecConfig.ubMte2BufferNum; + MxFp4NdScaleParams mxFp4NdScaleParams; + mxFp4NdScaleParams.antiQuantScaleBasePhyAddr = + (__local_mem__ uint8_t *) + ubAntiQuantScaleTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize] + .GetPhyAddr(); + + mxFp4NdScaleParams.antiQuantScaleF16PhyAddr0 = + (__local_mem__ xType *) + ubAntiQuantScaleAfterCastTotalBuffer_[ubMte2BufferIdx * + UB_BUFFER_INFO.antiQuantScaleAfterCastUbSingleBufferSize] + .GetPhyAddr(); + + mxFp4NdScaleParams.antiQuantScaleF16PhyAddr1 = + mxFp4NdScaleParams.antiQuantScaleF16PhyAddr0 + (VECTOR_REG_WIDTH >> 1); + if constexpr (wqmmConfig.bTrans) { + mxFp4NdScaleParams.ubLoopExternalAxis = + CeilDivide(ubMte2NSize, static_cast(4)); // 按照4行共128个数进行处理 + } else { + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + // nd场景scale采取对齐到默认值策略 + mxFp4NdScaleParams.ubLoopExternalAxis = CeilDivide(ubMte2KSize, MX_GROUPSIZE); + } else { + // nz场景，scale采取内轴128 element对齐的策略，vf一次做256个数，此处直接除以2向上对齐 + mxFp4NdScaleParams.ubLoopExternalAxis = CeilDivide(CeilDivide(ubMte2KSize, MX_GROUPSIZE), 2UL); + } + } + if constexpr (wqmmConfig.bTrans) { + AscendC::VF_CALL>(mxFp4NdScaleParams); + } else { + AscendC::VF_CALL>(mxFp4NdScaleParams); + } +} + +template +__aicore__ inline void BasicBlockLibVectorAntiQuantCompute::CopyKcScaleBiasGmToUb(uint64_t nRealL0Size, + uint64_t mRealL0Size, + uint64_t nOffset, + uint64_t mOffset) +{ + if (ubMte2AntiquantYLoopIdx_ >= UB_ANTI_QUANT_Y_BUFFER_NUM) { + WaitFlag( + vecEventIdAntiQuantYVToMte2_[ubMte2AntiquantYLoopIdx_ & (UB_ANTI_QUANT_Y_BUFFER_NUM - 1)]); + } + // nRealL0Size和mRealL0Size为实际MTE2搬运到UB的有效数据， + // 垃圾数据无需操作，搬出的时搬运有效数据即可。 + if (unlikely(mRealL0Size == 0)) { + ubMte2AntiquantYLoopIdx_++; // 避免当前核无任务时，SetVToMTE2()对同一个flagID重复SetFlag的问题 + return; + } + // nRealL0Size 需要 对齐 32B / sizeof(T) 32B是UB最小对齐单元 + uint64_t nCopyLenAlign = CeilAlign(nRealL0Size, static_cast(FP32_BLOCK_SIZE)); + + DataCopyPad2D(ubAntiQuantYPerChannelScaleTotalBuffer_[(ubMte2AntiquantYLoopIdx_ % UB_ANTI_QUANT_Y_BUFFER_NUM) * + ANTI_QUANT_Y_PER_CHANNEL_SCALE_SINGLE_BUFFER_SIZE], + antiQuantYPerChannelScaleGlobal_[nOffset], 1, nRealL0Size, nCopyLenAlign, nRealL0Size); + + uint64_t mCopyLenAlign = CeilAlign(mRealL0Size, static_cast(FP32_BLOCK_SIZE)); + DataCopyPad2D(ubAntiQuantYPerTokenScaleTotalBuffer_[(ubMte2AntiquantYLoopIdx_ % UB_ANTI_QUANT_Y_BUFFER_NUM) * + ANTI_QUANT_Y_PER_TOKEN_SCALE_SINGLE_BUFFER_SIZE], + antiQuantYPerTokenScaleGlobal_[mOffset], 1, mRealL0Size, mCopyLenAlign, mRealL0Size); + + if (hasBias_) { + DataCopyPad2D(ubAntiQuantYBiasTotalBuffer_[(ubMte2AntiquantYLoopIdx_ % UB_ANTI_QUANT_Y_BUFFER_NUM) * + ANTI_QUANT_Y_BIAS_SINGLE_BUFFER_SIZE], + antiQuantYBiasGlobal_[nOffset], 1, nRealL0Size, nCopyLenAlign, nRealL0Size); + } + event_t eventIdAntiquantYMTE2ToV = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdAntiquantYMTE2ToV); + WaitFlag(eventIdAntiquantYMTE2ToV); + ubMte2AntiquantYLoopIdx_++; +} + +/** +* @brief 该函数作用为对于搬运到UB的weight,scale, offst,按照标准VF计算单元(64,256)进行多次循环计算,总计算量为L1所需的大小 + 每一个VF计算单元的结果放置于ubHighBitTotalBuffer_中，其计算和MTE3搬运使用vecEventIdMte3ToV_控制同步 +* @param ubConsumeConfig 其中l1RequireVfComputeRealN,l1RequireVfComputeRealK表示L1上需要VEC计算的实际数据量 + nWeightLowBitUbOffset, kWeightLowBitUbOffset表示在搬运到UB上的weight上的偏移 +* @param weightHighBitL1 L1上的weight地址，用于反量化结束后存放 +* @param l1ConsumeConfig 其中l1RealExternalLen为L1上真实外轴长度, SplitTwoVecExternalOffset为L1上切分给两个VEC核的外轴偏 + 移大小，用于搬运到L1的dst地址偏移计算。 +*/ +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::WeightAntiQuantCompute(const UbConsumeConfig &ubConsumeConfig, + const LocalTensor &weightHighBitL1, + const L1ConsumeConfig &l1ConsumeConfig) +{ + uint64_t weightHighBitL1Offset; + uint64_t nRealLen; + uint64_t kRealLen; + TEventID vecEventIdMte3ToV[QUADRUPLE_BUFFER_NUM]; + + // 用临时变量接一下，优化编译的作用 + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + vecEventIdMte3ToV[0] = vecEventIdMte3ToV_[0]; + vecEventIdMte3ToV[1] = vecEventIdMte3ToV_[1]; + } else { + vecEventIdMte3ToV[0] = vecEventIdMte3ToV_[0]; + vecEventIdMte3ToV[1] = vecEventIdMte3ToV_[1]; + vecEventIdMte3ToV[2] = vecEventIdMte3ToV_[2]; + vecEventIdMte3ToV[3] = vecEventIdMte3ToV_[3]; + } + for (uint64_t antiQuantKOffset = 0; antiQuantKOffset < ubConsumeConfig.l1RequireVfComputeRealK; + antiQuantKOffset += VF_CONFIG.vfKStandardLen) { + for (uint64_t antiQuantNOffset = 0; antiQuantNOffset < ubConsumeConfig.l1RequireVfComputeRealN; + antiQuantNOffset += VF_CONFIG.vfNStandardLen) { + if (likely(ubComputeLoopIdx_ > UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) { + WaitFlag( + vecEventIdMte3ToV[ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)]); + } + nRealLen = antiQuantNOffset + VF_CONFIG.vfNStandardLen >= ubConsumeConfig.l1RequireVfComputeRealN + ? ubConsumeConfig.l1RequireVfComputeRealN - antiQuantNOffset + : VF_CONFIG.vfNStandardLen; + kRealLen = antiQuantKOffset + VF_CONFIG.vfKStandardLen >= ubConsumeConfig.l1RequireVfComputeRealK + ? ubConsumeConfig.l1RequireVfComputeRealK - antiQuantKOffset + : VF_CONFIG.vfKStandardLen; + WeightAntiQuantProcess(nRealLen, kRealLen, antiQuantNOffset, antiQuantKOffset, ubConsumeConfig); + + weightHighBitL1Offset = + ComputeWeightHighBitL1Offset(antiQuantNOffset, antiQuantKOffset, nRealLen, kRealLen, l1ConsumeConfig); + event_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); + WeightHighBitUbToL1(weightHighBitL1Offset, nRealLen, kRealLen, weightHighBitL1, + l1ConsumeConfig.l1RealExternalLen); + SetFlag( + vecEventIdMte3ToV[ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)]); + ubComputeLoopIdx_++; + } + } +} + +template +__aicore__ inline uint64_t +BasicBlockLibVectorAntiQuantCompute::ComputeWeightHighBitL1Offset(uint64_t antiQuantNOffset, + uint64_t antiQuantKOffset, + uint64_t nRealLen, uint64_t kRealLen, + const L1ConsumeConfig &l1ConsumeConfig) +{ + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + uint64_t l1RealExternalLenAlign = + CeilAlign(l1ConsumeConfig.l1RealExternalLen, static_cast(BLOCK_CUBE)); + if constexpr (!wqmmConfig.bTrans) { + return l1RealExternalLenAlign * antiQuantNOffset + + (antiQuantKOffset + l1ConsumeConfig.l1SplitTwoVecExternalOffset) * static_cast(BLOCK_CUBE); + } else { + return l1RealExternalLenAlign * antiQuantKOffset + + (antiQuantNOffset + l1ConsumeConfig.l1SplitTwoVecExternalOffset) * static_cast(BLOCK_CUBE); + } + } else { + if constexpr (!wqmmConfig.bTrans) { + uint64_t kRealLenAlign = CeilAlign(kRealLen, static_cast(BLOCK_CUBE)); + return kRealLenAlign * (antiQuantNOffset + l1ConsumeConfig.l1SplitTwoVecExternalOffset) + + antiQuantKOffset * static_cast(C0_SIZE); + } else { + uint64_t nRealLenAlign = CeilAlign(nRealLen, static_cast(BLOCK_CUBE)); + return nRealLenAlign * (antiQuantKOffset + l1ConsumeConfig.l1SplitTwoVecExternalOffset) + + antiQuantNOffset * static_cast(C0_SIZE); + } + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::WeightAntiQuantProcess(uint64_t nRealLen, uint64_t kRealLen, + uint64_t antiQuantNOffset, + uint64_t antiQuantKOffset, + const UbConsumeConfig &ubConsumeConfig) +{ + if constexpr ((wqmmConfig.weightFormat != CubeFormat::NZ && !wqmmConfig.bTrans) || + (wqmmConfig.weightFormat == CubeFormat::NZ && wqmmConfig.bTrans)) { + AntiQuantProcess(kRealLen, nRealLen, ubConsumeConfig.nWeightLowBitUbOffset + antiQuantNOffset, + ubConsumeConfig.kWeightLowBitUbOffset + antiQuantKOffset); + } else if constexpr ((wqmmConfig.weightFormat != CubeFormat::NZ && wqmmConfig.bTrans) || + (wqmmConfig.weightFormat == CubeFormat::NZ && !wqmmConfig.bTrans)) { + AntiQuantProcess(nRealLen, kRealLen, ubConsumeConfig.nWeightLowBitUbOffset + antiQuantNOffset, + ubConsumeConfig.kWeightLowBitUbOffset + antiQuantKOffset); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcess( + uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset) +{ + if constexpr (wqmmConfig.antiQuantType == QuantType::MX) { + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + AntiQuantProcessNdMx(vfExternalRealLen, vfInnerRealLen, nWeightLowBitUbOffset, kWeightLowBitUbOffset); + } else if constexpr (IsSameType::value) { + AntiQuantProcessNzMxA8W4(vfExternalRealLen, vfInnerRealLen, nWeightLowBitUbOffset, kWeightLowBitUbOffset); + } else { + AntiQuantProcessNzMx(vfExternalRealLen, vfInnerRealLen, nWeightLowBitUbOffset, kWeightLowBitUbOffset); + } + } else { + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + AntiQuantProcessNd(vfExternalRealLen, vfInnerRealLen, nWeightLowBitUbOffset, kWeightLowBitUbOffset); + } else { + AntiQuantProcessNz(vfExternalRealLen, vfInnerRealLen, nWeightLowBitUbOffset, kWeightLowBitUbOffset); + } + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcessNdMx(uint64_t vfExternalRealLen, + uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, + uint64_t kWeightLowBitUbOffset) +{ + uint64_t mte2BufIdx = (ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1); + uint64_t weightF16BufIdx = ubComputeLoopIdx_ & 1; + uint64_t weightLowBitBufOffset = mte2BufIdx * UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize; + uint64_t antiQuantScaleAfterCastBufOffset = mte2BufIdx * UB_BUFFER_INFO.antiQuantScaleAfterCastUbSingleBufferSize; + uint64_t weightLowBitOffset; + uint64_t antiQuantScaleAfterCastOffset; + + if constexpr (wqmmConfig.bTrans) { + weightLowBitOffset = nWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + kWeightLowBitUbOffset; + antiQuantScaleAfterCastOffset = (nWeightLowBitUbOffset * vecConfig.ubMte2InnerSize / MX_GROUPSIZE + + CeilDivide(kWeightLowBitUbOffset, MX_GROUPSIZE)) * + 2; // 在经过e8m0-->f16后，偏移需要乘2 + } else { + weightLowBitOffset = kWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + nWeightLowBitUbOffset; + antiQuantScaleAfterCastOffset = + CeilDivide(kWeightLowBitUbOffset, MX_GROUPSIZE) * vecConfig.ubMte2InnerSize + nWeightLowBitUbOffset; + } + + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value) { + weightLowBitOffset = weightLowBitOffset >> 1; + } + + MxFp4NdWeightParams mxFp4NdWeightParams; + mxFp4NdWeightParams.antiQuantScaleF16PhyAddr0 = + (__local_mem__ xType *)ubAntiQuantScaleAfterCastTotalBuffer_.GetPhyAddr(antiQuantScaleAfterCastBufOffset) + + antiQuantScaleAfterCastOffset; + mxFp4NdWeightParams.weightLowBitPhyAddr0 = + (__local_mem__ wType *)ubWeightInputLowBitTotalBuffer_.GetPhyAddr(weightLowBitBufOffset) + weightLowBitOffset; + mxFp4NdWeightParams.weightLowBitPhyAddr1 = mxFp4NdWeightParams.weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 2); + mxFp4NdWeightParams.weightF16PhyAddr0 = (__local_mem__ xType *)ubHighBitTotalBuffer_.GetPhyAddr( + weightF16BufIdx * UB_BUFFER_INFO.highBitDataUbSingleBufferSize); + mxFp4NdWeightParams.weightF16PhyAddr1 = + mxFp4NdWeightParams.weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); + if constexpr (wqmmConfig.bTrans) { + mxFp4NdWeightParams.ubLoopExternalAxis = vfExternalRealLen; + // 8个scale对应8/2 * 32=128个数 + mxFp4NdWeightParams.antiQuantScaleF16PhyAddr1 = mxFp4NdWeightParams.antiQuantScaleF16PhyAddr0 + 8; + AscendC::VF_CALL>(mxFp4NdWeightParams); + + } else { + mxFp4NdWeightParams.ubLoopExternalAxis = CeilDivide(vfExternalRealLen, MX_GROUPSIZE); + mxFp4NdWeightParams.antiQuantScaleF16PhyAddr1 = + mxFp4NdWeightParams.antiQuantScaleF16PhyAddr0 + (VECTOR_REG_WIDTH >> 1); + AscendC::VF_CALL>(mxFp4NdWeightParams); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcessNzMx(uint64_t vfExternalRealLen, + uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, + uint64_t kWeightLowBitUbOffset) +{ + if constexpr (!wqmmConfig.bTrans) { + Fp4NzParams fp4NzParams; + uint64_t ubMte2BufferIdx = (ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1); + + uint64_t antiQuantScaleAfterCastBufOffset = + ubMte2BufferIdx * UB_BUFFER_INFO.antiQuantScaleAfterCastUbSingleBufferSize + nWeightLowBitUbOffset; + fp4NzParams.antiQuantScaleBasePhyAddr = + (__local_mem__ xType *)ubAntiQuantScaleAfterCastTotalBuffer_[antiQuantScaleAfterCastBufOffset].GetPhyAddr(); + fp4NzParams.weightLowBitPhyAddr = + (__local_mem__ wType *) + ubWeightInputLowBitTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .GetPhyAddr() + + ((nWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + kWeightLowBitUbOffset * C0_SIZE) >> 1); + + fp4NzParams.weightHighBitPhyAddr = + (__local_mem__ xType *) + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) * + VEC_MAX_ELEM_B16] + .GetPhyAddr(); + + fp4NzParams.loopN1 = CeilDivide(vfExternalRealLen, static_cast(C0_SIZE)); + // 跳写UB避免bank冲突 + fp4NzParams.innerDstStride = VEC_MAX_ELEM_B16 * UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; + + fp4NzParams.loopGroupNum = CeilDivide(vfInnerRealLen, MX_GROUPSIZE); + // 小数据量vfInnerRealLen对齐到BLOCK_CUBE，再计算循环次数 + fp4NzParams.loopInnerNum = + vfInnerRealLen < MX_GROUPSIZE + ? CeilDivide(CeilAlign(vfInnerRealLen, static_cast(BLOCK_CUBE)) * C0_SIZE, VEC_MAX_ELEM_B16) + : CeilDivide(MX_GROUPSIZE * C0_SIZE, VEC_MAX_ELEM_B16); + fp4NzParams.groupDstStride = fp4NzParams.loopInnerNum * fp4NzParams.innerDstStride; + fp4NzParams.loopN1DstStride = fp4NzParams.loopGroupNum * fp4NzParams.groupDstStride; + AscendC::VF_CALL>(fp4NzParams); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcessNzMxA8W4(uint64_t vfExternalRealLen, + uint64_t vfInnerRealLen, + uint64_t nWeightLowBitUbOffset, + uint64_t kWeightLowBitUbOffset) +{ + MxA8W4NzParams mxA8W4NzParams; + uint64_t ubMte2BufferIdx = (ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1); + + mxA8W4NzParams.weightLowBitPhyAddr = + (__local_mem__ wType *) + ubWeightInputLowBitTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .GetPhyAddr() + + ((kWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + nWeightLowBitUbOffset * C0_SIZE) >> 1); + + mxA8W4NzParams.weightHighBitPhyAddr = + (__local_mem__ xType *) + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) * + VECTOR_REG_WIDTH] + .GetPhyAddr(); + + mxA8W4NzParams.loopKNum = CeilDivide(vfExternalRealLen, static_cast(C0_SIZE)); + mxA8W4NzParams.innerLoopNum = CeilDivide(CeilAlign(vfInnerRealLen, static_cast(BLOCK_CUBE)) * C0_SIZE, + static_cast(VECTOR_REG_WIDTH)); + // 跳写UB避免bank冲突 + mxA8W4NzParams.innerDstStride = VECTOR_REG_WIDTH * UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; + mxA8W4NzParams.loopKDstStride = mxA8W4NzParams.innerLoopNum * mxA8W4NzParams.innerDstStride; + AscendC::VF_CALL>(mxA8W4NzParams); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcessNd( + uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset) +{ + LocalAddressParam addressParam; + CalLocalAddrForVf(nWeightLowBitUbOffset, kWeightLowBitUbOffset, addressParam); + if constexpr (IsSameType::value || IsSameType::value) { + if constexpr (wqmmConfig.bTrans) { + AntiQuantB8CommonNdNk( + addressParam.antiQuantScaleBasePhyAddr, addressParam.antiQuantOffsetBasePhyAddr, + addressParam.weightLowBitPhyAddr0, addressParam.weightF16PhyAddr0, scaleValue_, offsetValue_, + vfExternalRealLen); + } else { + AntiQuantB8CommonNdKn( + addressParam.antiQuantScaleBasePhyAddr, addressParam.antiQuantOffsetBasePhyAddr, + addressParam.weightLowBitPhyAddr0, addressParam.weightF16PhyAddr0, scaleValue_, offsetValue_, + vfExternalRealLen); + } + } else if constexpr (IsSameType::value || IsSameType::value) { + CalculateParam calculateParam; + calculateParam.offsetValue = offsetValue_; + calculateParam.scaleValue = scaleValue_; + calculateParam.ubLoop = vfExternalRealLen; + if constexpr (wqmmConfig.bTrans) { + AscendC::VF_CALL>(addressParam, calculateParam); + } else { + AscendC::VF_CALL>(addressParam, calculateParam); + } + } else { + if constexpr (wqmmConfig.bTrans) { + AntiQuantInt4NdNk( + addressParam.antiQuantScaleBasePhyAddr, addressParam.antiQuantOffsetBasePhyAddr, + addressParam.weightLowBitPhyAddr0, addressParam.weightF16PhyAddr0, scaleValue_, offsetValue_, + vfExternalRealLen); + } else { + AntiQuantInt4NdKn( + addressParam.antiQuantScaleBasePhyAddr, addressParam.antiQuantOffsetBasePhyAddr, + addressParam.weightLowBitPhyAddr0, addressParam.weightF16PhyAddr0, scaleValue_, offsetValue_, + vfExternalRealLen); + } + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::AntiQuantProcessNz( + uint64_t vfExternalRealLen, uint64_t vfInnerRealLen, uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset) +{ + if constexpr (!wqmmConfig.bTrans) { + Int4NzParams int4NzParams; + CalInt4NzKnLocalAddrForVf(nWeightLowBitUbOffset, kWeightLowBitUbOffset, int4NzParams); + int4NzParams.loopN1 = CeilDivide(vfExternalRealLen, static_cast(C0_SIZE)); + // 跳写UB避免bank冲突，A16跳1024B，A8跳512B; MTE3对应跳读 + int4NzParams.innerDstStride = VEC_MAX_ELEM_B16 * UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; + + if constexpr (wqmmConfig.antiQuantType == QuantType::PER_GROUP) { + int4NzParams.antiQuantGroupSize = antiQuantGroupSize_; + int4NzParams.loopGroupNum = CeilDivide(vfInnerRealLen, antiQuantGroupSize_); + int4NzParams.loopInnerNum = + vfInnerRealLen < antiQuantGroupSize_ ? + CeilDivide(CeilAlign(vfInnerRealLen, static_cast(BLOCK_CUBE)) * C0_SIZE, + VEC_MAX_ELEM_B16) : + CeilDivide(antiQuantGroupSize_ * C0_SIZE, VEC_MAX_ELEM_B16); + int4NzParams.groupDstStride = int4NzParams.loopInnerNum * int4NzParams.innerDstStride; + int4NzParams.loopN1DstStride = int4NzParams.loopGroupNum * int4NzParams.groupDstStride; + AscendC::VF_CALL>(int4NzParams); + } else { + int4NzParams.loopInnerNum = CeilDivide( + static_cast(CeilAlign(vfInnerRealLen, static_cast(BLOCK_CUBE))) * C0_SIZE, + VEC_MAX_ELEM_B16); + int4NzParams.loopN1DstStride = int4NzParams.loopInnerNum * int4NzParams.innerDstStride; + + if (int4NzParams.loopN1 == 1) { + AscendC::VF_CALL>(int4NzParams); + } else { + AscendC::VF_CALL>(int4NzParams); + } + } + } +} + +template +__aicore__ inline void BasicBlockLibVectorAntiQuantCompute< + xType, wType, antiQuantScaleType, yType, wqmmConfig, + vecConfig>::CalInt4NzKnLocalAddrForVf(uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset, + Int4NzParams &int4NzParams) +{ + uint64_t ubMte2BufferIdx; + if constexpr (IsSameType::value) { + ubMte2BufferIdx = (ubMte2LoopIdx_ - 1) % vecConfig.ubMte2BufferNum; + int4NzParams.antiQuantScaleBasePhyAddr = + (__local_mem__ antiQuantScaleType *) + ubAntiQuantScaleTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize + + kWeightLowBitUbOffset / antiQuantGroupSize_ * VEC_MAX_ELEM_B16 + + nWeightLowBitUbOffset] + .GetPhyAddr(); + int4NzParams.antiQuantScaleMaskPhyAddr = (__local_mem__ uint8_t *)ubAntiQuantScaleMaskBuffer_.GetPhyAddr(); + } else { + ubMte2BufferIdx = (ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1); + int4NzParams.antiQuantScaleBasePhyAddr = + (__local_mem__ antiQuantScaleType *) + ubAntiQuantScaleTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize + + nWeightLowBitUbOffset] + .GetPhyAddr(); + int4NzParams.antiQuantOffsetBasePhyAddr = + (__local_mem__ xType *) + ubAntiQuantOffsetTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize + + nWeightLowBitUbOffset] + .GetPhyAddr(); + } + + int4NzParams.weightLowBitPhyAddr = + (__local_mem__ wType *) + ubWeightInputLowBitTotalBuffer_[ubMte2BufferIdx * UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize] + .GetPhyAddr() + + ((nWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + kWeightLowBitUbOffset * C0_SIZE) >> 1); + + int4NzParams.weightHighBitPhyAddr = + (__local_mem__ xType *) + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) * + VEC_MAX_ELEM_B16] + .GetPhyAddr(); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CalLocalAddrForVf( + uint64_t nWeightLowBitUbOffset, uint64_t kWeightLowBitUbOffset, LocalAddressParam &localAddressParam) +{ + uint64_t mte2BufIdx = (ubMte2LoopIdx_ - 1) & (vecConfig.ubMte2BufferNum - 1); + uint64_t antiquantParamOffset = + mte2BufIdx * UB_BUFFER_INFO.antiQuantScaleUbSingleBufferSize + nWeightLowBitUbOffset; + localAddressParam.antiQuantScaleBasePhyAddr = + (__local_mem__ xType *)ubAntiQuantScaleTotalBuffer_.GetPhyAddr(antiquantParamOffset); + localAddressParam.antiQuantScaleBasePhyAddr1 = localAddressParam.antiQuantScaleBasePhyAddr + VEC_MAX_ELEM_B16; + localAddressParam.antiQuantOffsetBasePhyAddr = + (__local_mem__ xType *)ubAntiQuantOffsetTotalBuffer_.GetPhyAddr(antiquantParamOffset); + localAddressParam.antiQuantOffsetBasePhyAddr1 = localAddressParam.antiQuantOffsetBasePhyAddr + VEC_MAX_ELEM_B16; + + uint64_t weightLowBitOffset; + if constexpr (wqmmConfig.bTrans) { + weightLowBitOffset = nWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + kWeightLowBitUbOffset; + } else { + weightLowBitOffset = kWeightLowBitUbOffset * vecConfig.ubMte2InnerSize + nWeightLowBitUbOffset; + } + + if constexpr (IsSameType::value) { + weightLowBitOffset = weightLowBitOffset >> 1; + } + + uint64_t weightLowBitBufOffset = mte2BufIdx * UB_BUFFER_INFO.weightInputLowBitUbSingleBufferSize; + localAddressParam.weightLowBitPhyAddr0 = + (__local_mem__ wType *)ubWeightInputLowBitTotalBuffer_.GetPhyAddr(weightLowBitBufOffset) + weightLowBitOffset; + if constexpr (IsSameType::value) { + // int4每次处理128个数即为64B, 256>>2=64 + localAddressParam.weightLowBitPhyAddr1 = localAddressParam.weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 2); + } else { + localAddressParam.weightLowBitPhyAddr1 = localAddressParam.weightLowBitPhyAddr0 + (VECTOR_REG_WIDTH >> 1); + } + + uint64_t weightF16BufIdx = ubComputeLoopIdx_ & 1; + localAddressParam.weightF16PhyAddr0 = (__local_mem__ xType *)ubHighBitTotalBuffer_.GetPhyAddr( + weightF16BufIdx * UB_BUFFER_INFO.highBitDataUbSingleBufferSize); + localAddressParam.weightF16PhyAddr1 = + localAddressParam.weightF16PhyAddr0 + WEIGHT_F16_UB_NZ_STRIDE * (VECTOR_REG_WIDTH >> 1); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::WeightHighBitUbToL1(uint64_t weightHighBitL1Offset, + uint64_t antiQuantRealN, uint64_t antiQuantRealK, + const LocalTensor &weightHighBitL1, + uint64_t l1RealExternalLen) +{ + DataCopyParams params; + if constexpr (wqmmConfig.weightFormat != CubeFormat::NZ) { + if constexpr (wqmmConfig.bTrans) { + params.blockLen = antiQuantRealN; + params.blockCount = CeilDivide(antiQuantRealK, static_cast(BLOCK_CUBE)); + params.srcStride = WEIGHT_F16_UB_NZ_STRIDE - antiQuantRealN; + params.dstStride = CeilAlign(l1RealExternalLen, static_cast(BLOCK_CUBE)) - antiQuantRealN; + } else { + params.blockLen = antiQuantRealK; + params.blockCount = CeilDivide(antiQuantRealN, static_cast(BLOCK_CUBE)); + params.srcStride = WEIGHT_F16_UB_NZ_STRIDE - antiQuantRealK; + params.dstStride = CeilAlign(l1RealExternalLen, static_cast(BLOCK_CUBE)) - antiQuantRealK; + } + DataCopy(weightHighBitL1[weightHighBitL1Offset], + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & 1) * UB_BUFFER_INFO.highBitDataUbSingleBufferSize], params); + } else { + if constexpr (wqmmConfig.antiQuantType == QuantType::MX && !IsSameType::value) { + // 小数据量vfInnerRealLen或K与BLOCK_CUBE, MX_GROUPSIZE上对齐大小相同时, mte3对齐到BLOCK_CUBE + if (antiQuantRealK < MX_GROUPSIZE || CeilAlign(antiQuantRealK, static_cast(BLOCK_CUBE)) == + CeilAlign(antiQuantRealK, static_cast(MX_GROUPSIZE))) { + CopyWeightHighBitForAligned(weightHighBitL1Offset, antiQuantRealN, antiQuantRealK, weightHighBitL1); + } else { + // mte3对齐到MX_GROUPSIZE + CopyWeightHighBitForUnaligned(weightHighBitL1Offset, antiQuantRealN, antiQuantRealK, weightHighBitL1); + } + } else { + CopyWeightHighBitForAligned(weightHighBitL1Offset, antiQuantRealN, antiQuantRealK, weightHighBitL1); + } + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyWeightHighBitForAligned(uint64_t weightHighBitL1Offset, + uint64_t antiQuantRealN, + uint64_t antiQuantRealK, + const LocalTensor &weightHighBitL1) +{ + DataCopyParams params; + if constexpr (!wqmmConfig.bTrans) { + params.blockCount = CeilAlign(antiQuantRealK, static_cast(BLOCK_CUBE)) * + CeilAlign(antiQuantRealN, static_cast(C0_SIZE)) / VEC_REG_ELEM; + } else { + params.blockCount = CeilAlign(antiQuantRealK, static_cast(C0_SIZE)) * + CeilAlign(antiQuantRealN, static_cast(BLOCK_CUBE)) / VEC_REG_ELEM; + } + + params.blockLen = (IsSameType::value || IsSameType::value) + ? VEC_REG_ELEM / ONE_BLK_SIZE + : VEC_REG_ELEM / BLOCK_CUBE; + params.srcStride = (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1) * params.blockLen; + params.dstStride = 0; // dst地址连续 + DataCopy( + weightHighBitL1[weightHighBitL1Offset], + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) * VEC_REG_ELEM], + params); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyWeightHighBitForUnaligned(uint64_t weightHighBitL1Offset, + uint64_t antiQuantRealN, + uint64_t antiQuantRealK, + const LocalTensor &weightHighBitL1) +{ + DataCopyParams params; + // 跳写UB避免bank冲突，A16跳1024B; MTE3对应跳读 + uint64_t innerDstStride = VEC_MAX_ELEM_B16 * UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; + uint64_t blockCubeAlignOfk = CeilAlign(antiQuantRealK, static_cast(BLOCK_CUBE)); + uint64_t mxGroupSizeAlignOfk = CeilAlign(antiQuantRealK, static_cast(MX_GROUPSIZE)); + for (int i = 0; i < CeilDivide(antiQuantRealN, static_cast(C0_SIZE)); i++) { + params.blockCount = + CeilAlign(antiQuantRealK, static_cast(BLOCK_CUBE)) / (VEC_MAX_ELEM_B16 / BLOCK_CUBE); + params.blockLen = VEC_MAX_ELEM_B16 / BLOCK_CUBE; + params.srcStride = (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1) * params.blockLen; + params.dstStride = 0; // dst地址连续 + weightHighBitL1Offset += i * blockCubeAlignOfk * BLOCK_CUBE; + // dst需要再加i * 32对齐后的blockCount * innerDstStride作为地址偏置 + DataCopy(weightHighBitL1[weightHighBitL1Offset], + ubHighBitTotalBuffer_[(ubComputeLoopIdx_ & (UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum - 1)) * + VEC_MAX_ELEM_B16 + + i * (mxGroupSizeAlignOfk / (VEC_MAX_ELEM_B16 / BLOCK_CUBE) * innerDstStride)], + params); + } +} + +template +__aicore__ inline void BasicBlockLibVectorAntiQuantCompute::AntiQuantYWithKc(uint64_t nRealL0Size, + uint64_t mRealL0Size) +{ + if (unlikely(mRealL0Size == 0)) { + SetFlag( + vecEventIdAntiQuantYVToMte2_[(ubMte2AntiquantYLoopIdx_ - 1) & (UB_ANTI_QUANT_Y_BUFFER_NUM - 1)]); + return; + } + + LocalAddressYParam addressParam; + CalLocalAddrForYVf(addressParam); + uint16_t loopN = CeilDivide(nRealL0Size, ANTIQUANT_Y_STANDARD_N_SIZE); + + if (hasBias_) { + AntiQuantYB32(addressParam, CeilAlign(nRealL0Size, ANTIQUANT_Y_STANDARD_N_SIZE), + ANTIQUANT_Y_STANDARD_N_SIZE, loopN, (uint16_t)mRealL0Size); + } else { + AntiQuantYB32(addressParam, CeilAlign(nRealL0Size, ANTIQUANT_Y_STANDARD_N_SIZE), + ANTIQUANT_Y_STANDARD_N_SIZE, loopN, (uint16_t)mRealL0Size); + } + + SetFlag( + vecEventIdAntiQuantYVToMte2_[(ubMte2AntiquantYLoopIdx_ - 1) & (UB_ANTI_QUANT_Y_BUFFER_NUM - 1)]); + event_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CalLocalAddrForYVf( + LocalAddressYParam &localAddressParam) +{ + uint64_t mte3BufIdx = (ubMte2AntiquantYLoopIdx_ - 1) & (UB_ANTI_QUANT_Y_BUFFER_NUM - 1); + uint64_t antiquantParamOffset = mte3BufIdx * ANTI_QUANT_Y_PER_CHANNEL_SCALE_SINGLE_BUFFER_SIZE; + localAddressParam.yOriginPhyAddr = (__local_mem__ int32_t *)ubHighBitTotalBuffer_.GetPhyAddr(); + localAddressParam.yPhyAddr = (__local_mem__ yType *)localAddressParam.yOriginPhyAddr; + localAddressParam.cScalePhyAddr = + (__local_mem__ float *)ubAntiQuantYPerChannelScaleTotalBuffer_.GetPhyAddr(antiquantParamOffset); + localAddressParam.kScalePhyAddr = + (__local_mem__ float *)ubAntiQuantYPerTokenScaleTotalBuffer_.GetPhyAddr(antiquantParamOffset); + if (hasBias_) { + localAddressParam.biasPhyAddr = + (__local_mem__ float *)ubAntiQuantYBiasTotalBuffer_.GetPhyAddr(antiquantParamOffset); + } +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::CopyYUbToGm( + uint64_t nRealL0Size, uint64_t mRealL0Size, __gm__ half *yGm, const BasicBlockOffsetParam &offsetParam, + uint64_t aivMOffset) +{ + if (unlikely(mRealL0Size == 0)) { + return; + } + antiQuantYF16Global_.SetGlobalBuffer(yGm); + uint64_t yGmAddrOffset = (offsetParam.mOffset + aivMOffset) * offsetParam.nSize + offsetParam.nOffset; + + DataCopyPad2D(antiQuantYF16Global_[yGmAddrOffset], ubHighBitTotalBuffer_.template ReinterpretCast(), + mRealL0Size, nRealL0Size, CeilAlign(nRealL0Size, ANTIQUANT_Y_STANDARD_N_SIZE) * 2, offsetParam.nSize); + + event_t eventIdMTE3ToV = static_cast(GetTPipePtr()->FetchEventID()); + SetFlag(eventIdMTE3ToV); + WaitFlag(eventIdMTE3ToV); +} + +template +__aicore__ inline void +BasicBlockLibVectorAntiQuantCompute::End() +{ + TEventID vecEventIdVToMte2[QUADRUPLE_BUFFER_NUM] = {vecEventIdVToMte2_[0], vecEventIdVToMte2_[1], + vecEventIdVToMte2_[2], vecEventIdVToMte2_[3]}; + TEventID vecEventIdMte3ToV[QUADRUPLE_BUFFER_NUM] = {vecEventIdMte3ToV_[0], vecEventIdMte3ToV_[1], + vecEventIdMte3ToV_[2], vecEventIdMte3ToV_[3]}; + TEventID vecEventIdAntiQuantYVToMte2[UB_ANTI_QUANT_Y_BUFFER_NUM] = {vecEventIdAntiQuantYVToMte2_[0], + vecEventIdAntiQuantYVToMte2_[1]}; + + for (uint16_t idx = 0; idx < ubComputeLoopIdx_ && idx < UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; idx++) { + WaitFlag(vecEventIdMte3ToV[idx]); + } + + for (uint16_t idx = 0; idx < ubMte2LoopIdx_ && idx < vecConfig.ubMte2BufferNum; idx++) { + WaitFlag(vecEventIdVToMte2[idx]); + } + + if constexpr (IsSameType::value) { + for (uint16_t idx = 0; idx < ubMte2AntiquantYLoopIdx_ && idx < UB_ANTI_QUANT_Y_BUFFER_NUM; idx++) { + WaitFlag(vecEventIdAntiQuantYVToMte2[idx]); + } + } + + for (uint16_t idx = 0; idx < vecConfig.ubMte2BufferNum; idx++) { + GetTPipePtr()->ReleaseEventID(vecEventIdVToMte2[idx]); + } + + for (uint16_t idx = 0; idx < UB_BUFFER_INFO.ubWeightOutputHighBitBufferNum; idx++) { + GetTPipePtr()->ReleaseEventID(vecEventIdMte3ToV[idx]); + } + if constexpr (IsSameType::value) { + for (uint16_t idx = 0; idx < UB_ANTI_QUANT_Y_BUFFER_NUM; idx++) { + GetTPipePtr()->ReleaseEventID(vecEventIdAntiQuantYVToMte2[idx]); + } + } +} +} // namespace WeightQuantBatchMatmulV2::Arch35 + +#endif // GROUPED_MATMUL_WEIGHT_QUANT_VEC_COMPUTE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/dlinfer_grouped_matmul_direct.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/dlinfer_grouped_matmul_direct.cpp new file mode 100644 index 00000000..75918858 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/dlinfer_grouped_matmul_direct.cpp @@ -0,0 +1,894 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul.cpp + * \brief + */ +#include "grouped_matmul_utils.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/grouped_matmul_tiling_data_apt.h" +using GMMWeightQuantTilingData = DlinferGroupedMatmulDirectTilingData::GMMWeightQuantTilingData; +using GMMNoQuantTilingData = DlinferGroupedMatmulDirectTilingData::GMMNoQuantTilingData; +using GMMQuantTilingData = DlinferGroupedMatmulDirectTilingData::GMMQuantTilingData; +#if defined(V310_GMM_QUANT) +#include "arch35/quant_adaptive_sliding_window_templates/gqmm_tiling_key.h" +#if defined(V310_GMM_QUANT_MX) || defined(V310_GMM_QUANT_CUBE) || defined(V310_GMM_QUANT_PERTENSOR_CUBE) +#include "arch35/quant_adaptive_sliding_window_templates/gqmm_cube_on_the_fly.h" +#endif +#if defined(V310_GMM_QUANT_MX) || defined(V310_GMM_QUANT_PERTENSOR_CUBE) +#include "arch35/quant_adaptive_sliding_window_templates/gqmm_init_output.h" +#endif +#if defined(V310_GMM_QUANT_MX) || defined(V310_GMM_QUANT_PERTENSOR_CUBE) +#include "arch35/quant_adaptive_sliding_window_templates/gqmm_mix_online_dynamic.h" +#endif +#if defined(V310_GMM_QUANT_PERTILE) +#include "arch35/quant_adaptive_sliding_window_templates/gqmm_act_pertile_kernel.h" +#endif +#elif defined(V310_GMM_ANTI_QUANT) +#include "arch35/weight_quant_basic_block/basic_block_config.h" +#include "arch35/weight_quant_basic_block/grouped_matmul_weight_quant_basic_controller.h" +#include "arch35/weight_quant_basic_block/grouped_matmul_weight_quant_resplit_controller.h" +#include "arch35/weight_quant_basic_block/weight_quant_basic_block.h" +#include "arch35/weight_quant_basic_block/weight_quant_vcv_basic_block.h" +#include "arch35/weight_quant_basic_block/weight_quant_tiling_key.h" +using WeightQuantBatchMatmulV2::Arch35::QuantType; +using WeightQuantBatchMatmulV2::Arch35::A16MXF4_NZKN; +using WeightQuantBatchMatmulV2::Arch35::MXA8W4_NZNK; +using WeightQuantBatchMatmulV2::Arch35::S8S4_NZKN_G; +using WeightQuantBatchMatmulV2::Arch35::WeightQuantMatmulBasicBlock; +using WeightQuantBatchMatmulV2::Arch35::WeightQuantVcvMatmulBasicBlock; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_0 = {2, 512}; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_1 = {4, 512}; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_2 = {2, 1024}; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_3 = {4, 256}; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_4 = {3, 512}; +static constexpr VecAntiQuantConfig VEC_ANTIQUANT_CONFIG_5 = {3, 384}; +#if defined(DT_FLOAT) && defined(ORIG_DTYPE_WEIGHT) && ORIG_DTYPE_WEIGHT == DT_FLOAT + #undef DTYPE_WEIGHT + #define DTYPE_WEIGHT fp4x2_e2m1_t +#endif +#if defined(DT_INT32) && defined(ORIG_DTYPE_WEIGHT) && ORIG_DTYPE_WEIGHT == DT_INT32 + #undef DTYPE_WEIGHT + #define DTYPE_WEIGHT AscendC::int4b_t + #undef ORIG_DTYPE_WEIGHT + #define ORIG_DTYPE_WEIGHT DT_INT4 +#endif +#else +#include "arch35/non_quant/grouped_matmul_basic_kernel.h" +#include "arch35/non_quant/grouped_matmul_tiling_key.h" +#endif +#else +#include "grouped_matmul_antiquant.h" +#include "grouped_matmul_vector.h" +#include "grouped_matmul_tiling_key.h" +#include "grouped_matmul.h" +#endif + +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + +#include "grouped_matmul_antiquant_a16w8_msd.h" +#include "grouped_matmul_antiquant_a8w4_msd_pre.h" +#include "grouped_matmul_antiquant_a8w4_msd.h" +#include "grouped_matmul_antiquant_a8w4_pre.h" +#include "grouped_matmul_antiquant_a8w4.h" +#include "grouped_matmul_antiquant_a8w4_msd_new.h" +#include "grouped_matmul_quant_mixcore.h" +#include "grouped_matmul_pre_tiling.h" +#include "grouped_matmul_a4w4.h" +#include "grouped_matmul_autotiling_a8w4.h" +#ifndef __CCE_KT_TEST__ +#include "grouped_matmul_fixaxismove_interface.cpp" +#endif +#endif + + +using namespace AscendC; +using namespace matmul; +using namespace GROUPED_MATMUL; + +#ifndef FORMAT_FRACTAL_NZ + #define FORMAT_FRACTAL_NZ +#endif + +namespace { +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ +constexpr CubeFormat wFormat = CubeFormat::NZ; +constexpr MatmulConfig matmulCFG = NZ_CFG_MDL; +#else +constexpr CubeFormat wFormat = CubeFormat::ND; +constexpr MatmulConfig matmulCFG = CFG_MDL; +#endif + +#if defined(GMM_ANTI_QUANT_A8W4_MSD) +constexpr MatmulConfig A8W4_GMM_CFG_MDL = GetNormalConfig(); +constexpr auto GetMmCFG() { + auto CFG = CFG_MDL; + CFG.isPartialOutput = true; + return CFG; +} +constexpr MatmulConfig A8W4_GMM_CFG_MDL_NEW = GetMmCFG(); +#endif +} + +template +using xType = MatmulType; + +template +using xTypeMSD = MatmulType; + +template +using weightType = MatmulType; + +template +using weightTypeMSD = MatmulType; + +using yType = MatmulType; + +using yTypeMSD = MatmulType; + +using biasType = MatmulType; + +namespace { + __aicore__ inline static constexpr MatmulApiStaticTiling GetGmmMatmulApiTiling(bool isND2NZ, bool transB) { + MatmulConfig conf = GenGmmConf(isND2NZ); + MatmulApiStaticTiling staticTilingTmp; + if (transB) { + staticTilingTmp = GetMatmulApiTiling, weightType, yType, biasType>(conf); + } else { + staticTilingTmp = GetMatmulApiTiling, weightType, yType, biasType>(conf); + } + staticTilingTmp.depthA1 = STATIC_TILING_DEPTH_A1_B1; + staticTilingTmp.depthB1 = STATIC_TILING_DEPTH_A1_B1; + staticTilingTmp.stepM = 1; + staticTilingTmp.stepN = 1; + staticTilingTmp.stepKa = STATIC_TILING_STEP_KA_KB; + staticTilingTmp.stepKb = STATIC_TILING_STEP_KA_KB; + staticTilingTmp.dbL0A = DOUBLE_BUFFER_L0A_L0B; + staticTilingTmp.dbL0B = DOUBLE_BUFFER_L0A_L0B; + staticTilingTmp.dbL0C = 1; + return staticTilingTmp; + } +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ + constexpr bool isWeightNZ = true; +#else + constexpr bool isWeightNZ = false; +#endif + constexpr static auto staticCFG = GetGmmMatmulApiTiling(isWeightNZ, false); + constexpr static auto staticCFGtransB = GetGmmMatmulApiTiling(isWeightNZ, true); +} // namespace + + +#define GMM_IMP(computeClass, processClass, transA, transB, sync, cfg) \ + do { \ + using matmulType = MMType, weightType, yType, biasType, cfg>; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + REGIST_MATMUL_OBJ(&tPipe, GetSysWorkSpacePtr(), mm, &mmTilingData_); \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, &mmTilingData_, gmmArrayAddr_, groupList, tiling); \ + op.Process(); \ + } while (0) + +#define GMM_CUBE_STATIC_TILING_IMP(processClass, transA, transB, sync, cfg) \ + do { \ + if ASCEND_IS_AIV { \ + return; \ + } \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + using matmulType = MMImplType, weightType, yType, biasType, cfg>; \ + matmulType::MT mm; \ + mm.SetSubBlockIdx(0); \ + mm.Init((TCubeTiling*)nullptr, &tPipe); \ + GMMCompute computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, nullptr, &tPipe); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, nullptr, 0, groupList, tiling); \ + op.InitStaticTiling((cfg).baseM, (cfg).baseN); \ + op.Process(); \ + } while (0) + +#define GMM_CV_SPLIT_STATIC_TILING_IMP(computeClass, processClass, transA, transB, sync, cfg, aType, bType, cType) \ + do { \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + using matmulType = MMImplType, bType, cType, biasType, cfg>; \ + matmulType::MT mm; \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init((TCubeTiling*)nullptr, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, nullptr, &tPipe); \ + computeOp.InitStaticTiling(&gmmBaseParams_, user1, (cfg).baseM, (cfg).baseN); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, nullptr, 0, groupList, tiling); \ + op.InitStaticTiling((cfg).baseM, (cfg).baseN); \ + op.Process(); \ + } while (0) + +#define GMM_CUBE_IMP(processClass, transA, transB, sync, cfg) \ + do { \ + if ASCEND_IS_AIV { \ + return; \ + } \ + using matmulType = MMImplType, weightType, yType, biasType, cfg>; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + GMMCompute computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, &mmTilingData_, gmmArrayAddr_, groupList, tiling); \ + op.Process(); \ + } while (0) + +#if defined(CONST_TILING) +#define GMM_CV_SPLIT_IMP(computeClass, processClass, transA, transB, sync, cfg, aType, bType, cType) \ + do { \ + using matmulType = MMImplType, bType, cType, biasType, cfg>; \ + matmulType::MT mm; \ + GMMTilingData gmmTilingData; \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, &mmTilingData_, gmmArrayAddr_, groupList, tiling); \ + op.Process(); \ + } while (0) +#else + #define GMM_CV_SPLIT_IMP(computeClass, processClass, transA, transB, sync, cfg, aType, bType, cType) \ + do { \ + using matmulType = MMImplType, bType, cType, biasType, cfg>; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + GMMPreTilingProcess preTiling; \ + preTiling.Init(groupList, gmmBaseParams_, mmTilingData_, &tPipe); \ + preTiling.Process(gmmBaseParams_, mmTilingData_); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + processClass op(computeOp); \ + op.Init(&gmmBaseParams_, &mmTilingData_, gmmArrayAddr_, groupList, tiling); \ + op.Process(); \ + } while (0) +#endif + +#define GMM_A4W4_IMP(computeClass, transA, transB, cfg, aType, bType, cType) \ + do { \ + using matmulType = MMImplType, bType, cType, biasType, cfg>; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, scale, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + computeOp.Process(); \ + } while (0) + +#define GMM_CV_SPLIT_IMP_A8W4_MSD(computeClass, cfg) \ + do { \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + if ASCEND_IS_AIV { \ + GMMA8W4PreProcess op1; \ + op1.Init(x, x, groupList, user1, gmmBaseParams_, &tPipe); \ + op1.Process(); \ + tPipe.Reset(); \ + tPipe.Destroy(); \ + tPipe.Init(); \ + } \ + using aT = MatmulType; \ + using bT = MatmulType; \ + using biasT = MatmulType; \ + using cT = MatmulType; \ + using matmulType = MMImplType; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass op(mm); \ + op.Init(x, weight, bias, groupList, scale, perTokenScale, offset, nullptr, nullptr, nullptr, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + op.Process(); \ + } while (0) + +#define GMM_CV_SPLIT_IMP_A8W4(computeClass, cfg) \ + do { \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + if ASCEND_IS_AIV { \ + GMMA8W4FakeQuantPreProcess op1; \ + op1.Init(weight, y, groupList, user1, gmmBaseParams_, &tPipe); \ + op1.Process(); \ + tPipe.Reset(); \ + tPipe.Destroy(); \ + tPipe.Init(); \ + } \ + SyncAll(); \ + using aT = MatmulType; \ + using bT = MatmulType; \ + using biasT = MatmulType; \ + using cT = MatmulType; \ + using matmulType = MMImplType; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass op(mm); \ + op.Init(x, weight, bias, groupList, scale, perTokenScale, offset, nullptr, nullptr, nullptr, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + op.Process(); \ + } while (0) + +#define GMM_CV_SPLIT_IMP_A8W4_FAKEA8W8(computeClass, cfg) \ + do { \ + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling); \ + if ASCEND_IS_AIV { \ + GMMA8W4FakeQuantPreProcess op1; \ + op1.Init(weight, y, scale, user1, gmmBaseParams_, &tPipe); \ + op1.Process(); \ + tPipe.Reset(); \ + tPipe.Destroy(); \ + tPipe.Init(); \ + } \ + SyncAll(); \ + GlobalTensor yGm; \ + yGm.SetGlobalBuffer((__gm__ int8_t *)workspace); \ + using aT = MatmulType; \ + using bT = MatmulType; \ + using biasT = MatmulType; \ + using cT = MatmulType; \ + using matmulType = MMImplType; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + GMMQuantMixCoreCompute computeOp(mm); \ + computeOp.isA8W4FakeQuant = true; \ + computeOp.Init(x, user1, bias, user1, offset, antiquantScale, antiquantOffset, groupList, perTokenScale, \ + y, user1, &gmmBaseParams_, &mmTilingData_, &tPipe); \ + GMMProcess op(computeOp); \ + op.Init(&gmmBaseParams_, &mmTilingData_, gmmArrayAddr_, groupList, tiling); \ + op.Process(); \ + } while (0) + +#define INVOKE_GMM_WEIGHT_QUANT_BASIC_CONTROLLER_OP_IMPL(templateClass, ...) \ + do { \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, gmmWeightQuantParam, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMWeightQuantTilingData, gmmArray, gmmArrayAddr_, tiling); \ + templateClass op; \ + op.Init(x, weight, antiquantScale, antiquantOffset, bias, groupList, y, &gmmBaseParams_, \ + &mmTilingData_, tiling, gmmArrayAddr_, &tPipe); \ + op.Process(); \ + } while (0) + +#define INVOKE_GMM_WEIGHT_QUANT_RESPLIT_CONTROLLER_OP_IMPL(templateClass, ...) \ + do { \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, gmmWeightQuantParam, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + templateClass op; \ + op.Init(x, weight, scale, antiquantScale, antiquantOffset, bias, groupList, perTokenScale, y, &gmmBaseParams_, \ + &mmTilingData_, tiling, &tPipe); \ + op.Process(); \ + } while (0) + +#define INVOKE_GMM_WEIGHT_QUANT_MXA8W4_CONTROLLER_OP_IMPL(templateClass, ...) \ + do { \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, gmmWeightQuantParam, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + templateClass op; \ + op.Init(x, weight, scale, antiquantScale, antiquantOffset, bias, groupList, perTokenScale, y, &gmmBaseParams_, \ + &mmTilingData_, tiling, &tPipe); \ + op.Process(); \ + } while (0) + +#define INVOKE_GMM_WEIGHT_QUANT_VCV_CONTROLLER_OP_IMPL(templateClass, ...) \ + do { \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, gmmWeightQuantParam, gmmBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + templateClass op; \ + op.Init(x, weight, scale, antiquantScale, antiquantOffset, bias, groupList, perTokenScale, y, &gmmBaseParams_, \ + &mmTilingData_, tiling, &tPipe); \ + op.Process(); \ + } while (0) + +#define GMM_QUANT_IMPL_CLASS(transposeX1, transposeX2, templateClass) \ + do { \ + templateClass \ + op; \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, gmmQuantParams, gmmQuantParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMQuantTilingData, gmmArray, gmmArrayAddr_, tiling); \ + op.Init(x, weight, bias, scale, groupList, perTokenScale, y, user1, &gmmQuantParams_, &mmTilingData_, \ + gmmArrayAddr_, &tPipe); \ + op.Process(); \ + } while (0) + +#define GMM_QUANT_MIX_IMPL_CLASS(transposeX1, transposeX2, templateClass) \ + do { \ + templateClass \ + op; \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, gmmQuantParams, gmmQuantParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMQuantTilingData, gmmArray, gmmArrayAddr_, tiling); \ + op.Init(x, weight, bias, scale, groupList, perTokenScale, y, user1, &gmmQuantParams_, &mmTilingData_, \ + gmmArrayAddr_, &tPipe); \ + op.Process(); \ + } while (0) + +#define GMM_QUANT_WITH_EMPTY_TENSOR_IMPL_CLASS(transposeX1, transposeX2, templateClass) \ + do { \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, gmmQuantParams, gmmQuantParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER_ADDR(GMMQuantTilingData, gmmArray, gmmArrayAddr_, tiling); \ + if ASCEND_IS_AIC { \ + templateClass \ + op; \ + op.Init(x, weight, bias, scale, groupList, perTokenScale, y, user1, &gmmQuantParams_, &mmTilingData_, \ + gmmArrayAddr_, &tPipe); \ + op.Process(); \ + } \ + if ASCEND_IS_AIV { \ + GQmmEmptyTensor(groupList, y, &gmmQuantParams_, gmmArrayAddr_, mmTilingData_.usedCoreNum, \ + &tPipe); \ + } \ + } while (0) + +#define GMM_QUANT_GB_IMPL_CLASS(xLayout, wLayout, yLayout) \ + do { \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, gmmQuantParams, gmmQuantParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GmmActPerTileKernel(x, weight, bias, scale, groupList, perTokenScale, y, user1, \ + &gmmQuantParams_, &mmTilingData_, &tPipe); \ + } while (0) + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#if defined(V310_GMM_QUANT) +template +#elif defined(V310_GMM_ANTI_QUANT) +template +#else +template +#endif +__global__ __aicore__ void dlinfer_grouped_matmul_direct(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR offset, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, + GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, GM_ADDR tiling) +#else +template +__global__ __aicore__ void dlinfer_grouped_matmul_direct(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR offset, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, + GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, GM_ADDR tiling) +#endif +{ + TPipe tPipe; + AscendCUtils::SetOverflow(1); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIC_ONLY); + GM_ADDR user1 = GetUserWorkspace(workspace); + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef __CCE_KT_TEST__ +#if defined(V310_GMM_QUANT) // Quant: A8W8 +REGISTER_TILING_DEFAULT(GMMQuantTilingData); +#if defined(V310_GMM_QUANT_MX) // mxfpx + if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_FIXP) { + GMM_QUANT_IMPL_CLASS(false, false, GmmASWKernel); + } else if constexpr (QUANT_B_TRANS == GMM_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_FIXP) { + GMM_QUANT_IMPL_CLASS(false, true, GmmASWKernel); + } +#endif +#if defined(V310_GMM_QUANT_CUBE) || defined(V310_GMM_QUANT_PERTENSOR_CUBE) // scale64/perTensor/double perTensor + if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_FIXP) { + GMM_QUANT_IMPL_CLASS(false, false, GmmASWKernel); + } else if constexpr (QUANT_B_TRANS == GMM_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_FIXP) { + GMM_QUANT_IMPL_CLASS(false, true, GmmASWKernel); + } +#endif +#if defined(V310_GMM_QUANT_MX) || defined(V310_GMM_QUANT_PERTENSOR_CUBE) // mx/perTensor/double perTensor + if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_TRANS + && KERNEL_TYPE == GMM_DEQUANT_FIXP) { + GMM_QUANT_WITH_EMPTY_TENSOR_IMPL_CLASS(true, false, GmmASWKernel); + } +#endif +#if defined(V310_GMM_QUANT_MIX) // perToken/SPLIT_K/scale bf16/fp32 + if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_VECTOR) { + GMM_QUANT_MIX_IMPL_CLASS(false, false, GQmmMixRegbaseKernel); + } else if constexpr (QUANT_B_TRANS == GMM_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_DEQUANT_VECTOR) { + GMM_QUANT_MIX_IMPL_CLASS(false, true, GQmmMixRegbaseKernel); + } else if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_TRANS + && KERNEL_TYPE == GMM_DEQUANT_VECTOR) { + GMM_QUANT_MIX_IMPL_CLASS(true, false, GQmmMixRegbaseKernel); + } +#endif +#if defined(V310_GMM_QUANT_PERTILE) + if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_PERGROUP_PERBLOCK) { + GMM_QUANT_GB_IMPL_CLASS(Act::Gemm::layout::RowMajor, Act::Gemm::layout::RowMajor, + Act::Gemm::layout::RowMajorAlign); + } else if constexpr (QUANT_B_TRANS == GMM_TRANS && QUANT_A_TRANS == GMM_NO_TRANS + && KERNEL_TYPE == GMM_PERGROUP_PERBLOCK) { + GMM_QUANT_GB_IMPL_CLASS(Act::Gemm::layout::RowMajor, Act::Gemm::layout::ColumnMajor, + Act::Gemm::layout::RowMajorAlign); + } else if constexpr (QUANT_B_TRANS == GMM_NO_TRANS && QUANT_A_TRANS == GMM_TRANS + && KERNEL_TYPE == GMM_PERGROUP_PERBLOCK) { + GMM_QUANT_GB_IMPL_CLASS(Act::Gemm::layout::ColumnMajor, Act::Gemm::layout::RowMajor, + Act::Gemm::layout::RowMajorAlign); + } +#endif +#elif defined(V310_GMM_ANTI_QUANT) + REGISTER_TILING_DEFAULT(GMMWeightQuantTilingData); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + #if ORIG_DTYPE_X == DT_INT8 + if constexpr (W_TYPE == WQGMM_FRACTAL_NZ && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_NO_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_512_BUF_NUM_DEFAULT + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + INVOKE_GMM_WEIGHT_QUANT_VCV_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, S8S4_NZKN_G, + VEC_ANTIQUANT_CONFIG_4); + } else if constexpr (W_TYPE == WQGMM_FRACTAL_NZ && + OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST && + C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && + WQ_B_TRANS == WQGMM_NO_TRANS && WQ_A_TRANS == WQGMM_NO_TRANS && + TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_384_BUF_NUM_3 && + ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && + ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + INVOKE_GMM_WEIGHT_QUANT_VCV_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, S8S4_NZKN_G, + VEC_ANTIQUANT_CONFIG_5); + } + #elif ORIG_DTYPE_X == DT_FLOAT8_E4M3FN + if constexpr (W_TYPE == WQGMM_FRACTAL_NZ && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_MX && WQ_B_TRANS == WQGMM_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + INVOKE_GMM_WEIGHT_QUANT_MXA8W4_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, MXA8W4_NZNK, + VEC_ANTIQUANT_CONFIG_3); + } + #elif ORIG_DTYPE_ANTIQUANT_SCALE == DT_FLOAT8_E8M0 + if constexpr (W_TYPE == WQGMM_FRACTAL_NZ && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_MX && WQ_B_TRANS == WQGMM_NO_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + INVOKE_GMM_WEIGHT_QUANT_RESPLIT_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, A16MXF4_NZKN, + VEC_ANTIQUANT_CONFIG_3); + } + #else + if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, true, QuantType::PER_CHANNEL, false, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_RESPLIT_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } else if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_TAIL_RESPLIT && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, true, QuantType::PER_CHANNEL, true, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_RESPLIT_CONTROLLER_OP_IMPL(GMMWeightQuantResplitController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } else if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_BASIC_BLOCK && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, true, QuantType::PER_CHANNEL, false, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_BASIC_CONTROLLER_OP_IMPL(GMMWeightQuantBasicController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } else if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_BASIC_BLOCK && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, true, QuantType::PER_CHANNEL, true, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_BASIC_CONTROLLER_OP_IMPL(GMMWeightQuantBasicController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } else if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_NOT_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_NO_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_BASIC_BLOCK && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, false, QuantType::PER_CHANNEL, false, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_BASIC_CONTROLLER_OP_IMPL(GMMWeightQuantBasicController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } else if constexpr (W_TYPE == WQGMM_ND && OFFSET_OR_BIAS_EXIT == WQGMM_ANTIQUANT_OFFSET_EXIST_BIAS_NOT_EXIST + && C_QUANT_TYPE == WQGMM_NONE && W_QUANT_TYPE == WQGMM_PER_CHANNEL && WQ_B_TRANS == WQGMM_NO_TRANS + && WQ_A_TRANS == WQGMM_NO_TRANS && TEMPLATE_CUSTOM_SC == WQGMM_MTE2_INNER_SIZE_256_BUF_NUM_4 + && ALGORITHM_SUB_CATEGORY == WQGMM_N_FIRST_BASIC_BLOCK && ALGORITHM_CATEGORY == WQGMM_VECTOR_ANTIQUANT) { + static constexpr WqmmConfig wqmmCfg = {false, false, QuantType::PER_CHANNEL, true, + QuantType::NONE, CubeFormat::ND}; + INVOKE_GMM_WEIGHT_QUANT_BASIC_CONTROLLER_OP_IMPL(GMMWeightQuantBasicController, wqmmCfg, + VEC_ANTIQUANT_CONFIG_3); + } + #endif +#else + REGISTER_TILING_DEFAULT(GMMNoQuantTilingData); + if constexpr (NO_QUANT_B_TRANS == GMM_NO_TRANS && NO_QUANT_A_TRANS == GMM_NO_TRANS) { + if constexpr (wFormat == CubeFormat::NZ) { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } else { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } + } else if constexpr (NO_QUANT_B_TRANS == GMM_NO_TRANS && NO_QUANT_A_TRANS == GMM_TRANS) { // x transposed + if ASCEND_IS_AIV { + EmptyTensor(groupList, y, tiling); + } + if ASCEND_IS_AIC { + if constexpr (wFormat == CubeFormat::NZ) { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } else { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } + } + } else if constexpr (NO_QUANT_B_TRANS == GMM_TRANS && NO_QUANT_A_TRANS == GMM_NO_TRANS) { // weight transposed + if constexpr (wFormat == CubeFormat::NZ) { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } else { + GmmNoQuantAswt(x, weight, bias, groupList, y, tiling); + } + } +#endif +#endif +#endif + +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) +#if defined(GMM_ANTI_QUANT_A8W4_MSD) + // ANTIQUANT_A8W4 + if constexpr (D_T_A == GMM_TPL_INT8 && D_T_B == GMM_TPL_INT4) { + if constexpr (A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_API_DEQUANT) { + GMM_CV_SPLIT_IMP_A8W4_MSD(GMMA8W4MSDCompute, A8W4_GMM_CFG_MDL); + } else if constexpr (A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_VECTOR_DEQUANT) { + GMM_CV_SPLIT_IMP_A8W4_MSD(GMMA8W4MSDComputeNew, A8W4_GMM_CFG_MDL_NEW); + } else if constexpr (A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERCHANNEL_ANTIQUANT) { + GMM_CV_SPLIT_IMP_A8W4_FAKEA8W8(GMMA8W4Compute, A8W4_GMM_CFG_MDL); + } else if constexpr (A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERGROUP_ANTIQUANT) { + GMM_CV_SPLIT_IMP_A8W4(GMMA8W4Compute, A8W4_GMM_CFG_MDL); + } else if constexpr (A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_AUTOTILING) { + GET_TILING_DATA_MEMBER(GMMTilingData, hpTilingData, tilingData, tiling); + GM_ADDR A = x; + GM_ADDR B = weight; + GM_ADDR C = y; + GM_ADDR groupListOptional = groupList; + GM_ADDR bias_ = bias; + GM_ADDR offset_ = offset; + GM_ADDR sa = perTokenScale; + GM_ADDR sw = scale; + GM_ADDR workspaceDevice = user1; + + GMMA4W8AutotilingCompute op(A, B, C, groupListOptional, bias_, offset_, sa, sw, workspaceDevice, + const_cast(&tilingData), &tPipe); + op.Init(); + op.Process(); + } + } +#elif defined(GMM_ANTI_QUANT) + // ANTIQUANT + if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) && + A16W8_KERNEL_TEMPLATE != GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) { + // ANTIQUANT_A16W4 & ANTIQUANT_A16W8_NOT_MSD + if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) { + GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, false, false, matmulCFG); + } else if constexpr (TRANS_B == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) { + GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, true, false, matmulCFG); + } else if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_2) { + GMM_IMP(GMMAntiquantComputePerformance, GMMAntiquantProcess, false, false, false, matmulCFG); + } + } +#if defined(ORIG_DTYPE_WEIGHT) && defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8 + // ANTIQUANT_A16W8_MSD + if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) && D_T_B == GMM_TPL_INT8 && + A16W8_KERNEL_TEMPLATE == GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) { + if constexpr (TRANS_B == 0) { + GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, false, false, matmulCFG, + xTypeMSD, weightTypeMSD, yTypeMSD); + } else if constexpr (TRANS_B == 1) { + GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, true, false, matmulCFG, + xTypeMSD, weightTypeMSD, yTypeMSD); + } + } +#endif + +#elif defined(GMM_QUANT_BF16) || defined(GMM_QUANT_FLOAT16) + // QUANT_A8W8O16 + if constexpr (D_T_A == GMM_TPL_INT8 && D_T_B == GMM_TPL_INT8 && (D_T_Y == GMM_TPL_BF16 || D_T_Y == GMM_TPL_FLOAT16) && + TRANS_A == 0 && A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE) { + if constexpr (IS_STATIC_TILING_API == 0) { + if constexpr (AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) { + if constexpr(IS_ENABLE_FIXED_AXIS == 0) { + if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMProcess, false, false, false, matmulCFG, xType, weightType, yType); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMProcess, false, true, false, matmulCFG, xType, weightType, yType); + } else if constexpr(TRANS_B == 0 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMGroupMSparseProcess, false, false, false, matmulCFG, xType, + weightType, yType); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMGroupMSparseProcess, false, true, false, matmulCFG, xType, + weightType, yType); + } + } else if constexpr(IS_ENABLE_FIXED_AXIS == 1 && TRANS_B == 0 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM) { + tPipe.Destroy(); + AscendC::SetMMLayoutTransform(true); + GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling) + using XDType = int8_t; + using WeightDType = int8_t; + using CDType = int32_t; + using ScaleDType = float; + using GrouplistDType = int64_t; + using PerTokenScaleDType = float; + using YDType = half; +#ifndef __CCE_KT_TEST__ + Catlass::grouped_matmul_fixaxismove( + gmmBaseParams_.m, gmmBaseParams_.k, gmmBaseParams_.n, gmmBaseParams_.groupNum, + x, weight, scale, groupList, perTokenScale, y, user1, gmmBaseParams_.coreNum); +#endif + } + } else if constexpr (AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_2) { + if constexpr (TRANS_B == 0) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMProcess, false, false, false, matmulCFG, xType, weightType, yType); + } else if constexpr (TRANS_B == 1) { + GMM_CV_SPLIT_IMP(GMMQuantMixCoreCompute, GMMProcess, false, true, false, matmulCFG, xType, weightType, yType); + } + } + } else if (IS_STATIC_TILING_API == 1) { + if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_STATIC_TILING_IMP(GMMQuantMixCoreCompute, GMMProcess, + false, false, false, staticCFG, xType, weightType, yType); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_STATIC_TILING_IMP(GMMQuantMixCoreCompute, GMMProcess, + false, true, false, staticCFGtransB, xType, weightType, yType); + } else if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_STATIC_TILING_IMP(GMMQuantMixCoreCompute, GMMGroupMSparseProcess, + false, false, false, staticCFG, xType, weightType, yType); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CV_SPLIT_STATIC_TILING_IMP(GMMQuantMixCoreCompute, GMMGroupMSparseProcess, + false, true, false, staticCFGtransB, xType, weightType, yType); + } + } + } +#elif defined(GMM_A4W4) + // QUANT_A4W4 + if constexpr (D_T_A == GMM_TPL_INT4 && D_T_B == GMM_TPL_INT4) { + GMM_A4W4_IMP(GMMA4W4Compute, false, false, matmulCFG, xType, weightType, yType); + } +#elif defined(GMM_QUANT_INT8) || defined(GMM_QUANT_INT32) + // QUANT_A8W8O8 & QUANT_A8W8O32 + if constexpr (D_T_A == GMM_TPL_INT8 && D_T_B == GMM_TPL_INT8 && (D_T_Y == GMM_TPL_INT8 || D_T_Y == GMM_TPL_INT32) && + TRANS_A == 0 && A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE && + AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) { + if constexpr (IS_STATIC_TILING_API == 0) { + if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_IMP(GMMGroupMSparseProcess, false, false, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_IMP(GMMGroupMSparseProcess, false, true, false, matmulCFGUnitFlag); + } + } else if constexpr (IS_STATIC_TILING_API == 1){ + if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_STATIC_TILING_IMP(GMMProcess, false, false, false, staticCFG); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_STATIC_TILING_IMP(GMMProcess, false, true, false, staticCFGtransB); + } else if constexpr (TRANS_B == 0 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_STATIC_TILING_IMP(GMMGroupMSparseProcess, false, false, false, staticCFG); + } else if constexpr (TRANS_B == 1 && GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) { + GMM_CUBE_STATIC_TILING_IMP(GMMGroupMSparseProcess, false, true, false, staticCFGtransB); + } + } + } +#elif defined(GMM_FLOAT) + // NO_QUANT + if (IS_STATIC_TILING_API == 0 && + A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE) { + if constexpr (TRANS_A == 0 && TRANS_B == 0 && + GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && + AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) { + GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_A == 0 && TRANS_B == 1 && + GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && + AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) { + GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_A == 0 && TRANS_B == 0 && + GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && + AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) { + GMM_CUBE_IMP(GMMGroupMSparseProcess, false, false, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_A == 0 && TRANS_B == 1 && + GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && + AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) { + GMM_CUBE_IMP(GMMGroupMSparseProcess, false, true, false, matmulCFGUnitFlag); + } else if constexpr (TRANS_A == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) { + if ASCEND_IS_AIV { + GET_TILING_DATA(tilingData, tiling); + EmptyTensorCompute(groupList, y, &tilingData); + } + if ASCEND_IS_AIC { + GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG); + } + } + } + +#endif +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 +#if defined(GMM_FLOAT) + if constexpr (TRANS_A == 0 && TRANS_B == 0) { + GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFG); + } else if constexpr (TRANS_A == 0 && TRANS_B == 1) { + GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFG); + } else if constexpr (TRANS_A == 1 && TRANS_B == 0) { + if ASCEND_IS_AIV { + GET_TILING_DATA(tilingData, tiling); + EmptyTensorCompute(groupList, y, &tilingData); + } + if ASCEND_IS_AIC { + GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG); + } + } +#endif +#endif +} diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/arch.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/arch.hpp new file mode 100644 index 00000000..0f9462fc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/arch.hpp @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_ARCH_ARCH_HPP +#define CATLASS_ARCH_ARCH_HPP + +#include "../../gmm_infra/base_defs.hpp" + +namespace Catlass::Arch { + +struct AtlasA2 { + static constexpr uint32_t BIAS_SIZE = 1024; + static constexpr uint32_t FIXBUF_SIZE = 7 * 1024; + static constexpr uint32_t UB_SIZE = 192 * 1024; + static constexpr uint32_t L1_SIZE = 512 * 1024; + static constexpr uint32_t L0A_SIZE = 64 * 1024; + static constexpr uint32_t L0B_SIZE = 64 * 1024; + static constexpr uint32_t L0C_SIZE = 128 * 1024; +}; + +struct PositionGM { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::GM; +}; + +struct PositionL1 { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::A1; +}; + +struct PositionL0A { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::A2; +}; + +struct PositionL0B { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::B2; +}; + +struct PositionL0C { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::CO1; +}; + +struct PositionUB { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::VECCALC; +}; + +} // namespace Catlass::Arch + +#endif // CATLASS_ARCH_ARCH_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/cross_core_sync.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/cross_core_sync.hpp new file mode 100644 index 00000000..ff9606cc --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/cross_core_sync.hpp @@ -0,0 +1,115 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_ARCH_CROSS_CORE_SYNC_HPP +#define CATLASS_ARCH_CROSS_CORE_SYNC_HPP + +#include "../../gmm_infra/base_defs.hpp" + +namespace Catlass::Arch { + +constexpr uint32_t MAX_REVERSE_DEPTH = 16; + +using FlagID = uint16_t; +constexpr FlagID AIV_INTER_BLOCK_BARRIER = 8; +constexpr FlagID AIC_INTER_BLOCK_BARRIER = 9; +constexpr FlagID AIV_INTER_SUBBLOCK_BARRIER = 10; +constexpr FlagID FFTS_MAX_FLAG = 7; + +struct CrossCoreFlag { + CATLASS_DEVICE + CrossCoreFlag() : id(0) {} + + CATLASS_DEVICE + CrossCoreFlag(FlagID id) : id(id) {} + + FlagID id; +}; + +template +struct CrossCoreFlagWithReverse { + CATLASS_DEVICE + CrossCoreFlagWithReverse() : id(0), reverseId(0) {} + + CATLASS_DEVICE + CrossCoreFlagWithReverse(FlagID id, FlagID reverseId) : id(id), reverseId(reverseId) {} + + FlagID id; + FlagID reverseId; + uint32_t count{ 0 }; +}; + +template +struct BarrierFlag { + static_assert(MODE != MODE, "Unsupported cross core barrier flag, can not find the specialization."); +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIC> { + static constexpr FlagID ID = AIC_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x1, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_SUBBLOCK_BARRIER; +}; + +template +CATLASS_DEVICE +void CrossCoreBarrier() +{ + constexpr FlagID flagId = BarrierFlag::ID; + AscendC::CrossCoreSetFlag(flagId); + AscendC::CrossCoreWaitFlag(flagId); +} + +template +CATLASS_DEVICE +void CrossCoreSetFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); +} + +CATLASS_DEVICE +void CrossCoreWaitFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); +} + +template +CATLASS_DEVICE +void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreWaitFlag(flag.reverseId); + flag.count = 0; + } +} + +template +CATLASS_DEVICE +void CrossCoreWaitFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreSetFlag(flag.reverseId); + flag.count = 0; + } +} + +} // namespace Catlass::Arch + +#endif // CATLASS_ARCH_CROSS_CORE_SYNC_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/local_tensor_buffer.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/local_tensor_buffer.hpp new file mode 100644 index 00000000..391ab78b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/local_tensor_buffer.hpp @@ -0,0 +1,233 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef INCLUDE_CATLASS_ARCH_MEMORY_H +#define INCLUDE_CATLASS_ARCH_MEMORY_H + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/arch/arch.hpp" + +namespace Catlass::Arch { + +struct LocalTensorBufferBase { +public: + template + CATLASS_DEVICE + AscendC::LocalTensor GetBufferByByte(const uint32_t offset) const + { + return tensor[offset].template ReinterpretCast(); + } + +protected: + CATLASS_DEVICE + LocalTensorBufferBase() = default; + + AscendC::LocalTensor tensor; +}; + +template < + class ArchTag, + AscendC::TPosition Position +> +struct LocalTensorBuffer { + static_assert(DEPENDENT_FALSE, "Unsupported local tensor buffer, can not find the specialization."); +}; + +/// Partial specialization for TPosition::A1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA1; + GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE); + tensor = tbufA1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::A2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA2; + GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE); + tensor = tbufA2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB1; + GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE); + tensor = tbufB1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::B2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB2; + GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE); + tensor = tbufB2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C1 +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC1; + GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE); + tensor = tbufC1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C2 +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2; + GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE); + tensor = tbufC2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::CO1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufCO1; + GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE); + tensor = tbufCO1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C2PIPE2GM +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2PIPE2GM; + GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE); + tensor = tbufC2PIPE2GM.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECIN +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECIN; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECIN; + GetTPipePtr()->InitBuffer(tbufVECIN, ArchTag::UB_SIZE); + tensor = tbufVECIN.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECOUT +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECOUT; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECOUT; + GetTPipePtr()->InitBuffer(tbufVECOUT, ArchTag::UB_SIZE); + tensor = tbufVECOUT.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECCALC +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECCALC; + GetTPipePtr()->InitBuffer(tbufVECCALC, ArchTag::UB_SIZE); + tensor = tbufVECCALC.Get(); + } +}; + +} // namespace Catlass::Arch + +#endif // INCLUDE_CATLASS_ARCH_MEMORY_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/resource.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/resource.hpp new file mode 100644 index 00000000..3a9f4304 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/arch/resource.hpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef INCLUDE_CATLASS_ARCH_RESOURCE_HPP +#define INCLUDE_CATLASS_ARCH_RESOURCE_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/arch/local_tensor_buffer.hpp" + +namespace Catlass::Arch { + +template +struct Resource { +public: + AscendC::TPipe pipe; + + LocalTensorBuffer l1Buf; + LocalTensorBuffer l0ABuf; + LocalTensorBuffer l0BBuf; + LocalTensorBuffer btBuf; + LocalTensorBuffer l0CBuf; + LocalTensorBuffer ubBuf; + + CATLASS_DEVICE + Resource() + { + // The initialization of AscendC::Tpipe will insert some synchronization interfaces, + // which may conflict with the usage by users. Therefore, the "destroy" interface is used for releasing. + pipe.Destroy(); + } +}; + +} // namespace Catlass::Arch + +#endif // INCLUDE_CATLASS_ARCH_RESOURCE_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/base_defs.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/base_defs.hpp new file mode 100644 index 00000000..02890ab5 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/base_defs.hpp @@ -0,0 +1,36 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_CATLASS_HPP +#define CATLASS_CATLASS_HPP + +#include + +#include "../gmm_infra/detail/alignment.hpp" +#include "../gmm_infra/detail/dependent_false.hpp" +#include "../gmm_infra/detail/macros.hpp" + +namespace Catlass { + +constexpr uint32_t BYTE_PER_C0 = 32; +constexpr uint32_t BYTE_PER_C2 = 64; +constexpr uint32_t C0_NUM_PER_FRACTAL = 16; +constexpr uint32_t BYTE_PER_FRACTAL = BYTE_PER_C0 * C0_NUM_PER_FRACTAL; + +constexpr uint32_t BYTE_PER_BLK = 32; +constexpr uint32_t BLK_NUM_PER_VECTOR_FRACTAL = 8; +constexpr uint32_t BYTE_PER_VECTOR_FRACTAL = BYTE_PER_BLK * BLK_NUM_PER_VECTOR_FRACTAL; + +constexpr uint64_t L2_OFFSET = 0; +constexpr uint32_t STRIDE_LIMIT = 65536; + +} // namespace Catlass + +#endif // CATLASS_CATLASS_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/coord.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/coord.hpp new file mode 100644 index 00000000..45548c52 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/coord.hpp @@ -0,0 +1,319 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_COORD_HPP +#define CATLASS_COORD_HPP + +#include "../gmm_infra/base_defs.hpp" + +namespace Catlass { + +/// Statically-sized array specifying Coords within a tensor +template < + int RANK_, ///< Logical rank of coordinate + class Index_ = uint32_t, ///< Index type used for each dimension + class LongIndex_ = int64_t ///< Long index type used for linear offsets +> +struct Coord { +public: + // Number of elements in Coord + static const int RANK = RANK_; + + // Index typen used to store elements + using Index = Index_; + + // Type used to represent linear offsets + using LongIndex = LongIndex_; + + // Default ctor initializes uniformly + CATLASS_HOST_DEVICE constexpr + explicit Coord(Index value = Index(0)) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = value; + } + } + + // Constructs from an array of integers + CATLASS_HOST_DEVICE constexpr + Coord(Index const (&idx_)[RANK]) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = idx_[i]; + } + } + + // Constructs frrom an array of integers + CATLASS_HOST_DEVICE + int Argmin() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] < idx[i]) { + i = j; + } + } + return i; + } + + // Returns the index of the dimension with greatest value + CATLASS_HOST_DEVICE + int Argmax() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] > idx[i]) { + i = j; + } + } + return i; + } + + // Returns true if Coord is non-zero + CATLASS_HOST_DEVICE + explicit operator bool() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return true; + } + } + return false; + } + + // Return true if Coord is uniformly zero. + CATLASS_HOST_DEVICE + bool operator!() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return false; + } + } + return true; + } + + // Element-wise addition + CATLASS_HOST_DEVICE + Coord operator+(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + b.idx[i]; + } + return c; + } + + // Add a scalar to each element + CATLASS_HOST_DEVICE + Coord operator+(const Index val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + val; + } + return c; + } + + // Element-wise subtraction + CATLASS_HOST_DEVICE + Coord operator-(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] - b.idx[i]; + } + return c; + } + + // Subtract a scalar from each element + CATLASS_HOST_DEVICE + Coord operator-(Index const val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] - val; + } + return c; + } + + // Element-wise multiply + CATLASS_HOST_DEVICE + Coord operator*(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] * b.idx[i]; + } + return c; + } + + // Element-wise division + CATLASS_HOST_DEVICE + Coord operator/(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] / b.idx[i]; + } + return c; + } + + // Element-wise mod + CATLASS_HOST_DEVICE + Coord operator%(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] % b.idx[i]; + } + return c; + } + + // In-place addition + CATLASS_HOST_DEVICE + Coord &operator+=(Coord const &b) + { + for (int i = 0; i < RANK; ++i) { + idx[i] += b.idx[i]; + } + return *this; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Coord const &b) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != b.idx[i]) { + return false; + } + } + return true; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Index const val) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != val) { + return false; + } + } + return true; + } + + // Member acces operator + CATLASS_HOST_DEVICE + Index &operator[](int dim) + { + return idx[dim]; + } + + // Member access operator + CATLASS_HOST_DEVICE + Index const &operator[](int dim) const + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE + Index &At() + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index &At(int dim) + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE + Index const &At() const + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index const &At(int dim) const + { + return idx[dim]; + } + + template + CATLASS_HOST_DEVICE + auto GetCoordByAxis() const + { + Index idx_[sizeof...(Is)]{idx[Is]...}; + return Coord{idx_}; + } + + CATLASS_HOST_DEVICE + static Coord Min(Coord const &a, Coord const &b) + { + Coord res; + for (int i = 0; i < RANK; ++i) { + res[i] = a[i] < b[i] ? a[i] : b[i]; + } + return res; + } + +private: + // Indices + Index idx[RANK]; +}; + +// Helper to make a 1-element coordinate +template +CATLASS_HOST_DEVICE constexpr +Coord<1, T> MakeCoord(T dim0) +{ + T values[1] = {dim0}; + return Coord<1, T>(values); +} + +/// Helper to make a 2-element coordinate +template +CATLASS_HOST_DEVICE constexpr +Coord<2, T> MakeCoord(T dim0, T dim1) +{ + T values[2] = {dim0, dim1}; + return Coord<2, T>(values); +} + +/// Helper to make a 3-element coordinate +template +CATLASS_HOST_DEVICE constexpr +Coord<3, T> MakeCoord(T dim0, T dim1, T dim2) +{ + T values[3] = {dim0, dim1, dim2}; + return Coord<3, T>(values); +} + +/// Helper to make a 4-element coordinate +template +CATLASS_HOST_DEVICE constexpr +Coord<4, T> MakeCoord(T dim0, T dim1, T dim2, T dim3) +{ + T values[4] = {dim0, dim1, dim2, dim3}; + return Coord<4, T>(values); +} + +} // namespace Catlass + +#endif // CATLASS_COORD_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/alignment.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/alignment.hpp new file mode 100644 index 00000000..132427b3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/alignment.hpp @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_ALIGNMENT_HPP +#define CATLASS_ALIGNMENT_HPP + +#include "../../gmm_infra/detail/macros.hpp" + +template +CATLASS_HOST_DEVICE +constexpr T RoundUp(const T &val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return (val + ALIGN - 1) / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE +constexpr T RoundUp(const T &val, const T align) +{ + return (val + align - 1) / align * align; +} + +template +CATLASS_HOST_DEVICE +constexpr T RoundDown(const T val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return val / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE +constexpr T RoundDown(const T val, const T align) +{ + return val / align * align; +} + +template +CATLASS_HOST_DEVICE +constexpr T CeilDiv(const T dividend) +{ + static_assert(DIVISOP != 0, "DIVISOP must not be 0"); + return (dividend + DIVISOP - 1) / DIVISOP; +} + +template +CATLASS_HOST_DEVICE +constexpr T CeilDiv(const T dividend, const T divisor) +{ + return (dividend + divisor - 1) / divisor; +} + +#endif // CATLASS_ALIGNMENT_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/callback.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/callback.hpp new file mode 100644 index 00000000..46e2cdc0 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/callback.hpp @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_DETAIL_CALLBACK_HPP +#define CATLASS_DETAIL_CALLBACK_HPP + +#include "../../gmm_infra/detail/macros.hpp" + +/// @brief Callback is an alternative to std::function, providing a general carrier +/// of callable structure with no parameters and no return value. Compared with function pointers +/// of type void (*)(), Callback can carry lambda expressions with captures, and does not need to +/// pay attention to the captured content. It should be noted that Callback itself does not store +/// the callable structure it carries like std::function, so it is necessary to ensure +/// that it is used within the life cycle of the callable structure. +struct Callback { + void const *func{nullptr}; + void (*caller)(void const *){nullptr}; + + Callback() = default; + + CATLASS_DEVICE + void operator()() const + { + if (func) { + caller(func); + } + } + + CATLASS_DEVICE + operator bool() const + { + return func != nullptr; + } +}; + +template +CATLASS_DEVICE +static void FuncWrapper(void const *func) +{ + (*static_cast(func))(); +} + +// Use this to make a callback +template +CATLASS_DEVICE +Callback MakeCallback(Func *func) +{ + Callback callback; + callback.func = func; + callback.caller = &FuncWrapper; + return callback; +} + +#endif // CATLASS_DETAIL_CALLBACK_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/dependent_false.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/dependent_false.hpp new file mode 100644 index 00000000..78543a8c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/dependent_false.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_DETAIL_DEPENDENT_FALSE_HPP +#define CATLASS_DETAIL_DEPENDENT_FALSE_HPP + +template +constexpr bool DEPENDENT_BOOL_VALUE = VALUE; + +template +constexpr bool DEPENDENT_FALSE = DEPENDENT_BOOL_VALUE; + +#endif // CATLASS_DETAIL_DEPENDENT_FALSE_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/macros.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/macros.hpp new file mode 100644 index 00000000..de0835e7 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/detail/macros.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_DETAIL_MACROS_HPP +#define CATLASS_DETAIL_MACROS_HPP + +#define CATLASS_DEVICE __forceinline__ [aicore] +#define CATLASS_HOST_DEVICE __forceinline__ [host, aicore] +#define CATLASS_GLOBAL __global__ [aicore] + +#endif // CATLASS_DETAIL_MACROS_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue.hpp new file mode 100644 index 00000000..74a8d350 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue.hpp @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP + +#include "../../../gmm_infra/base_defs.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class DispatchPolicy, + class... Args +> +class BlockEpilogue { + static_assert(DEPENDENT_FALSE, "Could not find an epilogue specialization"); +}; + +} // namespace Catlass::Epilogue::Block + +#include "../../../gmm_infra/epilogue/block/block_epilogue_per_token_dequant.hpp" +#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue_per_token_dequant.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue_per_token_dequant.hpp new file mode 100644 index 00000000..93645af0 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/block/block_epilogue_per_token_dequant.hpp @@ -0,0 +1,645 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/resource.hpp" +#include "../../../gmm_infra/epilogue/dispatch_policy.hpp" +#include "../../../gmm_infra/gemm_coord.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" +#include "../../../gmm_infra/matrix_coord.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/detail/callback.hpp" + +namespace Catlass::Epilogue::Block { + +template < + uint32_t UB_STAGES_, + class CType_, + class ScaleType_, + class PerTokenScaleType_, + class DType_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_ +> +class BlockEpilogue < + EpilogueAtlasA2PerTokenDequant, + CType_, + ScaleType_, + PerTokenScaleType_, + DType_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_ +> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementScale = typename ScaleType_::Element; + using LayoutScale = typename ScaleType_::Layout; + using ElementPerTokenScale = typename PerTokenScaleType_::Element; + using LayoutPerTokenScale = typename PerTokenScaleType_::Layout; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v) && + std::is_same_v && std::is_same_v, + "The element type template parameters of BlockEpilogue are wrong" + ); + static_assert( + std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong" + ); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops" + ); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + + (TileShape::COUNT + TileShape::COLUMN + TileShape::COUNT + TileShape::ROW) * sizeof(float) + + TileShape::ROW * BYTE_PER_BLK) + <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB" + ); + + struct Params { + __gm__ ElementScale *ptrScale{nullptr}; + LayoutScale layoutScale{}; + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params( + __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, + __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, LayoutD const &layoutD_ + ) : ptrScale(ptrScale_), layoutScale(layoutScale_), + ptrPerTokenScale(ptrPerTokenScale_), layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), layoutD(layoutD_) {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) : params(params) + { + size_t ubOffset = 0; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(float); + ubPerTokenScaleFp32Brcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubMul; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator() ( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC, Callback &&callback = Callback{} + ) + { + if (actualBlockShapeMNK.k() == 0) { + return; + } + callback(); + + // Calculate the offset of the current block + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + AscendC::GlobalTensor gmScale; + gmScale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gmPerTokenScale; + gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffsetInBlock = tileCoord * tileShape; + auto tileOffset = blockOffset + tileOffsetInBlock; + + auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; + auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); + + auto &ubC = ubCList[ubListId]; + LayoutC layoutUbC{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); + auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); + + auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; + auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); + + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); + + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); + auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); + + auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; + auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); + + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( + perTokenScaleTileShape); + + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, + layoutGmTilePerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + AscendC::Cast(ubScaleFp32, ubScale, AscendC::RoundMode::CAST_NONE, TileShape::COLUMN); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + AscendC::Cast(ubPerTokenScaleFp32, ubPerTokenScale, AscendC::RoundMode::CAST_NONE, TileShape::ROW); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + tileRowBroadcastMul(ubMul, ubCFp32, ubScaleFp32); + tileBroadcastOneBlk(ubPerTokenScaleFp32Brcb, ubPerTokenScaleFp32); + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleFp32Brcb); + AscendC::PipeBarrier(); + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubScaleFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleFp32; + AscendC::LocalTensor ubPerTokenScaleFp32Brcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; +}; + +template < + uint32_t UB_STAGES_, + class CType_, + class LayoutScale_, + class LayoutPerTokenScale_, + class DType_, + class TileRowBroadcastMul_, + class TileBroadcastOneBlk_, + class TileOneBlkColumnBroadcastMul_, + class TileCopy_, + class EpilogueTileSwizzle_ +> +class BlockEpilogue < + EpilogueAtlasA2PerTokenDequant, + CType_, + Gemm::GemmType, + Gemm::GemmType, + DType_, + TileRowBroadcastMul_, + TileBroadcastOneBlk_, + TileOneBlkColumnBroadcastMul_, + TileCopy_, + EpilogueTileSwizzle_ +> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementScale = float; + using LayoutScale = LayoutScale_; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = LayoutPerTokenScale_; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v), + "The element type template parameters of BlockEpilogue are wrong" + ); + static_assert( + std::is_same_v && std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong" + ); + + // Tile compute ops + using TileRowBroadcastMul = TileRowBroadcastMul_; + using TileBroadcastOneBlk = TileBroadcastOneBlk_; + using TileOneBlkColumnBroadcastMul = TileOneBlkColumnBroadcastMul_; + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyGmToUbScale = typename TileCopy_::CopyGmToUbX; + using CopyGmToUbPerTokenScale = typename TileCopy_::CopyGmToUbY; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + using EpilogueTileSwizzle = EpilogueTileSwizzle_; + + using TileShape = typename TileRowBroadcastMul::TileShape; + + static_assert( + TileShape::ROW == TileBroadcastOneBlk::COMPUTE_LENGTH && + std::is_same_v, + "TileShape must be consistent for all tile compute ops" + ); + + static_assert( + (UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + TileShape::COLUMN * sizeof(ElementScale) + + TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) + + (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + + TileShape::ROW * BYTE_PER_BLK) + <= ArchTag::UB_SIZE, + "TileShape is too large to fit in UB" + ); + + struct Params { + __gm__ ElementScale *ptrScale{nullptr}; + LayoutScale layoutScale{}; + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params( + __gm__ ElementScale *ptrScale_, LayoutScale const &layoutScale_, + __gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, LayoutD const &layoutD_ + ) : ptrScale(ptrScale_), layoutScale(layoutScale_), + ptrPerTokenScale(ptrPerTokenScale_), layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), layoutD(layoutD_) {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) : params(params) + { + size_t ubOffset = 0; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementC); + ubScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COLUMN * sizeof(ElementScale); + ubPerTokenScaleList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * sizeof(ElementPerTokenScale); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbScaleVMTE2List[i] = eventVMTE2++; + eventUbScaleMTE2VList[i] = eventMTE2V++; + eventUbPerTokenScaleVMTE2List[i] = eventVMTE2++; + eventUbPerTokenScaleMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbScaleVMTE2List[i]); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubCFp32 = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubMul = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::COUNT * sizeof(float); + ubPerTokenScaleBrcb = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += TileShape::ROW * BYTE_PER_BLK; + ubPerTokenMul = ubCFp32; + } + + CATLASS_DEVICE + ~BlockEpilogue() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator() ( + GemmCoord const &blockShapeMNK, + GemmCoord const &blockCoordMNK, + GemmCoord const &actualBlockShapeMNK, + AscendC::GlobalTensor const &gmBlockC, + LayoutC const &layoutBlockC, Callback &&callback = Callback{} + ) + { + if (actualBlockShapeMNK.k() == 0) { + return; + } + uint32_t coreIdx = AscendC::GetBlockIdx(); + callback(); + // Calculate the offset of the current block + MatrixCoord blockShape = blockShapeMNK.GetCoordMN(); + MatrixCoord blockCoord = blockCoordMNK.GetCoordMN(); + MatrixCoord actualBlockShape = actualBlockShapeMNK.GetCoordMN(); + MatrixCoord blockOffset = blockCoord * blockShape; + + AscendC::GlobalTensor gmScale; + gmScale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gmPerTokenScale; + gmPerTokenScale.SetGlobalBuffer(params.ptrPerTokenScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + auto ubTileStride = MakeCoord(static_cast(TileShape::COLUMN), 1L); + auto tileShape = TileShape::ToCoord(); + EpilogueTileSwizzle epilogueTileSwizzle(actualBlockShape, tileShape); + uint32_t tileLoops = epilogueTileSwizzle.GetLoops(); + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t subblockNum = AscendC::GetSubBlockNum(); + for (uint32_t loopIdx = subblockIdx; loopIdx < tileLoops; loopIdx += subblockNum) { + auto tileCoord = epilogueTileSwizzle.GetTileCoord(loopIdx); + auto actualTileShape = epilogueTileSwizzle.GetActualTileShape(tileCoord); + auto tileOffsetInBlock = tileCoord * tileShape; + auto tileOffset = blockOffset + tileOffsetInBlock; + + auto gmTileC = gmBlockC[layoutBlockC.GetOffset(tileOffsetInBlock)]; + auto layoutGmTileC = layoutBlockC.GetTileLayout(actualTileShape); + + auto &ubC = ubCList[ubListId]; + LayoutC layoutUbC{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + auto scaleTileOffset = tileOffset.template GetCoordByAxis<1>(); + auto scaleTileShape = actualTileShape.template GetCoordByAxis<1>(); + + auto gmTileScale = gmScale[params.layoutScale.GetOffset(scaleTileOffset)]; + auto layoutGmTileScale = params.layoutScale.GetTileLayout(scaleTileShape); + + auto &ubScale = ubScaleList[ubListId]; + auto layoutUbScale = LayoutScale::template MakeLayoutInUb(scaleTileShape); + + AscendC::WaitFlag(eventUbScaleVMTE2List[ubListId]); + copyGmToUbScale(ubScale, gmTileScale, layoutUbScale, layoutGmTileScale); + AscendC::SetFlag(eventUbScaleMTE2VList[ubListId]); + + auto perTokenScaleTileOffset = tileOffset.template GetCoordByAxis<0>(); + auto perTokenScaleTileShape = actualTileShape.template GetCoordByAxis<0>(); + + auto gmTilePerTokenScale = gmPerTokenScale[params.layoutPerTokenScale.GetOffset(perTokenScaleTileOffset)]; + auto layoutGmTilePerTokenScale = params.layoutPerTokenScale.GetTileLayout(perTokenScaleTileShape); + + auto &ubPerTokenScale = ubPerTokenScaleList[ubListId]; + auto layoutUbPerTokenScale = LayoutScale::template MakeLayoutInUb( + perTokenScaleTileShape); + + AscendC::WaitFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, layoutUbPerTokenScale, + layoutGmTilePerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbScaleMTE2VList[ubListId]); + tileRowBroadcastMul(ubMul, ubCFp32, ubScale); + AscendC::SetFlag(eventUbScaleVMTE2List[ubListId]); + + AscendC::WaitFlag(eventUbPerTokenScaleMTE2VList[ubListId]); + tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale); + AscendC::SetFlag(eventUbPerTokenScaleVMTE2List[ubListId]); + + AscendC::PipeBarrier(); + tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleBrcb); + AscendC::PipeBarrier(); + + auto &ubD = ubDList[ubListId]; + LayoutD layoutUbD{actualTileShape, ubTileStride}; + + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + AscendC::Cast(ubD, ubPerTokenMul, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + auto gmTileD = gmD[params.layoutD.GetOffset(tileOffset)]; + auto layoutGmTileD = params.layoutD.GetTileLayout(actualTileShape); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubScaleList[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbScaleVMTE2List[UB_STAGES]; + int32_t eventUbScaleMTE2VList[UB_STAGES]; + int32_t eventUbPerTokenScaleVMTE2List[UB_STAGES]; + int32_t eventUbPerTokenScaleMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32; + AscendC::LocalTensor ubMul; + AscendC::LocalTensor ubPerTokenScaleBrcb; + AscendC::LocalTensor ubPerTokenMul; + + TileRowBroadcastMul tileRowBroadcastMul; + TileBroadcastOneBlk tileBroadcastOneBlk; + TileOneBlkColumnBroadcastMul tileOneBlkColumnBroadcastMul; + + CopyGmToUbC copyGmToUbC; + CopyGmToUbScale copyGmToUbScale; + CopyGmToUbPerTokenScale copyGmToUbPerTokenScale; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_DEQUANT_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/dispatch_policy.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/dispatch_policy.hpp new file mode 100644 index 00000000..1718cfd5 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/dispatch_policy.hpp @@ -0,0 +1,39 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_DISPATCH_POLICY_HPP +#define CATLASS_EPILOGUE_DISPATCH_POLICY_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/arch/arch.hpp" + +namespace Catlass::Epilogue { + +// For AtlasA2, per token dequant +template +struct EpilogueAtlasA2PerTokenDequant { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; +}; +//////////////////////////// +/// new add +// For AtlasA2, GEMM +struct EpilogueAtlasA2Gemm { + using ArchTag = Arch::AtlasA2; +}; + +// For AtlasA2, GEMV +struct EpilogueAtlasA2Gemv { + using ArchTag = Arch::AtlasA2; +}; +/////////////////////////// +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_DISPATCH_POLICY_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_gm_to_ub.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_gm_to_ub.hpp new file mode 100644 index 00000000..ac57c363 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_gm_to_ub.hpp @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP +#define CATLASS_EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType +> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutSrc.shape(0), + layoutSrc.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) * sizeof(Element), + (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK, + 0 + ); + + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0 + ); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +/// @brief This copy instruction used to copy per token scale from GM to UB. +/// Copy the scale of shape (m,1) on GM to the first column of shape (m,n) on UB, +/// and pad the first block of each row (i.e. pad to shape (m,8) when element type is float). +/// @tparam ArchTag: Architecture tag. +/// @tparam GmType: Type of data on GM. +template < + class ArchTag, + class GmType +> +struct CopyPerTokenScale2Ub { + static_assert(std::is_same_v, + "Unsupported layout for CopyPerTokenScale2Ub."); + + using Element = typename GmType::Element; + using LayoutSrc = typename GmType::Layout; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyPerTokenScale2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, + LayoutSrc const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams; + AscendC::DataCopyPadExtParams padParams; + + dataCopyParams.blockCount = layoutSrc.shape(0); + dataCopyParams.blockLen = layoutSrc.shape(1) * sizeof(Element); // per token scale has only one column + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + // Pad the data to the complete block + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = 0; + + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + } +}; + +template < + class ArchTag, + class GmType +> +struct CopyGm2UbAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub aligned, can not find the specialization."); +}; + +template +struct CopyGm2UbAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyGm2UbAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // Catlass::Epilogue::Tile + +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_ub_to_gm.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_ub_to_gm.hpp new file mode 100644 index 00000000..686cbe70 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/copy_ub_to_gm.hpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP +#define CATLASS_EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType +> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0 + ); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + + +// new add vectorlayout version +template +struct CopyUb2Gm> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutDst.shape(0) * sizeof(Element), + 0, + 0, + 0 + ); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + }; +}; + + +template < + class ArchTag, + class GmType +> +struct CopyUb2GmAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm aligned, can not find the specialization."); +}; + +template +struct CopyUb2GmAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyUb2GmAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutDst.shape(0); + uint32_t cols = layoutDst.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // Catlass::Epilogue::Tile + +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_mul.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_mul.hpp new file mode 100644 index 00000000..fb184da4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_mul.hpp @@ -0,0 +1,137 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_BROADCAST_MUL_HPP +#define CATLASS_EPILOGUE_TILE_TILE_BROADCAST_MUL_HPP + +#include "../../../gmm_infra/base_defs.hpp" + +namespace Catlass::Epilogue::Tile { + +/// BroadcastMul computes the elementwise multiplication of a tensor of shape (m, n) and a tensor +/// of shape (m, n) after broadcasting. There are two broadcast modes: row-broadcast and +/// column-broadcast. + +/// @brief Computes the elementwise multiplication of a tensor with shape (m, n) and a tensor with +/// original shape (1, n) broadcast to (m, n). +/// @tparam ArchTag_ is the architecture tag. +/// @tparam ComputeType_ includes the element type and layout information. +/// @tparam TileShape_ is the shape (m, n). +template < + class ArchTag_, + class ComputeType_, + class TileShape_ +> +struct TileRowBroadcastMul { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + CATLASS_DEVICE + TileRowBroadcastMul() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + constexpr uint32_t maxRepeatTimes = 255; + constexpr uint32_t eleNumPerBlk = BYTE_PER_BLK / sizeof(ElementCompute); + + constexpr uint32_t blkNumPerColumn = TileShape::COLUMN / eleNumPerBlk; + AscendC::BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = blkNumPerColumn; + repeatParams.src0RepStride = blkNumPerColumn; + repeatParams.src1RepStride = 0; + + constexpr uint32_t rowNumPerCompute = maxRepeatTimes; + constexpr uint32_t colNumPerCompute = BYTE_PER_VECTOR_FRACTAL / sizeof(ElementCompute); + for (uint32_t rowOffset = 0; rowOffset < TileShape::ROW; rowOffset += rowNumPerCompute) { + uint32_t residueM = TileShape::ROW - rowOffset; + uint8_t repeatTimes = static_cast((residueM > rowNumPerCompute) ? rowNumPerCompute : residueM); + for (uint32_t colOffset = 0; colOffset < TileShape::COLUMN; colOffset += colNumPerCompute) { + uint32_t residueN = TileShape::COLUMN - colOffset; + uint64_t mask = (residueN > colNumPerCompute) ? colNumPerCompute : residueN; + AscendC::Mul( + ubOut[rowOffset * TileShape::COLUMN + colOffset], + ubIn0[rowOffset * TileShape::COLUMN + colOffset], + ubIn1[colOffset], + mask, repeatTimes, repeatParams + ); + } + } + } +}; + +/// @brief Compute the elementwise multiplication of a tensor of shape (m, n) and a tensor of shape +/// (m, eleNumPerBlk), which is broadcast from a tensor of shape (m, 1), broadcast to (m, n). +/// @tparam ArchTag_ is the architecture tag. +/// @tparam ComputeType_ includes the element type and layout information. +/// @tparam TileShape_ is the shape (m, n). +template < + class ArchTag_, + class ComputeType_, + class TileShape_ +> +struct TileOneBlkColumnBroadcastMul { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + CATLASS_DEVICE + TileOneBlkColumnBroadcastMul() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + constexpr uint32_t maxRepeatNum = 255; + constexpr uint32_t eleNumPerBlk = BYTE_PER_BLK / sizeof(ElementCompute); + + constexpr uint32_t blkNumPerColumn = TileShape::COLUMN / eleNumPerBlk; + AscendC::BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = blkNumPerColumn; + repeatParams.src0BlkStride = blkNumPerColumn; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 1; + repeatParams.src0RepStride = 1; + repeatParams.src1RepStride = 0; + + constexpr uint32_t rowNumPerCompute = BLK_NUM_PER_VECTOR_FRACTAL; + constexpr uint32_t colNumPerCompute = eleNumPerBlk * maxRepeatNum; + for (uint32_t rowOffset = 0; rowOffset < TileShape::ROW; rowOffset += rowNumPerCompute) { + uint32_t residueM = TileShape::ROW - rowOffset; + uint64_t mask = ((residueM > rowNumPerCompute) ? rowNumPerCompute : residueM) * eleNumPerBlk; + for (uint32_t colOffset = 0; colOffset < TileShape::COLUMN; colOffset += colNumPerCompute) { + uint32_t residueN = TileShape::COLUMN - colOffset; + uint8_t repeatTimes = static_cast( + ((residueN > colNumPerCompute) ? colNumPerCompute : residueN) / eleNumPerBlk); + AscendC::Mul( + ubOut[rowOffset * TileShape::COLUMN + colOffset], + ubIn0[rowOffset * TileShape::COLUMN + colOffset], + ubIn1[rowOffset * eleNumPerBlk], + mask, repeatTimes, repeatParams + ); + } + } + } +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_one_blk.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_one_blk.hpp new file mode 100644 index 00000000..9dec6de6 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_broadcast_one_blk.hpp @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_BROADCAST_ONE_BLK_HPP +#define CATLASS_EPILOGUE_TILE_TILE_BROADCAST_ONE_BLK_HPP + +#include "../../../gmm_infra/base_defs.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag_, + class ComputeType_, + uint32_t COMPUTE_LENGTH_ +> +struct TileBroadcastOneBlk { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; + + CATLASS_DEVICE + TileBroadcastOneBlk() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn + ) + { + constexpr uint32_t maxRepeatNum = 255; + constexpr uint32_t eleNumPerBlk = BYTE_PER_BLK / sizeof(ElementCompute); + + AscendC::BrcbRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.dstRepStride = BLK_NUM_PER_VECTOR_FRACTAL; + + constexpr uint32_t eleNumPerCompute = RoundDown(maxRepeatNum * BLK_NUM_PER_VECTOR_FRACTAL); + for (uint32_t offset = 0; offset < COMPUTE_LENGTH; offset += eleNumPerCompute) { + uint32_t residueM = COMPUTE_LENGTH - offset; + uint32_t computeM = (residueM > eleNumPerCompute) ? eleNumPerCompute : residueM; + uint8_t repeatTimes = static_cast(CeilDiv(computeM)); + AscendC::Brcb( + ubOut[offset * eleNumPerBlk], ubIn[offset], + repeatTimes, repeatParams + ); + } + } +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_copy.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_copy.hpp new file mode 100644 index 00000000..78e7b1d8 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_copy.hpp @@ -0,0 +1,107 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_COPY_HPP +#define CATLASS_EPILOGUE_TILE_TILE_COPY_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/epilogue/tile/copy_gm_to_ub.hpp" +#include "../../../gmm_infra/epilogue/tile/copy_ub_to_gm.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + /// Tag indicating architecture + class ArchTag, + class... Args +> +struct TileCopy { + static_assert(DEPENDENT_FALSE, "Unsupported tile copy, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for X matrix operand + class XType, + /// GemmType for D matrix operand + class DType +> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType +> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementY = typename YType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyGmToUbY = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType +> +struct TileCopyBf16 { + using ElementC = typename CType::Element; + using ElementX = bfloat16_t; + using ElementY = bfloat16_t; + using ElementD = bfloat16_t; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub>; + using CopyGmToUbY = CopyGm2Ub>; + using CopyUbToGmD = CopyUb2Gm>; +}; + +template < + class ArchTag, + class CType, + class ScaleType, + class PerTokenScaleType, + class DType +> +struct TileCopyPerTokenDequant { + using ElementC = typename CType::Element; + using ElementScale = typename ScaleType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbScale = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyPerTokenScale2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_EPILOGUE_TILE_TILE_COPY_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_swizzle.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_swizzle.hpp new file mode 100644 index 00000000..3d7f8cad --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/epilogue/tile/tile_swizzle.hpp @@ -0,0 +1,92 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_TILE_SWIZZLE_HPP +#define CATLASS_EPILOGUE_TILE_TILE_SWIZZLE_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/detail/alignment.hpp" +#include "../../../gmm_infra/matrix_coord.hpp" + +namespace Catlass::Epilogue::Tile { + +struct EpilogueIdentityTileSwizzle { + MatrixCoord blockShape; + MatrixCoord tileShape; + MatrixCoord loopsMN; + + CATLASS_DEVICE + EpilogueIdentityTileSwizzle() = default; + + CATLASS_DEVICE + EpilogueIdentityTileSwizzle(MatrixCoord const &blockShape, MatrixCoord const &tileShape) : + blockShape(blockShape), + tileShape(tileShape) + { + loopsMN = CeilDiv(blockShape, tileShape); + } + + CATLASS_DEVICE + uint32_t GetLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + CATLASS_DEVICE + MatrixCoord GetTileCoord(uint32_t loopIdx) const + { + return MatrixCoord{ loopIdx / loopsMN.column(), loopIdx % loopsMN.column() }; + } + + CATLASS_DEVICE + MatrixCoord GetActualTileShape(MatrixCoord const &tileCoord) const + { + return MatrixCoord::Min(tileShape, blockShape - tileCoord * tileShape); + } +}; + +struct EpilogueHorizontalTileSwizzle { + MatrixCoord blockShape; + MatrixCoord tileShape; + MatrixCoord loopsMN; + + CATLASS_DEVICE + EpilogueHorizontalTileSwizzle() = default; + + CATLASS_DEVICE + EpilogueHorizontalTileSwizzle(MatrixCoord const &blockShape, MatrixCoord const &tileShape) : + blockShape(blockShape), + tileShape(tileShape) + { + loopsMN = CeilDiv(blockShape, tileShape); + } + + CATLASS_DEVICE + uint32_t GetLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + CATLASS_DEVICE + MatrixCoord GetTileCoord(uint32_t loopIdx) const + { + return MatrixCoord{ loopIdx % loopsMN.row(), loopIdx / loopsMN.row() }; + } + + CATLASS_DEVICE + MatrixCoord GetActualTileShape(MatrixCoord const &tileCoord) const + { + return MatrixCoord::Min(tileShape, blockShape - tileCoord * tileShape); + } +}; + +} + +#endif // CATLASS_EPILOGUE_TILE_TILE_SWIZZLE_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad.hpp new file mode 100644 index 00000000..d58f83e6 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad.hpp @@ -0,0 +1,39 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy.hpp" +#include "../../../gmm_infra/gemm/tile/tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class AType, + class BType, + class CType, + class BiasType = void, + class TileCopy = Gemm::Tile::TileCopy, + class TileMmad = Gemm::Tile::TileMmad +> +struct BlockMmad { + static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); +}; + +} // namespace Catlass::Gemm::Block + +#include "../../../gmm_infra/gemm/block/block_mmad_preload_async_fixAxisMove_with_callback.hpp" + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad_preload_async_fixAxisMove_with_callback.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad_preload_async_fixAxisMove_with_callback.hpp new file mode 100644 index 00000000..fd546f13 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_mmad_preload_async_fixAxisMove_with_callback.hpp @@ -0,0 +1,495 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_FIXAXISMOVE_WITH_CALLBACK_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_FIXAXISMOVE_WITH_CALLBACK_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/resource.hpp" +#include "../../../gmm_infra/coord.hpp" +#include "../../../gmm_infra/detail/callback.hpp" +#include "../../../gmm_infra/gemm_coord.hpp" +#include "../../../gmm_infra/gemm/dispatch_policy.hpp" +#include "../../../gmm_infra/gemm/helper.hpp" + +namespace Catlass::Gemm::Block { + +template +struct BlockMmad, + L1TileShape_, L0TileShape_, AType_, BType_, CType_, BiasType_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = + MmadAtlasA2PreloadAsyncFixAxisMoveWithCallback; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1A_TILE_NUM = DispatchPolicy::L1A_TILE_NUM; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_RIFFLE_SHUFFLE = DispatchPolicy::ENABLE_RIFFLE_SHUFFLE; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA) * L1A_TILE_NUM; + static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + + static constexpr uint32_t L1A_K_ONE_STAGE = L1TileShape::K * L1A_TILE_NUM; + static constexpr uint32_t L1A_K_ONE_TIME = L1A_K_ONE_STAGE * L1A_STAGES; + + static constexpr uint32_t FLAG_ENABLE = 1; + static constexpr uint32_t MMTE1_FLAG_ENABLE = 1; + static constexpr uint32_t FIX_FLAG_ENABLE = 1; + + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + + static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout(L1TileShape::M, L1TileShape::K); + static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout(L1TileShape::K, L1TileShape::N); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + InitL1(resource, l1BufAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + } + + CATLASS_DEVICE + ~BlockMmad() + { + SynchronizeBlock(); + AscendC::WaitFlag(l1AEventList[0]); + for (uint32_t i = 0; i < L1B_STAGES; ++i) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + + // 左矩阵分四块copy + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + GemmCoord const &actualShape, int needCopyL1Left, int needAtomicAdd, uint8_t needLeftCopyFlag, + uint8_t isFisrtBlock, Callback const &callbackBeforeFixpipe, Callback const &callbackAfterFixpipe) + { + uint32_t kaTileCount = CeilDiv(actualShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_RIFFLE_SHUFFLE) { + startTileIdx = AscendC::GetBlockIdx() % kaTileCount; + } + // 左矩阵的切分 + for (uint32_t kaLoopIdx = 0; kaLoopIdx < kaTileCount; ++kaLoopIdx) { + // uint32_t kTileIdx = (kaLoopIdx < kaTileCount) ? kaLoopIdx : (kaLoopIdx - kaTileCount); + uint32_t kaTileIdx = (startTileIdx + kaLoopIdx < kaTileCount) ? (startTileIdx + kaLoopIdx) : + (startTileIdx + kaLoopIdx - kaTileCount); + + uint32_t kActualLeft = + (kaTileIdx < kaTileCount - 1) ? L1A_K_ONE_STAGE : (actualShape.k() - kaTileIdx * L1A_K_ONE_STAGE); + // Emission load instruction from GM to L1 + + // Load first matrix A tile from GM to L1 + if (needCopyL1Left && isFisrtBlock) { + MatrixCoord gmTileAOffset{0, kaTileIdx * L1A_K_ONE_STAGE}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + if (kaLoopIdx == 0) { + AscendC::WaitFlag(l1AEventList[0]); + } + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualLeft)); + copyGmToL1A(l1ATensorList[l1AListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[0]); + } + + // 右矩阵的切分 + uint32_t kbTileCount = 1; + for (uint32_t kbLoopIdx = 0; kbLoopIdx < kbTileCount; ++kbLoopIdx) { + uint32_t kActualRight = + (kbLoopIdx < kbTileCount - 1) ? L1TileShape::K : (kActualLeft - kbLoopIdx * L1TileShape::K); + + uint32_t kbIdx = (kaLoopIdx + startTileIdx) * kbTileCount + kbLoopIdx; + uint32_t kbTileIdx = (kbIdx < kaTileCount * kbTileCount) ? kbIdx : (kbIdx - kaTileCount * kbTileCount); + + MatrixCoord gmTileBOffset{kbTileIdx * L1TileShape::K, 0}; + + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActualRight, actualShape.n())); + + copyGmToL1B(l1BTensorList[l1BListId], gmTileB, L1B_LAYOUT, layoutTileB); + AscendC::SetFlag(l1BEventList[l1BListId]); + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? + (l1TileMmadParamsId + preloadCount) : + (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1AListId = l1AListId; + l1TileMmadParams.l1AOffset = kbLoopIdx; + l1TileMmadParams.l1BListId = l1BListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActualLeft = kActualLeft; + l1TileMmadParams.kActualRight = kActualRight; + l1TileMmadParams.isKLoopFirst = (kbLoopIdx == 0 && kaLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kbLoopIdx == kbTileCount - 1 && kaLoopIdx == kaTileCount - 1); + l1TileMmadParams.copyedL1Left = + needCopyL1Left && kbLoopIdx == 0 /*&& (blockIdx != 0 || blockIdx == blockCnt - 1)*/; + l1TileMmadParams.needAtomicAdd = needAtomicAdd; + l1TileMmadParams.isFisrtBlock = isFisrtBlock; + l1TileMmadParams.needSetFlag = (kaLoopIdx == kaTileCount - 1); + if (kbLoopIdx == kbTileCount - 1) { + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.callbackBeforeFixpipe = callbackBeforeFixpipe; + l1TileMmadParams.callbackAfterFixpipe = callbackAfterFixpipe; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1BListId = (l1BListId + 1) % L1B_STAGES; + } + if (needCopyL1Left && isFisrtBlock == 0) { + MatrixCoord gmTileAOffset{0, kaTileIdx * L1A_K_ONE_STAGE}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + if (kaLoopIdx == 0) { + AscendC::WaitFlag(l1AEventList[0]); + } + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActualLeft)); + copyGmToL1A(l1ATensorList[l1AListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[0]); + } + l1AListId = (l1AListId + 1 < kaTileCount) ? (l1AListId + 1) : 0; + } + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + +private: + struct L1TileMmadParams { + uint32_t l1AListId; + uint32_t l1AOffset; + uint32_t l1BListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActualLeft; + uint32_t kActualRight; + bool isKLoopFirst; + bool isKLoopLast; + int copyedL1Left; + int needAtomicAdd; + int needSetFlag; + uint8_t isFisrtBlock; + AscendC::GlobalTensor gmBlockC; + LayoutC layoutCInGm; + // Callback callback; + Callback callbackBeforeFixpipe; + Callback callbackAfterFixpipe; + + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + for (uint32_t i = 0; i < L1A_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + } + // 左矩阵不ping-ping，作为一块数据分布式copy + l1AEventList[0] = 0; + AscendC::SetFlag(l1AEventList[0]); + + for (uint32_t i = 0; i < L1B_STAGES; ++i) { + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1BEventList[i] = i + 1; + AscendC::SetFlag(l1BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound); + uint32_t nPartLoop = CeilDiv(params.nRound); + uint32_t kPartLoop = CeilDiv(params.kActualRight); + auto &l1ATensor = l1ATensorList[params.l1AListId]; + auto &l1BTensor = l1BTensorList[params.l1BListId]; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = + (mPartIdx < mPartLoop - 1) ? L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M); + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = + (kPartIdx < kPartLoop - 1) ? L0TileShape::K : (params.kActualRight - kPartIdx * L0TileShape::K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = + MakeCoord(mPartIdx, kPartIdx + params.l1AOffset * kPartLoop) * L0TileShape::ToCoordMK(); + auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; + if (params.isFisrtBlock) { + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0) && params.copyedL1Left) { + AscendC::WaitFlag(l1AEventList[0]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + // 左矩阵全部copy完成，才需要setFlag + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && params.needSetFlag && + params.copyedL1Left) { + AscendC::SetFlag(l1AEventList[0]); + } + } + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = + (nPartIdx < nPartLoop - 1) ? L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); + auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1BListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); + + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1BListId]); + } + if (nPartIdx == 0 && !params.isFisrtBlock) { + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0) && params.copyedL1Left) { + AscendC::WaitFlag(l1AEventList[0]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + // 左矩阵全部copy完成，才需要setFlag + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && params.needSetFlag && + params.copyedL1Left) { + AscendC::SetFlag(l1AEventList[0]); + } + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && + (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + + if (params.callbackBeforeFixpipe) { + params.callbackBeforeFixpipe(); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + if (params.needAtomicAdd) { + AscendC::SetAtomicAdd(); + } + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); + if (params.needAtomicAdd) { + AscendC::SetAtomicNone(); + } + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + if (params.needAtomicAdd) { + AscendC::SetAtomicAdd(); + } + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); + if (params.needAtomicAdd) { + AscendC::SetAtomicNone(); + } + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + + if (params.callbackAfterFixpipe) { + params.callbackAfterFixpipe(); + } + } + } + + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_ASYNC_FIXAXISMOVE_WITH_CALLBACK_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_swizzle.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_swizzle.hpp new file mode 100644 index 00000000..2bf7b86e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/block/block_swizzle.hpp @@ -0,0 +1,551 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_SWIZZLE_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_SWIZZLE_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/detail/alignment.hpp" +#include "../../../gmm_infra/gemm_coord.hpp" +#include "../../../gmm_infra/matrix_coord.hpp" + +namespace Catlass::Gemm::Block { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Block swizzling function for Gemms +template +struct GemmIdentityBlockSwizzle { + /// Data members + + GemmCoord problemShape; + MatrixCoord tileMN; + MatrixCoord loopsMN; + + /// Methods + + CATLASS_DEVICE + GemmIdentityBlockSwizzle() + { + } + + CATLASS_DEVICE + GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + : problemShape(problemShape_), tileMN(tileMN_) + { + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + } + + CATLASS_DEVICE + GemmIdentityBlockSwizzle(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) + : problemShape(problemShape_), tileMN(tileMN_), loopsMN(loopsMN_) + { + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + + loopsMN = CeilDiv(MatrixCoord(problemShape.GetCoordMN()), tileMN); + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, MatrixCoord const &tileMN_, MatrixCoord const &loopsMN_) + { + problemShape = problemShape_; + tileMN = tileMN_; + loopsMN = loopsMN_; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + CATLASS_DEVICE + uint32_t GetBatchIdx(uint32_t taskIdx) + { + return taskIdx / (GetCoreLoops()); + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t innerIdx = taskIdx % GetCoreLoops(); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMN.row(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.column()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.column()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMN.row() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMN.column() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMN.column(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMN.row()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMN.row()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMN.column() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMN.row() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, 0}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord) + { + uint32_t mActual = + (blockCoord.m() == (loopsMN.row() - 1)) ? (problemShape.m() - blockCoord.m() * tileMN.row()) : tileMN.row(); + uint32_t nActual = (blockCoord.n() == (loopsMN.column() - 1)) ? + (problemShape.n() - blockCoord.n() * tileMN.column()) : + tileMN.column(); + uint32_t kActual = problemShape.k(); + return GemmCoord{mActual, nActual, kActual}; + } +}; + +/// Block swizzling function for Splitk Gemms +template +struct SplitkGemmIdentityBlockSwizzle { + /// Data members + GemmCoord problemShape; + GemmCoord tileShape; + GemmCoord loopsMNK; + uint32_t splitkFactor = 1; // splite k dim into virtual cores + + + /// Methods + + CATLASS_DEVICE + SplitkGemmIdentityBlockSwizzle() + { + } + + CATLASS_DEVICE + SplitkGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, + uint32_t splitkFactor_ = 1) + : problemShape(problemShape_), tileShape(tileShape_), splitkFactor(splitkFactor_) + { + loopsMNK = CeilDiv(problemShape, tileShape); + } + + CATLASS_DEVICE + uint32_t GetKIdxBySplitkSliceIdx(uint32_t splitkSliceIdx) const + { + if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { + return (loopsMNK.k() / splitkFactor + 1) * splitkSliceIdx; + } else { + return splitkSliceIdx * (loopsMNK.k() / splitkFactor) + loopsMNK.k() % splitkFactor; + } + } + + CATLASS_DEVICE + uint32_t GetSplitkSliceIdx(uint32_t taskIdx) const + { + uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); + return taskIdx % GetCoreLoops() / mnLoops; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return loopsMNK.m() * loopsMNK.n() * splitkFactor; + } + + CATLASS_DEVICE + uint32_t GetBatchIdx(uint32_t taskIdx) + { + return taskIdx / GetCoreLoops(); + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); + uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); + + uint32_t innerIdx = taskIdx % (loopsMNK.m() * loopsMNK.n()); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord, uint32_t splitkSliceIdx) + { + uint32_t splitkSliceLen; + if (splitkSliceIdx < loopsMNK.k() % splitkFactor) { + splitkSliceLen = (loopsMNK.k() / splitkFactor + 1) * tileShape.k(); + } else { + splitkSliceLen = (loopsMNK.k() / splitkFactor) * tileShape.k(); + } + uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? + (problemShape.m() - blockCoord.m() * tileShape.m()) : + tileShape.m(); + uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? + (problemShape.n() - blockCoord.n() * tileShape.n()) : + tileShape.n(); + uint32_t kActual = (splitkSliceIdx == (splitkFactor - 1)) ? + (problemShape.k() - blockCoord.k() * tileShape.k()) : + splitkSliceLen; + return GemmCoord{mActual, nActual, kActual}; + } +}; + +/// Block swizzling function for Splitk Gemms +template +struct SplitkInOneCoreGemmIdentityBlockSwizzle { + /// Data members + + static constexpr uint32_t MAX_TILE_NUM = 5000; + static constexpr uint32_t MAX_LINE = 500; + + GemmCoord problemShape; + GemmCoord tileShape; + GemmCoord loopsMNK; + uint32_t splitkFactor = 1; + uint32_t maxKValue = 0; + uint32_t kPerLoop = 0; + uint32_t blockCntPerCore = 0; + uint32_t leftBlockCnt = 0; + uint32_t minPlusBlockIdx = 0; + uint32_t maxPlusBlockIdx = 0; + uint32_t blockIdexList[MAX_TILE_NUM]; + uint32_t blockCnt = 0; + uint32_t beginBlockIdex = 0; + + /// Methods + CATLASS_DEVICE + SplitkInOneCoreGemmIdentityBlockSwizzle() + { + } + + CATLASS_DEVICE + SplitkInOneCoreGemmIdentityBlockSwizzle(GemmCoord const &problemShape_, GemmCoord const &tileShape_, + uint32_t splitkFactor_ = 1) + : problemShape(problemShape_), tileShape(tileShape_), splitkFactor(splitkFactor_) + { + loopsMNK = CeilDiv(problemShape, tileShape); + } + + CATLASS_DEVICE + uint32_t GetKIdxBySplitkSliceIdx(uint32_t splitkSliceIdx) const + { + return kPerLoop * splitkSliceIdx; + } + + CATLASS_DEVICE + uint32_t GetSplitkSliceIdx(uint32_t taskIdx) const + { + uint32_t mnLoops = loopsMNK.m() * loopsMNK.n(); + return taskIdx % splitkFactor; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops(uint32_t coreIdx) const + { + uint32_t blockCntTmp = blockCntPerCore; + if (leftBlockCnt > 0 && + ((minPlusBlockIdx <= maxPlusBlockIdx && coreIdx >= minPlusBlockIdx && coreIdx <= maxPlusBlockIdx) || + (minPlusBlockIdx > maxPlusBlockIdx && (coreIdx >= minPlusBlockIdx || coreIdx <= maxPlusBlockIdx)))) { + blockCntTmp = blockCntPerCore + 1; + } + return blockCntTmp; + } + + CATLASS_DEVICE + uint32_t GetCoreLoops() const + { + return blockCnt; + } + + CATLASS_DEVICE + uint32_t GetBatchIdx(uint32_t taskIdx) + { + return taskIdx % splitkFactor; + } + + CATLASS_DEVICE + GemmCoord GetBlockCoord(uint32_t taskIdx) + { + uint32_t splitkSliceIdx = GetSplitkSliceIdx(taskIdx); + uint32_t kIdx = GetKIdxBySplitkSliceIdx(splitkSliceIdx); + + uint32_t innerIdx = taskIdx / splitkFactor % (loopsMNK.m() * loopsMNK.n()); + if constexpr (SwizzleDirection == 0) { // Zn + uint32_t tileBlockLoop = CeilDiv(loopsMNK.m(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.n()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.n()); + + uint32_t nRow = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nRow = loopsMNK.m() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nRow; + uint32_t nIdx = inTileBlockIdx / nRow; + if (tileBlockIdx % 2 == 1) { + nIdx = loopsMNK.n() - nIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } else if constexpr (SwizzleDirection == 1) { // Nz + uint32_t tileBlockLoop = CeilDiv(loopsMNK.n(), SwizzleOffset); + uint32_t tileBlockIdx = innerIdx / (SwizzleOffset * loopsMNK.m()); + uint32_t inTileBlockIdx = innerIdx % (SwizzleOffset * loopsMNK.m()); + + uint32_t nCol = SwizzleOffset; + if (tileBlockIdx == tileBlockLoop - 1) { + nCol = loopsMNK.n() - SwizzleOffset * tileBlockIdx; + } + uint32_t mIdx = inTileBlockIdx / nCol; + uint32_t nIdx = tileBlockIdx * SwizzleOffset + inTileBlockIdx % nCol; + if (tileBlockIdx % 2 == 1) { + mIdx = loopsMNK.m() - mIdx - 1; + } + return GemmCoord{mIdx, nIdx, kIdx}; + } + } + + CATLASS_DEVICE + GemmCoord GetActualBlockShape(GemmCoord blockCoord, uint32_t splitkSliceIdx) + { + uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? + (problemShape.m() - blockCoord.m() * tileShape.m()) : + tileShape.m(); + uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? + (problemShape.n() - blockCoord.n() * tileShape.n()) : + tileShape.n(); + uint32_t kActual = + (splitkSliceIdx == (splitkFactor - 1)) ? (problemShape.k() - blockCoord.k() * tileShape.k()) : maxKValue; + return GemmCoord{mActual, nActual, kActual}; + } + + CATLASS_DEVICE + GemmCoord GetActualBlockTailBefore(GemmCoord blockCoord, uint32_t splitkSliceIdx) + { + uint32_t mActual = (blockCoord.m() == (loopsMNK.m() - 1)) ? + (problemShape.m() - blockCoord.m() * tileShape.m()) : + tileShape.m(); + uint32_t nActual = (blockCoord.n() == (loopsMNK.n() - 1)) ? + (problemShape.n() - blockCoord.n() * tileShape.n()) : + tileShape.n(); + uint32_t kActual; + if (splitkFactor > 1) { + kActual = (splitkSliceIdx == (splitkFactor - 2)) ? + (problemShape.k() - blockCoord.k() * tileShape.k() - maxKValue) : + maxKValue; + } else { + kActual = (splitkSliceIdx == (splitkFactor - 1)) ? (problemShape.k() - blockCoord.k() * tileShape.k()) : + maxKValue; + } + + return GemmCoord{mActual, nActual, kActual}; + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, GemmCoord const &tileShape_, uint32_t splitkFactor_, + uint32_t maxKValue_) + { + problemShape = problemShape_; + tileShape = tileShape_; + splitkFactor = splitkFactor_; + maxKValue = maxKValue_; + + loopsMNK = CeilDiv(problemShape, tileShape); + kPerLoop = maxKValue / tileShape.k(); + } + + CATLASS_DEVICE + void SetBlockIdx(uint32_t coreNum, uint32_t coreIdx) + { + uint32_t blockDimM = loopsMNK.m(); + uint32_t blockDimN = loopsMNK.n(); + uint32_t lastNIdx[MAX_LINE] /*= {0}*/; + for (uint32_t i = 0; i < blockDimM; i++) { + lastNIdx[i] = 0; + } + for (uint32_t i = 0; i <= coreIdx; i++) { + uint32_t line = i % blockDimM; + uint32_t col = i / blockDimM; + if (col % 2 != 0) { + line = (col + 1) * blockDimM - i - 1; + } + uint32_t offset = coreIdx / 4; + uint32_t blockCntTmp = blockCnt; + if (i != coreIdx) { + blockCntTmp = GetCoreLoops(i); + } + for (uint32_t j = 0; j < blockCntTmp; j++) { + if (col % 2 == 0) { + while (lastNIdx[line] >= blockDimN) { + line++; + } + } else { + while (lastNIdx[line] >= blockDimN) { + line--; + } + } + if (i == coreIdx) { + uint32_t jOffset = (j + offset) % blockCnt; + blockIdexList[jOffset] = lastNIdx[line] + line * blockDimN; + } + lastNIdx[line]++; + } + } + } + + CATLASS_DEVICE + void SetBeginBlockIdx(uint32_t coreNum, uint32_t coreIdx) + { + for (uint32_t i = 0; i < coreIdx; i++) { + beginBlockIdex += GetCoreLoops(i); + } + } + + CATLASS_DEVICE + uint32_t GetCoreBlockIdx(uint32_t index) + { + return blockIdexList[index]; + } + + CATLASS_DEVICE + uint32_t GetCoreBeginBlockIdx() + { + return beginBlockIdex; + } + + CATLASS_DEVICE + uint8_t GetNeedLeftCopyFlag(uint32_t index) + { + uint32_t curBlockLine = 0; + uint32_t blockDimN = loopsMNK.n(); + if (index == blockCnt - 1) { + return 1; + } + if (index > 1) { + uint32_t lastBlockLine = blockIdexList[index - 1] / blockDimN; + uint32_t curBlockLine = blockIdexList[index] / blockDimN; + if (lastBlockLine != curBlockLine) { + return 1; + } + } + return 0; + } + + CATLASS_DEVICE + uint8_t GetNeedLeftCopyFlag(uint32_t index, uint32_t curBlockIdex) + { + uint32_t curBlockLine = 0; + uint32_t blockDimN = loopsMNK.n(); + if (index == blockCnt - 1) { + return 1; + } + if (index > 1) { + uint32_t lastBlockLine = (curBlockIdex - 1) / blockDimN; + uint32_t curBlockLine = curBlockIdex / blockDimN; + if (lastBlockLine != curBlockLine) { + return 1; + } + } + return 0; + } + + CATLASS_DEVICE + void Update(GemmCoord const &problemShape_, GemmCoord const &tileShape_, uint32_t splitkFactor_, + uint32_t maxKValue_, uint32_t groupIdx, uint32_t coreNum, uint32_t coreIdx) + { + problemShape = problemShape_; + tileShape = tileShape_; + splitkFactor = splitkFactor_; + maxKValue = maxKValue_; + + loopsMNK = CeilDiv(problemShape, tileShape); + kPerLoop = maxKValue / tileShape.k(); + + uint32_t totalBlockCnt = loopsMNK.m() * loopsMNK.n(); + blockCntPerCore = totalBlockCnt / coreNum; + leftBlockCnt = totalBlockCnt - blockCntPerCore * coreNum; + minPlusBlockIdx = (groupIdx * leftBlockCnt) % coreNum; + maxPlusBlockIdx = (groupIdx * leftBlockCnt + leftBlockCnt - 1) % coreNum; + + blockCnt = GetCoreLoops(coreIdx); + SetBlockIdx(coreNum, coreIdx); + } + + CATLASS_DEVICE + void UpdateEx(GemmCoord const &problemShape_, GemmCoord const &tileShape_, uint32_t splitkFactor_, + uint32_t maxKValue_, uint32_t groupIdx, uint32_t coreNum, uint32_t coreIdx) + { + problemShape = problemShape_; + tileShape = tileShape_; + splitkFactor = splitkFactor_; + maxKValue = maxKValue_; + + loopsMNK = CeilDiv(problemShape, tileShape); + kPerLoop = maxKValue / tileShape.k(); + + uint32_t totalBlockCnt = loopsMNK.m() * loopsMNK.n(); + blockCntPerCore = totalBlockCnt / coreNum; + leftBlockCnt = totalBlockCnt - blockCntPerCore * coreNum; + minPlusBlockIdx = (groupIdx * leftBlockCnt) % coreNum; + maxPlusBlockIdx = (groupIdx * leftBlockCnt + leftBlockCnt - 1) % coreNum; + + blockCnt = GetCoreLoops(coreIdx); + SetBeginBlockIdx(coreNum, coreIdx); + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_SWIZZLE_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/dispatch_policy.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/dispatch_policy.hpp new file mode 100644 index 00000000..735526e3 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/dispatch_policy.hpp @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_DISPATCH_POLICY_HPP +#define CATLASS_GEMM_DISPATCH_POLICY_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/arch/arch.hpp" + +namespace Catlass::Gemm { + +// Block Mmad Policies + +template +struct MmadAtlasA2Base { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t ASYNC = ASYNC_; +}; + +using MmadAtlasA2 = MmadAtlasA2Base; +using MmadAtlasA2Async = MmadAtlasA2Base; + +template +struct MmadAtlasA2PreloadAsyncFixAxisMoveWithCallback : public MmadAtlasA2Async { + static constexpr uint32_t PRELOAD_STAGES = PRELOAD_STAGES_; // Stages of emitting load instruction in advance + static constexpr uint32_t L1_STAGES = L1_STAGES_; + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1A_TILE_NUM = L1A_TILE_NUM_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_RIFFLE_SHUFFLE = ENABLE_RIFFLE_SHUFFLE_; +}; + + +//////////////////// +// new add +template +struct GemmAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_RIFFLE_SHUFFLE = ENABLE_RIFFLE_SHUFFLE_; + static constexpr bool ENABLE_ABBA = ENABLE_ABBA_; +}; + +struct GemvAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; +//////////////////// + +template +struct MmadAtlasA2PingpongBias : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; +} // namespace Catlass::Gemm + +#endif // CATLASS_GEMM_DISPATCH_POLICY_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/gemm_type.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/gemm_type.hpp new file mode 100644 index 00000000..e006c5b5 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/gemm_type.hpp @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_GEMM_TYPE_HPP +#define CATLASS_GEMM_GEMM_TYPE_HPP + +#include "../../gmm_infra/base_defs.hpp" + +namespace Catlass::Gemm { + +//////////////////////////////////////////////////////////////////// + +template +struct GemmType { + using Element = Element_; + using Layout = Layout_; + static constexpr AscendC::TPosition POSITION = POSITION_; +}; + +} // namespace Catlass::Gemm + +#endif // CATLASS_GEMM_GEMM_TYPE_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/helper.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/helper.hpp new file mode 100644 index 00000000..a439620e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/helper.hpp @@ -0,0 +1,254 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_HELPER_HPP +#define CATLASS_GEMM_HELPER_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/layout/layout.hpp" +#include "../../gmm_infra/gemm/gemm_type.hpp" + +namespace Catlass::Gemm::helper { + +template +struct L1AlignHelper { + static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct ElementAccumulatorSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported element accumulator selector, can not find the specialization."); +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template +struct L1ATypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BiasTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BiasTypeSelector { + using GMBiasType = void; + using L1BiasType = void; + using L0BiasType = void; +}; + +template +struct L1BiasTypeSelector, ElementAccumulator> { + using GMBiasType = Gemm::GemmType; + using L1BiasType = Gemm::GemmType; + using L0BiasType = Gemm::GemmType; +}; + +/////////////////////////////////////// +// new add +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1AndL0TypeSelectorGemm{ + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; +/////////////////////////////////////// +} // namespace Catlass::Gemm::helper + +#endif // CATLASS_GEMM_HELPER_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_l1.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_l1.hpp new file mode 100644 index 00000000..7d60210b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_l1.hpp @@ -0,0 +1,1067 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_GM_TO_L1_HPP +#define CATLASS_GEMM_TILE_COPY_GM_TO_L1_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy_tla.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void +> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void +> +struct CopyGmToL1IntervalDataCopy { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +//////////////////////////////////////// +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +/// Partial specialization for AtlasA2, half, RowMajor in and zN out. +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingRowMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, ColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(2)], dataCopyParams); + } + } +}; + +/// new add gemm +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(0); + uint32_t ndNum = layoutSrc.shape(0) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(0) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(1); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; //` + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(0) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(1)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(1) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(0); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; +//////////////////////////////////////// + +/////////////////////////////////////// +/// new add gemv, VectorLayout -> zN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + intriParams.nValue = 1; + intriParams.srcDValue = layoutSrc.shape(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + + + +/////////////////////////////////////// +/// new add gemv, ColumnMajor -> nN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +/// Partial specialization for zN in and zN out. +template < + class ArchTag, + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(1)); + uint32_t blockLen = RoundUp(layoutSrc.orgShape(0)); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(3) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(3); + uint64_t srcOffset = i * layoutSrc.stride(3); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for nZ in and nZ out. +template < + class ArchTag, + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(0)); + uint32_t blockLen = RoundUp(layoutSrc.orgShape(1)); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(1) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(1); + uint64_t srcOffset = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + } +}; + +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = layoutDst.shape(0) / ELE_NUM_PER_C0; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_GM_TO_L1_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_ub.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_ub.hpp new file mode 100644 index 00000000..1cd561b4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_gm_to_ub.hpp @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_GM_TO_UB_HPP +#define CATLASS_GEMM_TILE_COPY_GM_TO_UB_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy_tla.hpp" + +namespace Catlass::Gemm::Tile { + +} // Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_GM_TO_UB_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l0c_to_gm.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l0c_to_gm.hpp new file mode 100644 index 00000000..b606e43b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l0c_to_gm.hpp @@ -0,0 +1,210 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L0C_TO_GM_HPP +#define CATLASS_GEMM_TILE_COPY_L0C_TO_GM_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" + +namespace Catlass::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT +> +struct CopyL0CToGmQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +// CopyL0CToGm cast fp32 to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, half, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToGm cast fp32 to bf16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToGm output fp32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, float, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm output int32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int32_t, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast int32_t to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_TENSOR +> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_CHANNEL +> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false +> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_ +> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> +{ + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_ +> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> +{ + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::zN; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void +> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_L0C_TO_GM_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_bt.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_bt.hpp new file mode 100644 index 00000000..0cd174f8 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_bt.hpp @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L1_TO_BT_HPP +#define CATLASS_GEMM_TILE_COPY_L1_TO_BT_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" + +namespace Catlass::Gemm::Tile { + + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_L1_TO_BT_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0a.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0a.hpp new file mode 100644 index 00000000..a6c0a7a6 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0a.hpp @@ -0,0 +1,349 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L1_TO_L0A_HPP +#define CATLASS_GEMM_TILE_COPY_L1_TO_L0A_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy_tla.hpp" + + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void +> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0A(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0A(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0)));; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for(uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++){ + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType>{ + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0A(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0))); + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + for(uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++){ + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1) * 2], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType>{ + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0A(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; +////////////////////////////////////////// + +/// Partial specialization for zN in and zZ out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for float, zN in and zZ out. +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for int8_t, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +/// Partial specialization for float, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = RoundUp(l1M); + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l0MAlign = RoundUp(l0M); + uint16_t l0KAlign = RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_L1_TO_L0A_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0b.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0b.hpp new file mode 100644 index 00000000..d4291bff --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_l1_to_l0b.hpp @@ -0,0 +1,479 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L1_TO_L0B_HPP +#define CATLASS_GEMM_TILE_COPY_L1_TO_L0B_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/layout/layout.hpp" +#include "../../../gmm_infra/gemm/gemm_type.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy_tla.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void +> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0B(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for(uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++){ // K N + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0B(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = static_cast(CeilDiv(layoutDst.orgShape(1))) - 1; + for(uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++){ // K N + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + CATLASS_DEVICE + CopyL1ToL0B(){} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; +///////////////////////////////////////////// + +//////////////////////////////////////////// +/// new add gemv +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1)); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(3); i++) + { + AscendC::LoadData(dstTensor[i * layoutDst.stride(3)], srcTensor[i * layoutSrc.stride(3)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> +{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = layoutDst.shape(1) * layoutDst.shape(3); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + using Element = float; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(0)) - 1; + + for (uint32_t i = 0; i < CeilDiv<2 * ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++) + { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3) * 2], + srcTensor[i * layoutSrc.stride(3)], + loadDataParams); + } + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + using Element = int8_t; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(1)); i++) + { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3)], + srcTensor[i * layoutSrc.stride(3) * 2], + loadDataParams); + } + } +}; +//////////////////////////////////////////// + +/// Partial specialization for int8_t, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1N = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0K = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0N = layoutDst.shape(2) * layoutDst.shape(3); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l1NAlign = RoundUp(l1N); + uint16_t l0KAlign = RoundUp(l0K); + uint16_t l0NAlign = RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for zN in and nZ out. +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + if (layoutSrc.shape(3) == layoutDst.shape(3)) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1) * layoutDst.shape(3)); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_L1_TO_L0B_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_ub_to_gm.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_ub_to_gm.hpp new file mode 100644 index 00000000..43419f5c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/copy_ub_to_gm.hpp @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_UB_TO_GM_HPP +#define CATLASS_GEMM_TILE_COPY_UB_TO_GM_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/arch/arch.hpp" +#include "../../../gmm_infra/gemm/tile/tile_copy_tla.hpp" + +namespace Catlass::Gemm::Tile { + + +} // Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_COPY_UB_TO_GM_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy.hpp new file mode 100644 index 00000000..c94cb22b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy.hpp @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_TILE_COPY_HPP +#define CATLASS_GEMM_TILE_TILE_COPY_HPP + +#include +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/gemm/tile/copy_gm_to_l1.hpp" +#include "../../../gmm_infra/gemm/tile/copy_l0c_to_gm.hpp" +#include "../../../gmm_infra/gemm/tile/copy_l1_to_l0a.hpp" +#include "../../../gmm_infra/gemm/tile/copy_l1_to_l0b.hpp" +#include "../../../gmm_infra/gemm/tile/copy_l1_to_bt.hpp" +#include "../../../gmm_infra/gemm/tile/copy_gm_to_ub.hpp" +#include "../../../gmm_infra/gemm/tile/copy_ub_to_gm.hpp" +#include "../../../gmm_infra/gemm/helper.hpp" + + +namespace Catlass::Gemm::Tile { + +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void +> +struct TileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A< + ArchTag, typename helper::L1ATypeSelector::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B< + ArchTag, typename helper::L1BTypeSelector::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t, + void, + Gemm::Tile::CopyGmToL1>; +}; + +////////////////////////////// +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_TILE_COPY_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy_tla.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy_tla.hpp new file mode 100644 index 00000000..05049b4a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_copy_tla.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_TILE_COPY_TLA_HPP +#define CATLASS_GEMM_TILE_TILE_COPY_TLA_HPP + +#include "../../../gmm_infra/base_defs.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void +> +struct TileCopyTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTla, can not find the specialization."); +}; + +// Extended template for TileCopyTla that supports manually specifying LayoutTagSrc and LayoutTagDst. +// Users can specialize the copy class by LayoutTagSrc and LayoutTagDst. +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class LayoutTagSrc, + class LayoutTagDst +> +struct TileCopyTlaExt { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTlaExt, can not find the specialization."); +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_TILE_COPY_TLA_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_mmad.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_mmad.hpp new file mode 100644 index 00000000..a7947853 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm/tile/tile_mmad.hpp @@ -0,0 +1,104 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_TILE_MMAD_HPP +#define CATLASS_GEMM_TILE_TILE_MMAD_HPP + +#include "../../../gmm_infra/base_defs.hpp" +#include "../../../gmm_infra/gemm/helper.hpp" +namespace Catlass::Gemm::Tile { + +/////////////////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag_, + /// GemmType for A matrix operand + class AType_, + /// GemmType type for B matrix operand + class BType_, + /// GemmType type for Bias operand + class BiasType_ +> +struct TileMmad { + using ElementA = typename AType_::Element; + using ElementB = typename BType_::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + // Methods + + CATLASS_DEVICE + TileMmad() {} + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + AscendC::LocalTensor const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + l0BiasTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_TILE_MMAD_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm_coord.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm_coord.hpp new file mode 100644 index 00000000..d1398169 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/gemm_coord.hpp @@ -0,0 +1,158 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_COORD_HPP +#define CATLASS_GEMM_COORD_HPP + +#include "../gmm_infra/coord.hpp" + +namespace Catlass { + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + uint32_t M_ = 1, + /// Columns of matrix product + uint32_t N_ = 1, + /// Inner dimension of matrix product + uint32_t K_ = 1 +> +struct GemmShape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + static constexpr int64_t MN = M * N; + static constexpr int64_t MK = M * K; + static constexpr int64_t KN = N * K; + static constexpr int64_t MNK = M * N * K; + + static constexpr int64_t COUNT = MNK; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMN() + { + return MakeCoord(M, N); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMK() + { + return MakeCoord(M, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordKN() + { + return MakeCoord(K, N); + } +}; + +/// GemmCoord is a structure derived from Coord<3> that specifies a location within the +/// coordinate space of a Gemm problem. +struct GemmCoord : public Coord<3, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// Gemm M dimension - rows of the output C matrix + static constexpr int M_INDEX = 0; + + /// Gemm N dimension - columns of the output C matrix + static constexpr int N_INDEX = 1; + + /// Gemm K dimension - inner dimension of the Gemm problem + static constexpr int K_INDEX = 2; + + /// Default ctor + CATLASS_HOST_DEVICE + GemmCoord() {} + + /// Constructs from Coord<3> and a batch + CATLASS_HOST_DEVICE + GemmCoord(Coord<3, Index> const &coord) : Base(coord) {} + + /// Helper to construct from a K, N, M, batch variables + CATLASS_HOST_DEVICE + GemmCoord(Index m, Index n, Index k) : Base(MakeCoord(m, n, k)) {} + + /// Returns the Gemm M coordinate + CATLASS_HOST_DEVICE + Index const &m() const + { + return this->At(M_INDEX); + } + + /// Returns reference to the Gemm M coordinate + CATLASS_HOST_DEVICE + Index &m() + { + return this->At(M_INDEX); + } + + /// Returns the Gemm N coordinate + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + + /// Returns reference to the Gemm N coordinate + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + /// Returns the Gemm K coordinate + CATLASS_HOST_DEVICE + Index const &k() const + { + return this->At(K_INDEX); + } + + /// Returns reference to the Gemm K coordinate + CATLASS_HOST_DEVICE + Index &k() + { + return this->At(K_INDEX); + } + + CATLASS_HOST_DEVICE + auto GetCoordMN() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordMK() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordKN() const + { + return this->GetCoordByAxis(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_GEMM_COORD_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/layout.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/layout.hpp new file mode 100644 index 00000000..1649063e --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/layout.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_LAYOUT_LAYOUT_HPP +#define CATLASS_LAYOUT_LAYOUT_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/layout/matrix.hpp" +#include "../../gmm_infra/layout/vector.hpp" + +#endif // CATLASS_LAYOUT_LAYOUT_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/matrix.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/matrix.hpp new file mode 100644 index 00000000..1dbfd60a --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/matrix.hpp @@ -0,0 +1,1204 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_LAYOUT_MATRIX_HPP +#define CATLASS_LAYOUT_MATRIX_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/coord.hpp" +#include "../../gmm_infra/detail/alignment.hpp" +#include "../../gmm_infra/matrix_coord.hpp" + +namespace Catlass::layout { + +/// Mapping function for row-major matrices +struct RowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(cols), LongIndex(1))) {} + + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(ldm, LongIndex(1))) {} + + /// Ctor + CATLASS_HOST_DEVICE + RowMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + template + CATLASS_HOST_DEVICE + static RowMajor MakeLayoutInUb(MatrixCoord const &shape) + { + return RowMajor(shape.row(), shape.column(), RoundUp(shape.column())); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) * stride_[0] + LongIndex(coord.column()); + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + RowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return RowMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for col-major matrices +struct ColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), LongIndex(rows))) {} + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), ldm)) {} + + /// Ctor + CATLASS_HOST_DEVICE + ColumnMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) + LongIndex(coord.column()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + ColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return ColumnMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal +struct nZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr + nZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr + nZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr + static nZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nZ(orgRows, + orgCols, + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + 1, + colsRound * ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + nZ GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2)) + ); + return nZ(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and col-major between fractal +struct zN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr + zN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr + zN(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr + static zN MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zN(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * ELE_NUM_PER_C0); + } + + CATLASS_HOST_DEVICE + static zN MakeLayoutInL0C(MatrixCoord const &shape) + { + return zN(shape.row(), + shape.column(), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.row()), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.column()), + C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL * C0_NUM_PER_FRACTAL, + 1, + RoundUp(shape.row()) * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + zN GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2)) + ); + return zN(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and row-major between fractal +struct zZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr + zZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr + zZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr + static zZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zZ(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + colsRound * C0_NUM_PER_FRACTAL, + 1, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding rowmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is row-major within blocks and also row-major between blocks. +struct PaddingRowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingRowMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) : + orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)blockCols, (LongIndex)blockRows * (LongIndex)RoundUp(orgCols, blockCols), + (LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols)) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + + (LongIndex)coord.column() / blockCols * stride_[3] + + (LongIndex)coord.row() % blockRows * stride_[0] + + (LongIndex)coord.column() % blockCols; + } + + CATLASS_HOST_DEVICE + PaddingRowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingRowMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding columnmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is column-major within blocks and also column-major between blocks. +struct PaddingColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingColumnMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) : + orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols, (LongIndex)blockRows, + (LongIndex)RoundUp(orgRows, blockRows) * (LongIndex)blockCols)) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + + (LongIndex)coord.column() / blockCols * stride_[3] + + (LongIndex)coord.row() % blockRows + + (LongIndex)coord.column() % blockCols * stride_[2]; + } + + CATLASS_HOST_DEVICE + PaddingColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingColumnMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/////////////////////// +// new add layout nN +// nN layout +struct nN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + nN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) { + } + + /// Ctor + CATLASS_HOST_DEVICE + nN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE static nN MakeLayout(Index orgRows, Index orgCols) { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nN(orgRows, + orgCols, + + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + + 1, + ELE_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const& coord) const { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index& orgShape(int idx) { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape& shape() { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index& shape(int idx) { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride& stride() { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index& stride(int idx) { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; +} // namespace Catlass::layout + +#endif // CATLASS_LAYOUT_MATRIX_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/vector.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/vector.hpp new file mode 100644 index 00000000..5eba66e1 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/layout/vector.hpp @@ -0,0 +1,132 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_LAYOUT_VECTOR_HPP +#define CATLASS_LAYOUT_VECTOR_HPP + +#include "../../gmm_infra/base_defs.hpp" +#include "../../gmm_infra/coord.hpp" + +namespace Catlass::layout { + +struct VectorLayout { +public: + /// Logical rank of tensor + static constexpr int RANK = 1; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Shape vector + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + + /// Logical coordinate + using TensorCoord = Coord; + +public: + // Methods + + CATLASS_HOST_DEVICE + VectorLayout(Index size = 0) : shape_(MakeCoord(size)), stride_(MakeCoord(LongIndex(1))) {} + + CATLASS_HOST_DEVICE + VectorLayout(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + template + CATLASS_HOST_DEVICE + static VectorLayout MakeLayoutInUb(TensorCoord const &tileShape) + { + return VectorLayout{RoundUp(tileShape[0])}; + } + + CATLASS_HOST_DEVICE + LongIndex GetOffset(TensorCoord const &coord) const + { + return stride_[0] * coord[0]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + VectorLayout GetTileLayout(TensorCoord const &tileShape) const + { + return VectorLayout(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Stride data member + Shape shape_; + Stride stride_; +}; + +} // namespace Catlass::layout + +#endif // CATLASS_LAYOUT_VECTOR_HPP \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/matrix_coord.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/matrix_coord.hpp new file mode 100644 index 00000000..30682ade --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/matrix_coord.hpp @@ -0,0 +1,103 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_MATRIX_COORD_HPP +#define CATLASS_MATRIX_COORD_HPP + +#include "../gmm_infra/coord.hpp" + +namespace Catlass { + +template < + uint32_t ROW_ = 1, + uint32_t COLUMN_ = 1 +> +struct MatrixShape { + static constexpr uint32_t ROW = ROW_; + static constexpr uint32_t COLUMN = COLUMN_; + + static constexpr int64_t COUNT = ROW * COLUMN; + + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(ROW, COLUMN); + } +}; + +/// MatrixCoord wraps Coord<2, uint32_t> to provide a helper for accessing named dimensions. Classes +/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord. +struct MatrixCoord : public Coord<2, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Rows dimension + static constexpr uint32_t ROW_INDEX = 0; + + /// Columns dimension + static constexpr uint32_t COLUMN_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + MatrixCoord() {} + + /// Constructs from Coord<2> + CATLASS_HOST_DEVICE + MatrixCoord(Coord<2, Index> const &coord) : Base(coord) {} + + /// Helper to construct from a row and column + CATLASS_HOST_DEVICE + MatrixCoord(Index row, Index column) : Base(MakeCoord(row, column)) {} + + /// Helper to construct from a row and column, which are LongIndex based + CATLASS_HOST_DEVICE + MatrixCoord(LongIndex row, LongIndex column) : Base(MakeCoord(Index(row), Index(column))) {} + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index const &row() const { return this->At(ROW_INDEX); } + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index &row() { return this->At(ROW_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index const &column() const { return this->At(COLUMN_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index &column() { return this->At(COLUMN_INDEX); } + + /// Element-wise addition + CATLASS_HOST_DEVICE + MatrixCoord operator+(Base const &b) const + { + return MatrixCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + MatrixCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +} // namespace Catlass + +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/status.hpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/status.hpp new file mode 100644 index 00000000..449b6b56 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/gmm_infra/status.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_STATUS_HPP +#define CATLASS_STATUS_HPP + +namespace Catlass{ + +enum class Status{ kSuccess, kInvalid }; + +} // namespace Catlass + +#endif // CATLASS_STATUS_HPP diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul.h new file mode 100644 index 00000000..f903e8fb --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul.h @@ -0,0 +1,559 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_H +#define ASCENDC_GROUPED_MATMUL_H + +#include "grouped_matmul_utils.h" + +namespace GROUPED_MATMUL { + +constexpr uint32_t thresholdBlockNum = 8; // 8 is obtained by tests, indicating the threshold of basic block numbers + // in both directions when assigning data blocks to cube cores when using + // diagnal strategy +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 +constexpr uint32_t thresholdDimM = 1; // not needs any special strategies +#else +constexpr uint32_t thresholdDimM = 5; // 5 is obtained by tests, indicating the threshold for distinguishing + // strategies for large/small shapes +#endif + +/*@brief store variables for core split configuration +*/ +struct MNConfig { + uint32_t m = 0; + uint32_t k = 0; + uint32_t n = 0; + uint32_t baseM = 0; + uint32_t baseN = 0; + uint32_t mIdx = 0; + uint32_t nIdx = 0; + uint32_t vecNIdx = 0; // for A8W4 MSD NEW + uint32_t blockDimM = 0; + uint32_t blockDimN = 0; + uint32_t vecBlockDimN = 0; // for A8W4 MSD NEW + uint32_t singleM = 0; + uint32_t singleN = 0; + uint32_t vecSingleN = 0; // for A8W4 MSD NEW + uint32_t offsetM = 0; // for A8W4 MSD + uint64_t wBaseOffset = 0; + uint64_t nAxisBaseOffset = 0; + uint64_t mAxisBaseOffset = 0; + uint64_t xBaseOffset = 0; + uint64_t yBaseOffset = 0; + uint64_t wOutOffset = 0; + uint64_t workSpaceOffset = 0; + int64_t scaleIndex = -1; +}; + +template +__aicore__ inline void DataCopyPad2D(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t fullDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (fullDim0 - dim0) * sizeof(T); + params.dstStride = Ceil(dim0 * sizeof(T), UB_BLOCK_DOUBLE_UNIT_SIZE) * 2 - \ + Ceil(dim0 * sizeof(T), UB_BLOCK_UNIT_SIZE); + + DataCopyPadExtParams padParams; + padParams.isPad = true; + padParams.rightPadding = 0; + padParams.leftPadding = 0; + padParams.paddingValue = 0; + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2D(const GlobalTensor dst, const LocalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcFullDim0, uint32_t dstFullDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = static_cast((srcFullDim0 - dim0) * sizeof(T) / UB_BLOCK_UNIT_SIZE); + params.dstStride = (dstFullDim0 - dim0) * sizeof(T); + DataCopyPad(dst, src, params); +} + +__aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN) { + if (mnConfig.blockDimM <= thresholdDimM || thresholdDimM == 1) { + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; + } else { + uint32_t relativeBlock = curBlock - count; + uint32_t curThresholdM = relativeBlock >= AlignDown(mnConfig.blockDimM * mnConfig.blockDimN, thresholdM_dimN) ? + mnConfig.blockDimM % thresholdBlockNum : thresholdBlockNum; + uint32_t curThresholdM_thresholdN = curThresholdM * thresholdBlockNum; + uint32_t curThresholdN = relativeBlock % thresholdM_dimN >= AlignDown(curThresholdM * mnConfig.blockDimN, + curThresholdM_thresholdN) ? mnConfig.blockDimN % thresholdBlockNum : thresholdBlockNum; + + uint32_t localRelativeBlock = relativeBlock % thresholdM_dimN % curThresholdM_thresholdN; + mnConfig.mIdx = localRelativeBlock % curThresholdM + relativeBlock / thresholdM_dimN * thresholdBlockNum; + mnConfig.nIdx = (localRelativeBlock + localRelativeBlock / + LeastCommonMultiple(curThresholdM, curThresholdN)) % curThresholdN + relativeBlock % + thresholdM_dimN / curThresholdM_thresholdN * thresholdBlockNum; + } +} + +/** @brief GroupMatmul operator Class +*/ +template +class GMMProcess { + protected: + using B = typename ComputeType::B; + ComputeType& computeOp; // inernal computation operator + const GMMBaseParams* __restrict gmmBaseParams; + const TCubeTiling* __restrict mmTilingData; + + uint32_t blockIdx; + uint32_t coreIdx; + uint32_t groupNum; + int32_t preOffset = 0; + GM_ADDR groupListPtr; + GlobalTensor groupListGm; + TILING_TYPE* mListGm; + TILING_TYPE* kListGm; + TILING_TYPE* nListGm; + uint32_t baseM_ = 0; + uint32_t baseN_ = 0; + + public: + /** @brief constructor */ + __aicore__ inline GMMProcess(ComputeType& computeOp_) : computeOp(computeOp_) {} + + __aicore__ inline void Init(const GMMBaseParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn, + GM_ADDR groupList, GM_ADDR tiling); + + __aicore__ inline void InitStaticTiling(int32_t baseM, int32_t baseN); + + __aicore__ inline void Process(); + + bool isA8W4FakeQuant = false; + + protected: + __aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); +}; + +template +__aicore__ inline void GMMProcess::Init(const GMMBaseParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn, GM_ADDR groupList, GM_ADDR tiling) { + blockIdx = GetBlockIdx(); + coreIdx = blockIdx; + int64_t coreRation = GetTaskRation(); + if (coreRation > 1) { + coreIdx /= coreRation; + } + gmmBaseParams = gmmBaseParamsIn; + mmTilingData = mmTilingDataIn; + groupNum = gmmBaseParams->groupNum; + groupListPtr = groupList; + if (groupListPtr != nullptr) { + groupListGm.SetGlobalBuffer((__gm__ int64_t*)groupList); + } + mListGm = gmmArrayAddrIn; + kListGm = gmmArrayAddrIn + MKN_LIST_LEN; + nListGm = gmmArrayAddrIn + MKN_LIST_LEN * 2; +} + +template +__aicore__ inline void GMMProcess::InitStaticTiling(int32_t baseM, int32_t baseN) { + baseM_ = static_cast(baseM); + baseN_ = static_cast(baseN); +} + +template +__aicore__ inline void GMMProcess::SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig) { + SetMKN(splitValue, groupIdx, mnConfig); + if (mmTilingData != nullptr) { + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + } else { + mnConfig.baseM = baseM_; + mnConfig.baseN = baseN_; + } + mnConfig.singleM = mnConfig.baseM; + mnConfig.singleN = mnConfig.baseN; +#if defined(GMM_QUANT_BF16) || defined(GMM_QUANT_FLOAT16) || defined(GMM_FLOAT) + if (gmmBaseParams->singleN > 0) { // not sequential write + mnConfig.singleN = gmmBaseParams->singleN; + } +#endif +} + +template +__aicore__ inline void GMMProcess::SetMKN(const int32_t splitValue, const uint32_t groupIdx, + MNConfig &mnConfig) { + uint32_t singleWeight = gmmBaseParams->singleWeight; + uint32_t singleX = gmmBaseParams->singleX; + uint32_t singleY = gmmBaseParams->singleY; + bool isAllSingleTensor = singleWeight == 1 && singleX == 1 && singleY == 1; + uint32_t valueIdx = isAllSingleTensor ? 0 : groupIdx; + if (mmTilingData == nullptr) { + mnConfig.m = splitValue; + mnConfig.k = gmmBaseParams->k; + mnConfig.n = gmmBaseParams->n; + return; + } + + if (gmmBaseParams->groupType == 0) { + mnConfig.m = splitValue; + mnConfig.k = kListGm[valueIdx]; + mnConfig.n = nListGm[valueIdx]; + return; + } + + if (gmmBaseParams->groupType == 2) { + mnConfig.m = mListGm[valueIdx]; + mnConfig.k = splitValue; + mnConfig.n = nListGm[valueIdx]; + return; + } + + mnConfig.m = mListGm[groupIdx]; + mnConfig.k = kListGm[groupIdx]; + mnConfig.n = nListGm[groupIdx]; + return; +} + +template +__aicore__ inline void GMMProcess::UpdateMnConfig(MNConfig &mnConfig) { + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.scaleIndex++; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMProcess::Process() { + MNConfig mnConfig; + if (gmmBaseParams->groupType != -1) { // -1: no split + if (unlikely(groupListPtr == nullptr)) { + return; + } + preOffset = 0; + } + AscendC::WaitPreTaskEnd(); + for (uint32_t groupIdx = 0, count = 0; groupIdx < groupNum; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm); + SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = coreIdx >= count ? coreIdx : coreIdx + gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + computeOp.MMCompute(groupIdx, mnConfig, coreIdx); + computeOp.VectorCompute(mnConfig); + curBlock += gmmBaseParams->coreNum; + } + count = curCount % gmmBaseParams->coreNum; + } + computeOp.PostCompute(); + AscendC::SetNextTaskStart(); +} + +/** @brief GroupMatmul GroupType M sparse operator Class +*/ +template +class GMMGroupMSparseProcess : public GMMProcess { +public: + /** @brief constructor */ + __aicore__ inline GMMGroupMSparseProcess(ComputeType& computeOp_) : GMMProcess(computeOp_) {} + + __aicore__ inline void Process() + { + if (this->gmmBaseParams->groupType != -1) { // -1: no split + if (unlikely(this->groupListPtr == nullptr)) { + return; + } + } + + MNConfig mnConfig; + uint32_t groupListInnerShape = 2u; // shape: [e, 2] + uint32_t groupListShapeSize = this->groupNum * groupListInnerShape; + AscendC::WaitPreTaskEnd(); + for (uint32_t loop = 0, count = 0; loop < groupListShapeSize; loop += groupListInnerShape) { + int32_t splitValue = static_cast(this->groupListGm.GetValue(loop + 1)); + if (splitValue <= 0) { + break; + } + + uint32_t groupIdx = static_cast(this->groupListGm.GetValue(loop)); + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; + this->SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.nAxisBaseOffset = groupIdx * mnConfig.n; + if constexpr (GMMProcess::B::format == CubeFormat::NZ) { + // 16: nz format last two dim size + mnConfig.wBaseOffset = AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.nAxisBaseOffset); + } else { + mnConfig.wBaseOffset = mnConfig.k * mnConfig.nAxisBaseOffset; + } + + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = this->coreIdx >= count ? this->coreIdx : this->coreIdx + this->gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx); + this->computeOp.VectorCompute(mnConfig); + curBlock += this->gmmBaseParams->coreNum; + } + count = curCount % this->gmmBaseParams->coreNum; + } + + this->computeOp.PostCompute(); + AscendC::SetNextTaskStart(); + } +}; + + +/** @brief intenal computation class +*/ +template +class GMMCompute { + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = DTYPE_WEIGHT; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + bool isA8W4FakeQuant = false; + + /** @brief constructor */ + __aicore__ inline GMMCompute(typename mmType::MT& mm_) : mm(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR offset, + GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe); + + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx); + + __aicore__ inline void VectorCompute(MNConfig& mnConfig) {} + + __aicore__ inline void PostCompute() {} + + protected: + __aicore__ inline void SetGlobalBufferBias(uint32_t groupIdx, uint32_t tailN, const MNConfig mnConfig); + + __aicore__ inline GlobalTensor SetGlobalBufferW(uint32_t groupIdx, uint32_t tailN, MNConfig& mnConfig); + + __aicore__ inline uint64_t SetWOffset(uint32_t tailN, uint32_t k); + + protected: + TPipe* pipe; + typename mmType::MT& mm; // matmul operator + bool hasBias = false; + GM_ADDR xTensorPtr; + GM_ADDR weightTensorPtr; + GM_ADDR biasTensorPtr; + GM_ADDR yTensorPtr; + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor biasGm; + GlobalTensor yGm; +#if defined(GMM_QUANT_INT8) + GM_ADDR scaleTensorPtr; + GlobalTensor scaleGm; +#endif + uint32_t ubBaseN; + uint32_t ubBaseK; + uint32_t ubCalSize; + uint32_t singleWeight; + uint32_t singleX; + uint32_t singleY; + uint32_t coreNum; + uint32_t subBlockIdx; + bool mmWaitStatus; + uint32_t activeType; +}; + +template +__aicore__ inline void GMMCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR offset, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, + GM_ADDR groupList, GM_ADDR perTokenScale, GM_ADDR y, + GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, + TPipe* tPipe) { + xTensorPtr = x; + weightTensorPtr = weight; + biasTensorPtr = bias; + yTensorPtr = y; + pipe = tPipe; + ubBaseN = gmmBaseParams->ubBaseN; + ubBaseK = gmmBaseParams->ubBaseK; + ubCalSize = gmmBaseParams->ubCalSize; + singleWeight = gmmBaseParams->singleWeight; + singleX = gmmBaseParams->singleX; + singleY = gmmBaseParams->singleY; + coreNum = gmmBaseParams->coreNum; + subBlockIdx = GetSubBlockIdx(); + if (mmTilingData != nullptr) { + hasBias = mmTilingData->isBias != 0; + } + activeType = gmmBaseParams->activeType; + mmWaitStatus = false; +#if defined(GMM_QUANT_INT8) + scaleTensorPtr = scale; +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 + TBuf<> ubBuf; + pipe->InitBuffer(ubBuf, TOTAL_UB_SIZE / 2); + LocalTensor buf = ubBuf.template Get(); + mm.SetLocalWorkspace(buf); +#endif +} + +template +__aicore__ inline void GMMCompute::SetGlobalBufferBias(uint32_t groupIdx, + uint32_t tailN, const MNConfig mnConfig) { + if (hasBias) { + if (singleWeight == 0) { + biasGm.SetGlobalBuffer(GetTensorAddr(groupIdx, biasTensorPtr)); + } else { + biasGm.SetGlobalBuffer(GetTensorAddr(0, biasTensorPtr) + mnConfig.nAxisBaseOffset); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + constexpr bool isBiasEpilogue_ = + AscendC::IsSameType::value && + (AscendC::IsSameType::value || AscendC::IsSameType::value); + if constexpr (!isBiasEpilogue_) { + mm.SetBias(biasGm[tailN]); + } +#else + mm.SetBias(biasGm[tailN]); +#endif + } +} + +template +__aicore__ inline uint64_t GMMCompute::SetWOffset(uint32_t tailN, uint32_t k) { + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ && transposeW) { + wOffset = tailN * (UB_BLOCK_UNIT_SIZE / sizeof(BT)); // 32: quant is 32, float16 is 16 + } else if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else if constexpr (transposeW) { + wOffset = tailN * k; + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline GlobalTensor GMMCompute::SetGlobalBufferW( + uint32_t groupIdx, uint32_t tailN, MNConfig& mnConfig) { + uint64_t wOffset = SetWOffset(tailN, mnConfig.k); +#if defined(GMM_ANTI_QUANT) && !defined(GMM_ANTI_QUANT_A8W4_MSD) + return weightGm[transposeW ? mnConfig.workSpaceOffset - tailN + wOffset : mnConfig.workSpaceOffset]; +#else + GlobalTensor weightGmLocal; + if (singleWeight == 0) { + weightGmLocal.SetGlobalBuffer(GetTensorAddr(groupIdx, weightTensorPtr) + wOffset); + } else if (isA8W4FakeQuant) { + weightGmLocal.SetGlobalBuffer(reinterpret_cast<__gm__ BT *>(weightTensorPtr) + mnConfig.wBaseOffset + wOffset); + } else { + weightGmLocal.SetGlobalBuffer(GetTensorAddr(0, weightTensorPtr) + mnConfig.wBaseOffset + wOffset); + } + #if !(defined(ASCENDC_OOM) && ASCENDC_OOM == 1) + if (mnConfig.blockDimM == 1) { + weightGmLocal.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + #endif + return weightGmLocal; +#endif +} + +template +__aicore__ inline void GMMCompute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx) { + if (subBlockIdx != 0) { + return; + } + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM + : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + // init global buffer + if (singleX == 0) { + xGm.SetGlobalBuffer(GetTensorAddr(groupIdx, xTensorPtr)); + } else { + xGm.SetGlobalBuffer(GetTensorAddr(0, xTensorPtr) + mnConfig.xBaseOffset); + } + GlobalTensor weightGmLocal = SetGlobalBufferW(groupIdx, tailN, mnConfig); + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(xGm[xOffset], transposeX); + mm.SetTensorB(weightGmLocal, transposeW); +#if defined(GMM_QUANT_INT8) + if (singleWeight == 0) { + scaleGm.SetGlobalBuffer(GetTensorAddr(groupIdx, scaleTensorPtr)); + } else { + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scaleTensorPtr) + mnConfig.nAxisBaseOffset); + } + mm.SetQuantVector(scaleGm[tailN]); +#endif + SetGlobalBufferBias(groupIdx, tailN, mnConfig); + if (singleY == 0) { + yGm.SetGlobalBuffer(GetTensorAddr(groupIdx, yTensorPtr)); + } else { + yGm.SetGlobalBuffer(GetTensorAddr(0, yTensorPtr) + mnConfig.yBaseOffset); + } + #if defined(GMM_ANTI_QUANT) + mm.template IterateAll(yGm[outOffset], 0, false, true); + mmWaitStatus = true; + #else + mm.template IterateAll(yGm[outOffset], 0); + #endif +} + +} // namespace GROUPED_MATMUL +#endif // ASCENDC_GROUPED_MATMUL_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_a4w4.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_a4w4.h new file mode 100644 index 00000000..719e7a2b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_a4w4.h @@ -0,0 +1,354 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_a4w4.h + * \brief + */ + + #ifndef ASCENDC_GROUPED_MATMUL_A4W4_H + #define ASCENDC_GROUPED_MATMUL_A4W4_H + + #include "grouped_matmul_utils.h" + #include "grouped_matmul.h" + + #ifdef GMM_A4W4 + namespace GROUPED_MATMUL{ + using namespace matmul; + using namespace AscendC; + + using DTYPE_PERTOKEN_SCALE_A4W4 = float; + + #ifdef GMM_A4W4_BF16 + using DTYPE_Y_A4W4 = bfloat16_t; + #else + using DTYPE_Y_A4W4 = half; + #endif + + using DTYPE_X_A4W4 = int4b_t; + using DTYPE_WEIGHT_A4W4 = int4b_t; + using DTYPE_SCALE_A4W4 = uint64_t; + + constexpr uint64_t SYNC_AIV_TO_AIC = 3; + constexpr uint64_t SYNC_AIC_TO_AIV = 5; + constexpr uint32_t BUFFER_NUM = 2; + constexpr int32_t PARALL_NUM = 4; + + template +__aicore__ inline void DataCopyPad2DA4W4(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + + template + class GMMA4W4Compute { + public: + using aT = MatmulType; + using bT = typename mmType::BT; + using biasT = MatmulType; + using cT = MatmulType; + using DTYPE_OUT = DTYPE_Y_A4W4; + + public: + __aicore__ inline GMMA4W4Compute(typename mmType::MT &matmul) : mm(matmul) {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe); + __aicore__ inline void Process(); + private: + __aicore__ inline void InitUbBuffer(); + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorTilingCalc(MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, uint32_t& vecBaseM); + __aicore__ inline void ComputeDequantAndActivate(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t alignBaseN, + uint32_t curVecBaseN, uint32_t offsetM); + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM); + + private: + typename mmType::MT& mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor mmOutGm; + GlobalTensor scaleGm; + GlobalTensor perTokenScaleGm; + GlobalTensor groupListGm; + GlobalTensor yGm; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue perTokenScaleInQueue; + TBuf tmpBuff; + LocalTensor mmOutFp32Buf; + LocalTensor pertokenBrcbLocal; + LocalTensor perTokenResBuf; + LocalTensor calcTmpBuf; + uint32_t subBlockIdx; + uint32_t coreIdx; + uint32_t quantGroupSize_; + uint32_t cubeCount = 0; + uint32_t vecCount_ = 0; + uint32_t mmBaseBlockOffset_ = 0; + TPipe *pipe; + const GMMBaseParams *tiling; + const TCubeTiling* mmTilingData; + }; + + template + __aicore__ inline void GMMA4W4Compute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe) + { + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + weightGm.SetGlobalBuffer(GetTensorAddr(0, weight)); + mmOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ cT::T *>(workspace)); + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scale)); + perTokenScaleGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PERTOKEN_SCALE_A4W4 *>(perTokenScale)); + groupListGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(groupList)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + + this->tiling = gmmBaseParams; + this->mmTilingData = mmTilingData; + mmBaseBlockOffset_ = mmTilingData->baseM * mmTilingData->baseN; + + quantGroupSize_ = tiling->k / tiling->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); // 0/1 + coreIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + if (GetTaskRation() != 0) { + coreIdx /= GetTaskRation(); + } + } + pipe = tPipe; + InitUbBuffer(); + } + + template + __aicore__ inline void GMMA4W4Compute::InitUbBuffer() + { + if ASCEND_IS_AIC { + return; + } + pipe->InitBuffer(perTokenScaleInQueue, BUFFER_NUM, mmTilingData->baseM * sizeof(float)); + pipe->InitBuffer(vecInQueue, BUFFER_NUM, tiling->ubCalSize * sizeof(cT::T)); + + pipe->InitBuffer(vecOutQueue, BUFFER_NUM, tiling->ubCalSize * sizeof(DTYPE_OUT)); + + pipe->InitBuffer(tmpBuff, tiling->ubRestBytes); + + uint32_t ubCalSizeFloat = tiling->ubCalSize * sizeof(float); + mmOutFp32Buf = tmpBuff.GetWithOffset(tiling->ubCalSize, 0); + uint32_t offset = ubCalSizeFloat; + pertokenBrcbLocal = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + offset += ubCalSizeFloat; + perTokenResBuf = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + offset += ubCalSizeFloat; + calcTmpBuf = tmpBuff.GetWithOffset(ubCalSizeFloat, offset); + offset += ubCalSizeFloat; + } + + template + __aicore__ inline void GMMA4W4Compute::Process() + { + MNConfig mnConfig; + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + mnConfig.singleM = mnConfig.baseM; + mnConfig.singleN = mnConfig.baseN; + mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN); + int32_t preOffset = 0; + for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) { + int32_t m = GetSplitValueFromGroupList(groupIdx, preOffset, tiling, groupListGm); + if (m <= 0) { + continue; + } + mnConfig.m = static_cast(m); + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + tiling->coreNum; + uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, preCount, thresholdM_dimN); + MMCompute(groupIdx, mnConfig); + if ASCEND_IS_AIV { + VectorCompute(groupIdx, mnConfig); + } + curBlock += tiling->coreNum; + } + preCount = curCount % tiling->coreNum; + mnConfig.offsetM += mnConfig.m; + } + } + + template +__aicore__ inline void GMMA4W4Compute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + mnConfig.workSpaceOffset = mmBaseBlockOffset_ * \ + (coreIdx + (cubeCount % PARALL_NUM) * tiling->coreNum); + if ASCEND_IS_AIC { + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = tiling->n - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + + uint64_t xOffset = (mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM) * tiling->k; + uint64_t weightOffset; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = groupIdx * tiling->n * tiling->k + tailN * tiling->k; + } else { + weightOffset = groupIdx * tiling->n * tiling->k + tailN; + } + if (cubeCount >= PARALL_NUM) { + CrossCoreWaitFlag(SYNC_AIV_TO_AIC); + } + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize_); + GlobalTensor weightSlice; + for (uint32_t loopK = 0; loopK < tiling->quantGroupNum; loopK++) { + mm.SetTensorA(xGm[xOffset + loopK * quantGroupSize_]); + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize_ * 64]; + } else { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize_ * tiling->n]; + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + mm.SetQuantVector(scaleGm[groupIdx * tiling->n * tiling->quantGroupNum + loopK * tiling->n + tailN]); + uint64_t worskspaceOffset = mnConfig.workSpaceOffset; + #ifndef __CCE_KT_TEST__ + mm.Iterate(); + mm.GetTensorC(mmOutGm[worskspaceOffset], loopK == 0 ? 0 : 1, true); + #endif + worskspaceOffset += mmBaseBlockOffset_; + } + CrossCoreSetFlag<2, PIPE_FIX>(SYNC_AIC_TO_AIV); // 2: mode为2, group内同步 + } + cubeCount++; +} + +template +__aicore__ inline void GMMA4W4Compute::VectorTilingCalc( + MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, uint32_t& vecBaseM) +{ + curCubeSingleN = mnConfig.nIdx == mnConfig.blockDimN - 1 ? + tiling->n - mnConfig.nIdx * mnConfig.singleN : mnConfig.singleN; + curCubeSingleM = mnConfig.mIdx == mnConfig.blockDimM - 1 ? + mnConfig.m - mnConfig.mIdx * mnConfig.singleM : mnConfig.singleM; + vecBaseM = tiling->ubCalSize / (Ceil(mnConfig.baseN, uint32_t(8)) * 8); // 8: num int32_t in 32B ub block 32*256/256 + vecBaseM = vecBaseM < curCubeSingleM ? vecBaseM : curCubeSingleM; +} + +template +__aicore__ inline void GMMA4W4Compute::VectorCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + uint32_t curCubeSingleN; + uint32_t curCubeSingleM; + uint32_t vecBaseM; + VectorTilingCalc(mnConfig, curCubeSingleN, curCubeSingleM, vecBaseM); + uint32_t mGlobalOffset = mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM; // 2: 2 lines int4 to 1 line int8 + + uint64_t outOffset = mGlobalOffset * tiling->n + mnConfig.nIdx * mnConfig.singleN; + uint32_t curVecBaseN = mnConfig.baseN; + uint32_t taskRation = GetTaskRation(); // 2 + CrossCoreWaitFlag(SYNC_AIC_TO_AIV); + uint32_t nCount = 0; + for (uint32_t offsetN = 0; offsetN < curCubeSingleN; offsetN += mnConfig.baseN) { + if (unlikely(offsetN + mnConfig.baseN >= curCubeSingleN)) curVecBaseN = curCubeSingleN - offsetN; + uint32_t alignBaseN = Ceil(curVecBaseN, uint32_t(16)) * 16; // 16: fp16 num per 32B + uint32_t curVecBaseM = vecBaseM; + uint64_t mmOutOffset = mnConfig.workSpaceOffset + offsetN * mnConfig.baseM; + uint32_t mCount = 0; + for (uint32_t offsetM = 0; offsetM < curCubeSingleM; offsetM += vecBaseM) { + vecCount_++; + if (taskRation != 0 && vecCount_ % taskRation != subBlockIdx) { continue; } + if (unlikely(offsetM + vecBaseM >= curCubeSingleM)) { curVecBaseM = curCubeSingleM - offsetM; } + LocalTensor mmOutLocal = vecInQueue.AllocTensor(); + DataCopyPad2DA4W4(mmOutLocal, mmOutGm[mmOutOffset + offsetM * curVecBaseN], curVecBaseM, curVecBaseN, curVecBaseN); + vecInQueue.EnQue(mmOutLocal); + + ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM); + LocalTensor yLocal = vecOutQueue.DeQue(); + DataCopyPad2D(yGm[outOffset + offsetM * tiling->n + offsetN], yLocal, + curVecBaseM, curVecBaseN, alignBaseN, tiling->n); + vecOutQueue.FreeTensor(yLocal); + } + } + CrossCoreSetFlag<2, PIPE_MTE2>(SYNC_AIV_TO_AIC); // 2: mode为2, group内同步 +} + +template +__aicore__ inline void GMMA4W4Compute::ComputeDequantAndActivate(MNConfig& mnConfig, + uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, uint32_t offsetM) +{ + uint32_t computeSize = curVecBaseM * alignBaseN; + LocalTensor mmOutInUb = vecInQueue.DeQue(); + uint32_t castSize = 0; + if constexpr (mmType::BT::format == CubeFormat::ND) { + castSize = (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + castSize = computeSize; + } + Cast(mmOutFp32Buf, mmOutInUb, RoundMode::CAST_NONE, castSize); + PipeBarrier(); + vecInQueue.FreeTensor(mmOutInUb); + + DataCopyPerTokenScaleAndBrcb(mnConfig, curVecBaseM, alignBaseN, offsetM); + + Mul(perTokenResBuf, mmOutFp32Buf, pertokenBrcbLocal, computeSize); + PipeBarrier(); + LocalTensor yLocalInUb = vecOutQueue.AllocTensor(); + // Cast后获得最终输出 + Cast(yLocalInUb, perTokenResBuf, RoundMode::CAST_RINT, computeSize); + PipeBarrier(); + vecOutQueue.EnQue(yLocalInUb); +} + +template +__aicore__ inline void GMMA4W4Compute::DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, + uint32_t curBaseM, uint32_t alignBaseN, uint32_t offsetM) +{ + uint64_t perTokenScaleOffset = mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM + offsetM; + DataCopyPadExtParams padParams; + DataCopyExtParams perTokenScaleParams{1, static_cast(curBaseM * sizeof(float)), 0, 0, 0}; + + LocalTensor perTokenScaleLocal = perTokenScaleInQueue.AllocTensor(); + DataCopyPad(perTokenScaleLocal, perTokenScaleGm[perTokenScaleOffset], perTokenScaleParams, padParams); + perTokenScaleInQueue.EnQue(perTokenScaleLocal); + + perTokenScaleLocal = perTokenScaleInQueue.DeQue(); + + const uint32_t broadCastDst[2] = {curBaseM, alignBaseN}; + const uint32_t broadCastSrc[2] = {curBaseM, 1}; + BroadCast(pertokenBrcbLocal, perTokenScaleLocal, broadCastDst, broadCastSrc, calcTmpBuf); + perTokenScaleInQueue.FreeTensor(perTokenScaleLocal); +} + + } // namespace GROUPED_MATMUL + #endif + #endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant.h new file mode 100644 index 00000000..58e5f60c --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant.h @@ -0,0 +1,544 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_H + +#include "grouped_matmul.h" + +#ifdef GMM_ANTI_QUANT +namespace GROUPED_MATMUL { + +constexpr uint32_t CAST_THRESHOLD_CACHE_BIG = 16 * 1024 * 1024; // 16M is obtained by tests +constexpr uint32_t CAST_THRESHOLD_CACHE_SMALL = 10 * 1024 * 1024; // 10M is obtained by tests +constexpr uint32_t CAST_PERFORMANCE_MAX_N = 5120; +constexpr uint32_t CAST_MIN_SINGLE_K = 8; +constexpr int32_t BEST_UB_BASEN = 512; + +/*@brief store variables for core split configuration +*/ +struct CastWeightConfig { + uint32_t coreNum = 0; + uint32_t nUsedCore = 0; + uint32_t curDimN = 0; + uint32_t castRoundIdx = 0; + uint32_t workSpaceIdx = 0; + uint64_t wInNOffset = 0; + uint32_t wInKOffset = 0; + uint32_t curSingleN = 0; + uint32_t curSingleK = 0; + uint32_t tailN = 0; +}; + +/** @brief GroupMatmul Antiquant operator Class +*/ +template +class GMMAntiquantProcess : public GMMProcess{ + protected: + constexpr static bool antiquantPerformance = ComputeType::antiquantPerformanceFlag; + public: + /** @brief constructor */ + __aicore__ inline GMMAntiquantProcess(ComputeType& computeOp_) : GMMProcess(computeOp_) {} + + __aicore__ inline void Process(); + + private: + __aicore__ inline void SetAntiquantMNConfig(const uint64_t singleWorkSpaceSize, const uint32_t curBlock, bool& validCore, + CastWeightConfig& castConfig, MNConfig &mnConfig); + + __aicore__ inline void SetAntiquantCastConfig(uint32_t& curCount, MNConfig mnConfig, + CastWeightConfig& castConfig); + __aicore__ inline void AntiquantUpdateSingleM(MNConfig& mnConfig, uint32_t& dimM, uint32_t dimN); +}; + +template +__aicore__ inline void GMMAntiquantProcess::SetAntiquantMNConfig(const uint64_t singleWorkSpaceSize, + const uint32_t curBlock, bool& validCore, CastWeightConfig& castConfig, MNConfig &mnConfig) { + mnConfig.workSpaceOffset = castConfig.workSpaceIdx * singleWorkSpaceSize; + castConfig.workSpaceIdx = castConfig.workSpaceIdx == 0 ? 1 : 0; // next round use another workspace + castConfig.castRoundIdx = Ceil(curBlock + 1, castConfig.coreNum) - 1; // +1: let curBlock start from 1,-1: castRoundIdx start from 0 + castConfig.curDimN = castConfig.nUsedCore; + if (castConfig.castRoundIdx == Ceil(mnConfig.blockDimN, castConfig.nUsedCore) - 1) { // -1 last round + castConfig.curDimN = mnConfig.blockDimN - castConfig.castRoundIdx * castConfig.nUsedCore; + } + // compute dimM + uint32_t dimM = Max(castConfig.coreNum / castConfig.curDimN, 1); // 1: The minimum value of dimM is 1 + dimM = Min(Ceil(mnConfig.m, this->mmTilingData->baseM), dimM); + mnConfig.singleM = Ceil(mnConfig.m, dimM); + mnConfig.blockDimM = dimM; + mnConfig.mIdx = this->coreIdx / castConfig.curDimN; + mnConfig.nIdx = this->coreIdx % castConfig.curDimN; + validCore = this->coreIdx < dimM * castConfig.curDimN; +} + +template +__aicore__ inline void GMMAntiquantProcess::SetAntiquantCastConfig(uint32_t& curCount, + MNConfig mnConfig, + CastWeightConfig& castConfig) { + if (mnConfig.blockDimM > 0 && mnConfig.blockDimN > 0) { + // 16M and 10M is obtained by tests. When N is greater than 5120, the cache uses 10 MB for better performance + uint32_t cacheThreshold = mnConfig.n > CAST_PERFORMANCE_MAX_N ? CAST_THRESHOLD_CACHE_SMALL : CAST_THRESHOLD_CACHE_BIG; + // 16M/k is the length of N that needs to be calculated for single round. + // 16M/k/baseN is the coreNum required for single round calculation of the N-axis. + castConfig.nUsedCore = Min(Ceil(cacheThreshold, mnConfig.k * this->mmTilingData->baseN), castConfig.coreNum); + castConfig.nUsedCore = Min(castConfig.nUsedCore, mnConfig.blockDimN); + curCount = Ceil(mnConfig.blockDimN, castConfig.nUsedCore) * castConfig.coreNum; + } +} + +template +__aicore__ inline void GMMAntiquantProcess::AntiquantUpdateSingleM(MNConfig& mnConfig, + uint32_t& dimM, uint32_t dimN) { + if (dimM > 1 && dimN < this->gmmBaseParams->coreNum) { + uint32_t restCores = this->gmmBaseParams->coreNum / dimN; + if (dimM > restCores) { + mnConfig.singleM = Ceil(mnConfig.m, restCores); + dimM = Ceil(mnConfig.m, mnConfig.singleM); + } + } +} + +template +__aicore__ inline void GMMAntiquantProcess::Process() { + MNConfig mnConfig; + CastWeightConfig castConfig; + castConfig.coreNum = this->gmmBaseParams->coreNum; + bool validCore = true; + uint64_t singleWorkSpaceSize = this->gmmBaseParams->workspaceSize / 2; // 2: antiQuantNormal use 2 block workspace + if (this->gmmBaseParams->groupType != -1) { // -1: no need to split + this->preOffset = 0; + if (unlikely(this->groupListPtr == nullptr)) {this->groupNum = 0;} // not continue Process + } + for (uint32_t groupIdx = 0, count = 0; groupIdx < this->groupNum; ++groupIdx) { + int32_t splitValue = GetSplitValueFromGroupList(groupIdx, this->preOffset, this->gmmBaseParams, this->groupListGm); + this->SetMNConfig(splitValue, groupIdx, mnConfig); + uint32_t dimM = Ceil(mnConfig.m, mnConfig.singleM); + uint32_t dimN = Ceil(mnConfig.n, mnConfig.singleN); + if constexpr (!antiquantPerformance) { + AntiquantUpdateSingleM(mnConfig, dimM, dimN); + } + mnConfig.blockDimM = dimM; + mnConfig.blockDimN = dimN; + uint32_t curCount = count + dimM * dimN; + uint32_t curBlock = this->coreIdx >= count ? this->coreIdx : this->coreIdx + this->gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = thresholdBlockNum * dimN; + + if constexpr (antiquantPerformance) { + SetAntiquantCastConfig(curCount, mnConfig, castConfig); + } + + while (curBlock < curCount) { + if constexpr (antiquantPerformance) { // performance verison, will split dimN + SetAntiquantMNConfig(singleWorkSpaceSize, curBlock, validCore, castConfig, mnConfig); + } else { + mnConfig.workSpaceOffset = mnConfig.wBaseOffset; + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + } + this->computeOp.PreCompute(groupIdx, this->coreIdx, mnConfig, castConfig); + this->computeOp.MMSync(); + if (validCore) { + mnConfig.workSpaceOffset += mnConfig.nIdx * mnConfig.singleN; + if constexpr (antiquantPerformance) { + mnConfig.nIdx += castConfig.castRoundIdx * castConfig.nUsedCore; + } + this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx); + } + curBlock += this->gmmBaseParams->coreNum; + } + this->UpdateMnConfig(mnConfig); + count = curCount % this->gmmBaseParams->coreNum; + } +} + + +/** @brief intenal computation class +*/ +template +class GMMAntiquantCompute : public GMMCompute { + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = DTYPE_WEIGHT; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + constexpr static bool antiquantPerformanceFlag = antiquantPerformance; + + __aicore__ inline GMMAntiquantCompute(typename mmType::MT& mm_) : GMMCompute(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR offset, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR groupList, GM_ADDR perTokenScale, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe); + + __aicore__ inline void PreCompute(uint32_t groupIdx, + uint32_t coreIdx, MNConfig& mnConfig, CastWeightConfig& castConfig); + + __aicore__ inline void MMSync(); + + private: + + __aicore__ inline void CastWeightProcess(MNConfig& mnConfig, CastWeightConfig& castConfig); + __aicore__ inline void SetAntiQuantGlobalBuffer(uint32_t groupIdx, const MNConfig mnConfig); + __aicore__ inline void SetGmToUbDataCopyParams(const uint32_t curBaseN, const uint32_t curBaseK, + const MNConfig& mnConfig, DataCopyExtParams& intriParams); + __aicore__ inline void SetUbToGmDataCopyParams(const uint32_t curBaseN, const uint32_t alignRowLen, + const uint32_t curBaseK, const MNConfig& mnConfig, + DataCopyExtParams& intriParams); + __aicore__ inline void CastWeightCompute(uint32_t curCalcK, uint32_t curCalcAlignN); + __aicore__ inline void DataCopyScaleAndOffset(uint32_t curBaseN, uint32_t alignBaseN, + uint64_t realScaleOffset); + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyPerTokenScale(uint32_t curBaseM, uint64_t perTokenScaleOffset); + __aicore__ inline void PerTokenDequant(uint32_t curBaseM, uint32_t alignBaseN); + __aicore__ inline void SetPerTokenQuantRefreshedBuffer(const MNConfig mnConfig); + __aicore__ inline void ComputeUbBaseK(uint32_t curSingleK, uint32_t offsetK, uint32_t newBaseK, + uint32_t& curUsedGroupSize, uint32_t& curBaseK); + __aicore__ inline void FreeScaleAndOffset(bool& firstLoop); + + GlobalTensor weightAntiQuantGm; + GM_ADDR antiScaleTensorPtr; + GM_ADDR antiOffsetTensorPtr; + LocalTensor scaleInUb; + LocalTensor offsetInUb; + GlobalTensor antiScaleGM; + GlobalTensor antiOffsetGM; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue offsetInQueue; + TBuf tmpBuff; + LocalTensor tmpUb; + bool isPerGroup = false; + uint32_t perGroupSize; +}; + +template +__aicore__ inline void +GMMAntiquantCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, + GM_ADDR offset, GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR groupList, GM_ADDR perTokenScale, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe) { + this->GMMCompute::Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, + perTokenScale, y, workspace, gmmBaseParams, mmTilingData, tPipe); + antiScaleTensorPtr = antiquantScale; + antiOffsetTensorPtr = antiquantOffset; + perGroupSize = gmmBaseParams->quantParam; + isPerGroup = perGroupSize > 0; + this->weightGm.SetGlobalBuffer((__gm__ BT*)workspace); + uint32_t maxUbBaseN = BEST_UB_BASEN; + if constexpr (transposeW) { + maxUbBaseN = this->ubBaseN; + } + // scale should bigger than singleN, 32 alignment is required + this->pipe->InitBuffer(scaleInQueue, 2, maxUbBaseN * sizeof(BT)); + this->pipe->InitBuffer(offsetInQueue, 2, maxUbBaseN * sizeof(BT)); + this->pipe->InitBuffer(vecInQueue, 2, this->ubCalSize * GetTypeBits() / INT8_BITS); + this->pipe->InitBuffer(vecOutQueue, 2, this->ubCalSize * sizeof(BT)); + this->pipe->InitBuffer(tmpBuff, gmmBaseParams->ubRestBytes); + tmpUb = tmpBuff.Get(); +} + +template +__aicore__ inline void GMMAntiquantCompute::PreCompute(uint32_t groupIdx, + uint32_t coreIdx, MNConfig& mnConfig, CastWeightConfig& castConfig) { + if constexpr (!antiquantPerformance) { + if (this->subBlockIdx != 0) { + return; + } + } + castConfig.curSingleN = 0; + castConfig.curSingleK = 0; + castConfig.wInKOffset = 0; + castConfig.wInNOffset = 0; + mnConfig.wOutOffset = mnConfig.workSpaceOffset; + castConfig.tailN = 0; + if constexpr (antiquantPerformance) { // antiquant normal version + uint32_t blockDimK = Min(this->coreNum, Ceil(mnConfig.k, CAST_MIN_SINGLE_K)); + if (coreIdx >= blockDimK) { return; } + castConfig.curSingleK = Ceil(mnConfig.k, blockDimK); + castConfig.tailN = castConfig.castRoundIdx * castConfig.nUsedCore * mnConfig.singleN; + castConfig.wInNOffset = castConfig.tailN; + castConfig.wInKOffset = coreIdx * castConfig.curSingleK; + if (coreIdx == blockDimK - 1) { // -1: last dimK + castConfig.curSingleK = mnConfig.k - castConfig.curSingleK * coreIdx; + } + mnConfig.wOutOffset += castConfig.wInKOffset * mnConfig.n; + castConfig.curSingleN = castConfig.curDimN * mnConfig.singleN; + if (castConfig.castRoundIdx == Ceil(mnConfig.blockDimN, castConfig.nUsedCore) - 1) { // -1: last round + castConfig.curSingleN = mnConfig.n - castConfig.castRoundIdx * castConfig.nUsedCore * mnConfig.singleN; + } + } else { // antiquant generalized version + castConfig.curSingleN = mnConfig.singleN; + castConfig.curSingleK = mnConfig.k; + castConfig.tailN = mnConfig.nIdx * mnConfig.singleN; + castConfig.wInNOffset = this->transposeW ? castConfig.tailN * mnConfig.k : castConfig.tailN; + mnConfig.wOutOffset += castConfig.wInNOffset; + if (mnConfig.nIdx == mnConfig.blockDimN - 1) { + castConfig.curSingleN = mnConfig.n - mnConfig.nIdx * mnConfig.singleN; + } + } + SetAntiQuantGlobalBuffer(groupIdx, mnConfig); + CastWeightProcess(mnConfig, castConfig); +} + +template +__aicore__ inline void GMMAntiquantCompute::MMSync() { + if (this->mmWaitStatus) { + this->mm.WaitIterateAll(); + this->mmWaitStatus = false; + } + if constexpr (antiquantPerformance) { + SyncAll(); + } +} + +template +__aicore__ inline void +GMMAntiquantCompute::SetAntiQuantGlobalBuffer(uint32_t groupIdx, + const MNConfig mnConfig) { + if (this->singleWeight == 0) { + weightAntiQuantGm.SetGlobalBuffer(GetTensorAddr(groupIdx, this->weightTensorPtr)); + antiScaleGM.SetGlobalBuffer(GetTensorAddr(groupIdx, antiScaleTensorPtr)); + antiOffsetGM.SetGlobalBuffer(GetTensorAddr(groupIdx, antiOffsetTensorPtr)); + } else { + weightAntiQuantGm.SetGlobalBuffer(GetTensorAddr(0, this->weightTensorPtr) + mnConfig.wBaseOffset * GetTypeBits() / INT8_BITS); + uint64_t antiquantParamsOffset = mnConfig.nAxisBaseOffset; + if (isPerGroup) { + antiquantParamsOffset *= (mnConfig.k / perGroupSize); + } + antiScaleGM.SetGlobalBuffer(GetTensorAddr(0, antiScaleTensorPtr) + antiquantParamsOffset); + antiOffsetGM.SetGlobalBuffer(GetTensorAddr(0, antiOffsetTensorPtr) + antiquantParamsOffset); + } +} + + +template +__aicore__ inline void GMMAntiquantCompute::ComputeUbBaseK( + uint32_t curSingleK, uint32_t offsetK, uint32_t newBaseK, uint32_t& curUsedGroupSize, uint32_t& curBaseK) { + if (unlikely(offsetK + newBaseK >= curUsedGroupSize)) { + curBaseK = curUsedGroupSize - offsetK; + curUsedGroupSize += perGroupSize; + if (offsetK + curBaseK > curSingleK) { + curBaseK = curSingleK - offsetK; + } + } else if (unlikely(offsetK + newBaseK > curSingleK)) { + curBaseK = curSingleK - offsetK; + } else { + curBaseK = newBaseK; + } +} + + +template +__aicore__ inline void GMMAntiquantCompute::FreeScaleAndOffset(bool& firstLoop) { + if (firstLoop) { + firstLoop = false; + } else { + scaleInQueue.FreeTensor(scaleInUb); + offsetInQueue.FreeTensor(offsetInUb); + } +} + +template +__aicore__ inline void GMMAntiquantCompute::CastWeightProcess( + MNConfig& mnConfig, CastWeightConfig& castConfig) { + uint64_t wInOffset = castConfig.wInNOffset + static_cast(castConfig.wInKOffset) * mnConfig.n; + const uint32_t& curSingleK = castConfig.curSingleK; + const uint32_t& curSingleN = castConfig.curSingleN; + const uint32_t& scaleOffset = castConfig.tailN; + uint32_t newBaseK = this->ubBaseK; + uint32_t newBaseN = this->ubBaseN; + uint32_t usedGroupSize = mnConfig.k; + if (isPerGroup) { + newBaseK = Min(this->ubBaseK, perGroupSize); + if (!transposeW && newBaseK < perGroupSize && newBaseK > perGroupSize / 2 && mnConfig.n % newBaseN != 0) { + uint32_t tempUbBaseN = AlignDown(this->ubBaseK * this->ubBaseN / Ceil(perGroupSize, 2), 32); // 32:a factor + // ubBaseN cannot be larger than BEST_UB_BASEN, due to offset/scale queue size + if (tempUbBaseN <= BEST_UB_BASEN && mnConfig.n % tempUbBaseN == 0) { + newBaseK = Ceil(perGroupSize, 2); + newBaseN = tempUbBaseN; + } + } + usedGroupSize = perGroupSize + AlignDown(castConfig.wInKOffset, perGroupSize); + } + DataCopyPadExtParams padParams; + for (uint32_t offsetN(0), curBaseN(newBaseN), nCount(0); offsetN < curSingleN; offsetN += newBaseN) { + if (unlikely(offsetN + newBaseN > curSingleN)) { + curBaseN = curSingleN - offsetN; + } + uint32_t alignBaseN = AlignUp(curBaseN, UB_BLOCK_UNIT_SIZE * INT8_BITS / GetTypeBits()); + if (!isPerGroup) { + DataCopyScaleAndOffset(curBaseN, alignBaseN, scaleOffset + offsetN); + } + uint32_t curBaseK = newBaseK; + uint32_t curUsedGroupSize = usedGroupSize - castConfig.wInKOffset; + bool firstKLoop = true; + int32_t prePergroupIdx = -1; + int32_t curPergroupIdx = 0; + for (uint32_t offsetK(0), subCoreCount(nCount); offsetK < curSingleK; offsetK += curBaseK) { + ComputeUbBaseK(curSingleK, offsetK, newBaseK, curUsedGroupSize, curBaseK); + if constexpr (antiquantPerformance) { + if (this->subBlockIdx == (++subCoreCount) % 2) { // 2: two vectors + continue; + } + } + if (isPerGroup) { + curPergroupIdx = (offsetK + castConfig.wInKOffset) / perGroupSize; + if (firstKLoop || curPergroupIdx > prePergroupIdx) { // load new group + FreeScaleAndOffset(firstKLoop); + DataCopyScaleAndOffset(curBaseN, alignBaseN, scaleOffset + offsetN + curPergroupIdx * mnConfig.n); + prePergroupIdx = curPergroupIdx; + } + } + LocalTensor inLocal = vecInQueue.AllocTensor(); + DataCopyExtParams gmToUbIntriParams; + SetGmToUbDataCopyParams(curBaseN, curBaseK, mnConfig, gmToUbIntriParams); + uint64_t weightInOffset = transposeW ? offsetK + static_cast(offsetN) * mnConfig.k : + static_cast(offsetK) * mnConfig.n + offsetN; + DataCopyPad(inLocal, weightAntiQuantGm[(weightInOffset + wInOffset) * GetTypeBits() / INT8_BITS], gmToUbIntriParams, padParams); + vecInQueue.EnQue(inLocal); + + DataCopyExtParams ubToGmIntriParams; + if constexpr (transposeW) { + uint32_t alignBaseK = AlignUp(curBaseK, UB_BLOCK_UNIT_SIZE * INT8_BITS / GetTypeBits()); + CastWeightCompute(alignBaseK, alignBaseN); + SetUbToGmDataCopyParams(curBaseN, alignBaseK, curBaseK, mnConfig, ubToGmIntriParams); + } else { + CastWeightCompute(curBaseK, alignBaseN); + SetUbToGmDataCopyParams(curBaseN, alignBaseN, curBaseK, mnConfig, ubToGmIntriParams); + } + + // ResultCopy2GM + LocalTensor wResUb = vecOutQueue.DeQue(); + uint64_t weightOutOffset = transposeW ? mnConfig.wOutOffset + offsetK + offsetN * mnConfig.k : + mnConfig.wOutOffset + offsetK * mnConfig.n + offsetN; + DataCopyPad(this->weightGm[weightOutOffset], wResUb, ubToGmIntriParams); + vecOutQueue.FreeTensor(wResUb); + } + nCount = nCount == 0 ? 1: 0; + if (!(isPerGroup && firstKLoop)) { + scaleInQueue.FreeTensor(scaleInUb); + offsetInQueue.FreeTensor(offsetInUb); + } + } + + event_t eventIdMTE3ToS = static_cast(this->pipe->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventIdMTE3ToS); + WaitFlag(eventIdMTE3ToS); +} + +template +__aicore__ inline void +GMMAntiquantCompute::CastWeightCompute(uint32_t curCalcK, uint32_t curCalcAlignN) { + LocalTensor wInUb = vecInQueue.DeQue(); + wInUb.SetSize(curCalcK * curCalcAlignN); + LocalTensor wResUb = vecOutQueue.AllocTensor(); + LocalTensor tmpLocal = tmpUb.template ReinterpretCast(); + + AntiQuantShapeInfo shapeInfo; + if constexpr (transposeW) { + shapeInfo.offsetHeight = curCalcAlignN; + shapeInfo.offsetWidth = 1; + shapeInfo.scaleHeight = curCalcAlignN; + shapeInfo.scaleWidth = 1; + event_t eventId = static_cast(this->pipe->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventId); + WaitFlag(eventId); + } else { + shapeInfo.offsetHeight = 1; + shapeInfo.offsetWidth = curCalcAlignN; + shapeInfo.scaleHeight = 1; + shapeInfo.scaleWidth = curCalcAlignN; + } + // fp16 tempbuff is 0, bf16 tempbuff = offset.GetSize() * 2 * sizeof(float) + 64 * K * sizeof(float) + AscendAntiQuant(wResUb, wInUb, offsetInUb, scaleInUb, tmpLocal, curCalcK, shapeInfo); + + vecInQueue.FreeTensor(wInUb); + vecOutQueue.EnQue(wResUb); +} + +template +__aicore__ inline void +GMMAntiquantCompute::SetGmToUbDataCopyParams(const uint32_t curBaseN, + const uint32_t curBaseK, const MNConfig& mnConfig, DataCopyExtParams& intriParams) { + if constexpr (transposeW) { + intriParams.blockLen = Ceil(curBaseK * GetTypeBits(), INT8_BITS); + intriParams.blockCount = curBaseN; + intriParams.srcStride = Ceil((mnConfig.k - curBaseK) * GetTypeBits(), INT8_BITS); + intriParams.dstStride = 0; + } else { + intriParams.blockLen = Ceil(curBaseN * GetTypeBits(), INT8_BITS); + intriParams.blockCount = curBaseK; + intriParams.srcStride = Ceil((mnConfig.n - curBaseN) * GetTypeBits(), INT8_BITS); + intriParams.dstStride = 0; + } +} + +template +__aicore__ inline void +GMMAntiquantCompute::SetUbToGmDataCopyParams(const uint32_t curBaseN, + const uint32_t alignRowLen, const uint32_t curBaseK, const MNConfig& mnConfig, DataCopyExtParams& intriParams) { + if constexpr (transposeW) { + uint32_t alignBaseK = AlignUp(curBaseK, UB_BLOCK_UNIT_SIZE); + intriParams.blockLen = curBaseK * sizeof(BT); + intriParams.blockCount = curBaseN; + intriParams.srcStride = (alignRowLen - curBaseK) / (UB_BLOCK_UNIT_SIZE / sizeof(BT)); + intriParams.dstStride = (mnConfig.k - curBaseK) * sizeof(BT); + } else { + intriParams.blockLen = curBaseN * sizeof(BT); + intriParams.blockCount = curBaseK; + intriParams.srcStride = (alignRowLen - curBaseN) / (UB_BLOCK_UNIT_SIZE / sizeof(BT)); + intriParams.dstStride = (mnConfig.n - curBaseN) * sizeof(BT); + } +} + +template +__aicore__ inline void +GMMAntiquantCompute::DataCopyScaleAndOffset(uint32_t curBaseN, uint32_t alignBaseN, + uint64_t realScaleOffset) { + // copy scale and offset frome GM + DataCopyPadParams padParams; + DataCopyParams scaleParams; + scaleParams.blockLen = curBaseN * sizeof(BT); + scaleParams.blockCount = 1; + scaleParams.srcStride = 0; + scaleParams.dstStride = 0; + LocalTensor scaleLocal = scaleInQueue.AllocTensor(); + DataCopyPad(scaleLocal, antiScaleGM[realScaleOffset], scaleParams, padParams); + scaleInQueue.EnQue(scaleLocal); + + LocalTensor offsetLocal = offsetInQueue.AllocTensor(); + DataCopyPad(offsetLocal, antiOffsetGM[realScaleOffset], scaleParams, padParams); + offsetInQueue.EnQue(offsetLocal); + + scaleInUb = scaleInQueue.DeQue(); + scaleInUb.SetSize(alignBaseN); + offsetInUb = offsetInQueue.DeQue(); + offsetInUb.SetSize(alignBaseN); +} + +template +using GMMAntiquantComputePerformance = GMMAntiquantCompute; + +template +using GMMAntiquantComputeNorm = GMMAntiquantCompute; + +} // namespace GROUPED_MATMUL + +#endif // GMM_ANTI_QUANT +#endif // ASCENDC_GROUPED_MATMUL_ANTIQUANT_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a16w8_msd.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a16w8_msd.h new file mode 100644 index 00000000..400242b4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a16w8_msd.h @@ -0,0 +1,950 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a16w8_msd.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A16W8_MSD_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A16W8_MSD_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + + +#if defined(GMM_ANTI_QUANT) && defined(ORIG_DTYPE_WEIGHT) && defined(DT_INT8) && \ + ORIG_DTYPE_WEIGHT == DT_INT8 +namespace GROUPED_MATMUL { +static constexpr uint32_t A16W8_MSD_STEP = 2; +static constexpr uint32_t A16W8_MSD_PREPROCESS_MAX_GROUP = 12; +static constexpr uint32_t FACTOR_FOR_FLOAT_ALIGN_TO_32 = 8; +static constexpr uint32_t ALIGN_UB_BASE_K = 128; +static constexpr uint32_t POST_SKIP_ITER_NUM = 3; + +struct PreMNConfig { + uint32_t m = 0; + uint32_t k = 0; + uint32_t baseM = 0; + uint32_t baseK = 0; + uint32_t mIdx = 0; + uint32_t kIdx = 0; + uint32_t blockDimM = 0; + uint32_t blockDimK = 0; + uint32_t singleM = 0; + uint32_t singleMTail = 0; + uint64_t mAxisBaseOffset = 0; +}; + +struct PreBaseMNConfig { + uint32_t m = 0; + uint32_t k = 0; + uint64_t mAxisBaseOffset = 0; +}; + +/** @brief GroupMatmul operator Class +*/ +template +class GMMA16W8MSDProcess{ + protected: + using B = typename ComputeType::B; + ComputeType& computeOp; // internal computation operator + const GMMBaseParams* __restrict gmmBaseParams; + const TCubeTiling* __restrict mmTilingData; + + uint32_t blockIdx; + uint32_t coreIdx; + uint32_t groupNum; + uint32_t coreNum; + uint32_t ubCalSize; + int32_t preOffset; + int32_t preOffsetPre; + GM_ADDR groupListPtr; + GlobalTensor groupListGm; + TILING_TYPE* kListGm; + TILING_TYPE* nListGm; + + public: + /** @brief constructor */ + __aicore__ inline GMMA16W8MSDProcess(ComputeType& computeOp_) : computeOp(computeOp_) {} + + __aicore__ inline void Init(const GMMBaseParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn, + GM_ADDR groupList, GM_ADDR tiling); + + __aicore__ inline void Process(); + + private: + __aicore__ inline void PreProcess(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig, + uint32_t &preGroupIdx, uint32_t &preCoreCount, bool &isPreRequired); + + __aicore__ inline void TailProcess(MNConfig &mnConfig, uint32_t secondHalfIterCount); + + __aicore__ inline void SetMNConfigs(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); +}; + +template + __aicore__ inline void GMMA16W8MSDProcess::Init(const GMMBaseParams* __restrict gmmBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn, GM_ADDR groupList, GM_ADDR tiling) { + blockIdx = GetBlockIdx(); + coreIdx = blockIdx; + int64_t coreRation = GetTaskRation(); + if (coreRation > 1) { + coreIdx /= coreRation; + } + gmmBaseParams = gmmBaseParamsIn; + mmTilingData = mmTilingDataIn; + ubCalSize = gmmBaseParams->ubCalSize; + groupNum = gmmBaseParams->groupNum; + coreNum = gmmBaseParams->coreNum; + groupListPtr = groupList; + preOffset = 0; + preOffsetPre = 0; + if (groupListPtr != nullptr) { + groupListGm.SetGlobalBuffer((__gm__ int64_t*)groupList); + } + kListGm = gmmArrayAddrIn + MKN_LIST_LEN; + nListGm = gmmArrayAddrIn + MKN_LIST_LEN * 2; +} + +template +__aicore__ inline void GMMA16W8MSDProcess::SetMNConfigs( + PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig) { + preBaseMNConfig.k = kListGm[0]; + + mnConfig.k = preBaseMNConfig.k; + mnConfig.n = nListGm[0]; + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + mnConfig.singleM = mnConfig.baseM; + // 2048: least n for case enable singleN 1024; + // 2: according to experiments when n larger than or equal 2k, singleN 1024 has better performence + // 1024: larger singleN can reduce preprocess iter num than reduce syncall nums and has better performence + // 4: when satisfy reuqirements of different singleN, singleN should be quater of n align up to 1024 + mnConfig.singleN = mnConfig.n >= 2048 && mnConfig.n / mnConfig.k >= 2 ? + 1024 * Ceil(mnConfig.n / 4, 1024) : mnConfig.baseN; + mnConfig.singleN = mnConfig.singleN <= ubCalSize ? mnConfig.singleN : ubCalSize; +} + +template +__aicore__ inline void GMMA16W8MSDProcess::UpdateMnConfig(MNConfig &mnConfig) { + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMA16W8MSDProcess::PreProcess( + PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig, uint32_t &preGroupIdx, uint32_t &preCoreCount, + bool &isPreRequired) { + PreBaseMNConfig preBaseMNConfigs[A16W8_MSD_PREPROCESS_MAX_GROUP]; + uint32_t preValidGroupCount = 0; + while (preCoreCount < coreNum && preValidGroupCount < A16W8_MSD_PREPROCESS_MAX_GROUP && preGroupIdx < groupNum) { + preBaseMNConfig.mAxisBaseOffset += preBaseMNConfig.m; + preBaseMNConfig.m = GetSplitValueFromGroupList(preGroupIdx, preOffsetPre, gmmBaseParams, groupListGm); + preGroupIdx++; + if (preBaseMNConfig.m <= 0) { + continue; + } + preBaseMNConfigs[preValidGroupCount] = preBaseMNConfig; + preValidGroupCount++; + preCoreCount += Ceil(A16W8_MSD_STEP * preBaseMNConfig.m, mnConfig.singleM) * + Ceil(mnConfig.n, mnConfig.singleN); + } + if (preValidGroupCount == 0) { + isPreRequired = false; + return; + } + computeOp.PreProcess(preBaseMNConfigs, preValidGroupCount, mmTilingData->baseM / 2); + preCoreCount = preCoreCount % coreNum; +} + +template +__aicore__ inline void GMMA16W8MSDProcess::TailProcess(MNConfig &mnConfig, uint32_t secondHalfIterCount) { + uint32_t resPostLoop = POST_SKIP_ITER_NUM; + if (secondHalfIterCount < POST_SKIP_ITER_NUM) { + resPostLoop = secondHalfIterCount; + secondHalfIterCount = POST_SKIP_ITER_NUM; + } + for (uint32_t resIdx = 0; resIdx < resPostLoop; ++resIdx) { + computeOp.PostProcess(mnConfig, true, secondHalfIterCount); + secondHalfIterCount++; + } +} + +template +__aicore__ inline void GMMA16W8MSDProcess::Process() { + PreBaseMNConfig preBaseMNConfig; + MNConfig mnConfig; + uint32_t preValidGroupCount = 0; + uint32_t preGroupIdx = 0; + bool isPreRequired = false; + uint32_t secondHalfIterCount = 0; + SetMNConfigs(preBaseMNConfig, mnConfig); + if (mnConfig.k <= 0 || mnConfig.n <= 0) { + return; + } + for (uint32_t groupIdx(0), count(0), curBlock(0), curCount(0), preCoreCount(0); + groupIdx < groupNum; ++groupIdx) { + isPreRequired = preGroupIdx == groupIdx; + if (isPreRequired) { + PreProcess(preBaseMNConfig, mnConfig, preGroupIdx, preCoreCount, isPreRequired); + } + if (groupIdx > 0) { + UpdateMnConfig(mnConfig); + } + mnConfig.m = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm); + if ASCEND_IS_AIC { + if (isPreRequired) { + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); + } + } + if (mnConfig.m <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(A16W8_MSD_STEP * mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + curBlock = coreIdx >= count ? coreIdx : coreIdx + coreNum; + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; + computeOp.MMCompute(mnConfig); + computeOp.PostProcess(mnConfig, false, secondHalfIterCount); + secondHalfIterCount++; + curBlock += coreNum; + } + count = curCount % coreNum; + } + TailProcess(mnConfig, secondHalfIterCount); +} + +/** @brief intenal computation class +*/ +template +class GMMA16W8MSDCompute { + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using WT = DTYPE_WEIGHT; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + /** @brief constructor */ + __aicore__ inline GMMA16W8MSDCompute(typename mmType::MT& mm_) : mm(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR offset, + GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR group_list, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe); + + __aicore__ inline void MMCompute(MNConfig& mnConfig); + + __aicore__ inline void PreProcess(PreBaseMNConfig *preBaseMNConfigs, uint32_t preValidGroupCount, + uint32_t maxMPerGroup); + + __aicore__ inline void PostProcess(MNConfig& mnConfig, bool isLastGroup, uint32_t secondHalfIterCount); + + private: + __aicore__ inline void InitLocalTensor(); + + __aicore__ inline void InitWorkspace(GM_ADDR workspace); + + __aicore__ inline void PreProcessTiling(uint32_t m, uint32_t curCoreNum, uint32_t startCoreIdx, + PreBaseMNConfig &preBaseMNConfig, PreMNConfig &preMNConfig); + + __aicore__ inline void PreProcessCalc(uint32_t curCoreNum, uint32_t startCoreIdx, uint32_t &resSyncCount, + PreMNConfig &preMNConfig); + + __aicore__ inline void PreProcessSync(uint32_t preValidGroupCount, uint32_t &resSyncCount); + + __aicore__ inline void CopyOriginInput(uint32_t k, uint32_t curBaseM, uint32_t curBaseK, uint64_t xGmOffset); + + __aicore__ inline void CalcReduceSum(uint32_t curBaseM, uint32_t curBaseK, uint64_t gmReduceSumOffset); + + __aicore__ inline void CalcAMax(uint32_t curBaseM, uint32_t curBaseK, uint64_t gmReduceMaxOffset); + + __aicore__ inline void CopyInAmax(uint32_t curBaseM, uint64_t gmReduceMaxOffset); + + __aicore__ inline void CalcAMatrix(PreMNConfig &preMNConfig, uint32_t curBaseM, uint32_t curBaseK, + uint64_t gmReduceMaxOffset, uint64_t aOffsetGm); + + __aicore__ inline void CalcASum(MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM, + uint64_t offsetAndScaleOffset); + + __aicore__ inline void ProcessScaleAndBias(uint32_t n, uint32_t curBaseN, uint64_t offsetAndScaleOffset); + + __aicore__ inline void ProcessC1C2(MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM, + uint32_t curSingleM); + + __aicore__ inline void CalcCMatrix(MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM); + + __aicore__ inline void CopyOutFinalResult(uint32_t n, uint32_t curBaseM, uint32_t curBaseN, uint64_t yOffset); + + __aicore__ inline GlobalTensor SetGlobalBufferW(uint32_t tailN, MNConfig& mnConfig); + + __aicore__ inline uint64_t SetWOffset(uint32_t tailN, uint32_t k); + + TPipe* pipe; + typename mmType::MT& mm; // matmul operator + bool hasBias = false; + GM_ADDR weightTensorPtr; + GlobalTensor xGm; + GlobalTensor biasGm; + GlobalTensor scaleGm; + GlobalTensor offsetGm; + GlobalTensor mmOutGm; + GlobalTensor aMatrixGm; + GlobalTensor globalMaxGm; + GlobalTensor localSumGm; + GlobalTensor yGm; + + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue ReduceResultInQueue; + TBuf tmpBuff; + // LocalTensor used in stage1 (preprocess) + LocalTensor s1MiddleResult1; + LocalTensor s1MiddleResult2; + LocalTensor s1TmpBuf; + LocalTensor s1A1A2FP16; + // LocalTensor used in stage2 and stage3 (postprocess) + LocalTensor s23MiddleResult1; + LocalTensor s23MiddleResult2; + LocalTensor s23MiddleResult3; + LocalTensor cTmp; + LocalTensor processedScale; + LocalTensor processedBias; + LocalTensor globalReduceSum; + LocalTensor s23TmpBuf; + LocalTensor aMaxInUb; + + uint32_t cubeBaseM; + uint32_t ubCalSizeS1; + uint32_t ubCalSizeS2; + uint32_t coreNum; + uint32_t totalM; + uint32_t aicIdx; + uint32_t aivIdx; + uint32_t ubRestBytes; + uint32_t aMatrixSize; + MNConfig mnConfigs[POST_SKIP_ITER_NUM + 1]; +}; + +template +__aicore__ inline void GMMA16W8MSDCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR scale, GM_ADDR offset, GM_ADDR antiquantScale, + GM_ADDR antiquantOffset, GM_ADDR group_list, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, + TPipe* tPipe) { + weightTensorPtr = weight; + pipe = tPipe; + cubeBaseM = mmTilingData->baseM; + ubCalSizeS1 = 2 * gmmBaseParams->ubCalSize; + ubCalSizeS2 = gmmBaseParams->ubCalSize; + totalM = gmmBaseParams->m; + coreNum = gmmBaseParams->coreNum; + ubRestBytes = gmmBaseParams->ubRestBytes; + aMatrixSize = gmmBaseParams->workspaceSize; + hasBias = gmmBaseParams->hasBias == 1; + aicIdx = GetBlockIdx() / GetTaskRation(); + aivIdx = GetBlockIdx(); + + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + scaleGm.SetGlobalBuffer(GetTensorAddr(0, antiquantScale)); + offsetGm.SetGlobalBuffer(GetTensorAddr(0, antiquantOffset)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + if (hasBias) { + biasGm.SetGlobalBuffer(GetTensorAddr(0, bias)); + } + InitLocalTensor(); + InitWorkspace(workspace); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::InitLocalTensor() { + if ASCEND_IS_AIC { + return; + } + uint32_t alignedCoreNum = AlignUp<8>(coreNum); + pipe->InitBuffer(vecInQueue, 1, ubCalSizeS1 * sizeof(half)); + pipe->InitBuffer(vecOutQueue, 1, ubCalSizeS1 * sizeof(int8_t)); + pipe->InitBuffer(ReduceResultInQueue, 1, cubeBaseM / 2 * alignedCoreNum * sizeof(float)); + pipe->InitBuffer(tmpBuff, ubRestBytes); + uint32_t s1TmpUbOffset = 0; + uint32_t s23TmpUbOffset = 0; + // local tensor for stage1 + s1MiddleResult1 = tmpBuff.GetWithOffset(ubCalSizeS1, 0); + s1TmpUbOffset += ubCalSizeS1 * sizeof(float); + s1MiddleResult2 = tmpBuff.GetWithOffset(ubCalSizeS1, s1TmpUbOffset); + s1TmpUbOffset += ubCalSizeS1 * sizeof(float); + s1TmpBuf = tmpBuff.GetWithOffset(ubCalSizeS1, s1TmpUbOffset); + s1A1A2FP16 = tmpBuff.GetWithOffset(ubCalSizeS1, s1TmpUbOffset); + // local tensor for stage2 and stage3 + s23MiddleResult1 = tmpBuff.GetWithOffset(ubCalSizeS2, 0); + s23TmpUbOffset += ubCalSizeS2 * sizeof(float); + cTmp = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); + s23TmpUbOffset += ubCalSizeS2 * sizeof(float); + processedScale = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); + s23TmpUbOffset += ubCalSizeS2 * sizeof(float); + processedBias = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); + s23TmpUbOffset += ubCalSizeS2 * sizeof(float); + globalReduceSum = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); + s23MiddleResult2 = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); + s23TmpUbOffset += ubCalSizeS2 * sizeof(float); + s23MiddleResult3 = tmpBuff.GetWithOffset(ubCalSizeS2, s23TmpUbOffset); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::InitWorkspace(GM_ADDR workspace) { + uint32_t usedWorkspaceSize = 0; + globalMaxGm.SetGlobalBuffer((__gm__ float *)(workspace)); + if (aivIdx == 0) { // 0: use aiv 0 to init gm for amax + InitOutput(globalMaxGm, totalM * FACTOR_FOR_FLOAT_ALIGN_TO_32); + } + usedWorkspaceSize += totalM * sizeof(float) * FACTOR_FOR_FLOAT_ALIGN_TO_32; + localSumGm.SetGlobalBuffer((__gm__ float *)(workspace + usedWorkspaceSize)); + if (aivIdx == 1) { // 1: use aiv 1 to init gm for asum + InitOutput(localSumGm, totalM * coreNum); + } + usedWorkspaceSize += totalM * coreNum * sizeof(float); + aMatrixGm.SetGlobalBuffer((__gm__ int8_t *)(workspace + usedWorkspaceSize)); + usedWorkspaceSize += aMatrixSize * sizeof(int8_t); + mmOutGm.SetGlobalBuffer((__gm__ int32_t *)(workspace + usedWorkspaceSize)); + if ASCEND_IS_AIV { + SyncAll(); + } +} + +template +__aicore__ inline uint64_t GMMA16W8MSDCompute::SetWOffset(uint32_t tailN, uint32_t k) { + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ && transposeW) { + wOffset = tailN * (UB_BLOCK_UNIT_SIZE / sizeof(BT)); // 32: quant is 32, float16 is 16 + } else if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else if constexpr (transposeW) { + wOffset = k * tailN; + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline GlobalTensor GMMA16W8MSDCompute::SetGlobalBufferW( + uint32_t tailN, MNConfig& mnConfig) { + uint64_t wOffset = SetWOffset(tailN, mnConfig.k); + GlobalTensor weightGmLocal; + weightGmLocal.SetGlobalBuffer(GetTensorAddr(0, weightTensorPtr) + mnConfig.wBaseOffset + wOffset); + if (mnConfig.blockDimM == 1) { + weightGmLocal.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + return weightGmLocal; +} + +template +__aicore__ inline void GMMA16W8MSDCompute::PreProcess( + PreBaseMNConfig *preBaseMNConfigs, uint32_t preValidGroupCount, uint32_t maxMPerGroup) { + if ASCEND_IS_AIC { + return; + } + PreMNConfig preMNConfig; + uint32_t usedCoreNum = 0; // num of core not used + uint32_t curCoreNum = 0; + uint32_t tokenNumEachPreIter = 0; + uint32_t splitGroupNum = 0; + // since limitation on matmul baseM, data block of preprocess cannot only split by groupIdx. If m of a group + // larger than half of matmul baseM, this group should splited in preprocess, and should get group num after + // split first. + for (uint32_t gIdx = 0; gIdx < preValidGroupCount; ++gIdx) { + // ensure each group after splited has at least one core when preprocessing. + splitGroupNum += Ceil(preBaseMNConfigs[gIdx].m, maxMPerGroup); + tokenNumEachPreIter += preBaseMNConfigs[gIdx].m; + } + // num of core need to allocate to different group + uint32_t unAllocatedCoreNum = splitGroupNum <= coreNum ? coreNum - splitGroupNum : 0; + uint32_t resSyncCount = Ceil(splitGroupNum, coreNum); + uint32_t resTokenNum = tokenNumEachPreIter; // num of tokens not have corresponding core + for (uint32_t gIdx = 0; gIdx < preValidGroupCount; ++gIdx) { + uint32_t curGroupResTokenNum = preBaseMNConfigs[gIdx].m; + // if m of the group larger than half of matmul baseM, preprocess data step by step, + // each step accepts half of matmul baseM * k size of data. + while (curGroupResTokenNum > maxMPerGroup) { + curCoreNum = Ceil(unAllocatedCoreNum * maxMPerGroup, resTokenNum) + 1; + PreProcessTiling(maxMPerGroup, curCoreNum, usedCoreNum, preBaseMNConfigs[gIdx], preMNConfig); + PreProcessCalc(curCoreNum, usedCoreNum, resSyncCount, preMNConfig); + unAllocatedCoreNum -= (curCoreNum - 1); + resTokenNum -= maxMPerGroup; + usedCoreNum = (usedCoreNum + curCoreNum) % coreNum; + curGroupResTokenNum -= maxMPerGroup; + preBaseMNConfigs[gIdx].mAxisBaseOffset += maxMPerGroup; + } + curCoreNum = (unAllocatedCoreNum * curGroupResTokenNum / resTokenNum) + 1; + PreProcessTiling(curGroupResTokenNum, curCoreNum, usedCoreNum, preBaseMNConfigs[gIdx], preMNConfig); + PreProcessCalc(curCoreNum, usedCoreNum, resSyncCount, preMNConfig); + unAllocatedCoreNum -= (curCoreNum - 1); + resTokenNum -= curGroupResTokenNum; + usedCoreNum = (usedCoreNum + curCoreNum) % coreNum; + } + PreProcessSync(preValidGroupCount, resSyncCount); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::PreProcessSync( + uint32_t preValidGroupCount, uint32_t &resSyncCount) { + while (resSyncCount > 0) { + SyncAll(); + resSyncCount -= 1; + } + SyncAll(); + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::PreProcessTiling( + uint32_t m, uint32_t curCoreNum, uint32_t startCoreIdx, + PreBaseMNConfig &preBaseMNConfig, PreMNConfig &preMNConfig) { + if (aivIdx < startCoreIdx || aivIdx >= curCoreNum + startCoreIdx) { + return; + } + preMNConfig.m = m; + preMNConfig.k = preBaseMNConfig.k; + preMNConfig.mAxisBaseOffset = preBaseMNConfig.mAxisBaseOffset; + if (m < curCoreNum) { + preMNConfig.baseK = preMNConfig.k / curCoreNum; + preMNConfig.baseK = AlignUp(preMNConfig.baseK, ALIGN_UB_BASE_K); + preMNConfig.blockDimK = Ceil(preMNConfig.k, preMNConfig.baseK); + preMNConfig.blockDimM = curCoreNum / preMNConfig.blockDimK; + } else { + preMNConfig.baseK = preMNConfig.k ; + preMNConfig.blockDimK = 1; + preMNConfig.blockDimM = curCoreNum; + } + preMNConfig.singleM = Ceil(m, preMNConfig.blockDimM); + preMNConfig.blockDimM = Ceil(m, preMNConfig.singleM); // prevent wrapping when calc singleMTail + preMNConfig.singleMTail = m - (preMNConfig.blockDimM - 1) * preMNConfig.singleM; + preMNConfig.baseM = ubCalSizeS1 / preMNConfig.baseK; + preMNConfig.baseM = preMNConfig.baseM < preMNConfig.singleM ? preMNConfig.baseM : preMNConfig.singleM; + preMNConfig.mIdx = (aivIdx - startCoreIdx) / preMNConfig.blockDimK; + preMNConfig.kIdx = (aivIdx - startCoreIdx) % preMNConfig.blockDimK; +} + +template +__aicore__ inline void GMMA16W8MSDCompute::PreProcessCalc( + uint32_t curCoreNum, uint32_t startCoreIdx, uint32_t &resSyncCount, PreMNConfig &preMNConfig) { + if (aivIdx < startCoreIdx || aivIdx >= startCoreIdx + curCoreNum) { + return; + } + if (aivIdx < startCoreIdx + curCoreNum && aivIdx >= startCoreIdx + preMNConfig.blockDimM * preMNConfig.blockDimK) { + return; + } + uint32_t curBaseK = preMNConfig.kIdx < preMNConfig.blockDimK - 1 ? + preMNConfig.baseK : preMNConfig.k - preMNConfig.kIdx * preMNConfig.baseK; + uint32_t curBaseM = preMNConfig.baseM; + uint32_t curSingleM = preMNConfig.mIdx < preMNConfig.blockDimM - 1 ? + preMNConfig.singleM : preMNConfig.m - preMNConfig.mIdx * preMNConfig.singleM; + for (uint32_t offsetM = 0; offsetM < curSingleM; offsetM += preMNConfig.baseM) { + if (offsetM + preMNConfig.baseM >= curSingleM) { + curBaseM = curSingleM - offsetM; + } + uint64_t offsetBase = preMNConfig.mAxisBaseOffset + preMNConfig.mIdx * preMNConfig.singleM + offsetM; + uint64_t xGmOffset = offsetBase * preMNConfig.k + preMNConfig.kIdx * preMNConfig.baseK; + CopyOriginInput(preMNConfig.k, curBaseM, curBaseK, xGmOffset); + uint64_t gmReduceMaxOffset = offsetBase * FACTOR_FOR_FLOAT_ALIGN_TO_32; + uint64_t gmReduceSumOffset = offsetBase * coreNum + (aivIdx - startCoreIdx); + CalcAMax(curBaseM, curBaseK, gmReduceMaxOffset); + CalcReduceSum(curBaseM, curBaseK, gmReduceSumOffset); + } + SyncAll(); + resSyncCount -= 1; + curBaseM = preMNConfig.baseM; + for (uint32_t offsetM = 0; offsetM < curSingleM; offsetM += preMNConfig.baseM) { + if (offsetM + preMNConfig.baseM >= curSingleM) { + curBaseM = curSingleM - offsetM; + } + uint64_t offsetBase = preMNConfig.mAxisBaseOffset + preMNConfig.mIdx * preMNConfig.singleM + offsetM; + uint64_t xGmOffset = offsetBase * preMNConfig.k + preMNConfig.kIdx * preMNConfig.baseK; + uint64_t gmReduceMaxOffset = offsetBase * FACTOR_FOR_FLOAT_ALIGN_TO_32; + uint64_t aOffsetGm = + (preMNConfig.mAxisBaseOffset * 2 + preMNConfig.mIdx * preMNConfig.singleM + offsetM) * preMNConfig.k + + preMNConfig.kIdx * preMNConfig.baseK; + CopyOriginInput(preMNConfig.k, curBaseM, curBaseK, xGmOffset); + CalcAMatrix(preMNConfig, curBaseM, curBaseK, gmReduceMaxOffset, aOffsetGm); + } +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CopyOriginInput( + uint32_t k, uint32_t curBaseM, uint32_t curBaseK, uint64_t xGmOffset) { + uint32_t alignedBaseK = AlignUp<32>(curBaseK); + LocalTensor xLocal = vecInQueue.AllocTensor(); + DataCopyPad2D(xLocal, xGm[xGmOffset], curBaseM, curBaseK, k); + vecInQueue.EnQue(xLocal); + LocalTensor xFP16InUb = vecInQueue.DeQue(); + Cast(s1MiddleResult1, xFP16InUb, RoundMode::CAST_NONE, curBaseM * alignedBaseK); + PipeBarrier(); + vecInQueue.FreeTensor(xFP16InUb); + } + +template +__aicore__ inline void GMMA16W8MSDCompute::CalcReduceSum( + uint32_t curBaseM, uint32_t curBaseK, uint64_t gmReduceSumOffset) { + uint32_t alignedBaseK = AlignUp<32>(curBaseK); + LocalTensor blockReduceSumInUb = vecOutQueue.AllocTensor(); + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + ReduceSum(blockReduceSumInUb[idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32], s1MiddleResult1[idxM * alignedBaseK], + s1TmpBuf[idxM * alignedBaseK], curBaseK); + } + PipeBarrier(); + vecOutQueue.EnQue(blockReduceSumInUb); + LocalTensor blockReduceSum = vecOutQueue.DeQue(); + DataCopyExtParams aSumOutParams; + aSumOutParams.blockLen = sizeof(float); + aSumOutParams.blockCount = curBaseM; + aSumOutParams.srcStride = 0; + aSumOutParams.dstStride = (coreNum - 1) * sizeof(float); + DataCopyPad(localSumGm[gmReduceSumOffset], blockReduceSum, aSumOutParams); + vecOutQueue.FreeTensor(blockReduceSum); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CalcAMax( + uint32_t curBaseM, uint32_t curBaseK, uint64_t gmReduceMaxOffset) { + uint32_t alignedBaseK = AlignUp<32>(curBaseK); + Abs(s1MiddleResult2, s1MiddleResult1, curBaseM * alignedBaseK); + PipeBarrier(); + + // 计算ReduceMax + LocalTensor blockReduceMaxInUb = vecOutQueue.AllocTensor(); + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + ReduceMax(blockReduceMaxInUb[idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32], s1MiddleResult2[idxM * alignedBaseK], + s1TmpBuf[idxM * alignedBaseK], curBaseK, false); + } + PipeBarrier(); + vecOutQueue.EnQue(blockReduceMaxInUb); + LocalTensor blockReduceMax = vecOutQueue.DeQue(); + SetAtomicMax(); + DataCopyExtParams aMaxOutParams; + aMaxOutParams.blockLen = FACTOR_FOR_FLOAT_ALIGN_TO_32 * sizeof(float); + aMaxOutParams.blockCount = curBaseM; + aMaxOutParams.srcStride = 0; + aMaxOutParams.dstStride = 0; + DataCopyPad(globalMaxGm[gmReduceMaxOffset], blockReduceMax, aMaxOutParams); + SetAtomicNone(); + vecOutQueue.FreeTensor(blockReduceMax); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CopyInAmax(uint32_t curBaseM, uint64_t gmReduceMaxOffset) { + // copy amax from gm + LocalTensor aMaxLocal = ReduceResultInQueue.AllocTensor(); + DataCopyPadExtParams padParams; + DataCopyExtParams aMaxInParams; + aMaxInParams.blockLen = FACTOR_FOR_FLOAT_ALIGN_TO_32 * sizeof(float); + aMaxInParams.blockCount = curBaseM; + aMaxInParams.srcStride = 0; + aMaxInParams.dstStride = 0; + DataCopyPad(aMaxLocal, globalMaxGm[gmReduceMaxOffset], aMaxInParams, padParams); + ReduceResultInQueue.EnQue(aMaxLocal); + aMaxInUb = ReduceResultInQueue.DeQue(); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CalcAMatrix( + PreMNConfig &preMNConfig, uint32_t curBaseM, uint32_t curBaseK, uint64_t gmReduceMaxOffset, uint64_t aOffsetGm) { + uint32_t alignedBaseK = AlignUp<32>(curBaseK); + CopyInAmax(curBaseM, gmReduceMaxOffset); + event_t eventIdMTE2ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventIdMTE2ToS); + WaitFlag(eventIdMTE2ToS); + // calc a_tmp = 127 * x / amax for each row + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + float invertAMaxPerRow = 127.0f / aMaxInUb(idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(s1MiddleResult2[idxM * alignedBaseK], s1MiddleResult1[idxM * alignedBaseK], invertAMaxPerRow, + alignedBaseK); // a_tmp + } + PipeBarrier(); + ReduceResultInQueue.FreeTensor(aMaxInUb); + // calc a1 + LocalTensor a1Int8InUb = vecOutQueue.AllocTensor(); + + Cast(s1MiddleResult1, s1MiddleResult2, RoundMode::CAST_ROUND, curBaseM * alignedBaseK); // a1 + PipeBarrier(); + Cast(s1A1A2FP16, s1MiddleResult1, RoundMode::CAST_NONE, curBaseM * alignedBaseK); + PipeBarrier(); + Cast(a1Int8InUb, s1A1A2FP16, RoundMode::CAST_NONE, curBaseM * alignedBaseK); + vecOutQueue.EnQue(a1Int8InUb); + LocalTensor a1Int8 = vecOutQueue.DeQue(); + DataCopyPad2D(aMatrixGm[aOffsetGm], a1Int8, curBaseM, curBaseK, alignedBaseK, preMNConfig.k); + + // calc a2 + PipeBarrier(); + Sub(s1TmpBuf, s1MiddleResult2, s1MiddleResult1, curBaseM * alignedBaseK); // a_tmp - a1 + PipeBarrier(); + Muls(s1MiddleResult1, s1TmpBuf, static_cast(254), curBaseM * alignedBaseK); // 254 * (a_tmp - a1) + PipeBarrier(); + Cast(s1MiddleResult2, s1MiddleResult1, RoundMode::CAST_ROUND, curBaseM * alignedBaseK); // a2 + PipeBarrier(); + Cast(s1A1A2FP16, s1MiddleResult2, RoundMode::CAST_NONE, curBaseM * alignedBaseK); + vecOutQueue.FreeTensor(a1Int8); + LocalTensor a2Int8InUb = vecOutQueue.AllocTensor(); + PipeBarrier(); + Cast(a2Int8InUb, s1A1A2FP16, RoundMode::CAST_NONE, curBaseM * alignedBaseK); + vecOutQueue.EnQue(a2Int8InUb); + LocalTensor a2Int8 = vecOutQueue.DeQue(); + DataCopyPad2D(aMatrixGm[aOffsetGm + preMNConfig.m * preMNConfig.k], a2Int8, + curBaseM, curBaseK, alignedBaseK, preMNConfig.k); + vecOutQueue.FreeTensor(a2Int8); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::MMCompute(MNConfig& mnConfig) { + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + + mnConfig.workSpaceOffset = outOffset + A16W8_MSD_STEP * mnConfig.yBaseOffset; + if ASCEND_IS_AIC { + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : + A16W8_MSD_STEP * mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k + A16W8_MSD_STEP * mnConfig.xBaseOffset; + // init global buffer + GlobalTensor weightGm = SetGlobalBufferW(tailN, mnConfig); + mm.SetOrgShape(A16W8_MSD_STEP * mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(aMatrixGm[xOffset], transposeX); + mm.SetTensorB(weightGm, transposeW); + while (mm.Iterate()) { + mm.GetTensorC(mmOutGm[mnConfig.workSpaceOffset]); + } + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + } +} + +template +__aicore__ inline void GMMA16W8MSDCompute::PostProcess( + MNConfig& mnConfig, bool isLastGroup, uint32_t secondHalfIterCount) { + if ASCEND_IS_AIC { + return; + } + if (!isLastGroup) { + mnConfigs[secondHalfIterCount % (POST_SKIP_ITER_NUM + 1)] = mnConfig; + if (secondHalfIterCount < POST_SKIP_ITER_NUM) { + return; + } + } + MNConfig postMNConfig = mnConfigs[(secondHalfIterCount - POST_SKIP_ITER_NUM) % (POST_SKIP_ITER_NUM + 1)]; + uint32_t tailN = postMNConfig.nIdx * postMNConfig.singleN; + uint32_t curCubeSingleM = postMNConfig.mIdx < postMNConfig.blockDimM - 1 ? + postMNConfig.singleM : A16W8_MSD_STEP * postMNConfig.m - postMNConfig.mIdx * postMNConfig.singleM; + uint32_t curBaseN = postMNConfig.nIdx < postMNConfig.blockDimN - 1 ? + postMNConfig.singleN : postMNConfig.n - tailN; + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + uint32_t curBaseM = ubCalSizeS2 / alignedBaseN; + uint32_t curSingleM = curCubeSingleM / 2; + curBaseM = curBaseM < curSingleM ? curBaseM : curSingleM; + postMNConfig.singleM /= 2; + + uint64_t offsetAndScaleOffset = postMNConfig.nAxisBaseOffset + tailN; + ProcessScaleAndBias(postMNConfig.n, curBaseN, offsetAndScaleOffset); + for (uint32_t offsetM = 0; offsetM < curSingleM; offsetM += curBaseM) { + if (offsetM + curBaseM >= curSingleM) { + curBaseM = curSingleM - offsetM; + } + CalcASum(postMNConfig, curBaseM, curBaseN, offsetM, offsetAndScaleOffset); + if (offsetM == 0) { // only first iter need to wait for cube + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG); + } + ProcessC1C2(postMNConfig, curBaseM, curBaseN, offsetM, curSingleM); + CalcCMatrix(postMNConfig, curBaseM, curBaseN, offsetM); + uint64_t yOffset = (postMNConfig.mIdx * postMNConfig.singleM + offsetM) * postMNConfig.n + \ + postMNConfig.nIdx * postMNConfig.singleN + postMNConfig.yBaseOffset; + CopyOutFinalResult(postMNConfig.n, curBaseM, curBaseN, yOffset); + } +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CalcASum( + MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM, uint64_t offsetAndScaleOffset) { + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + uint32_t alignedCoreNum = AlignUp<8>(coreNum); + // process offset + LocalTensor offsetF16 = vecInQueue.AllocTensor(); + DataCopyPad2D(offsetF16, offsetGm[offsetAndScaleOffset], 1, curBaseN, postMNConfig.n); + vecInQueue.EnQue(offsetF16); + LocalTensor offsetF16InUb = vecInQueue.DeQue(); + Cast(s23MiddleResult1, offsetF16InUb, RoundMode::CAST_NONE, alignedBaseN); + PipeBarrier(); + vecInQueue.FreeTensor(offsetF16InUb); + + // calc global sum and mul with offset + LocalTensor localSum = ReduceResultInQueue.AllocTensor(); + DataCopyExtParams params; + params.blockCount = curBaseM; + params.blockLen = coreNum * sizeof(float); + params.srcStride = 0; + params.dstStride = 0; + + DataCopyPadExtParams padParams; + padParams.isPad = true; + padParams.rightPadding = alignedCoreNum - coreNum; + padParams.leftPadding = 0; + padParams.paddingValue = 0; + uint32_t gmReduceSumOffset = (postMNConfig.mAxisBaseOffset + postMNConfig.mIdx * postMNConfig.singleM + offsetM) * + coreNum; + DataCopyPad(localSum, localSumGm[gmReduceSumOffset], params, padParams); + ReduceResultInQueue.EnQue(localSum); + LocalTensor localSumInUb = ReduceResultInQueue.DeQue(); + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + ReduceSum(globalReduceSum[idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32], localSumInUb[idxM * alignedCoreNum], + s23MiddleResult3[idxM * alignedBaseN], alignedCoreNum); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float aSumPerRow = globalReduceSum.GetValue(idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(cTmp[idxM * alignedBaseN], s23MiddleResult1, aSumPerRow, alignedBaseN); // c_tmp = offset * asum + } + PipeBarrier(); + ReduceResultInQueue.FreeTensor(localSumInUb); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::ProcessScaleAndBias( + uint32_t n, uint32_t curBaseN, uint64_t offsetAndScaleOffset) { + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + LocalTensor scaleF16 = vecInQueue.AllocTensor(); + DataCopyPad2D(scaleF16, scaleGm[offsetAndScaleOffset], 1, curBaseN, n); + vecInQueue.EnQue(scaleF16); + LocalTensor scaleF16InUb = vecInQueue.DeQue(); + Cast(processedScale, scaleF16InUb, RoundMode::CAST_NONE, alignedBaseN); + PipeBarrier(); + vecInQueue.FreeTensor(scaleF16InUb); + if (hasBias) { + #if ORIG_DTYPE_X == DT_FLOAT16 + LocalTensor biasF16 = vecInQueue.AllocTensor(); + DataCopyPad2D(biasF16, biasGm[offsetAndScaleOffset], 1, curBaseN, n); + vecInQueue.EnQue(biasF16); + LocalTensor biasInUb = vecInQueue.DeQue(); + Cast(processedBias, biasInUb, RoundMode::CAST_NONE, alignedBaseN); + vecInQueue.FreeTensor(biasInUb); + #else + event_t eventIdVToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMte2); + WaitFlag(eventIdVToMte2); + DataCopyPad2D(processedBias, biasGm[offsetAndScaleOffset], 1, curBaseN, n); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + #endif + } +} + +template +__aicore__ inline void GMMA16W8MSDCompute::ProcessC1C2( + MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM, uint32_t curSingleM) { + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + LocalTensor c2S32 = vecInQueue.AllocTensor(); + uint64_t c2Offset = postMNConfig.workSpaceOffset + (curSingleM + offsetM) * postMNConfig.n; + DataCopyPad2D(c2S32, mmOutGm[c2Offset], curBaseM, curBaseN, postMNConfig.n); + vecInQueue.EnQue(c2S32); + LocalTensor c2S32InUb = vecInQueue.DeQue(); + Cast(s23MiddleResult2, c2S32InUb, RoundMode::CAST_NONE, curBaseM * alignedBaseN); // c2 + PipeBarrier(); + vecInQueue.FreeTensor(c2S32InUb); + + LocalTensor c1S32 = vecInQueue.AllocTensor(); + uint64_t c1Offset = postMNConfig.workSpaceOffset + offsetM * postMNConfig.n; + DataCopyPad2D(c1S32, mmOutGm[c1Offset], curBaseM, curBaseN, postMNConfig.n); + vecInQueue.EnQue(c1S32); + Muls(s23MiddleResult1, s23MiddleResult2, static_cast(1.0 / 254), curBaseM * alignedBaseN); // c2 / 254 + PipeBarrier(); + LocalTensor c1S32InUb = vecInQueue.DeQue(); + Cast(s23MiddleResult2, c1S32InUb, RoundMode::CAST_NONE, curBaseM * alignedBaseN); // c1 + PipeBarrier(); + vecInQueue.FreeTensor(c1S32InUb); + } + +template +__aicore__ inline void GMMA16W8MSDCompute::CalcCMatrix( + MNConfig& postMNConfig, uint32_t curBaseM, uint32_t curBaseN, uint32_t offsetM) { + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + + Add(s23MiddleResult3, s23MiddleResult2, s23MiddleResult1, curBaseM * alignedBaseN); // c1 + c2 / 254 + PipeBarrier(); + uint32_t gmReduceMaxOffset = (postMNConfig.mAxisBaseOffset + postMNConfig.mIdx * postMNConfig.singleM + offsetM) * + FACTOR_FOR_FLOAT_ALIGN_TO_32; + CopyInAmax(curBaseM, gmReduceMaxOffset); + event_t eventIdMTE2ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventIdMTE2ToS); + WaitFlag(eventIdMTE2ToS); + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + float aMaxPerRow = + aMaxInUb(idxM * FACTOR_FOR_FLOAT_ALIGN_TO_32) / 127.0f; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + // c = (c1 + c2 / 254) * amax / 127 + Muls(s23MiddleResult2[idxM * alignedBaseN], s23MiddleResult3[idxM * alignedBaseN], aMaxPerRow, alignedBaseN); + } + PipeBarrier(); + ReduceResultInQueue.FreeTensor(aMaxInUb); + Add(s23MiddleResult1, s23MiddleResult2, cTmp, curBaseM * alignedBaseN); // c + c_tmp + PipeBarrier(); + for (uint32_t idxM = 0; idxM < curBaseM; ++idxM) { + Mul(s23MiddleResult2[idxM * alignedBaseN], s23MiddleResult1[idxM * alignedBaseN], processedScale, + alignedBaseN); // (c + c_tmp) * scale + PipeBarrier(); + if (hasBias) { + Add(s23MiddleResult2[idxM * alignedBaseN], s23MiddleResult2[idxM * alignedBaseN], processedBias, + alignedBaseN); + } + } + PipeBarrier(); +} + +template +__aicore__ inline void GMMA16W8MSDCompute::CopyOutFinalResult( + uint32_t n, uint32_t curBaseM, uint32_t curBaseN, uint64_t yOffset) { + uint32_t alignedBaseN = AlignUp<32>(curBaseN); + LocalTensor outputInUb = vecOutQueue.AllocTensor(); + #if ORIG_DTYPE_X == DT_FLOAT16 + Cast(outputInUb, s23MiddleResult2, RoundMode::CAST_NONE, curBaseM * alignedBaseN); + #else + Cast(outputInUb, s23MiddleResult2, RoundMode::CAST_RINT, curBaseM * alignedBaseN); + #endif + vecOutQueue.EnQue(outputInUb); + LocalTensor output = vecOutQueue.DeQue(); + DataCopyPad2D(yGm[yOffset], output, curBaseM, curBaseN, alignedBaseN, n); + vecOutQueue.FreeTensor(output); +} + +} // namespace GROUPED_MATMUL + +#endif +#endif // ASCENDC_GROUPED_MATMUL_ANTIQUANT_A16W8_MSD_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4.h new file mode 100644 index 00000000..952fc9ec --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4.h @@ -0,0 +1,377 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + +#ifdef GMM_ANTI_QUANT_A8W4 +namespace GROUPED_MATMUL{ +using namespace matmul; +using namespace AscendC; + +using DTYPE_PERTOKEN_SCALE_A8W4 = float; +using DTYPE_BIAS_A8W4 = float; +using DTYPE_OFFSET_A8W4 = float; + +#ifdef GMM_ANTI_QUANT_A8W4_MSD_OUT_BF16 + using DTYPE_Y_A8W4 = bfloat16_t; + #define GMM_QUANT_BF16 +#else + using DTYPE_Y_A8W4 = half; + #define GMM_QUANT_FLOAT16 +#endif + +using DTYPE_X_DEV_A8W4 = int8_t; +using DTYPE_WEIGHT_DEV_A8W4 = int8_t; +using DTYPE_SCALE_DEV_A8W4 = uint64_t; + + template +__aicore__ inline void DataCopyPad2DA8W4NOMSD(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + +template +class GMMA8W4Compute { +public: + using aT = MatmulType; + using bT = typename mmType::BT; + using cT = MatmulType; + using DTYPE_OUT = DTYPE_Y_A8W4; + +public: + __aicore__ inline GMMA8W4Compute(typename mmType::MT &matmul) : mm(matmul) {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn); + __aicore__ inline void Process(); +private: + __aicore__ inline void InitUbBuffer(); + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorTilingCalc(MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, uint32_t& vecBaseM); + __aicore__ inline void ComputeDequantAndActivate(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t alignBaseN, + uint32_t curVecBaseN, uint32_t offsetM); + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyOffset(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM); + +private: + typename mmType::MT& mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor biasGm; // for 8 * weight + GlobalTensor mmOutGm; + GlobalTensor scaleGm; + GlobalTensor scaleGmF32; + GlobalTensor perTokenScaleGm; + GlobalTensor offsetGm; + GlobalTensor groupTokensGm; + GlobalTensor logitsGm; + GlobalTensor residualGm; + GlobalTensor tokenRanksGm; + GlobalTensor yGm; + GlobalTensor xRowSumGm; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue offsetInQueue; + TQue perTokenScaleInQueue; + TQue xRowSumInQueue; + TBuf tmpBuff; + LocalTensor scaleInUb; + LocalTensor offsetInUb; + LocalTensor mmOutFp32Buf; + LocalTensor pertokenBrcbLocal; + LocalTensor perTokenResBuf; + LocalTensor calcTmpBuf; + uint32_t subBlockIdx = 0; + uint32_t coreIdx = 0; + uint32_t quantGroupSize = 0; + uint32_t cubeCount = 0; + uint32_t vecCount = 0; + uint32_t xRowSumCount = 0; + uint32_t withOffset = 0; + TPipe *pipe; + const GMMBaseParams *tiling; + const TCubeTiling* mmTilingData; + + const uint64_t SYNC_AIV_TO_AIC = 3; + const uint64_t SYNC_AIC_TO_AIV = 5; + const uint32_t BUFFER_NUM = 2; + const uint32_t MM_BASE_BLOCK_OFFSET = 16384; // baseM * baseN = 128 * 128 +}; + +template +__aicore__ inline void GMMA8W4Compute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn) +{ + tiling = tilingData; + xRowSumCount = tiling->m; + withOffset = false; + + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(workspace)); + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scale)); + scaleGmF32.SetGlobalBuffer(GetTensorAddr(0, scale)); + perTokenScaleGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PERTOKEN_SCALE_A8W4 *>(pertoken_scale)); + groupTokensGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(group_tokens)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + const size_t weightSize = tilingData->groupNum * tilingData->k * tilingData->n * sizeof(int8_t); + mmOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ cT::T *>(workspace + weightSize)); + + this->mmTilingData = mmTilingData; + quantGroupSize = tiling->k / tiling->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + if (GetTaskRation() != 0) { + coreIdx /= GetTaskRation(); + } + } + pipe = tPipeIn; + InitUbBuffer(); +} + +template +__aicore__ inline void GMMA8W4Compute::InitUbBuffer() +{ + if ASCEND_IS_AIC { + return; + } + pipe->InitBuffer(perTokenScaleInQueue, BUFFER_NUM, mmTilingData->baseM * sizeof(float)); + pipe->InitBuffer(vecInQueue, BUFFER_NUM, tiling->ubCalSize * sizeof(cT::T)); + pipe->InitBuffer(vecOutQueue, BUFFER_NUM, tiling->ubCalSize * sizeof(DTYPE_OUT)); + pipe->InitBuffer(tmpBuff, tiling->ubRestBytes); + + uint32_t ubCalSizeFloat = tiling->ubCalSize * sizeof(float); + mmOutFp32Buf = tmpBuff.GetWithOffset(tiling->ubCalSize, 0); + uint32_t offset = ubCalSizeFloat; + pertokenBrcbLocal = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + offset += ubCalSizeFloat; + perTokenResBuf = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + offset += ubCalSizeFloat; + calcTmpBuf = tmpBuff.GetWithOffset(ubCalSizeFloat, offset); +} + +template +__aicore__ inline void GMMA8W4Compute::Process() +{ + MNConfig mnConfig; + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + mnConfig.singleM = mnConfig.baseM; + mnConfig.singleN = mnConfig.baseN; + mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN); + for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) { + int32_t m = static_cast(groupTokensGm.GetValue(groupIdx)); + if (m <= 0) { + continue; + } + mnConfig.m = static_cast(m); + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + tiling->coreNum; + uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, preCount, thresholdM_dimN); + MMCompute(groupIdx, mnConfig); + if ASCEND_IS_AIV { + VectorCompute(groupIdx, mnConfig); + } + curBlock += tiling->coreNum; + } + preCount = curCount % tiling->coreNum; + mnConfig.offsetM += mnConfig.m; + } +} + +template +__aicore__ inline void GMMA8W4Compute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + mnConfig.workSpaceOffset = MM_BASE_BLOCK_OFFSET * \ + (coreIdx + (cubeCount % tiling->parallNum) * tiling->coreNum); + if ASCEND_IS_AIC { + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = tiling->n - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + + uint64_t xOffset = (mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM) * tiling->k; + uint64_t weightOffset; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = static_cast(groupIdx) * tiling->n * tiling->k + tailN * tiling->k; + } else { + weightOffset = static_cast(groupIdx) * tiling->n * tiling->k + tailN; + } + if (cubeCount >= tiling->parallNum) { + CrossCoreWaitFlag(SYNC_AIV_TO_AIC); + } + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); + GlobalTensor weightSlice; + for (uint32_t loopK = 0; loopK < tiling->quantGroupNum; loopK++) { + mm.SetTensorA(xGm[xOffset + loopK * quantGroupSize]); + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * 32]; // 32: NZ分型32对齐 + } else { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * tiling->n]; + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + mm.SetQuantVector(scaleGm[groupIdx * tiling->n * tiling->quantGroupNum + loopK * tiling->n + tailN]); + uint64_t worskspaceOffset = mnConfig.workSpaceOffset; + mm.Iterate(); + mm.GetTensorC(mmOutGm[worskspaceOffset], loopK == 0 ? 0 : 1, true); + worskspaceOffset += MM_BASE_BLOCK_OFFSET; + } + CrossCoreSetFlag<2, PIPE_FIX>(SYNC_AIC_TO_AIV); // 2: mode为2, group内同步 + } + cubeCount++; +} + +template +__aicore__ inline void GMMA8W4Compute::DataCopyOffset(uint32_t curBaseN, uint32_t alignBaseN, uint64_t offsetOffset) +{ + DataCopyExtParams offsetParams{1, static_cast(curBaseN * sizeof(float)), 1, 1, 0}; + DataCopyPadExtParams offsetPadParams; + LocalTensor offsetLocal = offsetInQueue.AllocTensor(); + DataCopyPad(offsetLocal, offsetGm[offsetOffset], offsetParams, offsetPadParams); + offsetInQueue.EnQue(offsetLocal); + offsetInUb = offsetInQueue.DeQue(); +} + +template +__aicore__ inline void GMMA8W4Compute::VectorTilingCalc( + MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, uint32_t& vecBaseM) +{ + curCubeSingleN = mnConfig.nIdx == mnConfig.blockDimN - 1 ? + tiling->n - mnConfig.nIdx * mnConfig.singleN : mnConfig.singleN; + curCubeSingleM = mnConfig.mIdx == mnConfig.blockDimM - 1 ? + mnConfig.m - mnConfig.mIdx * mnConfig.singleM : mnConfig.singleM; + vecBaseM = tiling->ubCalSize / (Ceil(mnConfig.baseN, 8U) * 8); // 8: num int32_t in 32B ub block + vecBaseM = vecBaseM < curCubeSingleM ? vecBaseM : curCubeSingleM; +} + +template +__aicore__ inline void GMMA8W4Compute::VectorCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + uint32_t curCubeSingleN; + uint32_t curCubeSingleM; + uint32_t vecBaseM; + VectorTilingCalc(mnConfig, curCubeSingleN, curCubeSingleM, vecBaseM); + uint32_t mGlobalOffset = mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM; // 2: 2 lines int4 to 1 line int8 + + uint64_t outOffset = mGlobalOffset * tiling->n + mnConfig.nIdx * mnConfig.singleN; + uint32_t curVecBaseN = mnConfig.baseN; + uint32_t taskRation = GetTaskRation(); // 2 + CrossCoreWaitFlag(SYNC_AIC_TO_AIV); + uint32_t nCount = 0; + for (uint32_t offsetN = 0; offsetN < curCubeSingleN; offsetN += mnConfig.baseN) { + if (unlikely(offsetN + mnConfig.baseN >= curCubeSingleN)) curVecBaseN = curCubeSingleN - offsetN; + uint32_t alignBaseN = Ceil(curVecBaseN, 16U) * 16U; // 16: fp16 num per 32B + uint32_t curVecBaseM = vecBaseM; + uint64_t mmOutOffset = mnConfig.workSpaceOffset + offsetN * mnConfig.baseM; + uint32_t mCount = 0; + for (uint32_t offsetM = 0; offsetM < curCubeSingleM; offsetM += vecBaseM) { + vecCount++; + if (taskRation != 0 && vecCount % taskRation != subBlockIdx) { continue; } + if (unlikely(offsetM + vecBaseM >= curCubeSingleM)) { curVecBaseM = curCubeSingleM - offsetM; } + LocalTensor mmOutLocal = vecInQueue.AllocTensor(); + DataCopyPad2DA8W4NOMSD(mmOutLocal, mmOutGm[mmOutOffset + offsetM * curVecBaseN], curVecBaseM, curVecBaseN, curVecBaseN); + vecInQueue.EnQue(mmOutLocal); + + ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM); + LocalTensor yLocal = vecOutQueue.DeQue(); + DataCopyPad2D(yGm[outOffset + offsetM * tiling->n + offsetN], yLocal, + curVecBaseM, curVecBaseN, alignBaseN, tiling->n); + vecOutQueue.FreeTensor(yLocal); + } + } + CrossCoreSetFlag<2, PIPE_MTE2>(SYNC_AIV_TO_AIC); // 2: mode为2, group内同步 +} + +template +__aicore__ inline void GMMA8W4Compute::ComputeDequantAndActivate(MNConfig& mnConfig, + uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, uint32_t offsetM) +{ + uint32_t computeSize = curVecBaseM * alignBaseN; + LocalTensor mmOutInUb = vecInQueue.DeQue(); + uint32_t castSize = 0; + if constexpr (mmType::BT::format == CubeFormat::ND) { + castSize = (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + castSize = computeSize; + } + Cast(mmOutFp32Buf, mmOutInUb, RoundMode::CAST_NONE, castSize); + PipeBarrier(); + vecInQueue.FreeTensor(mmOutInUb); + + DataCopyPerTokenScaleAndBrcb(mnConfig, curVecBaseM, alignBaseN, offsetM); + + Mul(perTokenResBuf, mmOutFp32Buf, pertokenBrcbLocal, computeSize); + PipeBarrier(); + LocalTensor yLocalInUb = vecOutQueue.AllocTensor(); + // Cast后获得最终输出 + Cast(yLocalInUb, perTokenResBuf, RoundMode::CAST_RINT, computeSize); + PipeBarrier(); + vecOutQueue.EnQue(yLocalInUb); +} + +template +__aicore__ inline void GMMA8W4Compute::DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, + uint32_t curBaseM, uint32_t alignBaseN, uint32_t offsetM) +{ + uint64_t perTokenScaleOffset = mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM + offsetM; + DataCopyPadExtParams padParams; + DataCopyExtParams perTokenScaleParams{1, static_cast(curBaseM * sizeof(float)), 0, 0, 0}; + + LocalTensor perTokenScaleLocal = perTokenScaleInQueue.AllocTensor(); + DataCopyPad(perTokenScaleLocal, perTokenScaleGm[perTokenScaleOffset], perTokenScaleParams, padParams); + perTokenScaleInQueue.EnQue(perTokenScaleLocal); + + perTokenScaleLocal = perTokenScaleInQueue.DeQue(); + + const uint32_t broadCastDst[2] = {curBaseM, alignBaseN}; + const uint32_t broadCastSrc[2] = {curBaseM, 1}; + BroadCast(pertokenBrcbLocal, perTokenScaleLocal, broadCastDst, broadCastSrc, calcTmpBuf); + perTokenScaleInQueue.FreeTensor(perTokenScaleLocal); +} +} // namespace GROUPED_MATMUL +#endif +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd.h new file mode 100644 index 00000000..883857d0 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd.h @@ -0,0 +1,531 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4_msd.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + +#ifdef GMM_ANTI_QUANT_A8W4_MSD +namespace GROUPED_MATMUL{ +using namespace matmul; +using namespace AscendC; + +using DTYPE_PERTOKEN_SCALE_A8W4MSD = float; +using DTYPE_BIAS_A8W4MSD = float; +using DTYPE_OFFSET_A8W4MSD = float; + +#ifdef GMM_ANTI_QUANT_A8W4_MSD_OUT_BF16 + using DTYPE_Y_A8W4MSD = bfloat16_t; +#else + using DTYPE_Y_A8W4MSD = half; +#endif + +using DTYPE_X_DEV_A8W4MSD = int4b_t; +using DTYPE_WEIGHT_DEV_A8W4MSD = int4b_t; +using DTYPE_SCALE_DEV_A8W4MSD = uint64_t; + +constexpr uint64_t SYNC_AIV_TO_AIC = 3; +constexpr uint64_t SYNC_AIC_TO_AIV = 5; +constexpr uint32_t BUFFER_NUM = 1; +constexpr uint32_t MM_BASE_BLOCK_OFFSET = 32768; // baseM * baseN = 128 * 256 + +template +__aicore__ inline void DataCopyPad2DA8W4(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + // 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐 + params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2DA8W4ND(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + +template +__aicore__ inline void DataCopyPad2DA8W4(const GlobalTensor dst, const LocalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0, uint32_t dstDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + // 32: ub访问粒度为32B + params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32; + params.dstStride = (dstDim0 - dim0) * sizeof(T); + DataCopyPad(dst, src, params); +} + +template +class GMMA8W4MSDCompute { +public: + using aT = MatmulType; + using bT = typename mmType::BT; + using biasT = MatmulType; + using cT = MatmulType; + using DTYPE_OUT = DTYPE_Y_A8W4MSD; + +public: + __aicore__ inline GMMA8W4MSDCompute(typename mmType::MT &matmul) : mm(matmul) {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn); + __aicore__ inline void Process(); +private: + __aicore__ inline void InitUbBuffer(); + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void ComputeDequantAndActivate(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t alignBaseN, + uint32_t curVecBaseN, uint32_t offsetM); + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyOffset(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM); + __aicore__ inline void DataCopyAndBrcbOfRowSum(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM); + +private: + typename mmType::MT& mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor biasGm; // for 8 * weight + GlobalTensor mmOutGm; + GlobalTensor scaleGm; + GlobalTensor perTokenScaleGm; + GlobalTensor offsetGm; + GlobalTensor groupTokensGm; + GlobalTensor logitsGm; + GlobalTensor residualGm; + GlobalTensor tokenRanksGm; + GlobalTensor yGm; + GlobalTensor xRowSumGm; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue offsetInQueue; + TQue perTokenScaleInQueue; + TQue xRowSumInQueue; + TBuf tmpBuff; + LocalTensor scaleInUb; + LocalTensor offsetInUb; + LocalTensor buffer1; + LocalTensor buffer2; + LocalTensor buffer3; + LocalTensor buffer4; + LocalTensor buffer5; + LocalTensor buffer6; + LocalTensor buffer7; + uint32_t subBlockIdx; + uint32_t coreIdx; + uint32_t quantGroupSize; + uint32_t cubeCount = 0; + uint32_t vecCount = 0; + uint32_t xRowSumCount; + uint32_t withOffset; + TPipe *pipe; + const GMMBaseParams *tiling; + const TCubeTiling* mmTilingData; +}; + +template +__aicore__ inline void GMMA8W4MSDCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn) +{ + tiling = tilingData; + xRowSumCount = tiling->m; + withOffset = tiling->withOffset; + + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + weightGm.SetGlobalBuffer(GetTensorAddr(0, weight)); + biasGm.SetGlobalBuffer(GetTensorAddr(0, bias)); + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scale)); + perTokenScaleGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PERTOKEN_SCALE_A8W4MSD *>(pertoken_scale)); + groupTokensGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(group_tokens)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + if (withOffset == WITH_OFFSET) { + xRowSumGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(workspace)); + mmOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ cT::T *>(workspace + xRowSumCount * sizeof(float))); + offsetGm.SetGlobalBuffer(GetTensorAddr(0, offset)); + } else { + mmOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ cT::T *>(workspace)); + } + + this->mmTilingData = mmTilingData; + quantGroupSize = tiling->k / tiling->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + if (GetTaskRation() != 0) { + coreIdx /= GetTaskRation(); + } + } + pipe = tPipeIn; + InitUbBuffer(); +} + +template +__aicore__ inline void GMMA8W4MSDCompute::InitUbBuffer() +{ + if ASCEND_IS_AIC { + return; + } + pipe->InitBuffer(scaleInQueue, BUFFER_NUM, mmTilingData->baseN * sizeof(DTYPE_BIAS_A8W4MSD)); // bias queue + if (withOffset == WITH_OFFSET) { + pipe->InitBuffer(offsetInQueue, BUFFER_NUM, mmTilingData->baseN * sizeof(DTYPE_OFFSET_A8W4MSD)); + pipe->InitBuffer(xRowSumInQueue, BUFFER_NUM, Ceil(tiling->vBaseM * sizeof(float), 32) * 32); + } + pipe->InitBuffer(perTokenScaleInQueue, BUFFER_NUM, Ceil(tiling->vBaseM * sizeof(float), 32) * 32); + pipe->InitBuffer(vecInQueue, BUFFER_NUM, tiling->ubCalSize * 2 * sizeof(cT::T)); + pipe->InitBuffer(vecOutQueue, BUFFER_NUM, tiling->ubCalSize * sizeof(DTYPE_OUT)); + pipe->InitBuffer(tmpBuff, tiling->ubRestBytes); + + uint32_t ubCalSizeFloat = tiling->ubCalSize * sizeof(float); + // ub分配，依次划分中间结果，划分方式参考设计文档 + buffer1 = tmpBuff.GetWithOffset(tiling->ubCalSize * 2, 0); + uint32_t offset = ubCalSizeFloat * 2; + buffer2 = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + offset += ubCalSizeFloat; + buffer3 = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + buffer4 = tmpBuff.GetWithOffset(tiling->ubCalSize, 0); + buffer5 = tmpBuff.GetWithOffset(tiling->ubCalSize, ubCalSizeFloat); + if (withOffset == WITH_OFFSET) { + buffer6 = tmpBuff.GetWithOffset(ubCalSizeFloat, ubCalSizeFloat); + buffer7 = tmpBuff.GetWithOffset(ubCalSizeFloat, 0); + } else { + buffer6 = tmpBuff.GetWithOffset(2 * ubCalSizeFloat, offset); + buffer7 = tmpBuff.GetWithOffset(2 * ubCalSizeFloat, 0); + } +} + +template +__aicore__ inline void GMMA8W4MSDCompute::Process() +{ + MNConfig mnConfig; + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + mnConfig.singleM = mnConfig.baseM; + mnConfig.singleN = mnConfig.baseN; + mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN); + if ASCEND_IS_AIC { + SyncAll(); + } + for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) { + int32_t m = static_cast(groupTokensGm.GetValue(groupIdx)); + if (m <= 0) { + continue; + } + mnConfig.m = static_cast(m) * 2; // 2: int8 has been split in 2 int4 + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + tiling->coreNum; + + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN; + MMCompute(groupIdx, mnConfig); + if ASCEND_IS_AIV { + VectorCompute(groupIdx, mnConfig); + } + curBlock += tiling->coreNum; + } + preCount = curCount % tiling->coreNum; + mnConfig.offsetM += mnConfig.m; + } +} + +template +__aicore__ inline void GMMA8W4MSDCompute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + mnConfig.workSpaceOffset = MM_BASE_BLOCK_OFFSET * \ + (coreIdx + (cubeCount % tiling->parallNum) * tiling->coreNum); + if ASCEND_IS_AIC { + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = tiling->n - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + + uint64_t xOffset = (mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM) * tiling->k; + uint64_t weightOffset; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = static_cast(groupIdx) * tiling->n * tiling->k + tailN * tiling->k; + } else { + weightOffset = static_cast(groupIdx) * tiling->n * tiling->k + tailN; + } + if (cubeCount >= tiling->parallNum) { + CrossCoreWaitFlag(SYNC_AIV_TO_AIC); + } + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); // 8, 256, 512 --> 514us + GlobalTensor weightSlice; + for (uint32_t loopK = 0; loopK < tiling->quantGroupNum; loopK++) { + mm.SetTensorA(xGm[xOffset + loopK * quantGroupSize]); + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * tiling->n]; + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + mm.SetQuantVector(scaleGm[groupIdx * tiling->n * tiling->quantGroupNum + loopK * tiling->n + tailN]); + uint64_t worskspaceOffset = mnConfig.workSpaceOffset; +#ifndef __CCE_KT_TEST__ + mm.Iterate(); + mm.GetTensorC(mmOutGm[worskspaceOffset], loopK == 0 ? 0 : 1, true); +#endif + worskspaceOffset += MM_BASE_BLOCK_OFFSET; + } + CrossCoreSetFlag<2, PIPE_FIX>(SYNC_AIC_TO_AIV); // 2: mode为2, group内同步 + } + cubeCount++; +} +template +__aicore__ inline void GMMA8W4MSDCompute::VectorCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + uint32_t curCubeSingleN = mnConfig.singleN; + if (mnConfig.nIdx == mnConfig.blockDimN - 1) { curCubeSingleN = tiling->n - mnConfig.nIdx * mnConfig.singleN; } + uint32_t curCubeSingleM = mnConfig.singleM / 2; // 2: 2 lines int4 to 1 line int8 + uint32_t mGlobalOffset = mnConfig.offsetM / 2 + mnConfig.mIdx * curCubeSingleM; // 2: 2 lines int4 to 1 line int8 + uint64_t outOffset = mGlobalOffset * tiling->n + mnConfig.nIdx * mnConfig.singleN; + // 2: 2 lines int4 to 1 line int8 + if (mnConfig.mIdx == mnConfig.blockDimM - 1) { curCubeSingleM = mnConfig.m / 2 - mnConfig.mIdx * curCubeSingleM; } + uint32_t vecBaseM = tiling->ubCalSize / (Ceil(mnConfig.baseN, uint32_t(8)) * 8); // 8: num int32_t in 32B ub block 32*256/256 + vecBaseM = vecBaseM < curCubeSingleM ? vecBaseM : curCubeSingleM; + uint32_t curVecBaseN = mnConfig.baseN; + uint64_t scaleOffset = groupIdx * tiling->n + mnConfig.nIdx * mnConfig.singleN; + uint64_t offsetOffset = scaleOffset; + uint32_t taskRation = GetTaskRation(); + CrossCoreWaitFlag(SYNC_AIC_TO_AIV); + for (uint32_t offsetN = 0; offsetN < curCubeSingleN; offsetN += mnConfig.baseN) { + if (unlikely(offsetN + mnConfig.baseN >= curCubeSingleN)) curVecBaseN = curCubeSingleN - offsetN; + uint32_t alignBaseN = Ceil(curVecBaseN, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block + DataCopyScale(curVecBaseN, alignBaseN, scaleOffset + offsetN); + if (withOffset == WITH_OFFSET) { + DataCopyOffset(curVecBaseN, alignBaseN, offsetOffset + offsetN); + } + + uint32_t curVecBaseM = vecBaseM; + uint64_t mmOutOffset = mnConfig.workSpaceOffset + offsetN * mnConfig.baseM; + for (uint32_t offsetM = 0; offsetM < curCubeSingleM; offsetM += vecBaseM) { + vecCount++; + if (taskRation != 0 && vecCount % taskRation != subBlockIdx) { continue; } + if (unlikely(offsetM + vecBaseM >= curCubeSingleM)) { curVecBaseM = curCubeSingleM - offsetM; } + LocalTensor mmOutLocal = vecInQueue.AllocTensor(); + if constexpr (mmType::BT::format == CubeFormat::ND) { + DataCopyPad2DA8W4ND(mmOutLocal, mmOutGm[mmOutOffset + offsetM * 2 * curVecBaseN], + curVecBaseM, curVecBaseN, curVecBaseN * 2); // 2: 2 lines int4 to 1 line int8 + } else { + DataCopyPad2DA8W4(mmOutLocal, mmOutGm[mmOutOffset + offsetM * 2 * curVecBaseN], + curVecBaseM, curVecBaseN, curVecBaseN * 2); // 2: 2 lines int4 to 1 line int8 + } + uint32_t targetAddr; + if constexpr (mmType::BT::format == CubeFormat::ND) { + alignBaseN = (alignBaseN + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } + targetAddr = curVecBaseM * alignBaseN; + uint64_t lowBitAddr = mmOutOffset + (offsetM * 2 + 1) * curVecBaseN; + if constexpr (mmType::BT::format == CubeFormat::ND) { + DataCopyPad2DA8W4ND(mmOutLocal[targetAddr], + mmOutGm[lowBitAddr], + curVecBaseM, curVecBaseN, curVecBaseN * 2); // 2: 2 lines int4 to 1 line int8 + } else { + DataCopyPad2DA8W4(mmOutLocal[targetAddr], + mmOutGm[lowBitAddr], + curVecBaseM, curVecBaseN, curVecBaseN * 2); // 2: 2 lines int4 to 1 line int8 + } + vecInQueue.EnQue(mmOutLocal); + ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM); + LocalTensor yLocal = vecOutQueue.DeQue(); + DataCopyPad2DA8W4(yGm[outOffset + offsetM * tiling->n + offsetN], yLocal, + curVecBaseM, curVecBaseN, alignBaseN, tiling->n); + vecOutQueue.FreeTensor(yLocal); + } + scaleInQueue.FreeTensor(scaleInUb); + if (withOffset == WITH_OFFSET) { + offsetInQueue.FreeTensor(offsetInUb); + } + } + CrossCoreSetFlag<2, PIPE_MTE2>(SYNC_AIV_TO_AIC); // 2: mode为2, group内同步 +} +template +__aicore__ inline void GMMA8W4MSDCompute::ComputeDequantAndActivate(MNConfig& mnConfig, + uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, uint32_t offsetM) +{ + uint32_t computeSize = curVecBaseM * alignBaseN; + LocalTensor mmOutInUb = vecInQueue.DeQue(); + uint32_t castSize; + if constexpr (mmType::BT::format == CubeFormat::ND) { + castSize = computeSize + (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + castSize = computeSize * 2; + } + Cast(buffer1, mmOutInUb, RoundMode::CAST_NONE, castSize); + PipeBarrier(); + vecInQueue.FreeTensor(mmOutInUb); + const float RIGHT_MOVE = 16.0f; // right move int4 to int8 + Muls(buffer2, buffer1, RIGHT_MOVE, computeSize); + PipeBarrier(); + uint32_t addStartAddr; + if constexpr (mmType::BT::format == CubeFormat::ND) { + addStartAddr = (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + addStartAddr = computeSize; + } + Add(buffer3, buffer1[addStartAddr], buffer2, computeSize); + PipeBarrier(); + + uint32_t loop = alignBaseN / 64; // 256B为64个float，alignBaseN需约束为64倍数 + uint8_t blkStride = static_cast(alignBaseN * sizeof(float) / 32); //32: 单位32B + BinaryRepeatParams param(1, 1, 1, blkStride, blkStride, 0); + uint64_t mask = 64; + uint64_t last = alignBaseN % 64; + for (uint32_t i = 0; i < loop; i++) { + uint32_t offset = i * 64; // 每次64个元素 + Add(buffer2[offset], buffer3[offset], scaleInUb[offset], mask, curVecBaseM, param); + } + PipeBarrier(); + if (unlikely(last > 0)) { + uint32_t offset = loop * 64; + Add(buffer2[offset], buffer3[offset], scaleInUb[offset], last, curVecBaseM, param); + } + PipeBarrier(); + + if (withOffset == WITH_OFFSET) { + DataCopyAndBrcbOfRowSum(mnConfig, curVecBaseM, alignBaseN, offsetM); + PipeBarrier(); + + for (uint32_t i = 0; i < loop; i++) { + uint32_t offset = i * 64; // 每次64个元素 + Mul(buffer4[offset], buffer3[offset], offsetInUb[offset], mask, curVecBaseM, param); + } + PipeBarrier(); + if (unlikely(last > 0)) { + uint32_t offset = loop * 64; + Mul(buffer4[offset], buffer3[offset], offsetInUb[offset], last, curVecBaseM, param); + } + PipeBarrier(); + + Add(buffer5, buffer4, buffer2, computeSize); + PipeBarrier(); + } + + + DataCopyPerTokenScaleAndBrcb(mnConfig, curVecBaseM, alignBaseN, offsetM); + PipeBarrier(); + + LocalTensor yLocalInUb = vecOutQueue.AllocTensor(); + if (withOffset == WITH_OFFSET) { + Mul(buffer4, buffer5, buffer3, computeSize); + } else { + Mul(buffer4, buffer2, buffer3, computeSize); + } + PipeBarrier(); + Cast(yLocalInUb, buffer4, RoundMode::CAST_RINT, computeSize); + vecOutQueue.EnQue(yLocalInUb); +} + +template +__aicore__ inline void GMMA8W4MSDCompute::DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset) +{ + // GM拷贝scale + DataCopyPadExtParams padParams; + DataCopyExtParams scaleParams{1, static_cast(curBaseN * sizeof(float)), 1, 1, 0}; + LocalTensor scaleLocal = scaleInQueue.AllocTensor(); + DataCopyPad(scaleLocal, biasGm[scaleOffset], scaleParams, padParams); + scaleInQueue.EnQue(scaleLocal); + scaleInUb = scaleInQueue.DeQue(); +} + +template +__aicore__ inline void GMMA8W4MSDCompute::DataCopyOffset(uint32_t curBaseN, uint32_t alignBaseN, uint64_t offsetOffset) +{ + DataCopyExtParams offsetParams{1, static_cast(curBaseN * sizeof(float)), 1, 1, 0}; + DataCopyPadExtParams offsetPadParams; + LocalTensor offsetLocal = offsetInQueue.AllocTensor(); + DataCopyPad(offsetLocal, offsetGm[offsetOffset], offsetParams, offsetPadParams); + offsetInQueue.EnQue(offsetLocal); + offsetInUb = offsetInQueue.DeQue(); +} + + +template +__aicore__ inline void GMMA8W4MSDCompute::DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, + uint32_t curBaseM, uint32_t alignBaseN, uint32_t offsetM) +{ + uint64_t perTokenScaleOffset = mnConfig.offsetM / 2 + mnConfig.mIdx * mnConfig.singleM / 2 + offsetM; //2: m方向两行合并为1行 + uint32_t alignBaseM = Ceil(curBaseM, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block + // GM拷贝per token scale + DataCopyPadExtParams padParams; + DataCopyExtParams perTokenScaleParams{1, static_cast(curBaseM * sizeof(float)), 0, 0, 0}; + LocalTensor perTokenScaleLocal = perTokenScaleInQueue.AllocTensor(); + DataCopyPad(perTokenScaleLocal, perTokenScaleGm[perTokenScaleOffset], perTokenScaleParams, padParams); + + perTokenScaleInQueue.EnQue(perTokenScaleLocal); + + perTokenScaleLocal = perTokenScaleInQueue.DeQue(); + auto scaleTmp = perTokenScaleLocal; + + const uint32_t broadCastDst[2] = {curBaseM, alignBaseN}; + const uint32_t broadCastSrc[2] = {curBaseM, 1}; + BroadCast(buffer3, scaleTmp, broadCastDst, broadCastSrc, buffer7); + perTokenScaleInQueue.FreeTensor(perTokenScaleLocal); +} + +template +__aicore__ inline void GMMA8W4MSDCompute::DataCopyAndBrcbOfRowSum(MNConfig& mnConfig, + uint32_t curBaseM, uint32_t alignBaseN, uint32_t offsetM) +{ + const uint64_t xRowSumOffset = mnConfig.offsetM / 2 + mnConfig.mIdx * mnConfig.singleM / 2 + offsetM; //2: m方向两行合并为1行 + + DataCopyPadExtParams padParams; + DataCopyExtParams xRowSumParams{1, static_cast(curBaseM * sizeof(float)), 0, 0, 0}; + LocalTensor xRowSumLocal = xRowSumInQueue.AllocTensor(); + DataCopyPad(xRowSumLocal, xRowSumGm[xRowSumOffset], xRowSumParams, padParams); + + xRowSumInQueue.EnQue(xRowSumLocal); + xRowSumLocal = xRowSumInQueue.DeQue(); + + const uint32_t xRowSumBroadCastDst[2] = {curBaseM, alignBaseN}; + const uint32_t xRowSumBroadCastSrc[2] = {curBaseM, 1}; + BroadCast(buffer3, xRowSumLocal, xRowSumBroadCastDst, xRowSumBroadCastSrc, buffer6); + xRowSumInQueue.FreeTensor(xRowSumLocal); +} +} // namespace GROUPED_MATMUL +#endif +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_new.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_new.h new file mode 100644 index 00000000..0cbdbacd --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_new.h @@ -0,0 +1,523 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4_msd_new.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_NEW_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_NEW_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + +#ifdef GMM_ANTI_QUANT_A8W4_MSD +namespace GROUPED_MATMUL{ +using namespace matmul; +using namespace AscendC; + +using DTYPE_PERTOKEN_SCALE_A8W4MSD_NEW = float; +using DTYPE_BIAS_A8W4MSD_NEW = float; +using DTYPE_SCALE_DEV_A8W4MSD_NEW = int64_t; + +#ifdef GMM_ANTI_QUANT_A8W4_MSD_OUT_BF16 + using DTYPE_Y_A8W4MSD = bfloat16_t; +#else + using DTYPE_Y_A8W4MSD = half; +#endif + +using DTYPE_X_DEV_A8W4MSD_NEW = int4b_t; +using DTYPE_WEIGHT_DEV_A8W4MSD_NEW = int4b_t; + +constexpr uint64_t SYNC_AIV_TO_AIC_NEW = 3; +constexpr uint64_t SYNC_AIC_TO_AIV_NEW = 5; +constexpr uint32_t BUFFER_NUM_NEW = 1; +constexpr uint32_t MM_BASE_BLOCK_OFFSET_NEW = 16384; // baseM * baseN = 32 * 512 + +template +__aicore__ inline void DataCopyPad2DA8W4New(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + // 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐 + params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2DA8W4NDNew(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + +template +__aicore__ inline void DataCopyPad2DA8W4New(const GlobalTensor dst, const LocalTensor src, uint32_t dim1, uint32_t dim0, + uint32_t srcDim0, uint32_t dstDim0) { + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + // 32: ub访问粒度为32B + params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32; + params.dstStride = (dstDim0 - dim0) * sizeof(T); + DataCopyPad(dst, src, params); +} + +template +class GMMA8W4MSDComputeNew { +public: + using aT = MatmulType; + using bT = typename mmType::BT; + using biasT = MatmulType; + using cT = MatmulType; + using DTYPE_OUT = DTYPE_Y_A8W4MSD; + +public: + __aicore__ inline GMMA8W4MSDComputeNew(typename mmType::MT &matmul) : mm(matmul) {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn); + __aicore__ inline void Process(); +private: + __aicore__ inline void InitUbBuffer(); + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig); + __aicore__ inline void VectorCompute(uint32_t groupIdx, MNConfig& mnConfig, int loopK, uint64_t tailN, uint64_t workspaceOffset); + __aicore__ inline void DataCopyScaleDequant(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void ComputeDequantAndActivate(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t alignBaseN, + uint32_t curVecBaseN, uint32_t offsetM, int loopK, bool isLastBlock, uint32_t totalVecBaseM, uint32_t offsetVec0M); + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM, uint32_t offsetVec0M); + +private: + typename mmType::MT& mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor biasGm; + GlobalTensor mmOutGm; + GlobalTensor scaleGm; + GlobalTensor perTokenScaleGm; + GlobalTensor groupTokensGm; + GlobalTensor logitsGm; + GlobalTensor residualGm; + GlobalTensor tokenRanksGm; + GlobalTensor yGm; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue scaleInQueueDeqScale; + TQue perTokenScaleInQueue; + TBuf tmpBuff; + LocalTensor scaleInUb; + LocalTensor scaleInUb2; + LocalTensor scaleInUb2_i32; + LocalTensor scaleInUb2_f32; + + LocalTensor buffer1; + LocalTensor bufferAdd; + LocalTensor buffer2; + LocalTensor buffer3; + LocalTensor buffer4; + LocalTensor buffer5; + LocalTensor buffer6; + LocalTensor buffer7; + uint32_t subBlockIdx; + uint32_t coreIdx; + uint32_t quantGroupSize; + uint32_t cubeCount = 0; + uint32_t cubeId = 0; + uint32_t vecCount = 0; + TPipe *pipe; + const GMMBaseParams *tiling; + const TCubeTiling* mmTilingData; +}; + +template +__aicore__ inline void GMMA8W4MSDComputeNew::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR group_tokens, GM_ADDR scale, GM_ADDR pertoken_scale, GM_ADDR offset, GM_ADDR logits, GM_ADDR token_ranks, GM_ADDR residual, + GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* tilingData, const TCubeTiling* mmTilingData, TPipe *tPipeIn) +{ + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + weightGm.SetGlobalBuffer(GetTensorAddr(0, weight)); + biasGm.SetGlobalBuffer(GetTensorAddr(0, bias)); + mmOutGm.SetGlobalBuffer(reinterpret_cast<__gm__ cT::T *>(workspace)); + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scale)); + perTokenScaleGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PERTOKEN_SCALE_A8W4MSD_NEW *>(pertoken_scale)); + groupTokensGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(group_tokens)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + + tiling = tilingData; + this->mmTilingData = mmTilingData; + quantGroupSize = tiling->k / tiling->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + if ASCEND_IS_AIV { + if (GetTaskRation() != 0) { + coreIdx /= GetTaskRation(); + } + } + pipe = tPipeIn; + InitUbBuffer(); +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::InitUbBuffer() +{ + if ASCEND_IS_AIC { + return; + } + pipe->InitBuffer(scaleInQueue, BUFFER_NUM_NEW, mmTilingData->baseN * sizeof(DTYPE_BIAS_A8W4MSD_NEW)); + pipe->InitBuffer(scaleInQueueDeqScale, BUFFER_NUM_NEW, mmTilingData->baseN * sizeof(uint64_t)); + pipe->InitBuffer(perTokenScaleInQueue, BUFFER_NUM_NEW, Ceil(tiling->vBaseM * sizeof(float), 32) * 32); + pipe->InitBuffer(vecInQueue, BUFFER_NUM_NEW, tiling->ubCalSize * 2 * sizeof(cT::T)); + pipe->InitBuffer(vecOutQueue, BUFFER_NUM_NEW, tiling->ubCalSize * sizeof(DTYPE_OUT)); + pipe->InitBuffer(tmpBuff, tiling->ubRestBytes); + uint32_t ubCalSizeFloat = tiling->ubCalSize * sizeof(float); + + buffer1 = tmpBuff.GetWithOffset(tiling->ubCalSize * 2, 0); + bufferAdd = tmpBuff.GetWithOffset(tiling->ubCalSize * 2, tiling->ubCalSize * 2 * sizeof(float)); + buffer7 = tmpBuff.GetWithOffset(2 * ubCalSizeFloat, 0); + uint32_t offset = ubCalSizeFloat * 2; + buffer2 = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + buffer6 = tmpBuff.GetWithOffset(2 * ubCalSizeFloat, offset); + offset += ubCalSizeFloat; + buffer3 = tmpBuff.GetWithOffset(tiling->ubCalSize, offset); + buffer4 = tmpBuff.GetWithOffset(tiling->ubCalSize, 0); + buffer5 = tmpBuff.GetWithOffset(tiling->ubCalSize, ubCalSizeFloat); +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::Process() +{ + MNConfig mnConfig; + mnConfig.baseM = mmTilingData->baseM; + mnConfig.baseN = mmTilingData->baseN; + mnConfig.singleM = mnConfig.baseM; + mnConfig.singleN = tiling->n; + mnConfig.vecSingleN = 512; + mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN); + mnConfig.vecBlockDimN = Ceil(tiling->n, mnConfig.baseN); + if ASCEND_IS_AIC { + SyncAll(); + } + for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) { + int32_t m = static_cast(groupTokensGm.GetValue(groupIdx)); + if (m <= 0) { + continue; + } + mnConfig.m = static_cast(m) * 2; + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + tiling->coreNum; + + + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN; + MMCompute(groupIdx, mnConfig); + curBlock += tiling->coreNum; + } + preCount = curCount % tiling->coreNum; + mnConfig.offsetM += mnConfig.m; + } +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::MMCompute(uint32_t groupIdx, MNConfig& mnConfig) +{ + uint32_t tailN; + uint32_t curSingleN; + uint32_t curSingleM; + uint64_t xOffset; + uint64_t weightOffset; + GlobalTensor weightSlice; + int loopK = 0; + + tailN = mnConfig.nIdx * mnConfig.singleN; + curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = tiling->n - tailN; + } + curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + if ASCEND_IS_AIC { + xOffset = (mnConfig.offsetM + mnConfig.mIdx * mnConfig.singleM) * tiling->k; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = groupIdx * tiling->n * tiling->k + tailN * tiling->k; + } else { + weightOffset = groupIdx * tiling->n * tiling->k + tailN; + } + + mm.SetSingleShape(curSingleM, curSingleN, tiling->k); + + float tmp = 1.0; + uint64_t ans = static_cast(*reinterpret_cast(&tmp)); + mm.SetQuantScalar(ans); + uint64_t scaleOffset = groupIdx * tiling->n * tiling->quantGroupNum + loopK * tiling->n + tailN; + mm.SetTensorA(xGm[xOffset ]); + + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightSlice = weightGm[weightOffset + loopK * quantGroupSize * tiling->n]; + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + } + for (int nId = 0; nId < (curSingleN + 512 - 1) / 512; nId++) { + mnConfig.vecNIdx = nId; + uint64_t vecTailN = mnConfig.vecNIdx * mnConfig.baseN; + uint64_t vecCurSingleN = 512; + if (unlikely(mnConfig.vecNIdx == mnConfig.vecBlockDimN - 1)) { + vecCurSingleN = tiling->n - vecTailN; + } + for (int kId = 0; kId < tiling->k / 256; kId++) { + mnConfig.workSpaceOffset = MM_BASE_BLOCK_OFFSET_NEW * (coreIdx + (cubeId % tiling->parallNum) * tiling->coreNum); + uint64_t workspaceOffset = mnConfig.workSpaceOffset; + if ASCEND_IS_AIC { + if (cubeId >= tiling->parallNum) { + CrossCoreWaitFlag(SYNC_AIV_TO_AIC_NEW); + } +#ifndef __CCE_KT_TEST__ + mm.Iterate(); + mm.GetTensorC(mmOutGm[workspaceOffset], 0, 1); +#endif + CrossCoreSetFlag<2, PIPE_FIX>(SYNC_AIC_TO_AIV_NEW); + } + if ASCEND_IS_AIV { + VectorCompute(groupIdx, mnConfig, kId, vecTailN, workspaceOffset); + } + cubeId++; + } + } +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::VectorCompute(uint32_t groupIdx, MNConfig& mnConfig, int loopK, size_t tailN, uint64_t workspaceOffset) +{ + bool isLastBlock = (loopK + 1) % tiling->quantGroupNum == 0; + uint32_t curCubeSingleN = mnConfig.baseN; + if (mnConfig.vecNIdx == mnConfig.vecBlockDimN - 1) { curCubeSingleN = tiling->n - mnConfig.vecNIdx * mnConfig.baseN; } + uint32_t curCubeSingleM = mnConfig.singleM / 2; + uint32_t mGlobalOffset = mnConfig.offsetM / 2 + mnConfig.mIdx * curCubeSingleM; + uint64_t outOffset = mGlobalOffset * tiling->n + mnConfig.vecNIdx * mnConfig.baseN; + if (mnConfig.mIdx == mnConfig.blockDimM - 1) { curCubeSingleM = mnConfig.m / 2 - mnConfig.mIdx * curCubeSingleM; } + uint32_t vecBaseM = tiling->ubCalSize / (Ceil(mnConfig.baseN, uint32_t(8)) * 8); + vecBaseM = vecBaseM < curCubeSingleM ? vecBaseM : curCubeSingleM; + uint32_t totalVecBaseM = vecBaseM; + uint32_t vec0offsetM; + if (vecBaseM % 2 == 0) { + vecBaseM /= 2; + vec0offsetM = vecBaseM; + } else { + vecBaseM = vecBaseM / 2 + subBlockIdx; + vec0offsetM = totalVecBaseM / 2; + } + if (subBlockIdx == 1) { + outOffset += (vec0offsetM) * tiling->n; + } + uint32_t curVecBaseN = mnConfig.baseN; + uint64_t scaleOffset = groupIdx * tiling->n + mnConfig.vecNIdx * mnConfig.baseN; + uint32_t taskRation = GetTaskRation() == 0 ? 1 : GetTaskRation(); + CrossCoreWaitFlag(SYNC_AIC_TO_AIV_NEW); + uint32_t offsetN = 0; + if (unlikely(offsetN + mnConfig.baseN >= curCubeSingleN)) curVecBaseN = curCubeSingleN - offsetN; + uint32_t alignBaseN = Ceil(curVecBaseN, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block + if (isLastBlock) { DataCopyScale(curVecBaseN, alignBaseN, scaleOffset + offsetN); } + DataCopyScaleDequant(curVecBaseN, alignBaseN, groupIdx * tiling->n * tiling->quantGroupNum + loopK * tiling->n + tailN); + uint32_t curVecBaseM = vecBaseM; + uint64_t mmOutOffset = workspaceOffset + offsetN * mnConfig.baseM + subBlockIdx * vec0offsetM * curVecBaseN * 2; + uint32_t offsetM = 0; + LocalTensor mmOutLocal = vecInQueue.AllocTensor(); + if constexpr (mmType::BT::format == CubeFormat::ND) { + DataCopyPad2DA8W4NDNew(mmOutLocal, mmOutGm[mmOutOffset + offsetM * 2 * curVecBaseN], + curVecBaseM, curVecBaseN, curVecBaseN * 2); + } else { + DataCopyPad2DA8W4New(mmOutLocal, mmOutGm[mmOutOffset + offsetM * 2 * curVecBaseN], + curVecBaseM, curVecBaseN, curVecBaseN * 2); + } + uint32_t targetAddr; + if constexpr (mmType::BT::format == CubeFormat::ND) { + alignBaseN = (alignBaseN + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } + targetAddr = curVecBaseM * alignBaseN; + uint64_t lowBitAddr = mmOutOffset + (offsetM * 2 + 1) * curVecBaseN; + if constexpr (mmType::BT::format == CubeFormat::ND) { + DataCopyPad2DA8W4NDNew(mmOutLocal[targetAddr], + mmOutGm[lowBitAddr], + curVecBaseM, curVecBaseN, curVecBaseN * 2); + } else { + DataCopyPad2DA8W4New(mmOutLocal[targetAddr], + mmOutGm[lowBitAddr], + curVecBaseM, curVecBaseN, curVecBaseN * 2); + } + vecInQueue.EnQue(mmOutLocal); + ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM, loopK, isLastBlock, totalVecBaseM, vec0offsetM); + if (isLastBlock) { + uint64_t coreOutOffset = 0; + LocalTensor yLocal = vecOutQueue.DeQue(); + DataCopyPad2DA8W4New(yGm[outOffset + offsetM * tiling->n + offsetN + coreOutOffset], yLocal, + curVecBaseM, curVecBaseN, alignBaseN, tiling->n); + vecOutQueue.FreeTensor(yLocal); + } + scaleInQueueDeqScale.FreeTensor(scaleInUb2); + if (isLastBlock) { scaleInQueue.FreeTensor(scaleInUb); } + CrossCoreSetFlag<2, PIPE_MTE2>(SYNC_AIV_TO_AIC_NEW); // 2: mode为2, group内同步 +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::ComputeDequantAndActivate(MNConfig& mnConfig, + uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, uint32_t offsetM, int loopK, bool isLastBlock, uint32_t totalVecBaseM, uint32_t offsetVec0M) +{ + uint32_t computeSize = curVecBaseM * alignBaseN; + LocalTensor mmOutInUb = vecInQueue.DeQue(); + uint32_t castSize; + if constexpr (mmType::BT::format == CubeFormat::ND) { + castSize = computeSize + (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + castSize = computeSize * 2; + } + Cast(bufferAdd, mmOutInUb, RoundMode::CAST_NONE, castSize); + PipeBarrier(); + + int32_t maskCast = 256 / sizeof(float); + int repeatCast = alignBaseN / 64 * 2; + AscendC::PairReduceSum(scaleInUb2.ReinterpretCast(), scaleInUb2.ReinterpretCast(), repeatCast, maskCast, 1, 1, 8); + + PipeBarrier(); + uint64_t maskAdd = 64; + uint32_t loopAdd = alignBaseN / 64; + uint8_t blkStrideAdd = static_cast(alignBaseN * sizeof(float) / 32); + BinaryRepeatParams paramAdd(1, 1, 1, blkStrideAdd, blkStrideAdd, 0); + for (uint32_t i = 0; i < loopAdd; i++) { + uint32_t offset = i * 64; + if (loopK == 0){ + Mul(buffer1[offset], bufferAdd[offset], scaleInUb2.ReinterpretCast()[offset], maskAdd, subBlockIdx == 0 ? offsetVec0M*2 : (totalVecBaseM-offsetVec0M)*2, paramAdd); + PipeBarrier(); + } else { + Mul(bufferAdd[offset], bufferAdd[offset], scaleInUb2.ReinterpretCast()[offset], maskAdd, subBlockIdx == 0 ? offsetVec0M*2 : (totalVecBaseM-offsetVec0M)*2, paramAdd); + PipeBarrier(); + } + } + PipeBarrier(); + if (loopK != 0){ + PipeBarrier(); + Add(buffer1, buffer1, bufferAdd, castSize); + PipeBarrier(); + } + vecInQueue.FreeTensor(mmOutInUb); + + if (!isLastBlock) { + return; + } + + const float RIGHT_MOVE = 16.0f; + Muls(buffer2, buffer1, RIGHT_MOVE, computeSize); + PipeBarrier(); + uint32_t addStartAddr; + if constexpr (mmType::BT::format == CubeFormat::ND) { + addStartAddr = (computeSize + HALF_ALIGN - 1) / HALF_ALIGN * HALF_ALIGN; + } else { + addStartAddr = computeSize; + } + Add(buffer3, buffer1[addStartAddr], buffer2, computeSize); + PipeBarrier(); + + uint32_t loop = alignBaseN / 64; + uint8_t blkStride = static_cast(alignBaseN * sizeof(float) / 32); + BinaryRepeatParams param(1, 1, 1, blkStride, blkStride, 0); + uint64_t mask = 64; + uint64_t last = alignBaseN % 64; + for (uint32_t i = 0; i < loop; i++) { + uint32_t offset = i * 64; + Add(buffer2[offset], buffer3[offset], scaleInUb[offset], mask, curVecBaseM, param); + } + PipeBarrier(); + if (unlikely(last > 0)) { + uint32_t offset = loop * 64; + Add(buffer2[offset], buffer3[offset], scaleInUb[offset], last, curVecBaseM, param); + } + PipeBarrier(); + DataCopyPerTokenScaleAndBrcb(mnConfig, curVecBaseM, alignBaseN, offsetM, offsetVec0M); + Mul(buffer4, buffer2, buffer3, computeSize); + PipeBarrier(); + auto out = buffer4; + LocalTensor yLocalInUb = vecOutQueue.AllocTensor(); + Cast(yLocalInUb, out, RoundMode::CAST_RINT, computeSize); + PipeBarrier(); + vecOutQueue.EnQue(yLocalInUb); +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset) +{ + DataCopyPadExtParams padParams; + DataCopyExtParams scaleParams{1, static_cast(curBaseN * sizeof(float)), 1, 1, 0}; + LocalTensor scaleLocal = scaleInQueue.AllocTensor(); + DataCopyPad(scaleLocal, biasGm[scaleOffset], scaleParams, padParams); + scaleInQueue.EnQue(scaleLocal); + scaleInUb = scaleInQueue.DeQue(); +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::DataCopyScaleDequant(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset) +{ + DataCopyPadExtParams padParams; + DataCopyExtParams scaleParams{1, static_cast(curBaseN * sizeof(DTYPE_SCALE_DEV_A8W4MSD_NEW)), 1, 1, 0}; + LocalTensor scaleLocal2 = scaleInQueueDeqScale.AllocTensor(); + DataCopyPad(scaleLocal2, scaleGm[scaleOffset], scaleParams, padParams); + scaleInQueueDeqScale.EnQue(scaleLocal2); + scaleInUb2 = scaleInQueueDeqScale.DeQue(); +} + +template +__aicore__ inline void GMMA8W4MSDComputeNew::DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, + uint32_t curBaseM, uint32_t alignBaseN, uint32_t offsetM, uint32_t offsetVec0M) +{ + uint64_t perTokenScaleOffset = mnConfig.offsetM / 2 + mnConfig.mIdx * mnConfig.singleM / 2 + offsetM + subBlockIdx * offsetVec0M; + uint32_t alignBaseM = Ceil(curBaseM, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block + // GM拷贝per token scale + DataCopyPadExtParams padParams; + DataCopyExtParams perTokenScaleParams{1, static_cast(curBaseM * sizeof(float)), 0, 0, 0}; + LocalTensor perTokenScaleLocal = perTokenScaleInQueue.AllocTensor(); + DataCopyPad(perTokenScaleLocal, perTokenScaleGm[perTokenScaleOffset], perTokenScaleParams, padParams); + + perTokenScaleInQueue.EnQue(perTokenScaleLocal); + + perTokenScaleLocal = perTokenScaleInQueue.DeQue(); + auto scaleTmp = perTokenScaleLocal; + + const uint32_t broadCastDst[2] = {curBaseM, alignBaseN}; + const uint32_t broadCastSrc[2] = {curBaseM, 1}; + BroadCast(buffer3, scaleTmp, broadCastDst, broadCastSrc, buffer7); + perTokenScaleInQueue.FreeTensor(perTokenScaleLocal); +} +} // namespace GROUPED_MATMUL +#endif +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_pre.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_pre.h new file mode 100644 index 00000000..84c9513f --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_msd_pre.h @@ -0,0 +1,219 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4_msd_pre.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_PRE_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_MSD_PRE_H + +#include "kernel_operator.h" +#ifdef GMM_ANTI_QUANT_A8W4_MSD +namespace GROUPED_MATMUL{ +using namespace AscendC; +#define BUFFER_NUM_A8W4_PRE 1 +constexpr int TWO = 2; +constexpr int EIGHT = 8; +constexpr size_t LEN_128 = 128; // 16bit operator +constexpr int DATA_BLOCK_SIZE_32 = 32; +constexpr uint32_t WITH_OFFSET = 1; + +class GMMA8W4PreProcess { +public: + __aicore__ inline GMMA8W4PreProcess(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR groupList, GM_ADDR workspace, const GMMBaseParams& tilingData, TPipe *pipe); + __aicore__ inline void Process(); +private: + TQue vecInQueueX, vecInQueueXBak; + TQue vecOutQueueA1; + TQue vecOutQueueA2; + TQue vecOutQueueA3; + TQue vecOutQueue0F; + TQue vecInQueueG; + TQue vecInQueueGF; + TQue vecInWork; + TQue vecOutQueueRowSum; + TBuf tempBuff; + + LocalTensor xTensor; + LocalTensor xHighHalfTensor; + LocalTensor xHighFloatTensor; + LocalTensor xLowHalfTensor; + LocalTensor xLowHalfTensor2; + LocalTensor xHighI4Tensor; + LocalTensor xLowI4Tensor; + LocalTensor xLowI16Tensor; + LocalTensor groupListTensor; + LocalTensor groupListFTensor; + LocalTensor workTensor; + LocalTensor xRowSumTensor; + + GlobalTensor xGm; + GlobalTensor yGm; + GlobalTensor groupListGm; + GlobalTensor xRowSumGm; + + uint32_t vK; + uint32_t vKAlign; + uint32_t totalGroup{0}; + uint32_t blockDim; + uint32_t coreId; + uint32_t currentNum; + uint32_t startNum; + uint32_t groupNum; + uint32_t xRowSumRepeatNum; + uint32_t withOffset; +}; + +__aicore__ inline void GMMA8W4PreProcess::Init(GM_ADDR x, GM_ADDR y, GM_ADDR groupList, GM_ADDR workspace, const GMMBaseParams& tilingData, TPipe *pipe){ + xGm.SetGlobalBuffer(GetTensorAddr(0, x)); + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + groupListGm.SetGlobalBuffer((__gm__ int64_t *)groupList); + + vK = tilingData.k; + withOffset = tilingData.withOffset; + pipe->InitBuffer(vecInQueueX, BUFFER_NUM_A8W4_PRE, vK * sizeof(int8_t)); // 2KB + pipe->InitBuffer(vecOutQueueA1, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // 1KB + pipe->InitBuffer(vecOutQueueA2, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // 1KB + pipe->InitBuffer(vecOutQueueA3, BUFFER_NUM_A8W4_PRE, vK * sizeof(half)); // 4 KB + // xLowHalfTensor,xLowHalfTensor2 and xHighFloatTensor share the same buffer + pipe->InitBuffer(tempBuff, vK * sizeof(float)); + + constexpr int BUFFER_SIZE_256B = 128 * sizeof(int16_t); + pipe->InitBuffer(vecOutQueue0F, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B); // 256 B + groupNum = static_cast(tilingData.groupNum); + pipe->InitBuffer(vecInQueueG, BUFFER_NUM_A8W4_PRE, Ceil(groupNum * sizeof(int64_t), DATA_BLOCK_SIZE_32)); + pipe->InitBuffer(vecInQueueGF, BUFFER_NUM_A8W4_PRE, Ceil(groupNum * sizeof(float), DATA_BLOCK_SIZE_32)); + pipe->InitBuffer(vecInWork, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B * sizeof(float)); + + if (withOffset == WITH_OFFSET) { + xRowSumGm.SetGlobalBuffer((__gm__ float *)workspace); + pipe->InitBuffer(vecOutQueueRowSum, BUFFER_NUM_A8W4_PRE, 1 * sizeof(float)); + } + + startNum = GetBlockIdx(); + blockDim = GetBlockNum() * GetTaskRation(); +} + +__aicore__ inline void GMMA8W4PreProcess::Process() +{ + constexpr int32_t MASK = 128; + xTensor = vecInQueueX.AllocTensor(); + xHighI4Tensor = vecOutQueueA1.AllocTensor(); + xLowI4Tensor = vecOutQueueA2.AllocTensor(); + xHighHalfTensor = vecOutQueueA3.AllocTensor(); + const uint32_t xLowHalfOffset = vK * sizeof(half); + xLowHalfTensor = tempBuff.GetWithOffset(xLowHalfOffset, 0); + xLowHalfTensor2 = tempBuff.GetWithOffset(xLowHalfOffset, xLowHalfOffset); + xLowI16Tensor = vecOutQueue0F.AllocTensor(); + + if (withOffset == WITH_OFFSET) { + xHighFloatTensor = tempBuff.GetWithOffset(vK * sizeof(float), 0); + xRowSumTensor = vecOutQueueRowSum.AllocTensor(); + } + + Duplicate(xLowI16Tensor, static_cast(0x0F0F), MASK); // get rid of high 4 bits in every int8 + + const size_t LEN_VK = (vK / 2) / 128; // align to 128 + const size_t LAST_LEN_VK = (vK % 256) / 2; + const half one_eight = static_cast(0.0625); + //先计算要处理的group数 + groupListTensor = vecInQueueG.AllocTensor(); + groupListFTensor = vecInQueueGF.AllocTensor(); + workTensor = vecInWork.AllocTensor(); + DataCopyParams dataCopyParams{1, static_cast(groupNum * sizeof(int64_t)), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(groupListTensor, groupListGm, dataCopyParams, padParams); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + Cast(groupListFTensor, groupListTensor, AscendC::RoundMode::CAST_ROUND, groupNum); + PipeBarrier(); + ReduceSum(groupListFTensor, groupListFTensor, workTensor, groupNum); + PipeBarrier(); + Cast(groupListTensor, groupListFTensor, AscendC::RoundMode::CAST_ROUND, 1); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + totalGroup = groupListTensor.GetValue(0); + + SetFlag(EVENT_ID0); + SetFlag(EVENT_ID0); + SetFlag(EVENT_ID1); + for(uint32_t xloop = startNum; xloop < totalGroup; xloop += blockDim){ + uint64_t startAddr = xloop * vK; + WaitFlag(EVENT_ID0); + DataCopy(xTensor, xGm[startAddr], vK); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + Cast(xHighHalfTensor, xTensor, AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + + if (withOffset == WITH_OFFSET) { + Cast(xHighFloatTensor, xHighHalfTensor, AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + + SetFlag(EVENT_ID2); + WaitFlag(EVENT_ID2); + ReduceSum(xRowSumTensor, xHighFloatTensor, workTensor, vK); + SetFlag(EVENT_ID2); + WaitFlag(EVENT_ID2); + + DataCopyExtParams xRowSumParams{1, static_cast(1 * sizeof(float)), 0, 0, 0}; + DataCopyPad(xRowSumGm[xloop], xRowSumTensor, xRowSumParams); + AscendC::PipeBarrier(); + } + + Muls(xHighHalfTensor, xHighHalfTensor, one_eight, vK); + PipeBarrier(); + WaitFlag(EVENT_ID1); + Cast(xHighI4Tensor, xHighHalfTensor, AscendC::RoundMode::CAST_FLOOR, vK); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + DataCopy(yGm[startAddr], xHighI4Tensor.ReinterpretCast(), vK / 2); // 2: size int8 -> int4 + SetFlag(EVENT_ID1); + And(xLowHalfTensor.ReinterpretCast(), xTensor.ReinterpretCast(), xLowI16Tensor, LEN_128, LEN_VK, {1,1,1,8,8,0}); + if (LAST_LEN_VK > 0) { + And(xLowHalfTensor[LEN_VK * LEN_128].ReinterpretCast(), xTensor[LEN_VK * LEN_128 * TWO].ReinterpretCast(), xLowI16Tensor, LAST_LEN_VK, 1, {1,1,1,8,8,0}); + } + PipeBarrier(); + SetFlag(EVENT_ID0); + Cast(xLowHalfTensor2.ReinterpretCast(), xLowHalfTensor.ReinterpretCast(), AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + const half MINUS_EIGHT = static_cast(-8); // shrink 0~15 to -8~7 + Adds(xHighHalfTensor, xLowHalfTensor2, MINUS_EIGHT, vK); + PipeBarrier(); + WaitFlag(EVENT_ID0); + Cast(xLowI4Tensor, xHighHalfTensor.ReinterpretCast(), AscendC::RoundMode::CAST_NONE, vK); + SetFlag(EVENT_ID1); + WaitFlag(EVENT_ID1); + DataCopy(yGm[startAddr + vK / TWO], xLowI4Tensor.ReinterpretCast(), vK / TWO); + SetFlag(EVENT_ID0); + } + WaitFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + WaitFlag(EVENT_ID1); + SyncAll(); + vecInQueueX.FreeTensor(xTensor); + vecOutQueueA1.FreeTensor(xHighI4Tensor); + vecOutQueueA2.FreeTensor(xLowI4Tensor); + vecOutQueueA3.FreeTensor(xHighHalfTensor); + vecOutQueue0F.FreeTensor(xLowI16Tensor); + vecInQueueG.FreeTensor(groupListTensor); + vecInQueueGF.FreeTensor(groupListFTensor); + vecInWork.FreeTensor(workTensor); + if (withOffset == WITH_OFFSET) { + vecOutQueueRowSum.FreeTensor(xRowSumTensor); + } +} + +} // namespace +#endif +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_pre.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_pre.h new file mode 100644 index 00000000..37436314 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_antiquant_a8w4_pre.h @@ -0,0 +1,253 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4_pre.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_PRE_H +#define ASCENDC_GROUPED_MATMUL_ANTIQUANT_A8W4_PRE_H + +#include "kernel_operator.h" +#ifdef GMM_ANTI_QUANT_A8W4_MSD +namespace GROUPED_MATMUL{ +using namespace AscendC; + +#if ORIG_DTYPE_Y == DT_FLOAT16 + using DTYPE_SCALE_OUT = float; +#else + using DTYPE_SCALE_OUT = bfloat16_t; +#endif + +template +class GMMA8W4FakeQuantPreProcess { +public: + __aicore__ inline GMMA8W4FakeQuantPreProcess(){}; + __aicore__ inline void Init(GM_ADDR weight, GM_ADDR y, GM_ADDR groupList, GM_ADDR workspace, const GMMBaseParams& tilingData, TPipe *pipe); + __aicore__ inline void Process(); +private: + __aicore__ inline void ScaleProcess(); + + GlobalTensor weightGm; + GlobalTensor yGm; + GlobalTensor scaleOutGm; + GlobalTensor scaleGm; + + uint32_t totalGroup{0}; + uint32_t blockDim; + uint32_t startNum; + uint32_t groupNum; + bool isPerchannel = false; + + const int TWO = 2; + const int EIGHT = 8; + const int DATA_BLOCK_SIZE_32 = 32; + const uint32_t WITH_OFFSET = 0; + const GMMBaseParams *tiling; +}; + +template +__aicore__ inline void GMMA8W4FakeQuantPreProcess::Init(GM_ADDR weight, GM_ADDR y, GM_ADDR scale, GM_ADDR workspace, const GMMBaseParams& tilingData, TPipe *pipe){ + this->tiling = &tilingData; + weightGm.SetGlobalBuffer(GetTensorAddr(0, weight)); + yGm.SetGlobalBuffer((__gm__ int8_t *)workspace); + if (tiling->quantGroupNum == 1) { //per channel + this->isPerchannel = true; + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scale)); + scaleOutGm.SetGlobalBuffer((__gm__ DTYPE_SCALE_OUT *)((__gm__ int8_t *)workspace + tiling->groupNum * tiling->n * tiling->k)); + } +} +template +__aicore__ inline void GMMA8W4FakeQuantPreProcess::Process() +{ + const size_t baseSize = 24 * 2 * 1024; + const size_t HALF_UB = 96 * 1024; + const size_t BOTTOM_LOOP = 72 * 1024; + const size_t BOTTOM_LOOP_NZ = 48 * 1024; + AscendC::LocalTensor ALocalI4(AscendC::TPosition::VECIN, BOTTOM_LOOP, baseSize); + AscendC::LocalTensor ALocalF16(AscendC::TPosition::VECIN, 0, baseSize); + AscendC::LocalTensor ALocalI8(AscendC::TPosition::VECOUT, 0, baseSize); + AscendC::LocalTensor ALocalI8NZ(AscendC::TPosition::VECOUT, BOTTOM_LOOP_NZ, baseSize); +#if ASCENDC_CPU_DEBUG + + AscendC::LocalTensor BLocalI4(AscendC::TPosition::VECIN, 0, baseSize); + AscendC::LocalTensor BLocalF16(AscendC::TPosition::VECIN, 0, baseSize); + AscendC::LocalTensor BLocalI8(AscendC::TPosition::VECOUT, 0, baseSize); + AscendC::LocalTensor BLocalI8NZ(AscendC::TPosition::VECOUT, 0, baseSize); +#else + AscendC::LocalTensor BLocalI4(AscendC::TPosition::VECIN, HALF_UB + BOTTOM_LOOP, baseSize); + AscendC::LocalTensor BLocalF16(AscendC::TPosition::VECIN, HALF_UB, baseSize); + AscendC::LocalTensor BLocalI8(AscendC::TPosition::VECOUT, HALF_UB, baseSize); + AscendC::LocalTensor BLocalI8NZ(AscendC::TPosition::VECOUT, HALF_UB + BOTTOM_LOOP_NZ, baseSize); +#endif + constexpr uint32_t K_PER_LOOP = 24 * 1024 * 2 / 64; //one buffer: 24k int4 + + startNum = GetBlockIdx(); + totalGroup = tiling->n / 64; //限制为64对齐 + uint64_t lastGroupRest = 0; + if (totalGroup * 64 != tiling->n) { + lastGroupRest = tiling->n - totalGroup * 64; + totalGroup++; + } + blockDim = GetBlockNum() * GetTaskRation(); + + for(uint32_t eStart = 0; eStart < tiling->groupNum; eStart++) { + const size_t eStartAddrGm = eStart * tiling->n * tiling->k; + for(uint32_t nStart = startNum; nStart < totalGroup; nStart += blockDim) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + + bool isLastNotAlignN = false; + size_t nBlockSize = 64; + if (unlikely(nStart == totalGroup - 1 && lastGroupRest > 0)) { //nd：n非64对齐 + isLastNotAlignN = true; + nBlockSize = lastGroupRest; + } + for(uint32_t kStart = 0; kStart < tiling->k; kStart += K_PER_LOOP * 2) { //2: double buffer + const size_t kLenThisTimeTotal = (kStart + K_PER_LOOP * 2 > tiling->k) ? (tiling->k - kStart) : K_PER_LOOP * 2; + const uint16_t kLen = kLenThisTimeTotal / 2; //2: per buffer + const size_t handleBytePerBuffer = kLen * nBlockSize; + const size_t startAddrGm = tiling->k * 64 * nStart + kStart * nBlockSize; +#if ASCENDC_CPU_DEBUG +#else + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +#endif + AscendC::WaitFlag(EVENT_ID0); + DataCopy(ALocalI4, weightGm[eStartAddrGm + startAddrGm], handleBytePerBuffer); + AscendC::SetFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0); + Cast(ALocalF16, ALocalI4, AscendC::RoundMode::CAST_NONE, handleBytePerBuffer); + PipeBarrier(); + Cast(ALocalI8, ALocalF16, AscendC::RoundMode::CAST_NONE, handleBytePerBuffer); + + if constexpr (wFormat == CubeFormat::NZ) { + //to NZ 1 + PipeBarrier(); + DataCopy(ALocalI8NZ, ALocalI8, {kLen, 1, 1, 0}); + PipeBarrier(); + //to NZ 2 + DataCopy(ALocalI8NZ[kLen * 32], ALocalI8[32], {kLen, 1, 1, 0}); + PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0); + if constexpr (wFormat == CubeFormat::ND) { + AscendC::WaitFlag(EVENT_ID0); + DataCopy(yGm[eStartAddrGm + startAddrGm], ALocalI8, handleBytePerBuffer); + AscendC::SetFlag(EVENT_ID0); + } +#if ASCENDC_CPU_DEBUG +#else + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); +#endif + AscendC::WaitFlag(EVENT_ID1); + DataCopy(BLocalI4, weightGm[eStartAddrGm + startAddrGm + handleBytePerBuffer], handleBytePerBuffer); + AscendC::SetFlag(EVENT_ID1); + + AscendC::WaitFlag(EVENT_ID1); + Cast(BLocalF16, BLocalI4, AscendC::RoundMode::CAST_NONE, handleBytePerBuffer); + PipeBarrier(); + Cast(BLocalI8, BLocalF16, AscendC::RoundMode::CAST_NONE, handleBytePerBuffer); + + if constexpr (wFormat == CubeFormat::NZ) { + //to NZ 1 + PipeBarrier(); + DataCopy(BLocalI8NZ, BLocalI8, {kLen, 1, 1, 0}); + PipeBarrier(); + //to NZ 2 + DataCopy(BLocalI8NZ[kLen * 32], BLocalI8[32], {kLen, 1, 1, 0}); + PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID1); + + if constexpr (wFormat == CubeFormat::NZ) { //nz -> nz + AscendC::WaitFlag(EVENT_ID0); + DataCopy(yGm[eStartAddrGm + tiling->k * 64 * nStart + kStart * 32], ALocalI8NZ, kLen * 32); + DataCopy(yGm[eStartAddrGm + tiling->k * 64 * nStart + tiling->k * 32 + kStart * 32], ALocalI8NZ[kLen * 32], kLen * 32); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + DataCopy(yGm[eStartAddrGm + tiling->k * 64 * nStart + kLen * 32 + kStart * 32], BLocalI8NZ, kLen * 32); + DataCopy(yGm[eStartAddrGm + tiling->k * 64 * nStart + kLen * 32 + tiling->k * 32 + kStart * 32], BLocalI8NZ[kLen * 32], kLen * 32); + AscendC::SetFlag(EVENT_ID1); + } else { + AscendC::WaitFlag(EVENT_ID1); + DataCopy(yGm[eStartAddrGm + startAddrGm + handleBytePerBuffer], BLocalI8, handleBytePerBuffer); + AscendC::SetFlag(EVENT_ID1); + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + } + if(isPerchannel) { + ScaleProcess(); + } +} +template +__aicore__ inline void GMMA8W4FakeQuantPreProcess::ScaleProcess() { + const size_t SCALE_SIZE = 192 * 1024 / sizeof(int64_t); //u64 +#if ASCENDC_CPU_DEBUG + AscendC::LocalTensor scaleU64(AscendC::TPosition::VECIN, 0, SCALE_SIZE - 32); + AscendC::LocalTensor scaleF32(AscendC::TPosition::VECIN, 0, SCALE_SIZE - 32); + AscendC::LocalTensor scaleBF16(AscendC::TPosition::VECIN, 0, SCALE_SIZE - 32); +#else + AscendC::LocalTensor scaleU64(AscendC::TPosition::VECIN, 0, SCALE_SIZE); + AscendC::LocalTensor scaleF32(AscendC::TPosition::VECIN, 0, SCALE_SIZE); + AscendC::LocalTensor scaleBF16(AscendC::TPosition::VECIN, 0, SCALE_SIZE); +#endif + const uint64_t scale_n = tiling->n * 2; + const uint64_t each_core_u64 = (tiling->groupNum * tiling->n / blockDim) / 16 * 16; // align to 8 + const uint64_t this_core_u64 = startNum == blockDim - 1 ? tiling->groupNum * tiling->n - each_core_u64 * (blockDim - 1) : each_core_u64; + const uint64_t start_element_u64 = each_core_u64 * startNum; + int loop_size_u64 = 0; + PipeBarrier(); + for (int start_loc = 0; start_loc < this_core_u64; start_loc += SCALE_SIZE) { + loop_size_u64 = SCALE_SIZE; + if(start_loc + SCALE_SIZE > this_core_u64) loop_size_u64 = this_core_u64 - start_loc; + int loop_size_f32 = loop_size_u64 * 2; + DataCopy(scaleU64, scaleGm[start_element_u64 + start_loc], loop_size_u64); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + const size_t PR_SIZE = 64 * 255; //maximum element number in one instruction + int pr_last = 0; + for(int pr_start = 0; pr_start < loop_size_f32 ; pr_start += PR_SIZE) { + size_t pr_this_time = PR_SIZE; + if (pr_start + PR_SIZE > loop_size_f32) { + pr_this_time = loop_size_f32 - pr_start; + pr_last = pr_this_time - pr_this_time / 64 * 64; + } + int repeatCast = pr_this_time / 64; + AscendC::PairReduceSum(scaleF32[pr_start / 2], scaleF32[pr_start], repeatCast, 64, 1, 1, 8); + PipeBarrier(); + } + if (pr_last > 0) { + AscendC::PairReduceSum(scaleF32[(loop_size_f32 - pr_last)/ 2], scaleF32[loop_size_f32 - pr_last], 1, pr_last, 1, 1, 8); + } + if constexpr (sizeof(DTYPE_SCALE_OUT) == 2) { // to bf16 + PipeBarrier(); + Cast(scaleBF16, scaleF32, AscendC::RoundMode::CAST_FLOOR, loop_size_u64); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + #if ORIG_DTYPE_Y == DT_FLOAT16 + DataCopy(scaleOutGm[start_element_u64+start_loc],scaleF32 , loop_size_u64); + #else + DataCopy(scaleOutGm[start_element_u64+start_loc], scaleBF16, (loop_size_u64 + 16 - 1) / 16 * 16); + #endif + PipeBarrier(); + } +} + +} // namespace +#endif +#endif \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_autotiling_a8w4.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_autotiling_a8w4.h new file mode 100644 index 00000000..491bed71 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_autotiling_a8w4.h @@ -0,0 +1,2715 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_autotiling_a8w4.h + * \brief + */ +#ifndef ASCENDC_GROUPED_AUTOTILING_A8W4_H +#define ASCENDC_GROUPED_AUTOTILING_A8W4_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + +//内存设置别名 +#define GlobalMem TPosition::GM +#define UBMem TPosition::VECCALC +#define L1AMem TPosition::A1 +#define L1BMem TPosition::B1 +#define L0AMem TPosition::A2 +#define L0BMem TPosition::B2 +#define L0CMem TPosition::CO1 +#define L0C2Mem TPosition::CO2 + +// 数据类型长度 +#define INT2 0.25 +#define INT4 0.5 +#define INT6 0.75 +#define INT8 1 +#define INT16 2 +#define INT32 4 +#define INT64 8 +#define FP4 0.5 +#define FP6 0.75 +#define FP8 1 +#define FP16 2 +#define FP32 4 +#define FP64 8 + +/* 计算过程中系数 */ +#define FACTOR_2 2 +#define FACTOR_4 4 +#define FACTOR_7 7 +#define FACTOR_8 8 +#define FACTOR_10 10 +#define FACTOR_12 12 +#define FACTOR_16 16 +#define FACTOR_22 22 +#define FACTOR_32 32 +#define FACTOR_64 64 +#define FACTOR_128 128 +#define FACTOR_256 256 +#define FACTOR_512 512 +#define FACTOR_4096 4096 +#define FACTOR_16384 16384 +/* 维度数值 */ +#define DIM_VAL1 1 +#define DIM_VAL2 2 +#define DIM_VAL3 3 +/* 数组下标 */ +#define ARRAY_IDX0 0 +#define ARRAY_IDX1 1 +#define ARRAY_IDX2 2 +/* 平台硬件相关参数 */ +#define ELEM_BASE_ALIGN 16 // mixmum alignment value for element number in one dimension of tensor +#define RATIO_AIV2AIC 2 // ratio of number between AIV and AIC +#define COPY_BLK_BYTES 32 // bytes of data for one data move +#define MUL_BYTES 256 // bytes of mul each +#define MAT_FRAC_BYTES 512 // bytes of matrix fractal +/* Flag ID */ +#define LOCAL_FLAGID0 0x0 +#define LOCAL_FLAGID1 0x1 +#define LOCAL_FLAGID2 0x2 +#define LOCAL_FLAGID3 0x3 +#define LOCAL_FLAGID4 0x4 +#define LOCAL_FLAGID5 0x5 +#define LOCAL_FLAGID6 0x6 +#define LOCAL_FLAGID7 0x7 + +#define L0C_FORMAT_SIZE 4 +#define MAX_GROUP_LEN 256 + +namespace GROUPED_MATMUL +{ +using namespace AscendC; +using namespace matmul; + +namespace GMMHighPerf +{ +__aicore__ inline uint32_t ceilINT(uint32_t x, uint32_t y) +{ + return Ceil(x, y); +} + +//8的倍数向上取整 +__aicore__ inline uint32_t ceilINT_8(uint32_t num) +{ + if (num % FACTOR_8 == 0) { + return num; + } else { + return (num / FACTOR_8) * FACTOR_8 + FACTOR_8; + } +} + +//16的倍数向上取整 +__aicore__ inline uint32_t ceilINT_16(uint32_t num) +{ + if (num % FACTOR_16 == 0) { + return num; + } else { + return (num / FACTOR_16) * FACTOR_16 + FACTOR_16; + } +} + +__aicore__ inline uint32_t ceilINT_64(uint32_t num) +{ + if (num % FACTOR_64 == 0) { + return num; + } else { + return (num / FACTOR_64) * FACTOR_64 + FACTOR_64; + } +} + +__aicore__ inline uint32_t splitBy_64(uint32_t length) +{ + if (length <= FACTOR_64) { + return length; + } + uint32_t temp = length / FACTOR_2; + temp = ceilINT_64(temp); + return temp; +} + +struct Dim0 { +public: + __aicore__ inline Dim0() + { + dim = 0; + } + + __aicore__ inline Dim0(uint32_t in) + { + dim = in; + } + + uint32_t dim; + +private: + uint32_t ele_size; +}; + +struct Dim1 : public Dim0 { +public: + __aicore__ inline Dim1() : Dim0(DIM_VAL1) + { + s[0] = 0; + } + + __aicore__ inline Dim1(uint32_t in) : Dim0(DIM_VAL1) + { + s[0] = in; + } + + __aicore__ inline Dim1(uint32_t in, char name) : Dim0(DIM_VAL1) + { + s[0] = in; + axis_name[0] = name; + } + + __aicore__ inline Dim1(const Dim1 &d): Dim0(DIM_VAL1) + { + s[0] = d.s[0]; + axis_name[0] = d.axis_name[0]; + } + + __aicore__ inline uint32_t size() + { + return s[0]; + } + + __aicore__ inline uint32_t size_16() + { + return ceilINT_16(s[0]); + } + + __aicore__ inline Dim1 Dim_16() + { + return Dim1(ceilINT_16(s[0])); + } + + __aicore__ inline void clear() + { + s[0] = 0; + } + + __aicore__ inline void setAxisName(char n0='0') + { + axis_name[0] = n0; + } + + __aicore__ inline void copyAxisName(Dim1 &dim) + { + axis_name[0] = dim.axis_name[0]; + } + + __aicore__ inline void setByAxisName(uint32_t in, char name) + { + if (axis_name[0] == name) { + s[0] = in; + } else { + } + } + + __aicore__ inline void set(uint32_t in) + { + s[0] = in; + } + + __aicore__ inline uint32_t getAxisByName(char name) const + { + if (axis_name[0] == name) { + return s[0]; + } + return -1; + } + + __aicore__ inline bool operator==(Dim1 &in) + { + return s[0] == in.s[0]; + } + + __aicore__ inline void posIterator(Dim1 ori_pos, Dim1 base_tiling, Dim1 ori_vec, Dim1 &curr_vec) + { + uint32_t temp0 = s[0] + base_tiling.s[0]; + if (temp0 < ori_pos.s[0] + ori_vec.s[0]) { + s[0] = temp0; + } else { + curr_vec.s[0] = ori_pos.s[0] + ori_vec.s[0] - s[0]; + s[0] = ori_pos.s[0]; + } + } + + uint32_t s[DIM_VAL1]; + char axis_name[DIM_VAL1]; +}; + +struct Dim2 { +public: + uint32_t dim; + __aicore__ inline Dim2() + { + s[0] = 0; + s[1] = 0; + } + + __aicore__ inline Dim2(uint32_t in0, uint32_t in1) + { + s[0] = in0; + s[1] = in1; + } + + __aicore__ inline Dim2(uint32_t in0, uint32_t in1, char name0, char name1) + { + s[0] = in0; + s[1] = in1; + axis_name[0] = name0; + axis_name[1] = name1; + } + + __aicore__ inline Dim2(const Dim2 &d) + { + s[0] = d.s[0]; + s[1] = d.s[1]; + axis_name[0] = d.axis_name[0]; + axis_name[1] = d.axis_name[1]; + } + + __aicore__ inline uint32_t size() + { + return s[0] * s[1]; + } + + __aicore__ inline uint32_t size_16() + { + return ceilINT_16(s[0]) * ceilINT_16(s[1]); + } + + __aicore__ inline Dim2 Dim_16() + { + return Dim2(ceilINT_16(s[0]), ceilINT_16(s[1])); + } + + __aicore__ inline void clear() + { + s[0] = 0; + s[1] = 0; + } + + __aicore__ inline void set(uint32_t in0, uint32_t in1) + { + s[0] = in0; + s[1] = in1; + } + + __aicore__ inline void transpose() + { + uint32_t temp = s[0]; + s[0] = s[1]; + s[1] = temp; + + char temp_name = axis_name[0]; + axis_name[0] = axis_name[1]; + axis_name[1] = temp_name; + } + + __aicore__ inline void setAxisName(char n0='0', char n1='1') + { + axis_name[0] = n0; + axis_name[1] = n1; + } + + __aicore__ inline void copyAxisName(Dim2 &dim) + { + axis_name[0] = dim.axis_name[0]; + axis_name[1] = dim.axis_name[1]; + } + + __aicore__ inline void setByAxisName(uint32_t in, char name) + { + if (axis_name[0] == name) { + s[0] = in; + } else if (axis_name[1] == name) { + s[1] = in; + } else { + } + } + + __aicore__ inline uint32_t getAxisByName(char name) + { + if (axis_name[0] == name) { + return s[0]; + } else if (axis_name[1] == name) { + return s[1]; + } + return -1; + } + + __aicore__ inline bool operator==(Dim2 &in) + { + return s[0] == in.s[0] && s[1] == in.s[1]; + } + + __aicore__ inline void posIterator(Dim2 ori_pos, Dim2 base_tiling, Dim2 ori_vec, Dim2 &curr_vec) + { + uint32_t temp0 = s[0] + base_tiling.s[0]; + if (temp0 < ori_pos.s[0] + ori_vec.s[0]) { + s[0] = temp0; + } else { + curr_vec.s[0] = ori_pos.s[0] + ori_vec.s[0] - s[0]; + s[0] = ori_pos.s[0]; + + uint32_t temp1 = s[1] + base_tiling.s[1]; + if (temp1 < ori_pos.s[1] + ori_vec.s[1]) { + s[1] = temp1; + } else { + curr_vec.s[1] = ori_pos.s[1] + ori_vec.s[1] - s[1]; + s[1] = ori_pos.s[1]; + } + } + } + + uint32_t s[DIM_VAL2]; + char axis_name[DIM_VAL2]; +}; + +struct Dim3 : public Dim0 { +public: + __aicore__ inline Dim3() : Dim0(DIM_VAL3) + { + s[ARRAY_IDX0] = 0; + s[ARRAY_IDX1] = 0; + s[ARRAY_IDX2] = 0; + } + + __aicore__ inline Dim3(uint32_t in0, uint32_t in1, uint32_t in2) : Dim0(DIM_VAL3) + { + s[ARRAY_IDX0] = in0; + s[ARRAY_IDX1] = in1; + s[ARRAY_IDX2] = in2; + } + + __aicore__ inline Dim3(uint32_t in0, uint32_t in1, uint32_t in2, char name0, char name1, char name2) : Dim0(DIM_VAL3) + { + s[ARRAY_IDX0] = in0; + s[ARRAY_IDX1] = in1; + s[ARRAY_IDX2] = in2; + axis_name[ARRAY_IDX0] = name0; + axis_name[ARRAY_IDX1] = name1; + axis_name[ARRAY_IDX2] = name2; + } + + __aicore__ inline Dim3(const Dim3 &d): Dim0(DIM_VAL3) + { + s[ARRAY_IDX0] = d.s[ARRAY_IDX0]; + s[ARRAY_IDX1] = d.s[ARRAY_IDX1]; + s[ARRAY_IDX2] = d.s[ARRAY_IDX2]; + axis_name[ARRAY_IDX0] = d.axis_name[ARRAY_IDX0]; + axis_name[ARRAY_IDX1] = d.axis_name[ARRAY_IDX1]; + axis_name[ARRAY_IDX2] = d.axis_name[ARRAY_IDX2]; + } + + __aicore__ inline uint32_t size() + { + return s[ARRAY_IDX0] * s[ARRAY_IDX1] * s[ARRAY_IDX2]; + } + + __aicore__ inline uint32_t size_16() + { + return ceilINT_16(s[ARRAY_IDX0]) * ceilINT_16(s[ARRAY_IDX1]) * ceilINT_16(s[ARRAY_IDX2]); + } + + __aicore__ inline Dim3 Dim_16() + { + return Dim3(ceilINT_16(s[ARRAY_IDX0]), ceilINT_16(s[ARRAY_IDX1]), ceilINT_16(s[ARRAY_IDX2])); + } + + __aicore__ inline void clear() + { + s[ARRAY_IDX0] = 0; + s[ARRAY_IDX1] = 0; + s[ARRAY_IDX2] = 0; + } + + __aicore__ inline void set(uint32_t in0, uint32_t in1, uint32_t in2) + { + s[ARRAY_IDX0] = in0; + s[ARRAY_IDX1] = in1; + s[ARRAY_IDX2] = in2; + } + + __aicore__ inline void setAxisName(char n0='0', char n1='1', char n2='2') + { + axis_name[ARRAY_IDX0] = n0; + axis_name[ARRAY_IDX1] = n1; + axis_name[ARRAY_IDX2] = n2; + } + + __aicore__ inline void copyAxisName(Dim3 &dim) + { + axis_name[ARRAY_IDX0] = dim.axis_name[ARRAY_IDX0]; + axis_name[ARRAY_IDX1] = dim.axis_name[ARRAY_IDX1]; + axis_name[ARRAY_IDX2] = dim.axis_name[ARRAY_IDX2]; + } + + __aicore__ inline void setByAxisName(uint32_t in, char name) + { + if (axis_name[ARRAY_IDX0] == name) { + s[ARRAY_IDX0] = in; + } else if (axis_name[ARRAY_IDX1] == name) { + s[ARRAY_IDX1] = in; + } else if (axis_name[ARRAY_IDX2] == name) { + s[ARRAY_IDX2] = in; + } else { + } + } + + __aicore__ inline uint32_t getAxisByName(char name) + { + if (axis_name[ARRAY_IDX0] == name) { + return s[ARRAY_IDX0]; + } else if (axis_name[ARRAY_IDX1] == name) { + return s[ARRAY_IDX1]; + } else if (axis_name[ARRAY_IDX2] == name) { + return s[ARRAY_IDX2]; + } + return -1; + } + + __aicore__ inline bool operator==(Dim3 &in) + { + return s[ARRAY_IDX0] == in.s[ARRAY_IDX0] && s[ARRAY_IDX1] == in.s[ARRAY_IDX1] && s[ARRAY_IDX2] == in.s[ARRAY_IDX2]; + } + + __aicore__ inline void posIterator(Dim3 ori_pos, Dim3 base_tiling, Dim3 ori_vec, Dim3 &curr_vec) + { + //记录下初始的坐标 + uint32_t t_pos[DIM_VAL3]; + t_pos[ARRAY_IDX0] = s[ARRAY_IDX0]; + t_pos[ARRAY_IDX1] = s[ARRAY_IDX1]; + t_pos[ARRAY_IDX2] = s[ARRAY_IDX2]; + uint32_t temp0 = s[ARRAY_IDX0] + base_tiling.s[ARRAY_IDX0]; + if (temp0 < ori_pos.s[ARRAY_IDX0] + ori_vec.s[ARRAY_IDX0]) { + s[ARRAY_IDX0] = temp0; + } else { + s[ARRAY_IDX0] = ori_pos.s[ARRAY_IDX0]; + curr_vec.s[ARRAY_IDX0] = base_tiling.s[ARRAY_IDX0] - (temp0 - (ori_pos.s[ARRAY_IDX0] + ori_vec.s[ARRAY_IDX0])); + + uint32_t temp1 = s[ARRAY_IDX1] + base_tiling.s[ARRAY_IDX1]; + if (temp1 < ori_pos.s[ARRAY_IDX1] + ori_vec.s[ARRAY_IDX1]) { + s[ARRAY_IDX1] = temp1; + } else { + s[ARRAY_IDX1] = ori_pos.s[ARRAY_IDX1]; + curr_vec.s[ARRAY_IDX1] = base_tiling.s[ARRAY_IDX1] - (temp1 - (ori_pos.s[ARRAY_IDX1] + ori_vec.s[ARRAY_IDX1])); + + uint32_t temp2 = s[ARRAY_IDX2] + base_tiling.s[ARRAY_IDX2]; + if (temp2 < ori_pos.s[ARRAY_IDX2] + ori_vec.s[ARRAY_IDX2]) { + s[ARRAY_IDX2] = temp2; + } else { + s[ARRAY_IDX2] = ori_pos.s[ARRAY_IDX2]; + curr_vec.s[ARRAY_IDX2] = base_tiling.s[ARRAY_IDX2] - (temp2 - (ori_pos.s[ARRAY_IDX2] + ori_vec.s[ARRAY_IDX2])); + } + } + } + + uint32_t temp3 = t_pos[ARRAY_IDX1] + base_tiling.s[ARRAY_IDX1]; + if (temp3 > ori_pos.s[ARRAY_IDX1] + ori_vec.s[ARRAY_IDX1]) { + curr_vec.s[ARRAY_IDX1] = base_tiling.s[ARRAY_IDX1] - (temp3 - (ori_pos.s[ARRAY_IDX1] + ori_vec.s[ARRAY_IDX1])); + } + + uint32_t temp4 = t_pos[ARRAY_IDX2] + base_tiling.s[ARRAY_IDX2]; + if (temp4 > ori_pos.s[ARRAY_IDX2] + ori_vec.s[ARRAY_IDX2]) { + curr_vec.s[ARRAY_IDX2] = base_tiling.s[ARRAY_IDX2] - (temp4 - (ori_pos.s[ARRAY_IDX2] + ori_vec.s[ARRAY_IDX2])); + } + } + + uint32_t s[DIM_VAL3]; + char axis_name[DIM_VAL3]; +}; + +struct AxisTiling3Vector { + uint32_t s[DIM_VAL3]; +}; + +template +class UnisorShape { +public: + __aicore__ inline UnisorShape() + { + Dim dim; + vector_val = dim; + } + + __aicore__ inline UnisorShape(Dim vec) + { + vector_val = vec; + } + + __aicore__ inline unsigned size() + { + return vector_val.size(); + } + + Dim vector_val;// vector value to ensure the space of UnisorShape +}; + +template +class Cartesian { +public: + __aicore__ inline Cartesian() + { + Dim dim; + pos = dim; + vec = dim; + } + + __aicore__ inline Cartesian(const Cartesian &c) + { + pos = c.pos; + vec = c.vec; + } + + __aicore__ inline Cartesian(Dim v) + { + vec = v; + pos = v; + pos.clear(); + } + + __aicore__ inline Cartesian(Dim p, Dim v) + { + pos = p; + vec = v; + } + + __aicore__ inline unsigned size() + { + return vec.size(); + } + + Dim pos;//start position + Dim vec;// vector value to ensure the space of UnisorShape +}; + +template +class CartesianIterator +{ +public: + __aicore__ inline CartesianIterator(Cartesian cartesianC, Dim tiling) + { + iter_tiling = tiling; + ori_pos = cartesianC.pos; + ori_vec = cartesianC.vec; + offset_pos = cartesianC.pos; + active = false; + odd_even = false; + } + + __aicore__ inline bool posIterator(Cartesian &cartesianUnisorC) + { + odd_even = !odd_even; + cartesianUnisorC.pos = offset_pos; + Dim curr_vec = iter_tiling; + offset_pos.posIterator(ori_pos, iter_tiling, ori_vec, curr_vec); + cartesianUnisorC.vec = curr_vec; + if (cartesianUnisorC.pos == ori_pos && active) { + //用完了重置 + active = false; + offset_pos = ori_pos; + return false; + } + active = true; + return true; + } + + __aicore__ inline bool isLast() + { + if (offset_pos == ori_pos) { + return true; + } else { + return false; + } + } + + __aicore__ inline bool oddEven() + { + return odd_even; + } + + Dim iter_tiling; + Dim ori_pos; + Dim ori_vec; + Dim offset_pos; + bool active; + bool odd_even; +}; + +/* + * Uniform tensor for both global and local + * Unisor必须是一块联系的地址片段，但是shape可以指定Unisor的某一个Tiling地址空间 + */ +template +class Unisor : public UnisorShape , public Cartesian +{ +public: + float format_size; + GlobalTensor inputGlobal; + TBuf calcBuf; + + __aicore__ inline Unisor() : UnisorShape (), Cartesian () + { + } + + __aicore__ inline Unisor(const Unisor &t):UnisorShape (t.vector_val), Cartesian (t.pos, t.vec), format_size(t.format_size), inputGlobal(t.inputGlobal), calcBuf(t.calcBuf) + { + } + + //LocalTensor + __aicore__ inline Unisor(Dim vec, float format_size) : UnisorShape (vec), Cartesian (vec) + { + // LocalTensor要通过init()函数初始化 + this->format_size = format_size; + } + + //GlobalTensor + __aicore__ inline Unisor(GM_ADDR gm, Dim vec, float format_size) : UnisorShape (vec.Dim_16()), Cartesian (vec.Dim_16()) + { + this->format_size = format_size; + int64_t dataSize = format_size * vec.size_16(); + inputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(gm), dataSize); + } + + __aicore__ inline Unisor(GM_ADDR gm, Dim vec, float format_size, bool isReal) : UnisorShape (vec), Cartesian (vec) + { + this->format_size = format_size; + int64_t dataSize = format_size * vec.size(); + inputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(gm), dataSize); + } + + __aicore__ inline void setCartesian(Dim p, Dim v) + { + this->pos = p; + this->vec = v; + } + + __aicore__ inline void setCartesian(Cartesian cartesian) + { + this->pos = cartesian.pos; + this->vec = cartesian.vec; + } + + __aicore__ inline void setCartesianVec(Dim v) + { + this->vec = v; + } + + __aicore__ inline void transpose() + { + this->pos.transpose(); + this->vec.transpose(); + } + + __aicore__ inline void init(Dim vec, float format_size, TPipe *pipe) + { + // Shape和Cartesian初始化 + this->vector_val = vec.Dim_16(); + this->vector_val.copyAxisName(vec); + this->vec = vec; + this->pos.copyAxisName(vec); + this->format_size = format_size; + + int64_t size = format_size * vec.size_16(); + if (size > 0) { + pipe->InitBuffer(calcBuf, size); + } + } + + __aicore__ inline void init_RealShape(Dim vec, float format_size, TPipe *pipe) + { + //Shape和Cartesian初始化 + this->vector_val = vec; + this->vector_val.copyAxisName(vec); + this->vec = vec; + this->pos.copyAxisName(vec); + this->format_size = format_size; + + int64_t size = format_size * vec.size(); + if (size > 0) { + pipe->InitBuffer(calcBuf, size); + } + } + + template + __aicore__ inline LocalTensor get() + { + return calcBuf.template Get(); + } + + __aicore__ inline GlobalTensor getGM() + { + return inputGlobal; + } + + __aicore__ inline GlobalTensor getGM_32() + { + GlobalTensor tmp32tGlobalTensor; + tmp32tGlobalTensor.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(this->gm_addr), this->gm_datasize / FACTOR_4); + return tmp32tGlobalTensor; + } +}; + +class KernelBase { +public: + TPipe *pipe; + Dim3 splitDim; + Dim3 block; + uint32_t blockId; + uint32_t pattern = 1; + uint8_t output_type = 0; + bool offset_enable = false; + uint32_t nCoreAIC; + uint32_t nCoreAIV; + uint32_t szAvailUB; + uint32_t szAvailL0A; + uint32_t szAvailL0C; + +public: + __aicore__ inline void Init(TPipe* pipeIn) + { + pipe = pipeIn; + } + + template __aicore__ inline void NPU_Duplicate(Unisor &out, const T value) + { + LocalTensor srcLocal = out.get(); + uint32_t dstShape = out.vec.s[0] / FACTOR_2; + + Duplicate(srcLocal, value, dstShape); + PipeBarrier(); + } + + template __aicore__ inline void NPU_Adds(Unisor &out, Unisor &in, const T &value) + { + LocalTensor outLocal = out.get(); + LocalTensor inLocal = in.get(); + + uint32_t width = in.vector_val.s[1]; // in的整个R轴长度 + uint32_t height = in.vector_val.s[0]; // in的整个D轴长度 + + Dim2 posA = in.pos; + uint32_t posAHeight = posA.s[0]; + uint32_t posAWidth = posA.s[1]; + + Dim2 posOUT = out.pos; + uint32_t posOUTHeight = posOUT.s[0]; + uint32_t posOUTWidth = posOUT.s[1]; + + Dim2 vecOUT = out.vec; + uint32_t width_cal = vecOUT.s[1]; // 计算窗口 R 轴长度 + + AscendC::Adds(outLocal, inLocal, value, width_cal * height); + PipeBarrier(); + } + + template __aicore__ inline void NPU_Cast(Unisor &out, Unisor &in, const AscendC::RoundMode round_mode = AscendC::RoundMode::CAST_NONE) + { + LocalTensor srcLocal = in.get(); + LocalTensor dstLocal = out.get(); + Dim2 vec = in.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + uint32_t width = vec.s[1]; + uint32_t height = vec.s[0]; + uint32_t srcOffset = 0; + uint32_t dstOffset = 0; + AscendC::Cast(dstLocal, srcLocal, round_mode, width * height); + AscendC::PipeBarrier(); + } + + __aicore__ inline void NPU_Broadcast(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + LocalTensor outLocal = out.get(); + LocalTensor inLocal = in.get(); + + uint32_t width = vec.s[1]; + uint32_t height = vec.s[0]; + + uint32_t dstShape[ARRAY_IDX2] = {0, 0}; + uint32_t srcShape[ARRAY_IDX2] = {0, 0}; + dstShape[ARRAY_IDX0] = height; + dstShape[ARRAY_IDX1] = width; + srcShape[ARRAY_IDX0] = height; + srcShape[ARRAY_IDX1] = 1; + AscendC::BroadCast(outLocal, inLocal, dstShape, srcShape); + PipeBarrier(); + } + + // 向量乘：需要A的输入是dim1的 + __aicore__ inline void NPU_VecMul(Unisor &out, Unisor &A, Unisor &B) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + Dim2 vecB = B.vec; + LocalTensor outLocal = out.get(); + LocalTensor aLocal = A.get(); + LocalTensor bLocal = B.get(); + + uint32_t width = vec.s[1]; + uint32_t height = vec.s[0]; + + AscendC::Mul(outLocal, bLocal, aLocal, width * height); + PipeBarrier(); + } + + // 向量乘：需要A的输入是[1:]的 + __aicore__ inline void NPU_VecMuls(Unisor &out, Unisor &A, Unisor &B) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + Dim2 vecA = A.vec; + if (vecA.s[0] != 1) { + return; + } + + LocalTensor outLocal = out.get(); + LocalTensor aLocal = A.get(); + LocalTensor bLocal = B.get(); + + uint32_t width = vec.s[1]; + uint32_t height = vec.s[0]; + + // repeatTimes = 4, 一次迭代计算128个数, 共计算512个数 + // dstBlkStride, src0BlkStride, src1BlkStride = 1, 单次迭代内数据连续读取和写入 + // dstRepStride, src0RepStride, src1RepStride = 8, 相邻迭代间数据连续读取和写入 + const uint8_t dstRepStride = ceilINT(width, FACTOR_8); + const uint8_t src0RepStride = ceilINT(width, FACTOR_8); + + uint64_t mask = FACTOR_64; + uint16_t repeat = ceilINT(width, FACTOR_64); + for (int i = 0; i < repeat; i++) { + uint32_t dstOffset = FACTOR_64 * i; + uint32_t srcOffset = FACTOR_64 * i; + if (i == repeat - 1) { + mask = width - FACTOR_64 * i; + } + AscendC::Mul(outLocal[dstOffset], bLocal[dstOffset], aLocal[srcOffset], mask, height, {1, 1, 1, dstRepStride, src0RepStride, 0}); + } + PipeBarrier(); + } + + __aicore__ inline void NPU_Add_Bias(Unisor &out, Unisor &bias) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + LocalTensor outLocal = out.get(); + LocalTensor biasLocal = bias.get(); + + uint32_t width = vec.s[1]; + uint32_t height = vec.s[0]; + + uint32_t dstOffset = 0; + + uint64_t mask = FACTOR_64; + // repeatTimes = 4, 一次迭代计算128个数, 共计算512个数 + // dstBlkStride, src0BlkStride, src1BlkStride = 1, 单次迭代内数据连续读取和写入 + // dstRepStride, src0RepStride, src1RepStride = 8, 相邻迭代间数据连续读取和写入 + uint8_t dstRepStride = ceilINT(width, FACTOR_8); + uint8_t src0RepStride = ceilINT(width, FACTOR_8); + uint16_t repeat = ceilINT(width, FACTOR_64); + for (int i = 0; i < repeat; i++) { + uint32_t dstOffset = i * FACTOR_64; + uint32_t srcOffset = i * FACTOR_64; + if (i == repeat - 1) { + mask = width - i * FACTOR_64; + } + AscendC::Add(outLocal[dstOffset], outLocal[dstOffset], biasLocal[srcOffset], mask, height, {1, 1, 1, dstRepStride, src0RepStride, 0}); + } + PipeBarrier(); + } + + __aicore__ inline void NPU_Load(Unisor &out, Unisor &in) + { + Dim2 vec = in.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + GlobalTensor srcLocal = in.getGM(); + LocalTensor dstLocal = out.get(); + Dim2 pos = in.pos; + + float format_size = in.format_size; + int ele_num = COPY_BLK_BYTES; + int32_t in_width = in.vector_val.s[1] * format_size; + uint32_t in_height = in.vector_val.s[0]; + int32_t width = vec.s[1] * format_size; + uint32_t height = vec.s[0]; + uint32_t posHeight = pos.s[0]; + int32_t posWidth = pos.s[1] * format_size; + uint32_t widthBlock = ceilINT(width, ele_num); + + //ND->ND + uint16_t blockCount = height; + uint16_t blockLen = widthBlock; + uint16_t dstStride = 0; + uint16_t srcStride = ceilINT(in_width, ele_num) - widthBlock; + int64_t srcOffset = posHeight * in_width + posWidth; + int64_t dstOffset = 0; + DataCopy(dstLocal[dstOffset], srcLocal[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + __aicore__ void NPU_Store(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + LocalTensor src = in.get(); + GlobalTensor dst = out.getGM(); + Dim2 pos = out.pos; + + float format_size = out.format_size; + int ele_num = COPY_BLK_BYTES; + int32_t in_width = out.vector_val.s[1] * format_size; + uint32_t in_height = out.vector_val.s[0]; + int32_t width = vec.s[1] * format_size; + uint32_t height = vec.s[0]; + int32_t posWidth = pos.s[1] * format_size; + uint32_t posHeight = pos.s[0]; + uint32_t nBlocks = ceilINT(width, ele_num); + uint32_t mBlocks = ceilINT(height, ele_num); + uint32_t t_nBlocks = ceilINT(in_width, ele_num); + + //ND->ND + uint16_t blockCount = height; + uint16_t blockLen = nBlocks; + uint16_t srcStride = 0; + uint16_t dstStride = t_nBlocks - blockLen + t_nBlocks; + int64_t srcOffset = 0; + int64_t dstOffset = posHeight * in_width + posWidth; + + DataCopy(dst[dstOffset], src[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + __aicore__ void NPU_And(Unisor &out, Unisor &a, Unisor &b) + { + // 获取 UB 中的 Tensor + LocalTensor outLocal = out.get(); + LocalTensor aLocal = a.get(); + LocalTensor bLocal = b.get(); + uint32_t width = a.vec.s[1]; // a的整个R轴长度 + uint32_t height = a.vec.s[0]; // a的整个D轴长度 + + uint64_t mask = FACTOR_128; + uint8_t dstRepStride = FACTOR_8; + uint8_t src0RepStride = FACTOR_8; + uint8_t src1RepStride = 0; + + uint32_t dstOffset = 0; + uint32_t srcOffset = 0; + AscendC::And(outLocal[dstOffset], aLocal[dstOffset], bLocal[srcOffset], mask, height * width / FACTOR_128, {1, 1, 1, dstRepStride, src0RepStride, 0}); + + PipeBarrier(); + } + + // 底层API + // GM到UB + __aicore__ inline void NPU_Load_GMToUB(Unisor &out, Unisor &in, int stride = 1) + { + Dim2 vec = in.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + GlobalTensor srcLocal = in.getGM(); + LocalTensor dstLocal = out.get(); + Dim2 pos = in.pos; + + float format_size = in.format_size; + int64_t ele_num = COPY_BLK_BYTES; + int64_t in_width = in.vector_val.s[1] * format_size; + uint64_t in_height = in.vector_val.s[0]; + int64_t width = vec.s[1] * format_size; + uint64_t height = vec.s[0]; + uint64_t posHeight = pos.s[0]; + int64_t posWidth = pos.s[1] * format_size; + uint64_t widthBlock = ceilINT(width, ele_num); + + //ND->ND + uint16_t blockCount = height; + uint16_t blockLen = widthBlock; + uint16_t dstStride = 0; + uint16_t srcStride = stride * in_width / ele_num - widthBlock; + int64_t srcOffset = posHeight * in_width + posWidth; + int64_t dstOffset = 0; + DataCopy(dstLocal[dstOffset], srcLocal[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + // UB到GM + __aicore__ inline void NPU_Load_UBToGM(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + LocalTensor src = in.get(); + GlobalTensor dst = out.getGM(); + Dim2 pos = out.pos; + + float format_size = out.format_size; + int64_t ele_num = COPY_BLK_BYTES; + int64_t in_width = out.vector_val.s[1] * format_size; + uint64_t in_height = out.vector_val.s[0]; + int64_t width = vec.s[1] * format_size; + uint64_t height = vec.s[0]; + int64_t posWidth = pos.s[1] * format_size; + uint64_t posHeight = pos.s[0]; + uint64_t nBlocks = ceilINT(width, ele_num); + uint64_t mBlocks = ceilINT(height, ele_num); + uint64_t t_nBlocks = ceilINT(in_width, ele_num); + + //ND->ND + uint16_t blockCount = height; + uint16_t blockLen = nBlocks; + uint16_t srcStride = 0; + uint16_t dstStride = t_nBlocks - blockLen; + uint64_t srcOffset = 0; + uint64_t dstOffset = posHeight * in_width + posWidth; + DataCopy(dst[dstOffset], src[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + // 从 GM 到 L1A, ND -> Nz + __aicore__ inline void NPU_Load_GMToL1A(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + GlobalTensor src = in.getGM(); + LocalTensor dst = out.get(); + + float format_size = in.format_size; + int64_t ele_num = COPY_BLK_BYTES; + int64_t in_posWidth = format_size * in.pos.s[1]; + int64_t out_posWidth = format_size * out.pos.s[1]; + uint64_t in_posHeight = in.pos.s[0]; + uint64_t out_posHeight = out.pos.s[0]; + + int64_t in_width = in.vector_val.s[1] * format_size; + int64_t out_height = ceilINT_16(out.vector_val.s[0]); + int64_t height = vec.s[0]; + int64_t width = vec.s[1] * format_size; + + uint64_t srcOffset = in_posWidth + in_posHeight * in_width; + int64_t dstOffset = out_posHeight * ele_num + out_posWidth * out_height; + + AscendC::Nd2NzParams dataCopyA1Params; + dataCopyA1Params.ndNum = 1; + dataCopyA1Params.nValue = height; + dataCopyA1Params.dValue = width; + dataCopyA1Params.srcNdMatrixStride = 0; + dataCopyA1Params.srcDValue = in_width; + dataCopyA1Params.dstNzC0Stride = out_height; + dataCopyA1Params.dstNzNStride = 1; + dataCopyA1Params.dstNzMatrixStride = 0; + AscendC::DataCopy(dst[dstOffset], src[srcOffset], dataCopyA1Params); + } + + // 从 GM 到 L1B, ND -> Nz + __aicore__ inline void NPU_Load_GMToL1B(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + GlobalTensor src = in.getGM(); + LocalTensor dst = out.get(); + + float format_size = in.format_size; + int64_t ele_num = COPY_BLK_BYTES; + uint64_t in_posHeight = in.pos.s[0]; + int64_t in_posWidth = in.pos.s[1] * format_size; + uint64_t out_posHeight = out.pos.s[0]; + int64_t out_posWidth = out.pos.s[1] * format_size; + + int64_t in_height = in.vector_val.s[0]; + int64_t in_width = in.vector_val.s[1] * format_size; + int64_t out_height = out.vector_val.s[0]; + int64_t out_width = out.vector_val.s[1] * format_size; + int64_t height = vec.s[0]; + int64_t width = vec.s[1] * format_size; + + // NZ -> NZ 格式 + uint16_t blockCount = width / ele_num; + uint16_t blockLen = height; + uint16_t dstStride = out_height - height; + uint16_t srcStride = in_height - height; + + uint64_t srcOffset = in_posWidth * in_height + in_posHeight * ele_num; + uint64_t dstOffset = out_posWidth * out_height + out_posHeight * ele_num; + DataCopy(dst[dstOffset], src[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + // 从 L1 到 L0A, 转格式 + __aicore__ inline void NPU_Load_L1ToL0A(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + LocalTensor src = in.get(); + LocalTensor dst = out.get(); + float format_size = in.format_size; + int ele_num = COPY_BLK_BYTES; + + int32_t in_posHeight = in.pos.s[0]; + int32_t in_posWidth = in.pos.s[1] * format_size; + + uint32_t height = ceilINT_16(vec.s[0]); + int32_t width = vec.s[1] * format_size; + uint32_t in_height = ceilINT_16(in.vector_val.s[0]); + + uint32_t widthBlocks = width / FACTOR_32; + uint32_t heightBlocks = height / FACTOR_16; + + uint32_t srcStrideOffset = FACTOR_32 * in_height; + uint32_t dstStrideOffset = MAT_FRAC_BYTES; + + uint32_t srcOffset = in_posWidth * in_height + in_posHeight * ele_num; + uint32_t dstOffset = 0; + AscendC::LoadData2dParams loadL0AParams; + loadL0AParams.repeatTimes = heightBlocks; + loadL0AParams.srcStride = 1; + loadL0AParams.dstGap = widthBlocks - 1; + loadL0AParams.ifTranspose = false; + for (uint32_t i = 0; i < widthBlocks; i++) { + AscendC::LoadData(dst[dstOffset], src[srcOffset], loadL0AParams); + srcOffset = srcOffset + srcStrideOffset; + dstOffset = dstOffset + dstStrideOffset; + } + } + + // 从 L1 到 L0B, 转格式 + __aicore__ inline void NPU_Load_L1ToL0B(Unisor &out, Unisor &in) + { + Dim2 vec = out.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + float format_size = in.format_size; + + // INT4类型 + if (format_size == (float)INT4) { + LocalTensor src = in.get(); + LocalTensor dst = out.get(); + + uint32_t in_posHeight = in.pos.s[0]; + uint32_t in_posWidth = in.pos.s[1]; + + uint32_t height = vec.s[0]; + uint32_t width = vec.s[1]; + uint32_t in_height = in.vector_val.s[0]; + uint32_t widthBlocks = width / FACTOR_64; + uint32_t heightBlocks = height / FACTOR_64; + uint32_t in_heightBlocks = in_height / FACTOR_64; + + uint32_t srcOffset = in_posWidth * in_height + in_posHeight * FACTOR_64; + uint32_t dstOffset = 0; + + // 非转置 + if (vec.axis_name[0] == 'K') { + for (uint32_t i = 0; i < widthBlocks; i++) { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.srcStride = 1; + loadDataParams.repeatTimes = heightBlocks; + loadDataParams.dstGap = widthBlocks * FACTOR_4 - 1;//注意这是指原来相邻的两个方形搬到目的操作数之后的间隔大小 + loadDataParams.dstFracGap = 0; + AscendC::LoadDataWithTranspose(dst[dstOffset], src[srcOffset], loadDataParams); + srcOffset = srcOffset + FACTOR_64 * in_height; + dstOffset = dstOffset + MAT_FRAC_BYTES * FACTOR_8; + } + } else { //转置 + if (in_heightBlocks == heightBlocks) { + AscendC::LoadData2dParams loadL0BParams; + loadL0BParams.repeatTimes = heightBlocks * widthBlocks * FACTOR_4; + loadL0BParams.srcStride = 1; + loadL0BParams.dstGap = 0; + loadL0BParams.ifTranspose = false; + AscendC::LoadData(dst[dstOffset], src[srcOffset], loadL0BParams); + } else { + for (uint32_t i = 0; i < widthBlocks; i++) { + AscendC::LoadData2dParams loadL0BParams; + loadL0BParams.repeatTimes = heightBlocks * FACTOR_4; + loadL0BParams.srcStride = 1; + loadL0BParams.dstGap = 0; + loadL0BParams.ifTranspose = false; + AscendC::LoadData(dst[dstOffset], src[srcOffset], loadL0BParams); + srcOffset = srcOffset + in_height * FACTOR_64; + dstOffset = dstOffset + height * FACTOR_64; + } + } + } + } else if (format_size == (float)INT8) { // INT8类型 + LocalTensor src = in.get(); + LocalTensor dst = out.get(); + + uint32_t in_posHeight = in.pos.s[0]; + uint32_t in_posWidth = in.pos.s[1]; + + uint32_t height = vec.s[0]; + uint32_t width = vec.s[1]; + uint32_t in_height = in.vector_val.s[0]; + uint32_t widthBlocks = width / FACTOR_32; + uint32_t heightBlocks = height / FACTOR_32; + + uint32_t srcOffset = in_posWidth * in_height + in_posHeight * FACTOR_32; + uint32_t dstOffset = 0; + + // 非转置 + if (vec.axis_name[0]=='K') { + for (uint32_t i = 0; i < widthBlocks; i++) { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.srcStride = 1; + loadDataParams.repeatTimes = heightBlocks; + loadDataParams.dstGap = widthBlocks * FACTOR_2 - 1;// 注意这是指原来相邻的两个方形搬到目的操作数之后的间隔大小 + loadDataParams.dstFracGap = 0; + AscendC::LoadDataWithTranspose(dst[dstOffset], src[srcOffset], loadDataParams); + srcOffset = srcOffset + FACTOR_32 * in_height; + dstOffset = dstOffset + MAT_FRAC_BYTES * FACTOR_2; + } + } else { + AscendC::LoadData2dParams loadL0BParams; + loadL0BParams.repeatTimes = heightBlocks * widthBlocks * FACTOR_2; + loadL0BParams.srcStride = 1; + loadL0BParams.dstGap = 0; + loadL0BParams.ifTranspose = false; + AscendC::LoadData(dst[dstOffset], src[srcOffset], loadL0BParams); + } + } else if (format_size == (float)FP16) { // FP16类型 + LocalTensor src = in.get(); + LocalTensor dst = out.get(); + + uint32_t in_posHeight = in.pos.s[0]; + uint32_t in_posWidth = in.pos.s[1]; + + uint32_t height = vec.s[0]; + uint32_t width = vec.s[1]; + uint32_t in_height = in.vector_val.s[0]; + uint32_t widthBlocks = width / FACTOR_16; + uint32_t heightBlocks = height / FACTOR_16; + + uint32_t srcOffset = in_posWidth * in_height + in_posHeight * FACTOR_16; + uint32_t dstOffset = 0; + + // 非转置 + if (vec.axis_name[0]=='K') { + for (uint32_t i = 0; i < widthBlocks; i++) { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.srcStride = 1; + loadDataParams.repeatTimes = heightBlocks; + loadDataParams.dstGap = widthBlocks - 1;// 注意这是指原来相邻的两个方形搬到目的操作数之后的间隔大小 + AscendC::LoadDataWithTranspose(dst[dstOffset], src[srcOffset], loadDataParams); + srcOffset = srcOffset + (FACTOR_16 * in_height); + dstOffset = dstOffset + MUL_BYTES; + } + } else { // 转置 + AscendC::LoadData2dParams loadL0BParams; + loadL0BParams.repeatTimes = heightBlocks * widthBlocks; + loadL0BParams.srcStride = 1; + loadL0BParams.dstGap = 0; + loadL0BParams.ifTranspose = false; + AscendC::LoadData(dst[dstOffset], src[srcOffset], loadL0BParams); + } + } + } + + __aicore__ inline void NPU_Load_GMToL0B(Unisor &unisorB0, Unisor &unisorB1, Unisor &B, uint32_t MTE1_MTE2_EventID, uint32_t M_MTE1_EventID) + { + WaitFlag(MTE1_MTE2_EventID); + NPU_Load_GMToL1B(unisorB1, B); + SetFlag(MTE1_MTE2_EventID); + WaitFlag(MTE1_MTE2_EventID); + + WaitFlag(M_MTE1_EventID); + NPU_Load_L1ToL0B(unisorB0, unisorB1); + SetFlag(M_MTE1_EventID); + SetFlag(MTE1_MTE2_EventID); + } + + __aicore__ inline void NPU_Load_GMToL0A_CacheL1(Unisor &unisorA0, Unisor &unisorA1, + Unisor &A, uint32_t MTE1_MTE2_EventID_0, uint32_t MTE1_MTE2_EventID_1, uint32_t M_MTE1_EventID, bool NN_firstflag, bool NN_lastflag, bool firstNIter) + { + if (firstNIter) { + // ######## GM -> L1A + WaitFlag(MTE1_MTE2_EventID_1); + NPU_Load_GMToL1A(unisorA1, A); + SetFlag(MTE1_MTE2_EventID_1); + WaitFlag(MTE1_MTE2_EventID_1); + } + // ######## L1A -> L0A + WaitFlag(M_MTE1_EventID); + NPU_Load_L1ToL0A(unisorA0, unisorA1); + if (firstNIter) { + SetFlag(MTE1_MTE2_EventID_1); + } + SetFlag(M_MTE1_EventID); + } + + __aicore__ inline void NPU_Add_float(Unisor &out, Unisor &a, Unisor &b) + { + LocalTensor outLocal = out.get(); + LocalTensor aLocal = a.get(); + LocalTensor bLocal = b.get(); + + Dim2 vec = a.vec; + uint32_t size = vec.s[0] * vec.s[1]; + AscendC::Add(outLocal, aLocal, bLocal, size); + PipeBarrier(); + } + + __aicore__ inline void NPU_Add_float(Unisor &out, Unisor &a, Unisor &b) + { + LocalTensor outLocal = out.get(); + LocalTensor aLocal = a.get(); + LocalTensor bLocal = b.get(); + + Dim1 vec = a.vec; + uint32_t size = vec.s[0]; + AscendC::Add(outLocal, aLocal, bLocal, size); + PipeBarrier(); + } + + __aicore__ inline void NPU_Load_GMToUB(Unisor &out, Unisor &in) + { + Dim1 vec = in.vec; + if (vec.s[0] == 0) { + return; + } + GlobalTensor srcLocal = in.getGM(); + LocalTensor dstLocal = out.get(); + Dim1 pos = in.pos; + + float format_size = in.format_size; + int ele_num = COPY_BLK_BYTES; + int32_t in_width = in.vector_val.s[0] * format_size; + int32_t width = vec.s[0] * format_size; + int32_t posWidth = pos.s[0] * format_size; + uint32_t widthBlock = ceilINT(width, ele_num); + + //ND->ND + uint16_t blockCount = 1; + uint16_t blockLen = ceilINT(width, FACTOR_32); + uint16_t srcStride = 0; + uint16_t dstStride = 0; + int64_t srcOffset = posWidth; + int64_t dstOffset = 0; + DataCopy(dstLocal[dstOffset], srcLocal[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + __aicore__ inline void NPU_Load_UBToGM(Unisor &out, Unisor &in) + { + Dim1 vec = out.vec; + if (vec.s[0] == 0) { + return; + } + LocalTensor src = in.get(); + GlobalTensor dst = out.getGM(); + Dim1 pos = out.pos; + + float format_size = out.format_size; + int ele_num = COPY_BLK_BYTES; + int32_t in_width = out.vector_val.s[0] * format_size; + int32_t width = vec.s[0] * format_size; + int32_t posWidth = pos.s[0] * format_size; + + // ND->ND + uint16_t blockCount = 1; + uint16_t blockLen = ceilINT(width, FACTOR_32); + uint16_t srcStride = 0; + uint16_t dstStride = 0; + int64_t srcOffset = 0; + int64_t dstOffset = posWidth; + DataCopy(dst[dstOffset], src[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + __aicore__ inline void NPU_Load_GMToL1A_ND(Unisor &out, Unisor &in) + { + Dim2 vec = in.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + GlobalTensor srcLocal = in.getGM(); + LocalTensor dstLocal = out.get(); + Dim2 pos = in.pos; + + float format_size = in.format_size; + int ele_num = COPY_BLK_BYTES; + int32_t in_width = format_size * in.vector_val.s[1]; + int32_t width = format_size * vec.s[1]; + int32_t posWidth = format_size * pos.s[1]; + uint32_t in_height = in.vector_val.s[0]; + uint32_t posHeight = pos.s[0]; + uint32_t widthBlock = ceilINT(width, ele_num); + + //ND->ND + uint16_t dstStride = 0; + uint16_t srcStride = 0; + int64_t dstOffset = 0; + uint16_t blockCount = 1; + uint16_t blockLen = widthBlock; + int64_t srcOffset = posWidth + posHeight * in_width; + DataCopy(dstLocal[dstOffset], srcLocal[srcOffset], {blockCount, blockLen, srcStride, dstStride}); + } + + __aicore__ inline void NPU_Load_L0CToGM(GlobalTensor dst, TBuf &src, TBuf &deq, Cartesian dstCart, Cartesian deqCart, Dim2 dstShape, float dstFormat) + { + Dim2 vec = dstCart.vec; + if (vec.s[0] == 0 || vec.s[1] == 0) { + return; + } + + // int32_t 类型 + LocalTensor srcLocal = src.template Get(); + LocalTensor deqTensorLocal = deq.template Get(); + uint32_t deq_posWidth = deqCart.pos.s[1]; + + float format_size = dstFormat; + uint32_t out_posWidth = dstCart.pos.s[1]; + uint32_t out_posHeight = dstCart.pos.s[0]; + + uint32_t out_width = dstShape.s[1]; + uint32_t height = vec.s[0]; + uint32_t height_16 = ceilINT_16(vec.s[0]); + uint32_t width = vec.s[1]; + + uint64_t srcOffset = 0; + uint64_t dstOffset = (out_posHeight * out_width + out_posWidth) * (int)format_size; + AscendC::FixpipeParamsV220 fixpipeParams; + fixpipeParams.nSize = width; + fixpipeParams.mSize = height; + fixpipeParams.srcStride = height_16; + fixpipeParams.dstStride = out_width; + fixpipeParams.ndNum = 1; + fixpipeParams.srcNdStride = 0; + fixpipeParams.dstNdStride = 0; + fixpipeParams.quantPre = QuantMode_t::VDEQF16; + AscendC::Fixpipe(dst[dstOffset], srcLocal[srcOffset], deqTensorLocal[deq_posWidth], fixpipeParams); + } + + __aicore__ inline void NPU_Store_L0CToGM(Unisor &C, Unisor &unisorC0, Unisor unisorDeq, uint32_t eventID, bool atomic_flag) + { + // int32_t 类型 + if (atomic_flag) { + SetAtomicAdd(); + } + + SetFlag(eventID); + WaitFlag(eventID); + NPU_Load_L0CToGM(C.getGM(), unisorC0.calcBuf, unisorDeq.calcBuf, Cartesian(C.pos, C.vec), Cartesian(unisorDeq.pos, unisorDeq.vec), C.vector_val, C.format_size); + SetFlag(eventID); + if (atomic_flag) { + SetAtomicNone(); + } + } + + template __aicore__ inline void NPU_Muls(Unisor &out, Unisor &in, const T &value) + { + LocalTensor outLocal = out.get(); + LocalTensor inLocal = in.get(); + + Dim2 posOUT = out.pos; + uint32_t posOUTHeight = posOUT.s[0]; + uint32_t posOUTWidth = posOUT.s[1]; + + Dim2 vecOUT = out.vec; + uint32_t width_cal = vecOUT.s[1]; // 计算窗口 R 轴长度 + + uint32_t width = in.vector_val.s[1]; // in的整个R轴长度 + uint32_t height = in.vector_val.s[0]; // in的整个D轴长度 + + Dim2 posA = in.pos; + uint32_t posAHeight = posA.s[0]; + uint32_t posAWidth = posA.s[1]; + + AscendC::Muls(outLocal, inLocal, value, width_cal * height); + PipeBarrier(); + } +}; + +class KernelMatmul: public KernelBase { +public: + __aicore__ inline KernelMatmul(): KernelBase() {} + + __aicore__ inline void initBlockPos_3D(AxisTiling3Vector &splitRecord) + { + splitDim.s[ARRAY_IDX0] = splitRecord.s[ARRAY_IDX0]; + splitDim.s[ARRAY_IDX1] = splitRecord.s[ARRAY_IDX1]; + splitDim.s[ARRAY_IDX2] = splitRecord.s[ARRAY_IDX2]; + } + + __aicore__ inline uint32_t getMiniK_SplitBaseN(Dim2 baseCTiling, uint32_t total_k, float format_size) + { + int temp0 = 0; + int temp1 = 0; + + // split_N + temp0 = szAvailL0A / FACTOR_32 / baseCTiling.s[0] / format_size; + temp1 = szAvailL0A / baseCTiling.s[1] / format_size; + + uint32_t temp = temp0 > temp1 ? temp1 : temp0; + uint32_t result = 0; + if (temp > ceilINT_16(total_k)) { + result = total_k; + } else if (temp > FACTOR_16) { + if (temp % FACTOR_16 == 0) { + result = temp; + } else { + result = ceilINT_16(temp) - FACTOR_16; + } + } else { + result = FACTOR_16; + } + + return result; + } + + __aicore__ inline void NPU_Matmul(Unisor &out, Unisor &in0, Unisor &in1, bool isBias, Dim3 &realShape) + { + if (realShape.s[ARRAY_IDX0] == 0 || realShape.s[ARRAY_IDX1] == 0 || realShape.s[ARRAY_IDX2] == 0) { + return; + } + + float format_size = in0.format_size; + if (format_size == (float)INT4) { + // INT4类型 + LocalTensor src0 = in0.get(); + LocalTensor src1 = in1.get(); + LocalTensor dst = out.get(); + MmadParams mmadParams; + mmadParams.m = (uint16_t)ceilINT_16(realShape.s[ARRAY_IDX0]); + mmadParams.n = (uint16_t)realShape.s[ARRAY_IDX2]; + mmadParams.k = (uint16_t)realShape.s[ARRAY_IDX1]; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + mmadParams.isBias=isBias; + Mmad(dst, src0, src1, mmadParams); + } else if (format_size == (float)INT8) { + // INT8类型 + LocalTensor src0 = in0.get(); + LocalTensor src1 = in1.get(); + LocalTensor dst = out.get(); + MmadParams mmadParams; + mmadParams.m = (uint16_t)ceilINT_16(realShape.s[ARRAY_IDX0]); + mmadParams.n = (uint16_t)realShape.s[ARRAY_IDX2]; + mmadParams.k = (uint16_t)realShape.s[ARRAY_IDX1]; + mmadParams.cmatrixInitVal=true; + mmadParams.cmatrixSource=false; + mmadParams.isBias=isBias; + Mmad(dst, src0, src1, mmadParams); + } else if (format_size == (float)FP16) { + // FP16类型 + LocalTensor src0 = in0.get(); + LocalTensor src1 = in1.get(); + // 关键！！！ + // unitFlag = unitFlagIn; 预留参数，用户无需关心，使用默认值0即可 + // cmatrixSource = cmatrixSourceIn; 配置C矩阵初始值是否来源于C2（存放Bias的硬件缓存区）。默认值为false,Atlas 训练系列产品，仅支持配置为false。 + // cmatrixInitVal = cmatrixInitValIn; 配置C矩阵初始值是否为0。默认值true + // isBias: 配置是否需要累加初始矩阵，默认值为false + if (L0C_FORMAT_SIZE == FP16) { + // L0C是FP16类型 + LocalTensor dst = out.get(); + MmadParams mmadParams; + mmadParams.m = (uint16_t)realShape.s[ARRAY_IDX0]; + mmadParams.n = (uint16_t)realShape.s[ARRAY_IDX2]; + mmadParams.k = (uint16_t)realShape.s[ARRAY_IDX1]; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + mmadParams.isBias = isBias; + Mmad(dst, src0, src1, mmadParams); + } else { + // L0C是FP32类型 + LocalTensor dst = out.get(); + MmadParams mmadParams; + mmadParams.m = (uint16_t)ceilINT_16(realShape.s[ARRAY_IDX0]); + mmadParams.n = (uint16_t)realShape.s[ARRAY_IDX2]; + mmadParams.k = (uint16_t)realShape.s[ARRAY_IDX1]; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + mmadParams.isBias = isBias; + Mmad(dst, src0, src1, mmadParams); + } + } + } + + // loadEvenEventID是通过isEven决定的EventID！ + __aicore__ inline void npu_matmulUnisor_ping(Unisor &unisorC0, Unisor &unisorA0, Unisor &unisorB0, + Dim3 &realShape, uint32_t loadEvenEventID, uint32_t loadEventID, uint32_t storeCEventID, bool firstKIter) + { + if (firstKIter) { + WaitFlag(storeCEventID);// 与步骤4对称 + } + WaitFlag(loadEvenEventID); + WaitFlag(loadEventID); + NPU_Matmul(unisorC0, unisorA0, unisorB0, !firstKIter, realShape); + SetFlag(loadEventID); // 与步骤2对称 + } + + // loadEvenEventID是通过isEven决定的EventID！ + __aicore__ inline void npu_matmulUnisor_pong(Unisor &unisorC0, Unisor &unisorA0, Unisor &unisorB0, + Dim3 &realShape, uint32_t loadEvenEventID, uint32_t loadEventID, uint32_t storeCEventID, bool firstKIter) + { + if (firstKIter) { + WaitFlag(storeCEventID);// 与步骤4对称 + } + WaitFlag(loadEventID); + NPU_Matmul(unisorC0, unisorA0, unisorB0, !firstKIter, realShape); + SetFlag(loadEventID); + SetFlag(loadEvenEventID); + } + + __aicore__ inline void npu_user_defined_matmul_kernel_switch(Unisor &unisorA, Unisor &unisorB, Unisor &unisorC, Unisor &workUnisor, Unisor &biasUnisor, Unisor &saUnisor, Unisor &swUnisor, Dim2 &baseCTiling, uint32_t group_index, uint8_t kernel_index) + { + npu_user_defined_matmul_kernel_slave_CacheA_split_BaseN(unisorA, unisorB, unisorC, workUnisor, biasUnisor, saUnisor, swUnisor, baseCTiling, group_index); + } + + __aicore__ inline void npu_user_defined_matmul_kernel_slave_CacheA_split_BaseN(Unisor &A, Unisor &B, Unisor &C, Unisor &workUnisor, Unisor &biasUnisor, Unisor &saUnisor, Unisor &swUnisor, Dim2 &baseCTiling, uint32_t group_index) + { + if (g_coreType == AscendC::AIV) + { + return; + } + float format_size = A.format_size; + uint32_t split_k = A.vec.getAxisByName('K'); + uint32_t mk = FACTOR_512; + bool atomic_flag = (splitDim.s[0]>1); + uint32_t first_M = A.pos.getAxisByName('M'); + uint32_t Last_M = A.pos.getAxisByName('M') + A.vec.getAxisByName('M'); + uint32_t first_N = B.pos.getAxisByName('N'); + uint32_t Last_N = B.pos.getAxisByName('N') + B.vec.getAxisByName('N'); + uint32_t Base_M = baseCTiling.getAxisByName('M'); + uint32_t Base_N = baseCTiling.getAxisByName('N'); + Dim1 posN(B.pos.getAxisByName('N'), 'N'); + Dim1 vecN(B.vec.getAxisByName('N'), 'N'); + Dim1 posM(A.pos.getAxisByName('M'), 'M'); + Dim1 vecM(A.vec.getAxisByName('M'), 'M'); + + uint32_t B_posK = B.pos.getAxisByName('K'); + + // ping pong buffer + Unisor cacheUnisorA; + Unisor unisorBping; + Unisor unisorBpong; + + //L0A + Unisor unisorA0ping; + Unisor unisorA0pong; + + //L0B + Unisor unisorB0ping; + Unisor unisorB0pong; + + //L0C + Unisor unisorC0ping; + Unisor unisorC0pong; + + //取Base_N的一半（且是16的倍数） + uint32_t N_length11 = splitBy_64(Base_N); + uint32_t N_length22 = Base_N - N_length11; + + //BaseM x K + Dim2 vecAA = Dim2(Base_M, A.vec.getAxisByName('K'), 'M', 'K'); + + //miniK x BaseN + Dim2 vecMiniAA = Dim2(Base_M, mk, 'M', 'K'); + Dim2 vecMiniBB11 = Dim2(mk, N_length11, 'K', 'N'); + Dim2 vecMiniBB22 = Dim2(mk, N_length22, 'K', 'N'); + if (pattern != 1) { + vecMiniBB11.transpose(); + vecMiniBB22.transpose(); + } + Dim2 vecCC11 = Dim2(Base_M, N_length11, 'M', 'N'); + Dim2 vecCC22 = Dim2(Base_M, N_length22, 'M', 'N'); + + //L1A + cacheUnisorA.init(vecAA, A.format_size, this->pipe); + //L1B + unisorBping.init(vecMiniBB11, B.format_size, this->pipe); + unisorBpong.init(vecMiniBB22, B.format_size, this->pipe); + + //L0A + unisorA0ping.init(vecMiniAA, A.format_size, this->pipe); + unisorA0pong.init(vecMiniAA, A.format_size, this->pipe); + //L0B + unisorB0ping.init(vecMiniBB11, B.format_size, this->pipe); + unisorB0pong.init(vecMiniBB22, B.format_size, this->pipe); + //L0C + unisorC0ping.init(vecCC11, L0C_FORMAT_SIZE, this->pipe); //这里手动指定Float32 + unisorC0pong.init(vecCC22, L0C_FORMAT_SIZE, this->pipe); //这里手动指定Float32 + + Unisor unisorDeq; + Dim2 vecDeq = Dim2(1, vecN.s[0]); + unisorDeq.init_RealShape(vecDeq, INT64, this->pipe); + + Dim2 pos_sw = Dim2(group_index, posN.s[0]); + Cartesian cartesian_sw(pos_sw, vecDeq); + swUnisor.setCartesian(cartesian_sw); + + Dim2 pos_deq = Dim2(0, 0); + Cartesian cartesian_deq(pos_deq, vecDeq); + unisorDeq.setCartesian(cartesian_deq); + + bool first_flag = true; + + SetFlag(LOCAL_FLAGID0); + SetFlag(LOCAL_FLAGID1); + SetFlag(LOCAL_FLAGID2); + SetFlag(LOCAL_FLAGID3); + + SetFlag(LOCAL_FLAGID4); + SetFlag(LOCAL_FLAGID5); + SetFlag(LOCAL_FLAGID6); + SetFlag(LOCAL_FLAGID7); + SetFlag(LOCAL_FLAGID6); + SetFlag(LOCAL_FLAGID7); + + //K轴切分 + for (uint32_t i=0; i0) { + atomic_flag = true; + } else { + atomic_flag = false; + } + + bool isFirst_K = (i == 0); + bool isLast_K = (i == (splitDim.s[0]-1)); + uint32_t pos_K = i * split_k; + uint32_t current_K = split_k; + if (i == (splitDim.s[0]-1)) { + current_K = A.vector_val.s[1] - (i*split_k); + } + + uint32_t first_K = pos_K; + uint32_t last_K = pos_K + current_K; + + //N轴迭代器 + Cartesian cartesianN(posN, vecN); + Dim1 baseN(Base_N, 'N'); + CartesianIterator unisorIteratorN(cartesianN, baseN); + Cartesian cartesianUnisorN; + + //M轴迭代器 + Cartesian cartesianM(posM, vecM); + Dim1 baseM(Base_M, 'M'); + CartesianIterator unisorIteratorM(cartesianM, baseM); + Cartesian cartesianUnisorM; + + //K轴迭代器 + Dim1 pos_KK(pos_K, 'K'); + Dim1 vec_KK(current_K, 'K'); + Cartesian cartesianK(pos_KK, vec_KK); + Dim1 miniK(mk, 'K'); + CartesianIterator unisorIteratorK(cartesianK, miniK); + Cartesian cartesianUnisorK; + cartesianUnisorK.pos = pos_KK; + cartesianUnisorK.vec = vec_KK; + + uint32_t count = 0; + uint32_t write_point_count = 0; + while (unisorIteratorM.posIterator(cartesianUnisorM)) + { + while (unisorIteratorN.posIterator(cartesianUnisorN)) + { + Dim2 posC = Dim2(cartesianUnisorM.pos.getAxisByName('M'), cartesianUnisorN.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vecC = Dim2(cartesianUnisorM.vec.getAxisByName('M'), cartesianUnisorN.vec.getAxisByName('N'), 'M', 'N'); + + uint32_t N_length1 = splitBy_64(vecC.getAxisByName('N')); + uint32_t N_length2 = vecC.getAxisByName('N') - N_length1; + + //尾块处理 + if (vecC.getAxisByName('N')<=N_length11) { + N_length1=vecC.getAxisByName('N'); + N_length2=0; + } + + while (unisorIteratorK.posIterator(cartesianUnisorK)) + { + Dim1 posK = cartesianUnisorK.pos; + Dim1 vecK = cartesianUnisorK.vec; + Dim3 realShape1(vecC.getAxisByName('M'), vecK.getAxisByName('K'), N_length1, 'M', 'K', 'N'); + Dim3 realShape2(vecC.getAxisByName('M'), vecK.getAxisByName('K'), N_length2, 'M', 'K', 'N'); + + //GM_A矩阵坐标 + Dim2 posA_GM = Dim2(posC.getAxisByName('M'), posK.getAxisByName('K'), 'M', 'K'); + Dim2 vecA_GM = Dim2(vecC.getAxisByName('M'), vecK.getAxisByName('K'), 'M', 'K'); + Cartesian cartesianTempA_GM1(posA_GM, vecA_GM); + + //L1A矩阵坐标 + Dim2 posA_L11 = Dim2(0, posK.getAxisByName('K') - first_K, 'M', 'K'); + Cartesian cartesianTempA_L11(posA_L11, vecA_GM); + + //L0A矩阵坐标 + Dim2 posA_L01 = Dim2(0, 0, 'M', 'K'); + Cartesian cartesianTempA_L01(posA_L01, vecA_GM); + + //GM_B1矩阵坐标 + Dim2 posB_GM1 = Dim2(posK.getAxisByName('K'), posC.getAxisByName('N'), 'K', 'N'); + Dim2 vecB_GM1 = Dim2(vecK.getAxisByName('K'), N_length1, 'K', 'N'); + if (pattern != 1) { + posB_GM1.transpose(); + vecB_GM1.transpose(); + } + posB_GM1.s[1] = posB_GM1.s[1] + B_posK; + Cartesian cartesianTempB_GM1(posB_GM1, vecB_GM1); + + //GM_B2矩阵坐标 + Dim2 posB_GM2 = Dim2(posK.getAxisByName('K'), posC.getAxisByName('N') + N_length1, 'K', 'N'); + Dim2 vecB_GM2 = Dim2(vecK.getAxisByName('K'), N_length2, 'K', 'N'); + if (pattern != 1) { + posB_GM2.transpose(); + vecB_GM2.transpose(); + } + posB_GM2.s[1] = posB_GM2.s[1] + B_posK; + Cartesian cartesianTempB_GM2(posB_GM2, vecB_GM2); + + //L1_B1矩阵坐标 + Dim2 posB_L11 = Dim2(0, 0, 'K', 'N'); + if (pattern != 1) { + posB_L11.transpose(); + } + Cartesian cartesianTempB_L11(posB_L11, vecB_GM1); + + //L1_B2矩阵坐标 + Dim2 posB_L12 = Dim2(0, 0, 'K', 'N'); + if (pattern != 1) { + posB_L12.transpose(); + } + Cartesian cartesianTempB_L12(posB_L12, vecB_GM2); + + bool firstNIter = (cartesianUnisorN.pos.getAxisByName('N') == first_N); + bool LastNIter = (cartesianUnisorN.pos.getAxisByName('N') + Base_N >= Last_N); + bool firstMIter = (cartesianUnisorM.pos.getAxisByName('M') == first_M); + bool LastMIter = (cartesianUnisorM.pos.getAxisByName('M') + Base_M >= Last_M); + bool firstKIter = (posK.getAxisByName('K') == first_K); + bool lastKIter = (posK.getAxisByName('K') + mk >= last_K); + bool NN_firstflag = (firstNIter && firstKIter); + bool NN_lastflag = (LastNIter && lastKIter); + bool isFirstMN = (firstMIter && firstNIter); + bool isLastMN = (LastMIter && LastNIter); + bool postFlag = (isLast_K && lastKIter); + bool isEven = (count % 2 == 0); + + if (isEven) + { + A.setCartesian(cartesianTempA_GM1); + cacheUnisorA.setCartesian(cartesianTempA_L11); + unisorA0ping.setCartesian(cartesianTempA_L11); + NPU_Load_GMToL0A_CacheL1(unisorA0ping, cacheUnisorA, A, LOCAL_FLAGID4, LOCAL_FLAGID0, LOCAL_FLAGID4, NN_firstflag, NN_lastflag, firstNIter); + } + else + { + A.setCartesian(cartesianTempA_GM1); + cacheUnisorA.setCartesian(cartesianTempA_L11); + unisorA0pong.setCartesian(cartesianTempA_L11); + NPU_Load_GMToL0A_CacheL1(unisorA0pong, cacheUnisorA, A, LOCAL_FLAGID5, LOCAL_FLAGID1, LOCAL_FLAGID5, NN_firstflag, NN_lastflag, firstNIter); + } + + B.setCartesian(cartesianTempB_GM1); + unisorBping.setCartesian(cartesianTempB_L11); + unisorB0ping.setCartesian(cartesianTempB_L11); + NPU_Load_GMToL0B(unisorB0ping, unisorBping, B, LOCAL_FLAGID2, LOCAL_FLAGID6); + + if (isEven) + { + npu_matmulUnisor_ping(unisorC0ping, unisorA0ping, unisorB0ping, realShape1, LOCAL_FLAGID4, LOCAL_FLAGID6, LOCAL_FLAGID6, firstKIter); + } + else + { + npu_matmulUnisor_ping(unisorC0ping, unisorA0pong, unisorB0ping, realShape1, LOCAL_FLAGID5, LOCAL_FLAGID6, LOCAL_FLAGID6, firstKIter); + } + + if (lastKIter) + { + Dim2 posC1 = Dim2(cartesianUnisorM.pos.getAxisByName('M'), cartesianUnisorN.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vecC1 = Dim2(cartesianUnisorM.vec.getAxisByName('M'), N_length1, 'M', 'N'); + Cartesian cartesianTempC1(posC1, vecC1); + workUnisor.setCartesian(cartesianTempC1); + + if (first_flag) { + SetFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID0); + NPU_Load_GMToL1A_ND(unisorDeq, swUnisor); + SetFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID0); + first_flag = false; + } + + Dim2 pos_deq = Dim2(0, cartesianUnisorN.pos.getAxisByName('N')-posN.s[0]); + Dim2 vec_deq = Dim2(1, N_length2); + Cartesian cartesian_deq(pos_deq, vec_deq); + unisorDeq.setCartesian(cartesian_deq); + + NPU_Store_L0CToGM(workUnisor, unisorC0ping, unisorDeq, LOCAL_FLAGID6, atomic_flag); + } + + // ######### GM -> L0B + B.setCartesian(cartesianTempB_GM2); + unisorBpong.setCartesian(cartesianTempB_L12); + unisorB0pong.setCartesian(cartesianTempB_L12); + NPU_Load_GMToL0B(unisorB0pong, unisorBpong, B, LOCAL_FLAGID3, LOCAL_FLAGID7); + + if (isEven) + { + npu_matmulUnisor_pong(unisorC0pong, unisorA0ping, unisorB0pong, realShape2, LOCAL_FLAGID4, LOCAL_FLAGID7, LOCAL_FLAGID7, firstKIter); + } + else + { + npu_matmulUnisor_pong(unisorC0pong, unisorA0pong, unisorB0pong, realShape2, LOCAL_FLAGID5, LOCAL_FLAGID7, LOCAL_FLAGID7, firstKIter); + } + + // ######### L0C -> GM + if (lastKIter) + { + Dim2 posC2 = Dim2(cartesianUnisorM.pos.getAxisByName('M'), cartesianUnisorN.pos.getAxisByName('N') + N_length1, 'M', 'N'); + Dim2 vecC2 = Dim2(cartesianUnisorM.vec.getAxisByName('M'), N_length2, 'M', 'N'); + Cartesian cartesianTempC2(posC2, vecC2); + workUnisor.setCartesian(cartesianTempC2); + + Dim2 pos_deq = Dim2(0, cartesianUnisorN.pos.getAxisByName('N') + N_length1 -posN.s[0]); + Dim2 vec_deq = Dim2(1, N_length2); + Cartesian cartesian_deq(pos_deq, vec_deq); + unisorDeq.setCartesian(cartesian_deq); + NPU_Store_L0CToGM(workUnisor, unisorC0pong, unisorDeq, LOCAL_FLAGID7, atomic_flag); + } + count++; + } + } + } + } + + WaitFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID1); + WaitFlag(LOCAL_FLAGID2); + WaitFlag(LOCAL_FLAGID3); + + WaitFlag(LOCAL_FLAGID4); + WaitFlag(LOCAL_FLAGID5); + WaitFlag(LOCAL_FLAGID6); + WaitFlag(LOCAL_FLAGID7); + WaitFlag(LOCAL_FLAGID6); + WaitFlag(LOCAL_FLAGID7); + + this->pipe->Reset(); + } + + __aicore__ inline void npu_user_defined_matmul_kernel_slave_Post_Processing_off(Unisor &C, Unisor &workUnisor, + Unisor &biasUnisor, + Unisor &saUnisor, + Unisor &swUnisor) + { + //注意下面值不一定是16的倍数 + uint32_t pos_Work_M = workUnisor.pos.getAxisByName('M'); + uint32_t pos_Work_N = workUnisor.pos.getAxisByName('N'); + uint32_t write_point_count = pos_Work_N; + uint32_t work_N = workUnisor.vec.getAxisByName('N'); + float format_size = workUnisor.format_size; + + uint32_t startOffset1 = (szAvailL0C / FACTOR_2) * write_point_count; + uint32_t startOffset2 = (szAvailL0C / FACTOR_2) * write_point_count + (int)(work_N * format_size); + + uint32_t single_M = workUnisor.vec.getAxisByName('M'); + if (single_M == 0) { + return; + } + uint32_t current_M = ceilINT(single_M, FACTOR_2); //取一半 + bool single_aiv = (single_M < FACTOR_128); + if (single_aiv) { + current_M = single_M; + } + if (current_M % FACTOR_2 != 0) { + current_M = current_M +1; + } + uint32_t current_N = workUnisor.vec.getAxisByName('N'); + if (current_N == 0) { + return; + } + uint32_t pos_M = C.pos.getAxisByName('M'); + uint32_t pos_N = C.pos.getAxisByName('N'); + uint32_t pos_MM = pos_M; + uint32_t rowsum_m = ceilINT(current_M, FACTOR_2); + + uint32_t BlockIdx = GetBlockIdx(); + uint32_t SubBlockIdx = GetSubBlockIdx(); + //第二个V核，暂时不考虑尾块 + if (SubBlockIdx == 1) { + if (single_aiv) { + return; + } + pos_MM = pos_M + current_M; + current_M = single_M - current_M; + + if (single_M == FACTOR_2) { + return; + } + } + if (current_M == 0) { + return; + } + + //UB内存是192KB 需要三块Half大小的内存 每块不能超过64KB Base_M*Base_N <= 32*1024 + uint32_t ubSize = szAvailUB - FACTOR_4 * current_N - FACTOR_4 * current_M; + uint32_t Base_M = ubSize / FACTOR_12 / current_N; + if (Base_M % FACTOR_2 != 0) { + Base_M = Base_M -1; + } + + if (Base_M > current_M) { + Base_M = current_M; + } + uint32_t Base_N = current_N; + + //M轴迭代器 + Dim1 posM(pos_MM, 'M'); + Dim1 vecM(current_M, 'M'); + Cartesian cartesianM(posM, vecM); + Dim1 baseM(Base_M, 'M'); + CartesianIterator unisorIteratorM(cartesianM, baseM); + Cartesian cartesianUnisorM; + + //N轴迭代器 + Dim1 posN(pos_N, 'N'); + Dim1 vecN(current_N, 'N'); + Cartesian cartesianN(posN, vecN); + Dim1 baseN(Base_N, 'N'); + CartesianIterator unisorIteratorN(cartesianN, baseN); + Cartesian cartesianUnisorN; + + Dim2 pos_work(0, 0, 'M', 'N'); + Dim1 pos_work_d1(0, 'M'); + Dim2 vec_work(Base_M / FACTOR_2, Base_N, 'M', 'N'); + + //Bias初始化 + Dim2 vec_bias(1, current_N, 'M', 'N'); + + Unisor UBUnisor_Bias; + UBUnisor_Bias.init_RealShape(vec_bias, FACTOR_4, this->pipe); + UBUnisor_Bias.setCartesian(pos_work, vec_bias); + NPU_Load_GMToUB(UBUnisor_Bias, biasUnisor); + + //Sa初始化 + Dim1 vec_sa(Base_M / FACTOR_2, 'M'); + + Unisor UBUnisor_sa; + UBUnisor_sa.init(vec_sa, FACTOR_4, this->pipe); + + Unisor UBUnisor_sa_broadcast; + UBUnisor_sa_broadcast.init_RealShape(vec_work, FACTOR_4, this->pipe); + UBUnisor_sa_broadcast.setCartesian(pos_work, vec_work); + + Unisor UBUnisor_FP32_A1; + UBUnisor_FP32_A1.init_RealShape(vec_work, FACTOR_4, this->pipe); + + Unisor UBUnisor_FP32_A2; + UBUnisor_FP32_A2.init_RealShape(vec_work, FACTOR_4, this->pipe); + + Unisor UBUnisor_FP32_A1_FP16; + UBUnisor_FP32_A1_FP16.init_RealShape(vec_work, FACTOR_2, this->pipe); + + Unisor UBUnisor_FP32_A2_FP16; + UBUnisor_FP32_A2_FP16.init_RealShape(vec_work, FACTOR_2, this->pipe); + + Unisor UBUnisor_FP16; + UBUnisor_FP16.init_RealShape(vec_work, FACTOR_2, this->pipe); + + uint32_t count = 0; + AscendC::SetFlag(LOCAL_FLAGID0); + AscendC::SetFlag(LOCAL_FLAGID1); + while (unisorIteratorM.posIterator(cartesianUnisorM)) + { + while (unisorIteratorN.posIterator(cartesianUnisorN)) + { + bool isEven = (count % FACTOR_2 == 0); + Dim2 posC = Dim2(cartesianUnisorM.pos.getAxisByName('M'), cartesianUnisorN.pos.getAxisByName('N'), 'M', 'N'); + Dim2 out_posC = Dim2(cartesianUnisorM.pos.getAxisByName('M') / FACTOR_2, cartesianUnisorN.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vecC = Dim2(cartesianUnisorM.vec.getAxisByName('M') / FACTOR_2, cartesianUnisorN.vec.getAxisByName('N'), 'M', 'N'); + Dim1 pos_sa = Dim1(cartesianUnisorM.pos.getAxisByName('M') / FACTOR_2, 'M'); + Dim1 vec_sa = Dim1(cartesianUnisorM.vec.getAxisByName('M') / FACTOR_2, 'M'); + Dim2 pos_sw = Dim2(0, cartesianUnisorN.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vec_sw = Dim2(1, cartesianUnisorN.vec.getAxisByName('N'), 'M', 'N'); + + //1、WorkSpace->UB + workUnisor.setCartesian(posC, vecC); + UBUnisor_FP32_A1.setCartesian(pos_work, vecC); + UBUnisor_FP32_A2.setCartesian(pos_work, vecC); + UBUnisor_FP32_A1_FP16.setCartesian(pos_work, vecC); + UBUnisor_FP32_A2_FP16.setCartesian(pos_work, vecC); + UBUnisor_FP16.setCartesian(pos_work, vecC); + saUnisor.setCartesian(pos_sa, vec_sa); + UBUnisor_sa.setCartesian(pos_work_d1, vec_sa); + UBUnisor_sa_broadcast.setCartesian(pos_work, vecC); + + AscendC::WaitFlag(LOCAL_FLAGID0); + NPU_Load_GMToUB(UBUnisor_FP32_A1_FP16, workUnisor, FACTOR_2); + posC.s[0] = posC.s[0] + 1; + workUnisor.setCartesian(posC, vecC); + NPU_Load_GMToUB(UBUnisor_FP32_A2_FP16, workUnisor, FACTOR_2); + NPU_Load_GMToUB(UBUnisor_sa, saUnisor); + AscendC::SetFlag(LOCAL_FLAGID2); + AscendC::WaitFlag(LOCAL_FLAGID2); + + NPU_Cast(UBUnisor_FP32_A1, UBUnisor_FP32_A1_FP16); + NPU_Cast(UBUnisor_FP32_A2, UBUnisor_FP32_A2_FP16); + + NPU_Broadcast(UBUnisor_sa_broadcast, UBUnisor_sa); + + //A1乘16 + const float sixteen = 16.0; + NPU_Muls(UBUnisor_FP32_A1, UBUnisor_FP32_A1, sixteen); + + //A1+A2 + NPU_Add_float(UBUnisor_FP32_A1, UBUnisor_FP32_A1, UBUnisor_FP32_A2); + + NPU_Add_Bias(UBUnisor_FP32_A1, UBUnisor_Bias); + + // * Sa + NPU_VecMul(UBUnisor_FP32_A1, UBUnisor_sa_broadcast, UBUnisor_FP32_A1); + + //2、FP32->FP16 + AscendC::WaitFlag(LOCAL_FLAGID1); + if (this->output_type == 0) { + NPU_Cast(UBUnisor_FP16, UBUnisor_FP32_A1); + } else { + NPU_Cast(UBUnisor_FP16, UBUnisor_FP32_A1, AscendC::RoundMode::CAST_RINT); + } + + AscendC::SetFlag(LOCAL_FLAGID0); + AscendC::SetFlag(LOCAL_FLAGID3); + AscendC::WaitFlag(LOCAL_FLAGID3); + + //3、UB->GM + C.setCartesian(out_posC, vecC); + NPU_Load_UBToGM(C, UBUnisor_FP16); + AscendC::SetFlag(LOCAL_FLAGID1); + } + } + AscendC::WaitFlag(LOCAL_FLAGID0); + AscendC::WaitFlag(LOCAL_FLAGID1); + + this->pipe->Reset(); + } + + //A INT8->INT4 处理单核 + __aicore__ inline void npu_user_defined_matmul_kernel_slave_Single_A8ToA4_off(Unisor in, Unisor out) + { + //多核切分 + uint32_t BlockIdx = GetBlockIdx(); + uint32_t M_length = in.vec.s[0]; + uint32_t K_length = in.vec.s[1]; + + //切分成40个Vector核计算 + uint32_t singleM = ceilINT(M_length, nCoreAIV); + singleM = ceilINT_16(singleM); + if (singleM < FACTOR_32) { + singleM = FACTOR_32; + } + uint32_t singleK = K_length; + + //注意！！！BaseM要根据下面使用了多少个UB Unisor来计算 + uint32_t BaseM = (szAvailUB - singleK) / singleK / FACTOR_7; + if (BaseM > M_length) { + BaseM = M_length; + } + + Dim2 baseTiling(BaseM, singleK, 'M', 'K'); + + if (BlockIdx * singleM > M_length) { //已算完 + return ; + } + + uint32_t current_M = singleM; + if ((BlockIdx + 1) * singleM > M_length) { + current_M = M_length - (BlockIdx * singleM); + } + + if (current_M == 0) { + return; + } + + uint32_t posM_in = in.pos.s[0] + BlockIdx * singleM; + uint32_t posK_in = in.pos.s[1]; + uint32_t posM_out = out.pos.s[0] + BlockIdx * singleM; + uint32_t posK_out = out.pos.s[1]; + + in.pos.setAxisName('M', 'K'); + in.vec.setAxisName('M', 'K'); + out.pos.setAxisName('M', 'K'); + out.vec.setAxisName('M', 'K'); + + Unisor inputUnisor(in); + Unisor outputHighUnisor(out); + Unisor outputLowUnisor(out); + + Dim2 posIn = Dim2(posM_in, posK_in, 'M', 'K'); + Dim2 vecIn = Dim2(current_M, singleK, 'M', 'K'); + inputUnisor.setCartesian(posIn, vecIn); + + InternalRun_off(inputUnisor, outputHighUnisor, outputLowUnisor, baseTiling); + this->pipe->Reset(); + } + + __aicore__ inline void InternalRun_off(Unisor in, Unisor outHigh, Unisor outLow, Dim2 &baseTiling) + { + uint32_t posM_in = in.pos.getAxisByName('M'); + Dim1 posM(0, 'M'); + Dim1 vecM(in.vec.getAxisByName('M'), 'M'); + Cartesian cartesianM(posM, vecM); + Dim1 baseM(baseTiling.getAxisByName('M'), 'M'); + CartesianIterator unisorIteratorM(cartesianM, baseM); + Cartesian cartesianUnisorM; + cartesianUnisorM.pos.setAxisName('M'); + cartesianUnisorM.vec.setAxisName('M'); + + Dim1 posK(in.pos.getAxisByName('K'), 'K'); + Dim1 vecK(in.vec.getAxisByName('K'), 'K'); + Cartesian cartesianK(posK, vecK); + Dim1 baseK(baseTiling.getAxisByName('K'), 'K'); + CartesianIterator unisorIteratorK(cartesianK, baseK); + Cartesian cartesianUnisorK; + cartesianUnisorK.pos.setAxisName('K'); + cartesianUnisorK.vec.setAxisName('K'); + + Unisor maskUnisor; + Unisor unisorIn; + Unisor highUnisor; + Unisor highInt4Unisor; + Unisor lowInt8Unisor; + Unisor lowUnisor; + Unisor lowInt4Unisor; + { + Dim1 maskDim(MUL_BYTES / INT8, 'K'); + maskUnisor.init_RealShape(maskDim, INT8, this->pipe); + Dim2 vecIn(baseM.getAxisByName('M'), baseK.getAxisByName('K'), 'M', 'K'); + unisorIn.init_RealShape(vecIn, INT8, this->pipe); + highUnisor.init_RealShape(vecIn, FP16, this->pipe); + highInt4Unisor.init_RealShape(vecIn, INT4, this->pipe); + lowInt8Unisor.init_RealShape(vecIn, INT8, this->pipe); + lowUnisor.init_RealShape(vecIn, FP16, this->pipe); + lowInt4Unisor.init_RealShape(vecIn, INT4, this->pipe); + } + + // Step 0: Prepare a 1D mask filled with 0x0fU, which will be used to + // extract the lower 4 bits of each int8 of 'in'. + NPU_Duplicate(maskUnisor, 0x0f0fU); + + SetFlag(LOCAL_FLAGID0); + SetFlag(LOCAL_FLAGID0); + SetFlag(LOCAL_FLAGID1); + while (unisorIteratorM.posIterator(cartesianUnisorM)) + { + while (unisorIteratorK.posIterator(cartesianUnisorK)) + { + Dim2 posIn(cartesianUnisorM.pos.getAxisByName('M') + posM_in, + cartesianUnisorK.pos.getAxisByName('K'), + 'M', 'K'); + Dim2 vecIn(cartesianUnisorM.vec.getAxisByName('M'), + cartesianUnisorK.vec.getAxisByName('K'), + 'M', 'K'); + Dim2 posHighOut(cartesianUnisorM.pos.getAxisByName('M') * FACTOR_2 + posM_in * FACTOR_2, + cartesianUnisorK.pos.getAxisByName('K'), + 'M', 'K'); + Dim2 posLowOut(cartesianUnisorM.pos.getAxisByName('M') * FACTOR_2 + posM_in * FACTOR_2 + 1, + cartesianUnisorK.pos.getAxisByName('K'), + 'M', 'K'); + + Cartesian cartesianIn(posIn, vecIn); + Cartesian cartesianHighOut(posHighOut, vecIn); + Cartesian cartesianLowOut(posLowOut, vecIn); + in.setCartesian(cartesianIn); + outHigh.setCartesian(cartesianHighOut); + outLow.setCartesian(cartesianLowOut); + + // Step 1: Load a unisor of input data into 'unisorIn'. + WaitFlag(LOCAL_FLAGID0); + NPU_Load(unisorIn, in); + SetFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID0); + + // + // Higher 4 bits + // + + // Step 2: Cast 'unisorIn' to fp16, and save the results to 'highUnisor'. + NPU_Cast(highUnisor, unisorIn); + + // Step 6: Extract the lower 4 bits of each int8. + NPU_And(lowInt8Unisor, unisorIn, maskUnisor); + SetFlag(LOCAL_FLAGID0); + + // Step 3: Divide 'highUnisor' by 1/16 element-wisely to extract + // the higher 4 bits of each number. + const half one_over_sixteen = 0.0625; + NPU_Muls(highUnisor, highUnisor, one_over_sixteen); + + // Step 4: Cast 'highUnisor' to int4. + WaitFlag(LOCAL_FLAGID0); + NPU_Cast(highInt4Unisor, highUnisor, AscendC::RoundMode::CAST_FLOOR); + SetFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID0); + + // Step 5: Store 'highInt4Unisor' back to global memory. + NPU_Store(outHigh, highInt4Unisor); + SetFlag(LOCAL_FLAGID0); + + // + // Lower 4 bits + // + + // Step 7: Cast 'lowInt8Unisor' to half. + NPU_Cast(lowUnisor, lowInt8Unisor); + + // Step 8: Subtract 8 from 'lowUnisor'. + const half neg_eight = -8.0; + NPU_Adds(lowUnisor, lowUnisor, neg_eight); + + // Step 9: Cast 'lowUnisor' to int4. + WaitFlag(LOCAL_FLAGID1); + NPU_Cast(lowInt4Unisor, lowUnisor); + SetFlag(LOCAL_FLAGID1); + WaitFlag(LOCAL_FLAGID1); + + // Step 10: Store 'lowInt4Unisor' back to global memory. + NPU_Store(outLow, lowInt4Unisor); + SetFlag(LOCAL_FLAGID1); + } + } + WaitFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID1); + WaitFlag(LOCAL_FLAGID0); + } +}; + +__aicore__ inline void dynamic_unisor_programming(GM_ADDR gmA, GM_ADDR gmB, GM_ADDR gmC, GM_ADDR gmGrpList, + GM_ADDR gmBias, GM_ADDR gmOffset, GM_ADDR gmSa, GM_ADDR gmSw, + GM_ADDR gmWorkspaceDevice, A8W4HPTiling *tiling_data, + TPipe *pipe) +{ + auto ori_in0_shape = tiling_data->ori_in0_shape; + auto ori_in1_shape = tiling_data->ori_in1_shape; + auto ori_out_shape = tiling_data->ori_out_shape; + auto single_core_tiling = tiling_data->single_core_tiling; + auto single_core_base_tiling = tiling_data->single_core_base_tiling; + AxisTiling3Vector splitRecord{tiling_data->splitRecord[ARRAY_IDX0], tiling_data->splitRecord[ARRAY_IDX1], + tiling_data->splitRecord[ARRAY_IDX2]}; + + KernelMatmul op; + op.Init(pipe); + op.initBlockPos_3D(splitRecord); + uint32_t numAic = tiling_data->numAic; + uint32_t numAiv = tiling_data->numAiv; + uint32_t pattern = tiling_data->pattern; + op.pattern = pattern; + op.nCoreAIC = numAic; + op.nCoreAIV = numAiv; + op.szAvailUB = tiling_data->szUb; + op.szAvailL0A = tiling_data->szL0A; + op.szAvailL0C = tiling_data->szL0C; + op.pattern = pattern; + op.output_type = tiling_data->output_type; + + uint32_t blockIdx = GetBlockIdx();//当前核ID + uint32_t blockId = 0; + + if (g_coreType == AscendC::AIC) { + blockId = blockIdx; + } + + if (g_coreType == AscendC::AIV) { + blockId = blockIdx / FACTOR_2; + } + op.blockId = blockId; + + //GroupMatmul参数 + uint32_t group_type = tiling_data->group_type; + uint32_t group_num = tiling_data->group_num; + + //Group总长度 + uint32_t total_M = 0; // 16384 + uint32_t total_N = ori_out_shape[1]; // 4096 + uint32_t total_K_A = ori_in0_shape[1]; // 7168 + uint32_t total_K_B = ori_in0_shape[1] * group_num; // 512 + + //单个矩阵长度 + int64_t M_Lengths[MAX_GROUP_LEN]; + int64_t N_Length = total_N; + + //单核迭代长度 + uint32_t single_M = single_core_tiling[ARRAY_IDX0]; + uint32_t single_N = single_core_tiling[ARRAY_IDX1]; + uint32_t single_K = single_core_tiling[ARRAY_IDX2]; + + //Base迭代长度 + Dim2 baseCTiling(single_core_base_tiling[0], single_core_base_tiling[1], 'M', 'N'); + + float format_in = tiling_data->format_in; + float format_out = tiling_data->format_out; + + GlobalTensor groupListGm; + groupListGm.SetGlobalBuffer((__gm__ uint64_t *)gmGrpList); + + uint32_t true_group_num = 0; + for (uint32_t group_idx = 0; group_idx < group_num; group_idx++) { + if (group_type == 0) { // 从m轴方向切割 + int64_t mSizeInGroup = static_cast(groupListGm.GetValue(group_idx)); + if (mSizeInGroup != 0) { + true_group_num += 1; + } + M_Lengths[group_idx] = mSizeInGroup * FACTOR_2; + total_M += M_Lengths[group_idx]; + } + } + + constexpr uint32_t TOTAL_N_THRESHOLD_1024 = 1024; + constexpr uint32_t TOTAL_N_THRESHOLD_512 = 512; + constexpr float USAGE_RATE_THRESHOLD_1024 = 0.9f; + constexpr float USAGE_RATE_THRESHOLD_512 = 0.8f; + + uint32_t E_M_length = ceilINT(total_M, true_group_num); + int32_t UsageRate512 = ceilINT(total_N, TOTAL_N_THRESHOLD_512) * ceilINT(E_M_length, single_M) * true_group_num; + int32_t virtualCoreNum512 = ceilINT((uint32_t)UsageRate512, numAic) * numAic; + int32_t UsageRate1024 = ceilINT(total_N, TOTAL_N_THRESHOLD_1024) * ceilINT(E_M_length, single_M) * true_group_num; + int32_t virtualCoreNum1024 = ceilINT((uint32_t)UsageRate1024, numAic) * numAic; + if ((float)UsageRate1024/virtualCoreNum1024 >= USAGE_RATE_THRESHOLD_1024) { + single_N = TOTAL_N_THRESHOLD_1024;//single_N + } else if ((float)UsageRate512/virtualCoreNum512 >= USAGE_RATE_THRESHOLD_512) { + single_N = TOTAL_N_THRESHOLD_512;//single_N + } + Dim2 vecC(total_M / FACTOR_2, total_N); + Dim2 vecD(total_M, total_N); + Dim2 vecB; + + if (pattern == 1) { //B不转置 + vecB = Dim2(total_K_B, total_N); + } else { //B转置 + vecB = Dim2(total_N, total_K_B); + } + + Dim2 vecA(total_M / FACTOR_2, total_K_A); + Dim2 vecA_Work(total_M, total_K_A); + Dim2 vec_bias(group_num, total_N); + Dim2 vec_offset(group_num, total_N); + Dim1 vec_sa(total_M / FACTOR_2); + Dim2 vec_sw(group_num, total_N); + Dim2 vecWork(single_M, single_N); + Dim1 vecWorkRow(total_M / FACTOR_2); + + Unisor workUnisorPing(gmWorkspaceDevice, vecD, INT16); + uint32_t Temp_A_Offset = total_M * total_N * INT16; + Unisor inputA_workUnisor(gmWorkspaceDevice + Temp_A_Offset, vecA_Work, INT4); + + Unisor inputAUnisor(gmA, vecA, INT8); + Unisor inputBUnisor(gmB, vecB, format_in); + Unisor outputCUnisor(gmC, vecC, format_out); + + Unisor biasUnisor(gmBias, vec_bias, INT32, true); + Unisor saUnisor(gmSa, vec_sa, FP32, true); + Unisor swUnisor(gmSw, vec_sw, INT64, true); + + uint32_t work_count = 0; + + inputAUnisor.pos.setAxisName('M', 'K'); + inputAUnisor.vec.setAxisName('M', 'K'); + inputAUnisor.vector_val.setAxisName('M', 'K'); + + inputA_workUnisor.pos.setAxisName('M', 'K'); + inputA_workUnisor.vec.setAxisName('M', 'K'); + inputA_workUnisor.vector_val.setAxisName('M', 'K'); + + //B不转置 + if (pattern == 1) { + inputBUnisor.pos.setAxisName('K', 'N'); + inputBUnisor.vec.setAxisName('K', 'N'); + inputBUnisor.vector_val.setAxisName('K', 'N'); + } else { + //B转置 + inputBUnisor.pos.setAxisName('N', 'K'); + inputBUnisor.vec.setAxisName('N', 'K'); + inputBUnisor.vector_val.setAxisName('N', 'K'); + } + + outputCUnisor.pos.setAxisName('M', 'N'); + outputCUnisor.vec.setAxisName('M', 'N'); + outputCUnisor.vector_val.setAxisName('M', 'N'); + + workUnisorPing.pos.setAxisName('M', 'N'); + workUnisorPing.vec.setAxisName('M', 'N'); + workUnisorPing.vector_val.setAxisName('M', 'N'); + + uint8_t kernel_index = tiling_data->kernel_index; + uint32_t aicCoreId = 0; + + uint32_t pos_M_Offset = 0; + { + if (g_coreType == AscendC::AIV) { + uint32_t mIn = total_M / FACTOR_2; + uint32_t mOut = total_M; + Dim2 vecIn(mIn, total_K_A); + Dim2 vecOut(mOut, total_K_A); + Dim2 posIn(0, 0); + Dim2 posOut(0, 0); + + inputAUnisor.setCartesian(posIn, vecIn); + inputA_workUnisor.setCartesian(posOut, vecOut); //两倍关系 + + op.npu_user_defined_matmul_kernel_slave_Single_A8ToA4_off(inputAUnisor, inputA_workUnisor); + + SetFlag(LOCAL_FLAGID0); + WaitFlag(LOCAL_FLAGID0); + + AscendC::SyncAll(); + } + + if (g_coreType == AscendC::AIC){ + AscendC::SyncAll(); + } + + //Cube计算+后处理 + for (int32_t i=0; i multiCoreCartesianC0(C0pos, C0vec); + Dim1 singleCoreCTilingS0(single_M, 'M'); + CartesianIterator multiCoreUnisorIterator0(multiCoreCartesianC0, singleCoreCTilingS0); + Cartesian multiCoreCartesianUnisorC0; + + Dim1 C1pos(pos_N, 'N'), C1vec(N_Length, 'N'); + Cartesian multiCoreCartesianC1(C1pos, C1vec); + Dim1 singleCoreCTilingS1(single_N, 'N'); + CartesianIterator multiCoreUnisorIterator1(multiCoreCartesianC1, singleCoreCTilingS1); + Cartesian multiCoreCartesianUnisorC1; + + while (multiCoreUnisorIterator0.posIterator(multiCoreCartesianUnisorC0))//singleM + { + while (multiCoreUnisorIterator1.posIterator(multiCoreCartesianUnisorC1))//singleN + { + if ((aicCoreId%numAic) == blockId) + { + Dim2 posA = Dim2(multiCoreCartesianUnisorC0.pos.getAxisByName('M'), 0, 'M', 'K'); + Dim2 vecA = Dim2(multiCoreCartesianUnisorC0.vec.getAxisByName('M'), single_K, 'M', 'K'); + inputA_workUnisor.setCartesian(posA, vecA); + + //B不转置 + if (pattern == 1) { + Dim2 posB = Dim2(pos_N_group, multiCoreCartesianUnisorC1.pos.getAxisByName('N'), 'K', 'N'); + Dim2 vecB = Dim2(single_K, multiCoreCartesianUnisorC1.vec.getAxisByName('N'), 'K', 'N'); + inputBUnisor.vector_val.s[0] = total_K_A; + inputBUnisor.setCartesian(posB, vecB); + } else { + Dim2 posB = Dim2(multiCoreCartesianUnisorC1.pos.getAxisByName('N'), pos_K_group, 'N', 'K'); + Dim2 vecB = Dim2(multiCoreCartesianUnisorC1.vec.getAxisByName('N'), single_K, 'N', 'K'); + inputBUnisor.setCartesian(posB, vecB); + } + + Dim2 posC(multiCoreCartesianUnisorC0.pos.getAxisByName('M'), multiCoreCartesianUnisorC1.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vecC(multiCoreCartesianUnisorC0.vec.getAxisByName('M'), multiCoreCartesianUnisorC1.vec.getAxisByName('N'), 'M', 'N'); + outputCUnisor.setCartesian(posC, vecC); + workUnisorPing.setCartesian(posC, vecC); + uint32_t event_id = work_count % FACTOR_10; + + if (g_coreType == AscendC::AIC) { + op.npu_user_defined_matmul_kernel_switch(inputA_workUnisor, inputBUnisor, outputCUnisor, workUnisorPing, biasUnisor, saUnisor, swUnisor, baseCTiling, i, kernel_index); + AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(event_id); + } + + if (g_coreType == AscendC::AIV) { + AscendC::CrossCoreWaitFlag(event_id); + Dim2 pos_Bias(i, multiCoreCartesianUnisorC1.pos.getAxisByName('N'), 'M', 'N'); + Dim2 vec_Bias(1, multiCoreCartesianUnisorC1.vec.getAxisByName('N'), 'M', 'N'); + biasUnisor.setCartesian(pos_Bias, vec_Bias); + + op.npu_user_defined_matmul_kernel_slave_Post_Processing_off(outputCUnisor, workUnisorPing, biasUnisor, saUnisor, swUnisor); + } + work_count++; + } + aicCoreId++; + } + } + } + } +} + +class GMMA4W8AutotilingCompute { +private: + TPipe *pipe; + GM_ADDR gmA; + GM_ADDR gmB; + GM_ADDR gmC; + GM_ADDR gmGrpListOptional; + GM_ADDR gmBias; + GM_ADDR gmOffset; + GM_ADDR gmSa; + GM_ADDR gmSw; + GM_ADDR gmWorkspaceDevice; + A8W4HPTiling *tiling_data; + +public: + __aicore__ inline GMMA4W8AutotilingCompute(GM_ADDR addrA, GM_ADDR addrB, GM_ADDR addrC, GM_ADDR addrGrpList, + GM_ADDR addrBias, GM_ADDR addrOffset, GM_ADDR addrSa, GM_ADDR addrSw, + GM_ADDR addrWorkspace, A8W4HPTiling *tilingData, TPipe *p) + : pipe(p), gmA(addrA), gmB(addrB), gmC(addrC), gmGrpListOptional(addrGrpList), gmBias(addrBias), + gmOffset(addrOffset), gmSa(addrSa), gmSw(addrSw), gmWorkspaceDevice(addrWorkspace), tiling_data(tilingData) + { + } + + __aicore__ inline void Init() + { + gmA = GetTensorAddr(0, gmA); + gmB = GetTensorAddr(0, gmB); + gmC = GetTensorAddr(0, gmC); + gmBias = GetTensorAddr(0, gmBias); + gmOffset = GetTensorAddr(0, gmOffset); + gmSw = GetTensorAddr(0, gmSw); + } + + __aicore__ inline void Process() + { + dynamic_unisor_programming(gmA, gmB, gmC, gmGrpListOptional, gmBias, gmOffset, gmSa, gmSw, gmWorkspaceDevice, + tiling_data, pipe); + } +}; +} + +using GMMHighPerf::GMMA4W8AutotilingCompute; +} +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_interface.cpp b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_interface.cpp new file mode 100644 index 00000000..5e160f91 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_interface.cpp @@ -0,0 +1,121 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_fixaxismove_interface.cpp + * \brief + */ +#include "kernel_operator.h" +#include "grouped_matmul_fixaxismove_regular.h" + +namespace Catlass { + +template +CATLASS_DEVICE void grouped_matmul_fixaxismove(uint32_t m, uint32_t k, uint32_t n, uint32_t groupNum, + GM_ADDR gmA, GM_ADDR gmB, GM_ADDR gmScale, GM_ADDR group_list, + GM_ADDR per_token_scale, GM_ADDR y, GM_ADDR workspace, uint32_t aicCoreNum) { + using LayoutA = Catlass::layout::RowMajor; + using LayoutB = Catlass::layout::zN; + using LayoutD = Catlass::layout::RowMajor; + LayoutA layoutA{m, k}; + LayoutB layoutB = LayoutB::template MakeLayout(k, n); + Catlass::layout::VectorLayout layoutScale{n}; + Catlass::layout::VectorLayout layoutPerTokenScale{m}; + LayoutD layoutD{m, n}; + + using ArchTag = Arch::AtlasA2; + constexpr uint32_t preloadStages = 1; + // 左矩阵分四块copy + constexpr uint32_t l1AStages = 4; + constexpr uint32_t l1ABufferNum = 1; + constexpr uint32_t l1ATileNum = 1; + constexpr uint32_t l1BStages = 2; + constexpr uint32_t l0AStages = 2; + constexpr uint32_t l0BStages = 2; + constexpr uint32_t l0CStages = 1; + constexpr bool enableUnitFlag = true; + constexpr bool enableRiffleShuffle = true; + + constexpr uint32_t MAX_GROUP_NUM = 512; + constexpr uint32_t MAX_CORE_NUM = 24; + constexpr uint32_t MAX_TILE_NUM = 5000; + using DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncFixAxisMoveWithCallback< + preloadStages, + l1ABufferNum, l1AStages, l1ATileNum, l1BStages, l0AStages, l0BStages, l0CStages, + enableUnitFlag, enableRiffleShuffle + >; + using L1TileShape = GemmShape<128, 256, 512>; + using L0TileShape = GemmShape<128, 256, 128>; + + using AType = Gemm::GemmType; + using BType = Gemm::GemmType; + using CType = Gemm::GemmType; + + AscendC::GlobalTensor groupList; + groupList.SetGlobalBuffer((__gm__ GrouplistDType*)group_list); + uint32_t maxM = (uint32_t)(groupList.GetValue(0)); + for (int groupIdx = 1; groupIdx < groupNum; groupIdx++) { + uint32_t curM = uint32_t(groupList.GetValue(groupIdx) - groupList.GetValue(groupIdx - 1)); + if (curM > maxM) { + maxM = curM; + } + } + // kernel level + uint32_t blockNum = CeilDiv(n, L1TileShape::N) * CeilDiv(maxM, L1TileShape::M); + uint32_t blockNumPerCore1 = CeilDiv(blockNum, aicCoreNum); + uint32_t workspaceStages = 2; + if (blockNumPerCore1 > workspaceStages) { + workspaceStages = blockNumPerCore1; + } + + using BlockMmad = Gemm::Block::BlockMmad; + + constexpr uint32_t ubStages = 2; + using EpilogueDispatchPolicy = Epilogue::EpilogueAtlasA2PerTokenDequant; + using ScaleType = Gemm::GemmType; + using PerTokenScaleType = Gemm::GemmType; + using DType = Gemm::GemmType; + + using RowBroadcastMulType = Gemm::GemmType; + using BroadcastOneBlkType = Gemm::GemmType; + using OneBlkColumnBroadcastMulType = Gemm::GemmType; + + using EpilogueTileShape = MatrixShape<32, 256>; + using TileRowBroadcastMul = Epilogue::Tile::TileRowBroadcastMul; + using TileBroadcastOneBlk = Epilogue::Tile::TileBroadcastOneBlk; + using TileOneBlkColumnBroadcastMul = Epilogue::Tile::TileOneBlkColumnBroadcastMul; + using TileCopy = Epilogue::Tile::TileCopy; + using TileScheduler = Epilogue::Tile::EpilogueHorizontalTileSwizzle; + + using BlockEpilogue = Epilogue::Block::BlockEpilogue; + + using BlockScheduler = typename Gemm::Block::SplitkInOneCoreGemmIdentityBlockSwizzle<1, 0>; + + using MatmulKernel = DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove; + + typename MatmulKernel::Params params{ + m, k, n, groupNum, workspaceStages, + group_list, + gmA, layoutA, + gmB, layoutB, + gmScale, layoutScale, + per_token_scale, layoutPerTokenScale, + y, layoutD, workspace + }; + MatmulKernel matmul; + matmul(params); +} +} // namespace Catlass \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_regular.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_regular.h new file mode 100644 index 00000000..51532752 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_fixaxismove_regular.h @@ -0,0 +1,408 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_fixaxismove_regular.h + * \brief + */ +#ifndef GROUPED_MATMUL_FIXAXISMOVE_REGULAR_H +#define GROUPED_MATMUL_FIXAXISMOVE_REGULAR_H + +#include "gmm_infra/base_defs.hpp" +#include "gmm_infra/coord.hpp" +#include "gmm_infra/matrix_coord.hpp" +#include "gmm_infra/gemm_coord.hpp" +#include "gmm_infra/arch/cross_core_sync.hpp" +#include "gmm_infra/arch/resource.hpp" +#include "gmm_infra/arch/arch.hpp" +#include "gmm_infra/layout/layout.hpp" +#include "gmm_infra/detail/callback.hpp" +#include "gmm_infra/gemm/dispatch_policy.hpp" +#include "gmm_infra/gemm/gemm_type.hpp" +#include "gmm_infra/gemm/block/block_swizzle.hpp" +#include "gmm_infra/gemm/block/block_mmad.hpp" +#include "gmm_infra/epilogue/dispatch_policy.hpp" +#include "gmm_infra/epilogue/block/block_epilogue.hpp" +#include "gmm_infra/epilogue/tile/tile_broadcast_mul.hpp" +#include "gmm_infra/epilogue/tile/tile_broadcast_one_blk.hpp" +#include "gmm_infra/epilogue/tile/tile_swizzle.hpp" +#include "gmm_infra/epilogue/tile/tile_copy.hpp" + +constexpr uint32_t MAX_GROUP_NUM = 512; +constexpr uint32_t MAX_CORE_NUM = 24; +constexpr uint32_t MAX_TILE_NUM = 5000; +constexpr uint32_t MAX_WORKSPACE = 20; + +namespace Catlass { + +template +class DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove { +public: + using BlockMmad = BlockMmad_; + using ArchTag = typename BlockMmad::ArchTag; + using L1TileShape = typename BlockMmad::L1TileShape; + using ElementA = typename BlockMmad::ElementA; + using LayoutA = typename BlockMmad::LayoutA; + using ElementB = typename BlockMmad::ElementB; + using LayoutB = typename BlockMmad::LayoutB; + using ElementC = typename BlockMmad::ElementC; + using LayoutC = typename BlockMmad::LayoutC; + using ElementAccumulator = typename BlockMmad::ElementAccumulator; + + using BlockEpilogue = BlockEpilogue_; + using ElementScale = typename BlockEpilogue::ElementScale; + using LayoutScale = typename BlockEpilogue::LayoutScale; + using ElementPerTokenScale = typename BlockEpilogue::ElementPerTokenScale; + using LayoutPerTokenScale = typename BlockEpilogue::LayoutPerTokenScale; + using ElementD = typename BlockEpilogue::ElementD; + using LayoutD = typename BlockEpilogue::LayoutD; + using EpilogueParams = typename BlockEpilogue::Params; + + using ElementGroupList = ElementGroupList_; + + using BlockScheduler = BlockScheduler_; + + /// Parameters structure + struct Params { + // Data members + uint32_t m; + uint32_t k; + uint32_t n; + uint32_t problemCount; + uint32_t workspaceStages; + __gm__ ElementGroupList *ptrGroupList; + GM_ADDR ptrA; + LayoutA layoutA; + GM_ADDR ptrB; + LayoutB layoutB; + // __gm__ ElementScale *ptrScale; + GM_ADDR ptrScale; + LayoutScale layoutScale; + __gm__ ElementPerTokenScale *ptrPerTokenScale; + // GM_ADDR ptrPerTokenScale; + LayoutPerTokenScale layoutPerTokenScale; + GM_ADDR ptrD; + LayoutD layoutD; + GM_ADDR ptrWorkspace; + + // Methods + CATLASS_HOST_DEVICE + Params() + { + } + + CATLASS_HOST_DEVICE + Params(uint32_t m_, uint32_t k_, uint32_t n_, uint32_t problemCount_, uint32_t workspaceStages_, + GM_ADDR ptrGroupList_, GM_ADDR ptrA_, LayoutA layoutA_, GM_ADDR ptrB_, LayoutB layoutB_, + GM_ADDR ptrScale_, LayoutScale layoutScale_, GM_ADDR ptrPerTokenScale_, + LayoutPerTokenScale layoutPerTokenScale_, GM_ADDR ptrD_, LayoutD layoutD_, GM_ADDR ptrWorkspace_) + : m(m_), k(k_), n(n_), problemCount(problemCount_), workspaceStages(workspaceStages_), + ptrGroupList(reinterpret_cast<__gm__ ElementGroupList *>(ptrGroupList_)), ptrA((ptrA_)), + layoutA(layoutA_), ptrB((ptrB_)), layoutB(layoutB_), ptrScale((ptrScale_)), layoutScale(layoutScale_), + ptrPerTokenScale(reinterpret_cast<__gm__ ElementPerTokenScale *>(ptrPerTokenScale_)), + layoutPerTokenScale(layoutPerTokenScale_), ptrD((ptrD_)), layoutD(layoutD_), ptrWorkspace(ptrWorkspace_) + { + } + }; + + // Methods + CATLASS_DEVICE + DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove() + { + } + + CATLASS_DEVICE + ~DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove() + { + } + + template + CATLASS_DEVICE __gm__ T *GetTensorAddr(uint16_t index, GM_ADDR tensorPtr) + { + __gm__ uint64_t *dataAddr = reinterpret_cast<__gm__ uint64_t *>(tensorPtr); + uint64_t tensorPtrOffset = *dataAddr; // The offset of the data address from the first address. + // Moving 3 bits to the right means dividing by sizeof(uint64 t). + __gm__ uint64_t *retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T *>(*(retPtr + index)); + } + + template + CATLASS_DEVICE void operator()(Params const ¶ms); + + template <> + CATLASS_DEVICE void operator()(Params const ¶ms) + { + BlockScheduler blockScheduler; + BlockMmad blockMmad(resource); + + // Represent the full gm + AscendC::GlobalTensor gmC; + gmC.SetGlobalBuffer(reinterpret_cast<__gm__ ElementC *>(params.ptrWorkspace)); + + AscendC::GlobalTensor groupList; + groupList.SetGlobalBuffer(params.ptrGroupList); + + Arch::FlagID flagId = 0; + flagAivFinishComputeList = Arch::CrossCoreFlag(flagId++); + for (uint32_t stageId = 0; stageId < params.workspaceStages; ++stageId) { + flagAicFinishStoreList[stageId] = Arch::CrossCoreFlag(flagId++); + aicWaitFuncList[stageId] = {this, stageId}; + aicSetFuncList[stageId] = {this, stageId}; + } + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + int64_t gmGroupOffsetA = 0; + int64_t gmGroupOffsetB = 0; + auto layoutC = layout::RowMajor{L1TileShape::M * coreNum * params.workspaceStages, L1TileShape::N}; + + int32_t lastGmOffsetA = -1; + int32_t lastGroupIdx = -1; + uint8_t isFirstBlock = 1; + uint32_t stageId = 0; + uint32_t stageUsed = 0; + for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) { + uint32_t currentM = (groupIdx == 0) ? groupList.GetValue(groupIdx) : + (groupList.GetValue(groupIdx) - groupList.GetValue(groupIdx - 1)); + uint32_t kNum = CeilDiv(params.k, BlockMmad::L1A_K_ONE_TIME); + GemmCoord inGroupProblemShape{currentM, params.n, params.k}; + + LayoutA layoutA = params.layoutA.GetTileLayout(inGroupProblemShape.GetCoordMK()); + LayoutB layoutB = params.layoutB; + + blockScheduler.Update(inGroupProblemShape, GemmCoord(L1TileShape::M, L1TileShape::N, L1TileShape::K), kNum, + BlockMmad::L1A_K_ONE_TIME, groupIdx, coreNum, coreIdx); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + + uint32_t blockNum = CeilDiv(params.n, L1TileShape::N) * CeilDiv(currentM, L1TileShape::M); + uint32_t blockNumPerCore = CeilDiv(blockNum, coreNum); + + AscendC::GlobalTensor gmA; + gmA.SetGlobalBuffer(GetTensorAddr(0, params.ptrA) + gmGroupOffsetA); + + AscendC::GlobalTensor gmB; + gmB.SetGlobalBuffer(GetTensorAddr(0, params.ptrB) + gmGroupOffsetB); + if (CeilDiv(currentM, L1TileShape::M) == 1) { + gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE); + } + + uint32_t kOffset = 0; + + for (uint32_t kIdx = 0; kIdx < kNum; kIdx++) { + stageId = 0; + for (uint32_t i = 0; i < coreLoops; i++) { + Callback callbackBeforeFixpipe{}; + uint32_t blockIdx = blockScheduler.GetCoreBlockIdx(i); + uint32_t loopIdx = blockIdx * kNum + kIdx; + uint8_t needLeftCopyFlag = blockScheduler.GetNeedLeftCopyFlag(i); + + // Compute block location + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = blockScheduler.GetActualBlockTailBefore(blockCoord, kIdx); + + uint32_t kOffset = 0; + if (kNum > 1 && kIdx == kNum - 1) { + kOffset = params.k - BlockMmad::L1A_K_ONE_TIME; + } else { + kOffset = blockCoord.k() * L1TileShape::K; + } + MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, kOffset}; + MatrixCoord offsetB{kOffset, blockCoord.n() * L1TileShape::N}; + MatrixCoord offsetC{(coreIdx * params.workspaceStages + stageId) * L1TileShape::M, 0}; + int64_t gmOffsetA = layoutA.GetOffset(offsetA); + int64_t gmOffsetB = layoutB.GetOffset(offsetB); + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + int needCopyL1Left = 0; + if (gmOffsetA != lastGmOffsetA || groupIdx != lastGroupIdx) { + needCopyL1Left = 1; + } + lastGmOffsetA = gmOffsetA; + lastGroupIdx = groupIdx; + int needAtomicAdd = (kIdx > 0); + if (kIdx == kNum - 1) { + Callback callbackAfterFixpipe = MakeCallback(&aicSetFuncList[stageId]); + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, gmC[gmOffsetC], layoutC, + actualBlockShape, needCopyL1Left, needAtomicAdd, needLeftCopyFlag, isFirstBlock, + // callbackBeforeFixpipe, callbackAfterFixpipe + Callback{}, callbackAfterFixpipe); + } else { + blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, gmC[gmOffsetC], layoutC, + actualBlockShape, needCopyL1Left, needAtomicAdd, needLeftCopyFlag, isFirstBlock, + Callback{}, Callback{}); + callbackAfterFixpipe(); + } + } else { + if (groupIdx > 0 && kIdx == 0 && i == 0) { + callbackBeforeFixpipe = MakeCallback(&aicWaitFuncList[stageId]); + } + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, gmC[gmOffsetC], layoutC, + actualBlockShape, needCopyL1Left, needAtomicAdd, needLeftCopyFlag, isFirstBlock, + callbackBeforeFixpipe, Callback{}); + } else { + callbackBeforeFixpipe(); + blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, gmC[gmOffsetC], layoutC, + actualBlockShape, needCopyL1Left, needAtomicAdd, needLeftCopyFlag, isFirstBlock, + Callback{}, Callback{}); + } + } + stageId++; + if (isFirstBlock == 1) { + isFirstBlock = 0; + } + } + } + + gmGroupOffsetA += inGroupProblemShape.m() * inGroupProblemShape.k(); + gmGroupOffsetB += inGroupProblemShape.k() * inGroupProblemShape.n(); + } + + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad.SynchronizeBlock(); + } + Arch::CrossCoreWaitFlag(flagAivFinishComputeList); + } + + template <> + CATLASS_DEVICE void operator()(Params const ¶ms) + { + BlockScheduler blockScheduler; + BlockEpilogue blockEpilogue(resource); + + uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + Arch::FlagID flagId = 0; + flagAivFinishComputeList = Arch::CrossCoreFlag(flagId++); + for (uint32_t stageId = 0; stageId < params.workspaceStages; ++stageId) { + flagAicFinishStoreList[stageId] = Arch::CrossCoreFlag(flagId++); + } + + uint32_t coreNum = AscendC::GetBlockNum(); + int64_t gmGroupOffsetScale = 0; + int64_t gmGroupOffsetPerTokenScale = 0; + int64_t gmGroupOffsetD = 0; + + AscendC::GlobalTensor gmC; + gmC.SetGlobalBuffer(reinterpret_cast<__gm__ ElementC *>(params.ptrWorkspace)); + auto layoutC = layout::RowMajor{L1TileShape::M * coreNum * params.workspaceStages, L1TileShape::N}; + + AscendC::GlobalTensor groupList; + groupList.SetGlobalBuffer(params.ptrGroupList); + + uint32_t aicCoreNum = AscendC::GetBlockNum(); + uint32_t aivCoreIdx = AscendC::GetBlockIdx(); + uint32_t aivNumPerAic = AscendC::GetSubBlockNum(); + uint32_t aicCoreIdx = aivCoreIdx / aivNumPerAic; + uint32_t kNum = CeilDiv(params.k, BlockMmad::L1A_K_ONE_TIME); + uint32_t stageId = 0; + auto ptrScaleBegin = GetTensorAddr(0, params.ptrScale); + auto ptrDBegin = GetTensorAddr(0, params.ptrD); + for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) { + uint32_t currentM = (groupIdx == 0) ? groupList.GetValue(groupIdx) : + (groupList.GetValue(groupIdx) - groupList.GetValue(groupIdx - 1)); + GemmCoord inGroupProblemShape{currentM, params.n, params.k}; + + LayoutScale layoutScale = params.layoutScale; + LayoutPerTokenScale layoutPerTokenScale = + params.layoutPerTokenScale.GetTileLayout(inGroupProblemShape.template GetCoordByAxis<0>()); + LayoutD layoutD = params.layoutD.GetTileLayout(inGroupProblemShape.GetCoordMN()); + + EpilogueParams epilogueParams{ptrScaleBegin + gmGroupOffsetScale, + layoutScale, + params.ptrPerTokenScale + gmGroupOffsetPerTokenScale, + layoutPerTokenScale, + ptrDBegin + gmGroupOffsetD, + layoutD}; + + blockScheduler.Update(inGroupProblemShape, GemmCoord(L1TileShape::M, L1TileShape::N, L1TileShape::K), kNum, + BlockMmad::L1A_K_ONE_TIME, groupIdx, coreNum, coreIdx); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + blockEpilogue.UpdateParams(epilogueParams); + + uint32_t blockNum = CeilDiv(params.n, L1TileShape::N) * CeilDiv(currentM, L1TileShape::M); + uint32_t blockNumPerCore = CeilDiv(blockNum, aicCoreNum); + int blockCntIdx = groupIdx * aicCoreNum + aicCoreIdx; + uint32_t blockIdx = groupIdx * blockNumPerCore * aicCoreNum + aicCoreIdx * blockNumPerCore; + + GemmCoord blockShapeMNK = L1TileShape::ToCoord(); + + for (uint32_t i = 0; i < coreLoops; i++) { + uint32_t blockIdx = blockScheduler.GetCoreBlockIdx(i); + uint32_t loopIdx = blockIdx * kNum; + + // Compute block location + GemmCoord blockCoordMNK = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShapeMNK = blockScheduler.GetActualBlockShape(blockCoordMNK, 0); + + MatrixCoord offsetC{(coreIdx * params.workspaceStages + stageId) * L1TileShape::M, 0}; + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + auto gmBlockC = gmC[gmOffsetC]; + auto layoutBlockC = layoutC.GetTileLayout(actualBlockShapeMNK.GetCoordMN()); + + Arch::CrossCoreWaitFlag(flagAicFinishStoreList[stageId]); + blockEpilogue(blockShapeMNK, blockCoordMNK, actualBlockShapeMNK, gmBlockC, layoutBlockC); + stageId = (stageId + 1 < coreLoops) ? (stageId + 1) : 0; + blockIdx++; + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(flagAivFinishComputeList); + gmGroupOffsetScale += inGroupProblemShape.n(); + gmGroupOffsetPerTokenScale += inGroupProblemShape.m(); + gmGroupOffsetD += inGroupProblemShape.m() * inGroupProblemShape.n(); + } + } + +private: + friend struct AicWaitFunc; + friend struct AicSetFunc; + + struct AicWaitFunc { + using MatmulKernel = + DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove; + + CATLASS_DEVICE + AicWaitFunc() = default; + + CATLASS_DEVICE + void operator()() const + { + Arch::CrossCoreWaitFlag(ptr->flagAivFinishComputeList); + } + + MatmulKernel *ptr{nullptr}; + uint32_t stageId; + }; + + struct AicSetFunc { + using MatmulKernel = + DlinferGroupedMatmulDirectSliceMPerTokenDequantFixAxisMove; + + CATLASS_DEVICE + AicSetFunc() = default; + + CATLASS_DEVICE + void operator()() const + { + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(ptr->flagAicFinishStoreList[stageId]); + } + + MatmulKernel *ptr{nullptr}; + uint32_t stageId; + }; + + Arch::CrossCoreFlag flagAicFinishStoreList[MAX_WORKSPACE]; + Arch::CrossCoreFlag flagAivFinishComputeList; + + AicWaitFunc aicWaitFuncList[MAX_WORKSPACE]; + AicSetFunc aicSetFuncList[MAX_WORKSPACE]; + Arch::Resource resource; +}; + +} // namespace Catlass +#endif diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_pre_tiling.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_pre_tiling.h new file mode 100644 index 00000000..0cc32293 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_pre_tiling.h @@ -0,0 +1,160 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_pre_tiling.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_PRE_TILING_H +#define ASCENDC_GROUPED_MATMUL_PRE_TILING_H + +#include "kernel_operator.h" + +namespace GROUPED_MATMUL{ +using namespace AscendC; + +constexpr int32_t BASE_M_ALIGN_UP = 16; +constexpr int32_t BASE_SINGLE_N_ALIGN_UP = 256; +constexpr int32_t SINGLE_N_AVOID = 768; // by experiment +constexpr uint32_t L1_SIZE = 512 * 1024; +constexpr float EFFECTIVE_TASK_RATIO = 0.85; +constexpr uint32_t GROUP_LIST_SPARSE_M = 2U; + +class GMMPreTilingProcess { +public: + __aicore__ inline GMMPreTilingProcess(){}; + __aicore__ inline void Init(GM_ADDR groupList, GMMBaseParams& tilingData, TCubeTiling &mmTilingData , TPipe *pipe); + __aicore__ inline void Process(GMMBaseParams& tilingData,TCubeTiling &mmTilingData); +private: + __aicore__ inline void GetTokensPerGroup(); + __aicore__ inline void FullLoadA(const int32_t baseM, TCubeTiling &mmTilingData); + GlobalTensor groupListGm; + uint32_t groupListType = 0; + int32_t groupType = 0; + int32_t coreNum = 0; + uint32_t groupNum = 0; + int32_t totalTokenNum = 0; + int32_t tokenPerGroup = 0; + int32_t baseM = 0; + int32_t baseN = 0; + int32_t baseK = 0; + int32_t stepM = 0; + int32_t stepKa = 0; + int32_t m = 0; + int32_t n = 0; + int32_t k = 0; + int32_t depthB1 = 0; + int64_t isPreTiling = 0; +}; + +__aicore__ inline void GMMPreTilingProcess::GetTokensPerGroup() { + if(groupListType == 0) { + totalTokenNum = static_cast(groupListGm.GetValue(groupNum - 1)); + } else { + for(int i = 0; i < groupNum; i++) { + totalTokenNum += static_cast(groupListGm.GetValue(i)); + } + } + tokenPerGroup = totalTokenNum / groupNum; + return; +} + +__aicore__ inline void GMMPreTilingProcess::Init(GM_ADDR groupList, GMMBaseParams& tilingData, TCubeTiling& mmTilingData , TPipe *pipe) { +#if ORIG_DTYPE_X != DT_INT8 || ORIG_DTYPE_WEIGHT != DT_INT8 + return; +#else + groupListGm.SetGlobalBuffer((__gm__ int64_t *)groupList); + groupListType = tilingData.groupListType; + groupType = tilingData.groupType; + groupNum = tilingData.groupNum; + coreNum = tilingData.coreNum; + m = mmTilingData.M; + n = mmTilingData.N; + k = mmTilingData.Ka; + baseM = mmTilingData.baseM; + baseN = mmTilingData.baseN; + baseK = mmTilingData.baseK; + stepM = mmTilingData.stepM; + stepKa = mmTilingData.stepKa; + depthB1 = mmTilingData.depthB1; + isPreTiling = tilingData.isPreTiling; +#endif +} + +__aicore__ inline void GMMPreTilingProcess::Process(GMMBaseParams& tilingData, TCubeTiling& mmTilingData) { +#if ORIG_DTYPE_X != DT_INT8 || ORIG_DTYPE_WEIGHT != DT_INT8 + return; +#else + if (!isPreTiling || groupType != 0 || groupListType == GROUP_LIST_SPARSE_M) { + return; + } + + if (isPreTiling == 2) { // 2: white list pretiling key + tilingData.singleN = mmTilingData.singleCoreN; + FullLoadA(baseM, mmTilingData); + return; + } + // get token num in groupList + GetTokensPerGroup(); + + // modify baseM + int32_t mDim = Ceil(tokenPerGroup, baseM); + int32_t nDim = Ceil(n, baseN); + int32_t taskNum = mDim * nDim * groupNum; + int32_t taskNumPerCore = Ceil(taskNum, coreNum); + int32_t newBaseM = AlignUp(Ceil(tokenPerGroup, mDim), BASE_M_ALIGN_UP); + if(newBaseM < BASE_M_ALIGN_UP) { + return; + } else { + mmTilingData.baseM = newBaseM < baseM ? newBaseM : baseM; + } + + // modify singleN + if(taskNumPerCore >= 2 && n > baseN) { // 2 : taskNum < coreNum, singleN remains unchanged. + int32_t curNDim = 0; + int32_t curTaskNum = 0; + bool isModifySingleN = false; + int32_t bestSingleN = tilingData.singleN; + float ratio = 0; + // select the max singleN that meets the requirement of ratio. + for (int i = 1; i <= coreNum; ++i) { + bestSingleN = AlignUp(Ceil(n, i), BASE_SINGLE_N_ALIGN_UP); + curNDim = Ceil(n, bestSingleN); + curTaskNum = mDim * curNDim * groupNum; + ratio = static_cast(curTaskNum) / AlignUp(curTaskNum, coreNum); +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ + isModifySingleN = ratio >= EFFECTIVE_TASK_RATIO; +#else + isModifySingleN = ratio >= EFFECTIVE_TASK_RATIO && bestSingleN != SINGLE_N_AVOID; +#endif + if (!isModifySingleN) { + continue; + } + mmTilingData.singleCoreN = bestSingleN; + tilingData.singleN = bestSingleN; + // try full load A in L1 + FullLoadA(newBaseM, mmTilingData); + break; + } + return; + } +#endif +} + +__aicore__ inline void GMMPreTilingProcess::FullLoadA(const int32_t curBaseM, TCubeTiling &mmTilingData) { + int32_t A1FullLoadSize = curBaseM * AlignUp(k, baseK) + baseN * baseK * depthB1; + if (A1FullLoadSize < L1_SIZE) { + mmTilingData.stepKa = Ceil(k, baseK); + mmTilingData.depthA1 = mmTilingData.stepKa * stepM; + } +} + +} // GROUPED_MATMUL +#endif // ASCENDC_GROUPED_MATMUL_PRE_TILING_H \ No newline at end of file diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_quant_mixcore.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_quant_mixcore.h new file mode 100644 index 00000000..bb5b5dd4 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_quant_mixcore.h @@ -0,0 +1,549 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_quant_mixcore.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H +#define ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H + +#include "grouped_matmul_utils.h" +#include "grouped_matmul.h" + +#if defined(GMM_QUANT_BF16) || defined(GMM_QUANT_FLOAT16) +namespace GROUPED_MATMUL { +/*@brief store variables for core split configuration +*/ +constexpr int32_t PIPELINE_NUM = 4; +constexpr uint32_t BROADCAST_DIM = 2; +constexpr uint32_t FP32_PER_REPEAT = 64; + +/** @brief intenal computation class +*/ +template +class GMMQuantMixCoreCompute : public GMMCompute { + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = DTYPE_WEIGHT; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + + /** @brief constructor */ + __aicore__ inline GMMQuantMixCoreCompute(typename mmType::MT& mm_) : GMMCompute(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR scale, GM_ADDR offset, + GM_ADDR antiquantScale, GM_ADDR antiquantOffset, GM_ADDR group_list, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, TPipe* tPipe); + + __aicore__ inline void InitStaticTiling(const GMMBaseParams* __restrict gmmBaseParams, GM_ADDR workspace, + int32_t baseM, int32_t baseN); + + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx); + + __aicore__ inline void VectorCompute(MNConfig& mnConfig); + + __aicore__ inline void PostCompute(); + + private: + __aicore__ inline void Dequant(MNConfig& mnConfig); + + __aicore__ inline void SetPerTokenQuantStaticBuffer(const GMMBaseParams* __restrict gmmBaseParams, + GM_ADDR workspace); + + __aicore__ inline void DataCopyScale(uint32_t curBaseN, uint32_t alignBaseN, uint64_t scaleOffset); + + __aicore__ inline void DataCopyBias(uint32_t curBaseN, uint32_t alignBaseN, uint64_t biasOffset); + + __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN, + uint32_t offsetM); + + __aicore__ inline void SetPerTokenQuantRefreshedBuffer(const MNConfig mnConfig); + + __aicore__ inline void ActivationCompute(uint32_t computeSize, LocalTensor preResUb, + LocalTensor actTmpLocal); + + __aicore__ inline void ComputeDequantAndActivate(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, + uint32_t offsetM); + + __aicore__ inline void PerTokenQuant(uint32_t curVecBaseM, uint32_t alignBaseN); + + __aicore__ inline void DataCopyOut(MNConfig& mnConfig, uint32_t curVecBaseM, uint32_t curVecBaseN, + uint32_t alignBaseN, uint64_t outOffset); + + __aicore__ inline void VectorTilingCalc(MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, + uint32_t& vecBaseN, uint32_t& vecBaseM); + + GM_ADDR scaleTensorPtr; + GM_ADDR perTokenScaleTensorPtr; + GM_ADDR biasTensorPtr; + GlobalTensor scaleGm; + GlobalTensor perTokenScaleGm; + GlobalTensor mmOutGm; + GlobalTensor biasGm; + // define the que + TQue vecInQueue; + TQue vecOutQueue; + TQue scaleInQueue; + TQue perTokenScaleInQueue; + TQue biasInQueue; + TQue biasInQueueFp32; + TBuf tmpBuff; + LocalTensor mmOutInUb; + LocalTensor scaleInUb; + LocalTensor biasLocal; + LocalTensor biasLocalFp32; + LocalTensor perTokenScaleInUb; + LocalTensor dequantMiddleResult; + LocalTensor sharedTmpLocal; + LocalTensor mulsResultLocal; + LocalTensor pertokenBrcbLocal; + LocalTensor actResultLocal; + bool sequentialWrite = true; + bool isPerTokenQuant; + uint32_t cubeNum; // Matmul completions on the kernel + uint32_t nOffset; // antiquant n offset + uint32_t baseM_ = 0; + uint32_t baseN_ = 0; + uint32_t singleWeight = 0; + static constexpr bool isBiasEpilogue_ = + AscendC::IsSameType::value && + (AscendC::IsSameType::value || AscendC::IsSameType::value); +}; + +template +__aicore__ inline void GMMQuantMixCoreCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, + GM_ADDR scale, GM_ADDR offset, GM_ADDR antiquantScale, + GM_ADDR antiquantOffset, GM_ADDR groupList, + GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, + const GMMBaseParams* __restrict gmmBaseParams, + const TCubeTiling* __restrict mmTilingData, + TPipe* tPipe) { + this->GMMCompute::Init(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, groupList, + perTokenScale, y, workspace, gmmBaseParams, mmTilingData, tPipe); + isPerTokenQuant = gmmBaseParams->quantParam == 1; + singleWeight = gmmBaseParams->singleWeight; + scaleTensorPtr = scale; + perTokenScaleTensorPtr = perTokenScale; + biasTensorPtr = bias; + cubeNum = 0; + if (mmTilingData != nullptr) { + baseM_ = static_cast(mmTilingData->baseM); + baseN_ = static_cast(mmTilingData->baseN); + size_t workspace_offset = 0; + if (this->GMMCompute::isA8W4FakeQuant == true) { + workspace_offset = gmmBaseParams->groupNum * gmmBaseParams->k * gmmBaseParams->n * sizeof(int8_t) + + gmmBaseParams->groupNum * gmmBaseParams->n * sizeof(float); + scaleTensorPtr = scaleTensorPtr + gmmBaseParams->groupNum * gmmBaseParams->k * gmmBaseParams->n * sizeof(int8_t); + isPerTokenQuant = true; + } + SetPerTokenQuantStaticBuffer(gmmBaseParams, workspace + workspace_offset); + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::InitStaticTiling(const GMMBaseParams* __restrict gmmBaseParams, + GM_ADDR workspace, int32_t baseM, int32_t baseN) { + baseM_ = static_cast(baseM); + baseN_ = static_cast(baseN); + SetPerTokenQuantStaticBuffer(gmmBaseParams, workspace); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::PostCompute() { + if ASCEND_IS_AIC { + for (int32_t idx = 0; idx < Min(cubeNum, PIPELINE_NUM); ++idx) { + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); + } + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig, + uint32_t coreIdx) { + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM + : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + // init global buffer + if (this->singleX == 0) { + this->xGm.SetGlobalBuffer(GetTensorAddr(groupIdx, this->xTensorPtr)); + } else { + this->xGm.SetGlobalBuffer(GetTensorAddr(0, this->xTensorPtr) + mnConfig.xBaseOffset); + } + GlobalTensor weightGm = this->SetGlobalBufferW(groupIdx, tailN, mnConfig); + if ASCEND_IS_AIC { + this->mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + this->mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + this->mm.SetTensorA(this->xGm[xOffset], transposeX); + this->mm.SetTensorB(weightGm, transposeW); + this->SetGlobalBufferBias(groupIdx, tailN, mnConfig); + while (this->mm.Iterate()) { + if (sequentialWrite) { + mnConfig.workSpaceOffset = mnConfig.baseN * mnConfig.baseM * \ + (coreIdx + (cubeNum % PIPELINE_NUM) * this->coreNum); + } else { + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + } + if (this->cubeNum >= PIPELINE_NUM) { + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); + } + this->mm.GetTensorC(mmOutGm[mnConfig.workSpaceOffset], 0, sequentialWrite); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + cubeNum++; + } + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::VectorCompute(MNConfig& mnConfig) { + nOffset = 0; + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM + : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + int nDim = Ceil(curSingleN, mnConfig.baseN); + + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + if ASCEND_IS_AIV { + SetPerTokenQuantRefreshedBuffer(mnConfig); + for(int i = 0; i < nDim; i++) { + if (sequentialWrite) { + mnConfig.workSpaceOffset = mnConfig.baseN * mnConfig.baseM * \ + (GetBlockIdx() / GetTaskRation() + (cubeNum % PIPELINE_NUM) * this->coreNum); + } else { + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + } + cubeNum++; + Dequant(mnConfig); + nOffset += mnConfig.baseN; + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG); + } + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::ComputeDequantAndActivate(MNConfig& mnConfig, + uint32_t curVecBaseM, uint32_t alignBaseN, uint32_t curVecBaseN, uint32_t offsetM) { + DataCopyPerTokenScaleAndBrcb(mnConfig, curVecBaseM, alignBaseN, offsetM); + mmOutInUb = vecInQueue.DeQue(); + LocalTensor yLocalInUb = vecOutQueue.AllocTensor(); + + #if defined(GMM_QUANT_BF16) + if (!isPerTokenQuant && this->activeType == 0) { // BF16 static quantization without activation. + AscendDequant(yLocalInUb, mmOutInUb, scaleInUb, sharedTmpLocal, {curVecBaseM, alignBaseN, curVecBaseN}); + vecInQueue.FreeTensor(mmOutInUb); + vecOutQueue.EnQue(yLocalInUb); + return; + } + #endif + AscendDequant(dequantMiddleResult, mmOutInUb, scaleInUb, sharedTmpLocal, {curVecBaseM, alignBaseN, curVecBaseN}); + PipeBarrier(); + LocalTensor preResUb = dequantMiddleResult; + LocalTensor yFP32LocalInUb = dequantMiddleResult; + LocalTensor actTmpLocal = sharedTmpLocal; + // pertoken antiquant + if (isPerTokenQuant) { + PerTokenQuant(curVecBaseM, alignBaseN); + preResUb = mulsResultLocal; + yFP32LocalInUb = mulsResultLocal; + actTmpLocal = tmpBuff.GetWithOffset(2 * this->ubCalSize * sizeof(float), 0); + } + // activation function + if (this->activeType != 0) { + uint32_t computeSize = curVecBaseM * alignBaseN; + ActivationCompute(computeSize, preResUb, actTmpLocal); + yFP32LocalInUb = actResultLocal; + } + // get final output after Cast + #if defined(GMM_QUANT_BF16) + Cast(yLocalInUb, yFP32LocalInUb, RoundMode::CAST_RINT, curVecBaseM * alignBaseN); + #elif defined(GMM_QUANT_FLOAT16) + Cast(yLocalInUb, yFP32LocalInUb, RoundMode::CAST_NONE, curVecBaseM * alignBaseN); + #endif + PipeBarrier(); + vecInQueue.FreeTensor(mmOutInUb); + vecOutQueue.EnQue(yLocalInUb); + return; +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::PerTokenQuant(uint32_t curVecBaseM, uint32_t alignBaseN) +{ + uint32_t tailNum = alignBaseN % FP32_PER_REPEAT; + uint8_t repeatStride = alignBaseN * sizeof(float) / UB_BLOCK_UNIT_SIZE; + uint64_t perchannelResOffset = 0; + uint64_t alignedN = alignBaseN - tailNum; + while (perchannelResOffset < alignedN) { + Mul(mulsResultLocal[perchannelResOffset], dequantMiddleResult[perchannelResOffset], pertokenBrcbLocal, + FP32_PER_REPEAT, curVecBaseM, {1, 1, 0, repeatStride, repeatStride, 1}); + perchannelResOffset += FP32_PER_REPEAT; + } + if (tailNum != 0) { + Mul(mulsResultLocal[perchannelResOffset], dequantMiddleResult[perchannelResOffset], pertokenBrcbLocal, tailNum, + curVecBaseM, {1, 1, 0, repeatStride, repeatStride, 1}); + } + PipeBarrier(); + if constexpr (isBiasEpilogue_) { + for (int i = 0; i < curVecBaseM; i++) { + Add(mulsResultLocal[i * alignBaseN], mulsResultLocal[i * alignBaseN], biasLocalFp32, alignBaseN); + } + PipeBarrier(); + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::VectorTilingCalc( + MNConfig& mnConfig, uint32_t& curCubeSingleN, uint32_t& curCubeSingleM, uint32_t& vecBaseN, + uint32_t& vecBaseM) { + curCubeSingleN = mnConfig.nIdx == mnConfig.blockDimN - 1 ? + mnConfig.n - mnConfig.nIdx * mnConfig.singleN : mnConfig.singleN; + curCubeSingleM = mnConfig.mIdx == mnConfig.blockDimM - 1 ? + mnConfig.m - mnConfig.mIdx * mnConfig.singleM : mnConfig.singleM; + vecBaseN = mnConfig.baseN; + vecBaseM = this->ubCalSize / AlignUp(vecBaseN, static_cast(UB_BLOCK_DOUBLE_UNIT_SIZE / sizeof(int32_t))); + vecBaseM = Min(vecBaseM, curCubeSingleM); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::Dequant(MNConfig& mnConfig) { + uint32_t curCubeSingleN; + uint32_t curCubeSingleM; + uint32_t vecBaseN; + uint32_t vecBaseM; + VectorTilingCalc(mnConfig, curCubeSingleN, curCubeSingleM, vecBaseN, vecBaseM); + uint32_t curVecBaseN = vecBaseN; + uint32_t curVecBaseM; + uint32_t vecCount = 0; + if (nOffset + vecBaseN >= curCubeSingleN) { + curVecBaseN = curCubeSingleN - nOffset; + } + uint32_t rowLength = sequentialWrite ? curVecBaseN : mnConfig.n; + uint32_t taskRation = GetTaskRation(); + for (uint32_t offsetN = nOffset; offsetN < curCubeSingleN; offsetN += vecBaseN) { + uint32_t alignBaseN = AlignUp(curVecBaseN, static_cast(UB_BLOCK_DOUBLE_UNIT_SIZE / sizeof(int32_t))); + uint64_t scaleOffset = mnConfig.nIdx * mnConfig.singleN + offsetN; + DataCopyScale(curVecBaseN, alignBaseN, scaleOffset); + if constexpr (isBiasEpilogue_) { + DataCopyBias(curVecBaseN, alignBaseN, scaleOffset); + } + curVecBaseM = vecBaseM; + if (unlikely(offsetN == nOffset)) { + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG); + } + for (uint32_t offsetM = 0; offsetM < curCubeSingleM; offsetM += vecBaseM) { + vecCount++; + if (vecCount % taskRation != this->subBlockIdx) { + continue; + } + if (unlikely(offsetM + vecBaseM >= curCubeSingleM)) { + curVecBaseM = curCubeSingleM - offsetM; + } + // use AscendDequant interface to do perchannel dequant + uint64_t mmOutOffset = mnConfig.workSpaceOffset + offsetM * static_cast(rowLength); + LocalTensor mmOutLocal = vecInQueue.AllocTensor(); + DataCopyPad2D(mmOutLocal, mmOutGm[mmOutOffset], curVecBaseM, curVecBaseN, rowLength); + vecInQueue.EnQue(mmOutLocal); + ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM); + uint64_t outOffset = (mnConfig.mIdx * mnConfig.singleM + offsetM) * mnConfig.n + \ + mnConfig.nIdx * mnConfig.singleN + offsetN; + DataCopyOut(mnConfig, curVecBaseM, curVecBaseN, alignBaseN, outOffset); + } + scaleInQueue.FreeTensor(scaleInUb); + // once a base block + break; + } +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::DataCopyOut(MNConfig& mnConfig, uint32_t curVecBaseM, + uint32_t curVecBaseN, uint32_t alignBaseN, + uint64_t outOffset) { + // Copy the result of vector to yGm. + LocalTensor yLocal = vecOutQueue.DeQue(); + DataCopyPad2D(this->yGm[outOffset], yLocal, curVecBaseM, curVecBaseN, alignBaseN, mnConfig.n); + vecOutQueue.FreeTensor(yLocal); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::ActivationCompute(uint32_t computeSize, + LocalTensor preResUb, + LocalTensor actTmpLocal) { + ActiveType active = ActiveType(this->activeType); + if (active == ActiveType::FASTGELU) { + FasterGelu(actResultLocal, preResUb, actTmpLocal, computeSize); + } else if (active == ActiveType::RELU) { + Relu(actResultLocal, preResUb, computeSize); + } else if (active == ActiveType::SILU) { + Silu(actResultLocal, preResUb, computeSize); + } else if (active == ActiveType::GELU_TANH) { + Gelu(actResultLocal, preResUb, actTmpLocal, computeSize); + } + PipeBarrier(); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::SetPerTokenQuantStaticBuffer( + const GMMBaseParams* __restrict gmmBaseParams, GM_ADDR workspace) { + // Initialize ub and gm memories that do not need to be reinitialized due to changes in groupidx. + if ASCEND_IS_AIV { + // 2: enabling double buffer, occupying two buffer. + this->pipe->InitBuffer(scaleInQueue, 2, baseN_ * sizeof(DTYPE_SCALE)); + if (isPerTokenQuant) { + // 2: enabling double buffer, occupying two buffer. + this->pipe->InitBuffer(perTokenScaleInQueue, 2, baseM_ * sizeof(float)); + if constexpr (isBiasEpilogue_) { + this->pipe->InitBuffer(biasInQueue, 2, baseN_ * sizeof(DTYPE_BIAS)); + } + } + // 2: enabling double buffer, occupying two buffer. + this->pipe->InitBuffer(vecInQueue, 2, this->ubCalSize * sizeof(CT)); + // 2: enabling double buffer, occupying two buffer. + this->pipe->InitBuffer(vecOutQueue, 2, this->ubCalSize * sizeof(DTYPE_Y)); + this->pipe->InitBuffer(tmpBuff, gmmBaseParams->ubRestBytes); + dequantMiddleResult = tmpBuff.GetWithOffset(this->ubCalSize, 0); + #if defined(GMM_QUANT_FLOAT16) + uint32_t factor = 1; + #else + uint32_t factor = 0; + #endif + // 2: Indicates the first two blocks of ub are already occupied. + factor = !isPerTokenQuant && this->activeType == 0 ? factor : 2; + uint32_t ubCalSizeFloat = this->ubCalSize * sizeof(float); + uint32_t offset = factor * ubCalSizeFloat; + // 2: Indicates a temporary space twice the size is needed. + sharedTmpLocal = tmpBuff.GetWithOffset(2 * ubCalSizeFloat, offset); + if (isPerTokenQuant) { + // 2: Indicates the first two blocks of ub are already occupied. + mulsResultLocal = tmpBuff.GetWithOffset(this->ubCalSize, 2 * ubCalSizeFloat); + pertokenBrcbLocal = tmpBuff.GetWithOffset(this->ubCalSize, ubCalSizeFloat); + if constexpr (isBiasEpilogue_) { + biasLocalFp32 = tmpBuff.GetWithOffset(baseN_, 4 * ubCalSizeFloat); + } + } + if (this->activeType != 0) { + // 2: Indicates the first three blocks of ub are already occupied. + uint32_t offsetAct = !isPerTokenQuant ? ubCalSizeFloat : 3 * ubCalSizeFloat; + actResultLocal = tmpBuff.GetWithOffset(this->ubCalSize, offsetAct); + } + } + mmOutGm.SetGlobalBuffer((__gm__ MM_DTYPE_Y *)workspace); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::SetPerTokenQuantRefreshedBuffer(const MNConfig mnConfig) { + // Initialize gm memories that need to be reinitialized due to changes in groupidx. + // Currently, pertoken quant only supports single-tensor mode, + // hence set according to x and weight single-tensor mode. + // Add an if branch if multi-tensor mode for weght is required. + if (this->GMMCompute::isA8W4FakeQuant == true) { + scaleGm.SetGlobalBuffer((__gm__ DTYPE_SCALE *)scaleTensorPtr + mnConfig.nAxisBaseOffset); + } else if (singleWeight == 0) { + scaleGm.SetGlobalBuffer(GetTensorAddr(mnConfig.scaleIndex, scaleTensorPtr)); + } else { + scaleGm.SetGlobalBuffer(GetTensorAddr(0, scaleTensorPtr) + mnConfig.nAxisBaseOffset); + } + if (isPerTokenQuant) { + perTokenScaleGm.SetGlobalBuffer((__gm__ float *)perTokenScaleTensorPtr + mnConfig.mAxisBaseOffset); + if constexpr (isBiasEpilogue_) { + biasGm.SetGlobalBuffer(GetTensorAddr(0, biasTensorPtr) + mnConfig.nAxisBaseOffset); + } + } + // Add an if branch if multi-tensor mode for y is required. + this->yGm.SetGlobalBuffer(GetTensorAddr(0, this->yTensorPtr) + mnConfig.yBaseOffset); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::DataCopyScale(uint32_t curBaseN, + uint32_t alignBaseN, + uint64_t scaleOffset) +{ + // GM copy scale + DataCopyPadExtParams padParams; + DataCopyExtParams scaleParams; + scaleParams.blockLen = curBaseN * sizeof(DTYPE_SCALE); + scaleParams.blockCount = 1; + scaleParams.srcStride = 0; + scaleParams.dstStride = 0; + LocalTensor scaleLocal = scaleInQueue.AllocTensor(); + DataCopyPad(scaleLocal, scaleGm[scaleOffset], scaleParams, padParams); + scaleInQueue.EnQue(scaleLocal); + + scaleInUb = scaleInQueue.DeQue(); + scaleInUb.SetSize(alignBaseN); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::DataCopyBias(uint32_t curBaseN, uint32_t alignBaseN, + uint64_t biasOffset) +{ + DataCopyPadExtParams padParams; + DataCopyExtParams biasParams; + biasParams.blockLen = curBaseN * sizeof(DTYPE_BIAS); + biasParams.blockCount = 1; + biasParams.srcStride = 0; + biasParams.dstStride = 0; + LocalTensor biasLocal = biasInQueue.AllocTensor(); + DataCopyPad(biasLocal, biasGm[biasOffset], biasParams, padParams); + biasInQueue.EnQue(biasLocal); + biasLocal = biasInQueue.DeQue(); + biasLocal.SetSize(alignBaseN); + + Cast(biasLocalFp32, biasLocal, AscendC::RoundMode::CAST_NONE, alignBaseN); + PipeBarrier(); + biasInQueue.FreeTensor(biasLocal); +} + +template +__aicore__ inline void GMMQuantMixCoreCompute::DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, + uint32_t curBaseM, + uint32_t alignBaseN, + uint32_t offsetM) +{ + if (!isPerTokenQuant) { + return; + } + uint64_t perTokenScaleOffset = mnConfig.mIdx * mnConfig.singleM + offsetM; + // GM copy per token scale + DataCopyPadExtParams padParams; + DataCopyExtParams perTokenScaleParams; + perTokenScaleParams.blockLen = curBaseM * sizeof(float); + perTokenScaleParams.blockCount = 1; + perTokenScaleParams.srcStride = 0; + perTokenScaleParams.dstStride = 0; + LocalTensor perTokenScaleLocal = perTokenScaleInQueue.AllocTensor(); + DataCopyPad(perTokenScaleLocal, perTokenScaleGm[perTokenScaleOffset], perTokenScaleParams, padParams); + perTokenScaleInQueue.EnQue(perTokenScaleLocal); + + perTokenScaleInUb = perTokenScaleInQueue.DeQue(); + uint8_t repeatTimes = Ceil(curBaseM, 8); // curBaseM is 8 aligned; + Brcb(pertokenBrcbLocal, perTokenScaleInUb, repeatTimes, {1,8}); + perTokenScaleInQueue.FreeTensor(perTokenScaleInUb); +} + +} // namespace GROUPED_MATMUL + +#endif +#endif // ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_tiling_key.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_tiling_key.h new file mode 100644 index 00000000..173f3b1b --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_tiling_key.h @@ -0,0 +1,516 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file grouped_matmul_tiling_key.h + * \brief + */ + +#ifndef GROUPED_MATMUL_TILING_KEY_H +#define GROUPED_MATMUL_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +// Datatype definition reference to DT_*** +#define GMM_TPL_INVALID 0xFFFFFFFF +#define GMM_TPL_FLOAT 0 +#define GMM_TPL_FLOAT16 1 +#define GMM_TPL_INT8 2 +#define GMM_TPL_INT32 3 +#define GMM_TPL_BF16 27 +#define GMM_TPL_INT4 29 + +#define GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM 0 +#define GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT 1 +#define GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM 2 + +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE 0 +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_API_DEQUANT 1 +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_VECTOR_DEQUANT 2 +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERCHANNEL_ANTIQUANT 3 +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERGROUP_ANTIQUANT 4 +#define GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_AUTOTILING 5 + +#define GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE 0 +#define GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD 1 +#define GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT 2 + +#define GROUPED_MATMUL_CUBE_ONLY 0 +#define GROUPED_MATMUL_AIV_AIC_RATIO_1 1 +#define GROUPED_MATMUL_AIV_AIC_RATIO_2 2 + +#define SET_GMM_A8W4_TPL_ARGS(kernel_type, kernel_template, aiv_aic_ratio) \ + ASCENDC_TPL_KERNEL_TYPE_SEL(kernel_type), \ + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), \ + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT4), \ + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), \ + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \ + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), \ + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), \ + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \ + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_##kernel_template), \ + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, aiv_aic_ratio), \ + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + +#define SET_GMM_A16W4_TPL_ARGS(kernel_type, dtype, trans_b, aiv_aic_ratio) \ + ASCENDC_TPL_KERNEL_TYPE_SEL(kernel_type), \ + ASCENDC_TPL_DTYPE_SEL(D_T_A, dtype), \ + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT4), \ + ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \ + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \ + ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \ + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\ + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \ + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, aiv_aic_ratio), \ + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + +#define SET_GMM_A16W8_TPL_ARGS(dtype) \ + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), \ + ASCENDC_TPL_DTYPE_SEL(D_T_A, dtype), \ + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), \ + ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \ + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \ + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), \ + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\ + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \ + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), \ + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), \ + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + +#define SET_GMM_NO_QUANT_TPL_ARGS(kernel_type, dtype, trans_a, trans_b, aiv_aic_ratio) \ + ASCENDC_TPL_KERNEL_TYPE_SEL(kernel_type), \ + ASCENDC_TPL_DTYPE_SEL(D_T_A, dtype), \ + ASCENDC_TPL_DTYPE_SEL(D_T_B, dtype), \ + ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \ + ASCENDC_TPL_BOOL_SEL(TRANS_A, trans_a), \ + ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \ + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM),\ + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \ + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \ + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, aiv_aic_ratio), \ + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + +ASCENDC_TPL_ARGS_DECL( + DlinferGroupedMatmulDirect, + // feature map datatype + ASCENDC_TPL_DTYPE_DECL( + D_T_A, GMM_TPL_FLOAT, GMM_TPL_FLOAT16, GMM_TPL_BF16, GMM_TPL_INT8, GMM_TPL_INT4 + ), + // weight datatype + ASCENDC_TPL_DTYPE_DECL( + D_T_B, GMM_TPL_FLOAT, GMM_TPL_FLOAT16, GMM_TPL_BF16, GMM_TPL_INT8, GMM_TPL_INT4 + ), + // output datatype + ASCENDC_TPL_DTYPE_DECL( + D_T_Y, GMM_TPL_FLOAT, GMM_TPL_INT32, GMM_TPL_FLOAT16, GMM_TPL_BF16, GMM_TPL_INT8 + ), + // feature map transpose + ASCENDC_TPL_BOOL_DECL(TRANS_A, 0, 1), + // weight transpose + ASCENDC_TPL_BOOL_DECL(TRANS_B, 0, 1), + // group_list_type value + ASCENDC_TPL_UINT_DECL( + GROUP_LIST_TYPE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM + ), + // use of matmul static tiling high-level api + ASCENDC_TPL_BOOL_DECL(IS_STATIC_TILING_API, 0, 1), + // kernel template in A8W4 senario + ASCENDC_TPL_UINT_DECL( + A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE, // Not in A8W4 senario + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_API_DEQUANT, // MSD template, using API to dequant + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_MSD_VECTOR_DEQUANT, // MSD template, using vector to dequant + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERCHANNEL_ANTIQUANT, // Fake A8W8 template with perchannel quantization + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_PERGROUP_ANTIQUANT, // Antiquant template without MSD strategy + GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_AUTOTILING, // Autotiling template with tuning config + ), + // kernel template in A16W8 senario + ASCENDC_TPL_UINT_DECL( + A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE, // Not in A16W8 senario + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD, // A16W8 MSD template + GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT, // A16W8 antiquant template + ), + // ratio of AIV:AIC + ASCENDC_TPL_UINT_DECL( + AIV_AIC_RATIO, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_CUBE_ONLY, + GROUPED_MATMUL_AIV_AIC_RATIO_1, + GROUPED_MATMUL_AIV_AIC_RATIO_2 + ), + // IS_ENABLE_FIXED_AXIS_ALGORITHM + ASCENDC_TPL_BOOL_DECL(IS_ENABLE_FIXED_AXIS, 0, 1), +); + + +ASCENDC_TPL_SEL( +#if defined(__CCE_AICORE__) +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) +#if defined(GMM_ANTI_QUANT_A8W4_MSD) + // ANTIQUANT_A8W4 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, MSD_API_DEQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, MSD_VECTOR_DEQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, PERCHANNEL_ANTIQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, PERGROUP_ANTIQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, AUTOTILING, GROUPED_MATMUL_AIV_AIC_RATIO_2)), +#elif defined(GMM_ANTI_QUANT) + // ANTIQUANT_A16W4 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, GMM_TPL_BF16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, GMM_TPL_FLOAT16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 1, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 1, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + #if defined(ORIG_DTYPE_WEIGHT) && defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8 + // ANTIQUANT_A16W8 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W8_TPL_ARGS(GMM_TPL_BF16)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W8_TPL_ARGS(GMM_TPL_FLOAT16)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_FLOAT16), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + #endif +#elif defined(GMM_QUANT_BF16) || defined(GMM_QUANT_FLOAT16) + // QUANT_A8W8O16 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // QUANT_A8W8O16_ENABLE_FIXED_AXIS_ALGORITHM + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 1) + ), +#elif defined(GMM_A4W4) + // QUANT_A4W4 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT4), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT4), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), +#elif defined(GMM_QUANT_INT8) + // QUANT_A8W8O8 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_0), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_INT8), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_CUBE_ONLY), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), +#elif defined(GMM_QUANT_INT32) + // QUANT_A8W8O32 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_0), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_CUBE_ONLY), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), +#elif defined(GMM_FLOAT) + // NO_QUANT + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_BF16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_BF16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), +#endif +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 +#if defined(GMM_FLOAT) + // NO_QUANT + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_BF16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_FLOAT16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_FLOAT, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_BF16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_FLOAT16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_AICORE, GMM_TPL_FLOAT, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_VECTOR_CORE, GMM_TPL_BF16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_VECTOR_CORE, GMM_TPL_FLOAT16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_VECTOR_CORE, GMM_TPL_FLOAT, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)) +#endif +#endif + +#else // not defined(__CCE_AICORE__): compile tiling only + // ANTIQUANT_A8W4 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, MSD_API_DEQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, MSD_VECTOR_DEQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, PERCHANNEL_ANTIQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, PERGROUP_ANTIQUANT, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A8W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, AUTOTILING, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + // ANTIQUANT_A16W4 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, GMM_TPL_BF16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_2, GMM_TPL_FLOAT16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_2)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 1, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W4_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 1, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + // ANTIQUANT_A16W8 + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W8_TPL_ARGS(GMM_TPL_BF16)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_A16W8_TPL_ARGS(GMM_TPL_FLOAT16)), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_FLOAT16), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // QUANT_A8W8O16 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // QUANT_A8W8O16_ENABLE_FIXED_AXIS_ALGORITHM + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_1), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 1) + ), + // QUANT_A4W4 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT4), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT4), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_BF16, GMM_TPL_FLOAT16), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_2), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // QUANT_A8W8O8 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_0), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_INT8), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_CUBE_ONLY), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // QUANT_A8W8O32 + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_0), + ASCENDC_TPL_DTYPE_SEL(D_T_A, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_B, GMM_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(D_T_Y, GMM_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), + ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), + ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, + GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, + GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, + GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), + ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0, 1), + ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), + ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_CUBE_ONLY), + ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0) + ), + // NO_QUANT + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_BF16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT16, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT, 0, 0, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_BF16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT16, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_0, GMM_TPL_FLOAT, 0, 1, GROUPED_MATMUL_CUBE_ONLY)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_BF16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT16, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)), + ASCENDC_TPL_ARGS_SEL(SET_GMM_NO_QUANT_TPL_ARGS(ASCENDC_TPL_MIX_AIC_1_1, GMM_TPL_FLOAT, 1, 0, GROUPED_MATMUL_AIV_AIC_RATIO_1)) +#endif +); +#endif // GROUPED_MATMUL_TILING_KEY_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_utils.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_utils.h new file mode 100644 index 00000000..bac53c41 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_utils.h @@ -0,0 +1,372 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_utils.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_UTILS_H +#define ASCENDC_GROUPED_MATMUL_UTILS_H + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + #if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 + #define DTYPE_L0C_LOCAL int32_t + #else + #define DTYPE_L0C_LOCAL float + #endif + #if defined(ORIG_DTYPE_X) && defined(ORIG_DTYPE_WEIGHT) && defined(DT_FLOAT8_E5M2) && defined(DT_FLOAT8_E4M3FN) && \ + defined(DT_HIFLOAT8) && defined(DT_INT8) && defined(DT_FLOAT4_E2M1) && defined(DT_FLOAT4_E1M2) && \ + defined(DT_INT4) && \ + ((ORIG_DTYPE_X == DT_INT8 && ORIG_DTYPE_WEIGHT == DT_INT8) || \ + (ORIG_DTYPE_X == DT_HIFLOAT8 && ORIG_DTYPE_WEIGHT == DT_HIFLOAT8) || \ + ((ORIG_DTYPE_X == DT_FLOAT8_E5M2 || ORIG_DTYPE_X == DT_FLOAT8_E4M3FN) && \ + (ORIG_DTYPE_WEIGHT == DT_FLOAT8_E5M2 || ORIG_DTYPE_WEIGHT == DT_FLOAT8_E4M3FN)) || \ + ((ORIG_DTYPE_X == DT_FLOAT4_E2M1 || ORIG_DTYPE_X == DT_FLOAT4_E1M2) && \ + (ORIG_DTYPE_WEIGHT == DT_FLOAT4_E2M1 || ORIG_DTYPE_WEIGHT == DT_FLOAT4_E1M2)) || \ + (ORIG_DTYPE_X == DT_INT4 && ORIG_DTYPE_WEIGHT == DT_INT4)) + #define V310_GMM_QUANT + #if defined(ORIG_DTYPE_SCALE) && defined(DT_FLOAT8_E8M0) && ORIG_DTYPE_SCALE == DT_FLOAT8_E8M0 + #define V310_GMM_QUANT_MX + #elif defined(ORIG_DTYPE_SCALE) && defined(DT_UINT64) && defined(DT_INT64) && \ + (ORIG_DTYPE_SCALE != DT_UINT64 && ORIG_DTYPE_SCALE != DT_INT64) + #define V310_GMM_QUANT_MIX + #define V310_GMM_QUANT_PERTENSOR_CUBE + #if (ORIG_DTYPE_X != DT_INT8 && ORIG_DTYPE_SCALE == DT_FLOAT) + #define V310_GMM_QUANT_PERTILE + #endif + #else + #define V310_GMM_QUANT_CUBE + #endif + #endif + + #if defined(ORIG_DTYPE_X) && defined(ORIG_DTYPE_WEIGHT) && ORIG_DTYPE_X != ORIG_DTYPE_WEIGHT + #if ((ORIG_DTYPE_X == DT_FLOAT16 || ORIG_DTYPE_X == DT_BF16) && \ + (ORIG_DTYPE_WEIGHT == DT_FLOAT8_E5M2 || ORIG_DTYPE_WEIGHT == DT_FLOAT8_E4M3FN || \ + ORIG_DTYPE_WEIGHT == DT_HIFLOAT8 || ORIG_DTYPE_WEIGHT == DT_INT8 || ORIG_DTYPE_WEIGHT == DT_FLOAT4_E2M1 || \ + ORIG_DTYPE_WEIGHT == DT_FLOAT4_E1M2 || ORIG_DTYPE_WEIGHT == DT_FLOAT || ORIG_DTYPE_WEIGHT == DT_INT32 || \ + ORIG_DTYPE_WEIGHT == DT_INT4)) || \ + (ORIG_DTYPE_X == DT_INT8 && (ORIG_DTYPE_WEIGHT == DT_INT4 || ORIG_DTYPE_WEIGHT == DT_INT32)) || \ + (ORIG_DTYPE_X == DT_FLOAT8_E4M3FN && (ORIG_DTYPE_WEIGHT == DT_FLOAT4_E2M1 || ORIG_DTYPE_WEIGHT == DT_FLOAT)) + #define V310_GMM_ANTI_QUANT + #endif + #endif +#endif + +#if defined(ORIG_DTYPE_X) && defined(ORIG_DTYPE_WEIGHT) && defined(ORIG_DTYPE_Y) && defined(DT_INT8) && \ + defined(DT_BF16) && defined(DT_INT4) + #if ORIG_DTYPE_X == ORIG_DTYPE_WEIGHT + #if ORIG_DTYPE_X == DT_INT8 + #if ORIG_DTYPE_Y == DT_BF16 + #define GMM_QUANT_BF16 + #define MM_DTYPE_Y int32_t + #elif ORIG_DTYPE_Y == DT_FLOAT16 + #define GMM_QUANT_FLOAT16 + #define MM_DTYPE_Y int32_t + #elif ORIG_DTYPE_Y == DT_INT32 + #define GMM_QUANT_INT32 + #else + #define GMM_QUANT_INT8 + #endif + #elif ORIG_DTYPE_X == DT_INT4 + #define GMM_A4W4 + #define MM_DTYPE_Y half + #if ORIG_DTYPE_Y == DT_BF16 + #define GMM_A4W4_BF16 + #elif ORIG_DTYPE_Y == DT_FLOAT16 + #define GMM_A4W4_FP16 + #endif + #else + #define GMM_FLOAT + #endif + #else + #define GMM_ANTI_QUANT + #if ORIG_DTYPE_X == DT_INT8 && ORIG_DTYPE_WEIGHT == DT_INT4 + #define GMM_ANTI_QUANT_A8W4_MSD + #define GMM_ANTI_QUANT_A8W4 + #if ORIG_DTYPE_Y == DT_BF16 + #define GMM_ANTI_QUANT_A8W4_MSD_OUT_BF16 + #else + #define GMM_ANTI_QUANT_A8W4_MSD_OUT_FP16 + #endif + #define MM_DTYPE_Y int32_t + #else + #define GMM_ANTI_QUANT + #endif + #endif +#endif + +#if defined(DTYPE_Y) && !defined(MM_DTYPE_Y) + #define MM_DTYPE_Y DTYPE_Y +#endif + +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + #ifdef GMM_ANTI_QUANT_A8W4_MSD_OUT_BF16 + #undef DTYPE_SCALE + #define DTYPE_SCALE bfloat16_t + #elif defined(GMM_ANTI_QUANT_A8W4_MSD_OUT_FP16) + #undef DTYPE_SCALE + #define DTYPE_SCALE float + #endif +#endif + +#if defined(CONST_TILING) + #define TILING_TYPE const int32_t +#else + #define TILING_TYPE __gm__ int32_t +#endif + +#if defined(CONST_TILING) + #if defined(V310_GMM_ANTI_QUANT) + #define GET_TILING_DATA_MEMBER_ADDR(tilingType, member, var, tiling) \ + GET_TILING_DATA_MEMBER(GMMWeightQuantTilingData, member, obj, tiling); \ + const int32_t* (var) = (const int32_t*)((const uint8_t*)&obj); + #else + #define GET_TILING_DATA_MEMBER_ADDR(tilingType, member, var, tiling) \ + GET_TILING_DATA_MEMBER(tilingType, member, obj, tiling); \ + const int32_t* (var) = (const int32_t*)((const uint8_t*)&obj); + #endif +#else + #define GET_TILING_DATA_MEMBER_ADDR(tilingType, member, var, tiling) \ + size_t offset##var = (size_t)(&((tilingType*)0)->member); \ + __gm__ int32_t* (var) = (__gm__ int32_t*)((tiling) + (offset##var)); +#endif + +namespace GROUPED_MATMUL { +using namespace AscendC; + +constexpr uint32_t INT8_BITS = 8; // a int8 number has 8 bits +constexpr int32_t MKN_LIST_LEN = 128; // 128: predefined array legnth +constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data +constexpr uint32_t UB_BLOCK_DOUBLE_UNIT_SIZE = 64; // 64: a block has 64 bytes data +constexpr uint32_t HALF_UB_BLOCK_UNIT_SIZE = UB_BLOCK_UNIT_SIZE / 2; // 2: a float16 data has two bytes +#if ((defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003)) && \ + defined(ORIG_DTYPE_X) && defined(ORIG_DTYPE_WEIGHT) && \ + ORIG_DTYPE_X == DT_INT8 && ORIG_DTYPE_WEIGHT == DT_INT8 +constexpr MatmulConfig NZ_CFG_MDL = + GetMDLConfig(false, false, 0, true, false, false, true, true, true, false, false, true); +constexpr MatmulConfig matmulCFGUnitFlag{.doMultiDataLoad = true, .enUnitFlag = true, .enableKdimReorderLoad = true}; +#else +constexpr MatmulConfig NZ_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true); +constexpr MatmulConfig matmulCFGUnitFlag{false, false, true, 0, 0, 0, false, false, false, false, false, 0, 0, 0, + 0, 0, 0, 0, true}; +#endif + +constexpr uint64_t SYNC_AIV_AIC_FLAG = 3; +constexpr uint64_t SYNC_AIC_AIV_FLAG = 5; +constexpr uint64_t SYNC_MODE2 = 2; +// used for static tiling template +constexpr uint32_t BASIC_BLOCK_SIZE_128 = 128; +constexpr uint32_t BASIC_BLOCK_SIZE_256 = 256; +constexpr int32_t STATIC_TILING_DEPTH_A1_B1 = 8; +constexpr int32_t STATIC_TILING_STEP_KA_KB = 4; +constexpr uint64_t DOUBLE_BUFFER_L0A_L0B = 2; +constexpr uint32_t STATIC_TILING_MAX_K = 8192; +constexpr uint32_t STATIC_TILING_MAX_SINGLE_N = 1024; + +template +struct MMType { + using AT = AT_; + using BT = BT_; + using CT = CT_; + using BiasT = BiasT_; + using MT = matmul::Matmul; +}; + +template +struct MMImplType { + using AT = AT_; + using BT = BT_; + using CT = CT_; + using BiasT = BiasT_; + using MT = matmul::MatmulImpl; +}; + +enum class ActiveType : std::uint8_t { + INVALID_TYPE = 0, + RELU, + GELU_TANH, + GELU_ERR_FUNC, + FASTGELU, + SILU +}; + +template +__aicore__ inline T GreatestCommonDivisor(T a, T b) { + T c = a; + if (a < b) { + a = b; + b = c; + } + while (b != 0) { + c = a; + a = b; + b = c % b; + } + return a; +} + +template +__aicore__ inline T LeastCommonMultiple(T a, T b) { + return a * b / GreatestCommonDivisor(a, b); +} + +template +__aicore__ inline T Max(T a, T b) { + return a > b ? a : b; +} + +template +__aicore__ inline T Min(T a, T b) { + return a > b ? b : a; +} + +template +__aicore__ inline T AlignUp(T a) { + return (a + base - 1) / base * base; +} + +template +__aicore__ inline T AlignUp(T a, T base) { + return (a + base - 1) / base * base; +} + +template +__aicore__ inline T AlignDown(T a, T base) { + if (unlikely(base == 0)) { + return a; + } + return a / base * base; +} + +template <> +__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a) { + // to be Multiple of 4, result should be in a format of b(xxxx,x100). + // This means last two bits should be zero, requiring that + // result = num & b(1111,1100) = num & (~3). + // &(~3) operator may reduces num into the range [num, num - 3]. + // As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case. + // In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise, + // the result will not be least multiple of 4 for 3. In other cases like [num, num - 2], + // num = a + 3 also satisfies the goal condition. + return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a) { + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a) { + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a) { + // refer to the above comments. + return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero} +} + +template +__aicore__ inline __gm__ T* GetTensorAddr(uint16_t index, GM_ADDR tensorPtr) { + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr); + uint64_t tensorPtrOffset = *dataAddr; // The offset of the data address from the first address. + // Moving 3 bits to the right means dividing by sizeof(uint64 t). + __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T*>(*(retPtr + index)); +} + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ != 310 +__aicore__ inline int32_t GetSplitValueFromGroupList(uint32_t groupIdx, int32_t &preOffset, + const GMMBaseParams* __restrict &gmmBaseParams, + const GlobalTensor &groupListGm) { + int32_t splitValue = 0; + if (likely(gmmBaseParams->groupType != -1)) { // -1: no need to split + if (gmmBaseParams->groupListType == 0) { + int32_t offset = static_cast(groupListGm.GetValue(groupIdx)); + splitValue = offset - preOffset; + preOffset = offset; + } else { + splitValue = static_cast(groupListGm.GetValue(groupIdx)); + } + } + return splitValue; +} +#endif + +template +__aicore__ inline constexpr uint32_t GetTypeBits() { + if constexpr (IsSameType::value) { + return 4; // 4: int4 bits number + } + return sizeof(T) * INT8_BITS; +} + +__aicore__ static constexpr MatmulConfig GenGmmConf(bool isND2NZ) { + return { + .doNorm = false, + .doBasicBlock = false, + .doMultiDataLoad = true, + .basicM = BASIC_BLOCK_SIZE_128, + .basicN = BASIC_BLOCK_SIZE_256, + .basicK = BASIC_BLOCK_SIZE_128, + .intrinsicsCheck = false, + .isNBatch = false, + .enVecND2NZ = isND2NZ, + .doSpecialBasicBlock = false, + .doMTE2Preload = 0, + .singleCoreM = BASIC_BLOCK_SIZE_128, + .singleCoreN = STATIC_TILING_MAX_SINGLE_N, + .singleCoreK = STATIC_TILING_MAX_K, + .stepM = 0, + .stepN = 0, + .baseMN = 0, + .singleCoreMN = 0, + .enUnitFlag = true, + .isPerTensor = false, + .hasAntiQuantOffset = false, + .doIBShareNorm = false, + .doSpecialMDL = false, + .enableInit = false, + .batchMode = BatchMode::NONE, + .enableEnd = true, + .enableGetTensorC = true, + .enableSetOrgShape = true, + .enableSetBias = false, + .enableSetTail = true, + .enableQuantVector = false, + .enableSetDefineData = false, + .iterateMode = IterateMode::ITERATE_MODE_DEFAULT, + .enableReuse = true, + .enableUBReuse = true, + .enableL1CacheUB = false, + .intraBlockPartSum = false, + .iterateOrder = IterateOrder::UNDEF, + .scheduleType = ScheduleType::INNER_PRODUCT, + .enableDoubleCache = false, + .isBiasBatch = true, + .enableStaticPadZeros = false, + .isA2B2Shared = false, + .enableKdimReorderLoad = false, + .isCO1Shared = false, + }; +} + +} // namespace GROUPED_MATMUL + +#endif // ASCENDC_GROUPED_MATMUL_UTILS_H diff --git a/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_vector.h b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_vector.h new file mode 100644 index 00000000..caf243e6 --- /dev/null +++ b/dlinfer/vendor/ascend/csrc/grouped_matmul_direct/opp/op_kernel/grouped_matmul_vector.h @@ -0,0 +1,76 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_vector.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_VECTOR_H +#define ASCENDC_GROUPED_MATMUL_VECTOR_H + +#include "grouped_matmul_utils.h" + +namespace GROUPED_MATMUL { + +template +__aicore__ inline void EmptyTensorCompute(GM_ADDR groupListPtr, GM_ADDR y, const GMMTilingData* __restrict tiling) { + const GMMBaseParams* __restrict gmmBaseParams = &tiling->gmmBaseParams; + // In the V2 interface, grouptype is -1 after host is grouped. Thus, grouptype can be either -1 or 2. + if (groupListPtr == nullptr || gmmBaseParams->groupType == 0) { + return; + } + + GlobalTensor yGm; + GlobalTensor groupListGm; + yGm.SetGlobalBuffer(GetTensorAddr(0, y)); + if (groupListPtr != nullptr) { + groupListGm.SetGlobalBuffer((__gm__ int64_t*)groupListPtr); + } + uint64_t yBaseOffset = 0; + int32_t preOffset = 0; + uint32_t singleWeight = gmmBaseParams->singleWeight; + uint32_t singleX = gmmBaseParams->singleX; + uint32_t singleY = gmmBaseParams->singleY; + bool isAllSingleTensor = singleWeight == 1 && singleX == 1 && singleY == 1; + + const int32_t *ubM = tiling->gmmArray.mList; + const int32_t *ubK = tiling->gmmArray.kList; + const int32_t *ubN = tiling->gmmArray.nList; + int64_t coreIdx = GetBlockIdx(); + int64_t coreRation = GetTaskRation(); + if (coreRation > 1) { + coreIdx /= coreRation; + } + + for (uint32_t groupIdx = 0; groupIdx < gmmBaseParams->groupNum; ++groupIdx) { + int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm); + uint32_t m = isAllSingleTensor && gmmBaseParams->groupType == 2 ? *ubM : *(ubM + groupIdx); + uint32_t k = *ubK < 0 && gmmBaseParams->groupType == 2 ? splitValue : *(ubK + groupIdx); + uint32_t n = isAllSingleTensor ? *ubN : *(ubN + groupIdx); + + if (k == 0) { + uint32_t singleM = Ceil(m, gmmBaseParams->coreNum); + singleM = AlignUp(singleM); + uint32_t cursingleM = singleM; + if (singleM * coreIdx >= m) { + yBaseOffset += m * n; + continue; + } else if (m - singleM * coreIdx < singleM) { + cursingleM = m - singleM * coreIdx; + } + InitOutput(yGm[yBaseOffset + coreIdx * singleM * n], cursingleM * n, 0); + } + yBaseOffset += m * n; + } +} + +} // namespace GROUPED_MATMUL + +#endif // ASCENDC_GROUPED_MATMUL_VECTOR_H diff --git a/dlinfer/vendor/ascend/grouped_matmul_direct.py b/dlinfer/vendor/ascend/grouped_matmul_direct.py new file mode 100644 index 00000000..8d4d3458 --- /dev/null +++ b/dlinfer/vendor/ascend/grouped_matmul_direct.py @@ -0,0 +1,42 @@ +"""Bundled Ascend GroupedMatmul direct2560 extension.""" + +from typing import Optional + +import torch +import torch_npu + +_EXTENSION_IMPORT_ERROR: Optional[BaseException] = None + +try: + from . import _grouped_matmul_direct +except (ImportError, OSError) as exc: + _grouped_matmul_direct = None + _EXTENSION_IMPORT_ERROR = exc + + +def is_available() -> bool: + """Whether the bundled pybind extension and op-api loaded successfully.""" + return _grouped_matmul_direct is not None + + +def unavailable_reason() -> str: + if _EXTENSION_IMPORT_ERROR is None: + return "" + return str(_EXTENSION_IMPORT_ERROR) + + +def grouped_matmul( + x: torch.Tensor, + weight: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int = 1, +) -> torch.Tensor: + if _grouped_matmul_direct is None: + raise RuntimeError( + "DLInfer bundled GroupedMatmul extension is unavailable: " + f"{unavailable_reason()}" + ) + stream_handle = torch_npu.npu.current_stream(x.device).npu_stream + return _grouped_matmul_direct.grouped_matmul( + x, weight, group_list, group_list_type, stream_handle + ) diff --git a/dlinfer/vendor/ascend/moe.py b/dlinfer/vendor/ascend/moe.py index a0a66033..aaf5c8a3 100644 --- a/dlinfer/vendor/ascend/moe.py +++ b/dlinfer/vendor/ascend/moe.py @@ -2,6 +2,7 @@ import torch import torch.distributed as dist from dlinfer.utils.type_annotation import MoECommType +from dlinfer.vendor.ascend import grouped_matmul_direct # aclnnGroupedMatmulV5 requires the groupList tensor to have at most 1024 @@ -14,6 +15,13 @@ # True -> catch-all (identical scheme to graph capture) # False -> per-chunk row slicing (weight views, each row computed once) _MOE_PREFILL_USE_CATCHALL = os.environ.get("DLINFER_MOE_PREFILL_CATCHALL", "0") == "1" +_GMM_EXPERIMENT_MODE = os.environ.get("DLINFER_GMM_EXPERIMENT", "chunked") +_GMM_EXPERIMENT_ALLOWED_MODES = {"chunked", "direct2560"} +if _GMM_EXPERIMENT_MODE not in _GMM_EXPERIMENT_ALLOWED_MODES: + raise RuntimeError( + "DLINFER_GMM_EXPERIMENT must be set before Python starts to one of " + f"{sorted(_GMM_EXPERIMENT_ALLOWED_MODES)}, got {_GMM_EXPERIMENT_MODE!r}" + ) class ChunkedMoeWeightLayout: @@ -35,6 +43,23 @@ def build_chunked_moe_storage_layout(num_experts: int): if num_experts <= MAX_GROUP_LIST_SIZE: return num_experts, None + # Select the physical expert-weight layout before checkpoint loading and + # graph capture. Direct mode keeps exactly one continuous 2560-row tensor. + use_direct = _GMM_EXPERIMENT_MODE == "direct2560" + if use_direct: + if num_experts != 2560: + raise RuntimeError( + "DLINFER_GMM_EXPERIMENT=direct2560 is restricted to exactly " + f"2560 logical experts, got {num_experts}" + ) + print( + "DLInfer MoE2560 enabled: backend=bundled_pybind " + "logical_experts=2560 weight_rows=2560 packed=False " + "group_list_type=1", + flush=True, + ) + return num_experts, None + chunk_size = MAX_GROUP_LIST_SIZE - 2 num_chunks = (num_experts + chunk_size - 1) // chunk_size storage_num_experts = num_experts + 2 * num_chunks @@ -81,29 +106,43 @@ def _grouped_mlp( down_weights: torch.Tensor, group_list: torch.Tensor, group_list_type: int, + use_bundled_direct: bool = False, ): + grouped_matmul = ( + grouped_matmul_direct.grouped_matmul + if use_bundled_direct + else lambda x, weight, groups, list_type: torch.ops.npu.npu_grouped_matmul( + [x], + [weight], + group_list=groups, + split_item=2, + group_type=0, + group_list_type=list_type, + )[0] + ) + # up sample - up_proj = torch.ops.npu.npu_grouped_matmul( - [hidden_states], - [gate_up_weights.transpose(1, 2)], - group_list=group_list, - split_item=2, - group_type=0, - group_list_type=group_list_type, - )[0] + up_weight = ( + gate_up_weights if use_bundled_direct else gate_up_weights.transpose(1, 2) + ) + up_proj = grouped_matmul( + hidden_states, + up_weight, + group_list, + group_list_type, + ) # activation gate_cache = torch.ops.npu.npu_swiglu(up_proj, -1) # down sample - down_proj = torch.ops.npu.npu_grouped_matmul( - [gate_cache], - [down_weights.transpose(1, 2)], - group_list=group_list, - split_item=2, - group_type=0, - group_list_type=group_list_type, - )[0] + down_weight = down_weights if use_bundled_direct else down_weights.transpose(1, 2) + down_proj = grouped_matmul( + gate_cache, + down_weight, + group_list, + group_list_type, + ) return down_proj @@ -247,6 +286,38 @@ def apply_mlp( hidden_states, gate_up_weights, down_weights, group_list, group_list_type ) + use_direct = _GMM_EXPERIMENT_MODE == "direct2560" + if use_direct: + if not grouped_matmul_direct.is_available(): + raise RuntimeError( + "DLINFER_GMM_EXPERIMENT=direct2560 requires the bundled " + "GroupedMatmul extension: " + f"{grouped_matmul_direct.unavailable_reason()}" + ) + if num_experts != 2560: + raise RuntimeError(f"direct2560 requires 2560 experts, got {num_experts}") + if chunked_moe_layout is not None: + raise RuntimeError("direct2560 received a packed chunked weight layout") + if gate_up_weights.dim() != 3 or down_weights.dim() != 3: + raise RuntimeError("direct2560 requires rank-3 single-weight tensors") + if gate_up_weights.size(0) != 2560 or down_weights.size(0) != 2560: + raise RuntimeError( + "direct2560 requires continuous 2560-row weights, got " + f"gate_up={gate_up_weights.size(0)}, down={down_weights.size(0)}" + ) + if group_list_type != 1: + raise RuntimeError( + f"direct2560 requires group_list_type=1, got {group_list_type}" + ) + return _grouped_mlp( + hidden_states, + gate_up_weights, + down_weights, + group_list, + group_list_type=1, + use_bundled_direct=True, + ) + # More experts than aclnnGroupedMatmulV5 supports: split into chunks of at # most MAX_GROUP_LIST_SIZE groups. Work in per-expert token counts # (group_list_type=1) regardless of the incoming layout. diff --git a/tests/test_grouped_matmul_direct.py b/tests/test_grouped_matmul_direct.py new file mode 100644 index 00000000..6696b7cf --- /dev/null +++ b/tests/test_grouped_matmul_direct.py @@ -0,0 +1,56 @@ +import pytest + +from dlinfer.vendor.ascend import grouped_matmul_direct +import torch + + +torch_npu = pytest.importorskip("torch_npu") + + +@pytest.mark.skipif(not torch_npu.npu.is_available(), reason="requires Ascend NPU") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_grouped_matmul_direct2560_matches_chunked(dtype): + assert ( + grouped_matmul_direct.is_available() + ), grouped_matmul_direct.unavailable_reason() + + experts = 2560 + hidden_size = 16 + output_size = 16 + tokens = 8 + device = torch.device("npu:0") + + x = torch.randn(tokens, hidden_size, device=device, dtype=dtype) + # Production stores weights as [E, N, K] and passes a transposed view. + stored_weight = torch.randn( + experts, output_size, hidden_size, device=device, dtype=dtype + ) + weight = stored_weight.transpose(1, 2) + group_list = torch.zeros(experts, device=device, dtype=torch.int64) + group_list[:tokens] = 1 + + direct = grouped_matmul_direct.grouped_matmul(x, stored_weight, group_list, 1) + + references = [] + row_start = 0 + expert_start = 0 + while expert_start < experts: + expert_end = min(expert_start + 1024, experts) + chunk_groups = group_list[expert_start:expert_end] + row_end = row_start + int(chunk_groups.sum().cpu()) + if row_end > row_start: + references.append( + torch.ops.npu.npu_grouped_matmul( + [x[row_start:row_end]], + [weight[expert_start:expert_end]], + group_list=chunk_groups, + split_item=2, + group_type=0, + group_list_type=1, + )[0] + ) + row_start = row_end + expert_start = expert_end + + reference = torch.cat(references, dim=0) + torch.testing.assert_close(direct, reference, rtol=0, atol=0)