EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
csrc/moe/scatter_nd_update_v2/CMakeLists.txt Normal file
View File

@@ -0,0 +1,19 @@
#
# 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(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
if(NOT ENABLE_TEST AND NOT BENCHMARK)
list(REMOVE_ITEM CURRENT_DIRS tests)
endif()
foreach(SUB_DIR ${CURRENT_DIRS})
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
add_subdirectory(${SUB_DIR})
endif()
endforeach()

32
csrc/moe/scatter_nd_update_v2/op_host/CMakeLists.txt Normal file
View File

@@ -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.
#/
#optiling and opapi
#message(STATUS "=== Debug: start ops.index.scatter_nd_update_v2.CMakeLists.txt ")
#add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE scatter_nd_update_v2 ACLNNTYPE aclnn_exclude)
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnnExc PRIVATE
scatter_nd_update_v2_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME ScatterNdUpdateV2
OPTIONS --cce-auto-sync=on
-Wno-deprecated-declarations
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE scatter_nd_update_v2 ACLNNTYPE aclnn_exclude)
endif()

821
csrc/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_binary.json Normal file
View File

@@ -0,0 +1,821 @@
{
"op_type": "ScatterNdUpdateV2",
"op_list": [
{
"bin_filename": "ScatterNdUpdateV2_float16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float32_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float32_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bfloat16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bfloat16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int8_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int8_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int32_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int32_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int64_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int64_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bool_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bool_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
}
]
}

13
csrc/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_simplified_key.ini Normal file
View File

@@ -0,0 +1,13 @@
; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
; [某算子]
; default=xx
; ascendxx=xx
; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
; 2)如果仅有default配置：各个版本按default配置
; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
; 6)是否是AscendC算子，由 ops/built-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
[ScatterNdUpdateV2]
default=0

821
csrc/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_binary.json Normal file
View File

@@ -0,0 +1,821 @@
{
"op_type": "ScatterNdUpdateV2",
"op_list": [
{
"bin_filename": "ScatterNdUpdateV2_float16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float32_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_float32_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "float32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bfloat16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bfloat16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bfloat16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int8_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int8_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int8",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int16_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int16_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int16",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int32_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int32_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int64_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_int64_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bool_int32",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int32",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
},
{
"bin_filename": "ScatterNdUpdateV2_bool_int64",
"inputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "indices",
"index": 1,
"dtype": "int64",
"format": "ND",
"paramType": "required",
"shape": [-2]
},
{
"name": "updates",
"index": 2,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"outputs": [
{
"name": "var",
"index": 0,
"dtype": "bool",
"format": "ND",
"paramType": "required",
"shape": [-2]
}
],
"attrs": [
{
"name": "strides",
"dtype": "list_int",
"value": null
},
{
"name": "use_locking",
"dtype": "bool",
"value": null
}
]
}
]
}

13
csrc/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_simplified_key.ini Normal file
View File

@@ -0,0 +1,13 @@
; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
; [某算子]
; default=xx
; ascendxx=xx
; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
; 2)如果仅有default配置：各个版本按default配置
; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
; 6)是否是AscendC算子，由 ops/built-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
[ScatterNdUpdateV2]
default=0

142
csrc/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.cpp Normal file
View File

@@ -0,0 +1,142 @@
/**
* 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_scatter_nd_update_v2.cpp
* \brief
*/
#include "aclnn_scatter_nd_update_v2.h"
#include "scatter_nd_update_v2.h"
#include "level0/broadcast_to.h"
#include "aclnn_kernels/contiguous.h"
#include "level0/squeeze.h"
#include "level0/unsqueeze.h"
#include "aclnn_kernels/common/op_error_check.h"
#include "opdev/make_op_executor.h"
#include "opdev/platform.h"
#include "opdev/op_dfx.h"
#include "opdev/tensor_view_utils.h"
#include "opdev/op_executor.h"
using namespace op;
#ifdef __cplusplus
extern "C" {
#endif
// 根据API定义，需要列出所能支持的所有dtype
static const std::initializer_list<op::DataType> ASCEND910_DTYPE_DTYPE_SUPPORT_LIST = {
op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL};
static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST = {
op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL, op::DataType::DT_INT16,
op::DataType::DT_BF16, op::DataType::DT_INT64, op::DataType::DT_INT8, op::DataType::DT_INT32};
static const std::initializer_list<op::DataType> ASCEND950_DTYPE_DTYPE_SUPPORT_LIST = {
op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL,
op::DataType::DT_BF16, op::DataType::DT_INT64, op::DataType::DT_INT8,
op::DataType::DT_FLOAT8_E5M2, op::DataType::DT_FLOAT8_E4M3FN, op::DataType::DT_FLOAT8_E8M0};
static const std::initializer_list<op::DataType> INDEX_DTYPE_SUPPORT_LIST = {op::DataType::DT_INT64,
op::DataType::DT_INT32};
static bool CheckNotNull(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) {
OP_CHECK_NULL(varRef, return false);
OP_CHECK_NULL(indices, return false);
OP_CHECK_NULL(updates, return false);
return true;
}
static const std::initializer_list<DataType>& GetDtypeSupportList() {
// if (Ops::NN::AclnnUtil::IsRegbase()) {
// return ASCEND950_DTYPE_DTYPE_SUPPORT_LIST;
// }
if (GetCurrentPlatformInfo().GetSocVersion() >= SocVersion::ASCEND910B &&
GetCurrentPlatformInfo().GetSocVersion() <= SocVersion::ASCEND910E) {
return ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST;
} else {
return ASCEND910_DTYPE_DTYPE_SUPPORT_LIST;
}
}
static bool CheckDtypeValid(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) {
// 检查self的数据类型是否在算子的支持列表内
auto supportList = GetDtypeSupportList();
OP_CHECK_DTYPE_NOT_SUPPORT(varRef, supportList, return false);
// 检查index的数据类型是否在算子的支持列表内
OP_CHECK_DTYPE_NOT_SUPPORT(indices, INDEX_DTYPE_SUPPORT_LIST, return false);
// varRef和updates的数据类型要一致
if (varRef->GetDataType() != updates->GetDataType()) {
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "updates dtype %s should be in same with varRef dtype %s.",
op::ToString(updates->GetDataType()).GetString(), op::ToString(varRef->GetDataType()).GetString());
return false;
}
return true;
}
static aclnnStatus CheckParams(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) {
// 1. 检查参数是否为空指针
CHECK_RET(CheckNotNull(varRef, indices, updates), ACLNN_ERR_PARAM_NULLPTR);
// 2. 检查输入的数据类型是否在API支持的数据类型范围之内
CHECK_RET(CheckDtypeValid(varRef, indices, updates), ACLNN_ERR_PARAM_INVALID);
return ACLNN_SUCCESS;
}
aclnnStatus aclnnScatterNdUpdateV2GetWorkspaceSize(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates,
const aclIntArray *strides, uint64_t* workspaceSize, aclOpExecutor** executor) {
L2_DFX_PHASE_1(aclnnScatterNdUpdateV2, DFX_IN(varRef, indices, updates), DFX_OUT(varRef));
// 固定写法，创建OpExecutor
auto uniqueExecutor = CREATE_EXECUTOR();
CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
// 固定写法，参数检查
auto ret = CheckParams(varRef, indices, updates);
CHECK_RET(ret == ACLNN_SUCCESS, ret);
if (varRef->IsEmpty() || indices->IsEmpty() || updates->IsEmpty()) {
*workspaceSize = 0;
uniqueExecutor.ReleaseTo(executor);
return ACLNN_SUCCESS;
}
// 将输入varRef转换成连续的tensor
// auto varRefContiguous = l0op::Contiguous(varRef, uniqueExecutor.get());
// CHECK_RET(varRefContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
varRef->SetStorageShape(varRef->GetViewShape());
// 将输入indices转换成连续的tensor
auto indicesContiguous = l0op::Contiguous(indices, uniqueExecutor.get());
CHECK_RET(indicesContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
// 将输入updates转换成连续的tensor
auto updatesContiguous = l0op::Contiguous(updates, uniqueExecutor.get());
CHECK_RET(updatesContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
// 执行L0算子 (kernel直接原地更新varRef, 无需Contiguous和ViewCopy)
auto scatterUpdateRes = l0op::ScatterNdUpdateV2(varRef, indicesContiguous, updatesContiguous, strides, false, uniqueExecutor.get());
CHECK_RET(scatterUpdateRes != nullptr, ACLNN_ERR_INNER_NULLPTR);
// 获取计算过程中需要使用的workspace大小
*workspaceSize = uniqueExecutor->GetWorkspaceSize();
uniqueExecutor.ReleaseTo(executor);
return ACLNN_SUCCESS;
}
aclnnStatus aclnnScatterNdUpdateV2(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream) {
L2_DFX_PHASE_2(aclnnScatterNdUpdateV2);
// 固定写法，调用框架能力，完成计算
return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
}
#ifdef __cplusplus
}
#endif

56
csrc/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.h Normal file
View File

@@ -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 OP_API_INC_SCATTER_ND_UPDATE_V2_H_
#define OP_API_INC_SCATTER_ND_UPDATE_V2_H_
#include "aclnn/aclnn_base.h"
#include "aclnn_util.h"
#ifdef __cplusplus
extern "C" {
#endif
/**
* @brief aclnnScatterNdUpdateV2的第一段接口，根据具体的计算流程，计算workspace大小。
* @domain aclnnop_ops_infer
* @domain aclnnop_ops_train
* 算子功能： 将tensor updates中的值按指定的索引indices逐个更新tensor var中的值。
* @param [in] varRef: npu device侧的aclTensor, 数据类型支持FLOAT16, FLOAT32, BOOL
* INT64，BFLOAT16，支持非连续的Tensor，数据格式支持ND。
* @param [in] indices: npu device侧的aclTensor，数据类型支持INT32, INT64类型。支持非连续的Tensor，数据格式支持ND。
* @param [in] updates: npu device侧的aclTensor，数据类型支持FLOAT16, FLOAT32, BOOL
* INT64，BFLOAT16，支持非连续的Tensor，数据格式支持ND,
* @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
* @param [out] executor: 返回op执行器，包含算子计算流程。
* @return aclnnStatus: 返回状态码
*/
ACLNN_API aclnnStatus aclnnScatterNdUpdateV2GetWorkspaceSize(aclTensor* varRef, const aclTensor* indices,
const aclTensor* updates, const aclIntArray* strides, uint64_t* workspaceSize,
aclOpExecutor** executor);
/**
* @brief: aclnnScatterNdUpdateV2的第二段接口，用于执行计算
* @domain aclnnop_ops_infer
* @domain aclnnop_ops_train
* 算子功能： 将tensor updates中的值按指定的索引indices逐个更新tensor var中的值。
* @param [in] workspace: 在npu device侧申请的workspace内存起址。
* @param [in] workspaceSize: 在npu device侧申请的workspace大小，由第一段接口aclnnScatterNdUpdateV2GetWorkspaceSize获取。
* @param [in] stream: acl stream流。
* @param [in] executor: op执行器，包含了算子计算流程。
* @return aclnnStatus: 返回状态码。
*/
ACLNN_API aclnnStatus aclnnScatterNdUpdateV2(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
aclrtStream stream);
#ifdef __cplusplus
}
#endif
#endif // OP_API_INC_SCATTER_ND_UPDATE_V2_H_

81
csrc/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.cpp Normal file
View File

@@ -0,0 +1,81 @@
/**
* 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 scatter_nd_update_v2.cpp
* \brief
*/
#include "scatter_nd_update_v2.h"
#include "opdev/make_op_executor.h"
#include "opdev/op_dfx.h"
#include "opdev/op_log.h"
#include "opdev/aicpu/aicpu_task.h"
#include "opdev/op_def.h"
#include "opdev/op_executor.h"
// #include "op_api/aclnn_util.h"
#include "aclnn_kernels/common/op_error_check.h"
using namespace op;
namespace l0op {
OP_TYPE_REGISTER(ScatterNdUpdateV2);
// AiCore支持的ScatterUpdate类型
static const std::initializer_list<op::DataType> AICORE_DTYPE_SUPPORT_LIST = {
op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL};
static const std::initializer_list<op::DataType> ASCEND910B_AICORE_DTYPE_SUPPORT_LIST = {
op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL, op::DataType::DT_BF16,
op::DataType::DT_INT64, op::DataType::DT_INT8};
inline static bool IsAiCoreSupport(const aclTensor* self) {
// ScatterNdUpdateV2只需要判断self
if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B ||
GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910_93) {
return CheckType(self->GetDataType(), ASCEND910B_AICORE_DTYPE_SUPPORT_LIST);
}
return CheckType(self->GetDataType(), AICORE_DTYPE_SUPPORT_LIST);
}
// AiCore的执行逻辑
inline static const aclTensor* ScatterNdUpdateV2AiCore(const aclTensor* self, const aclTensor* indices,
const aclTensor* updates, const aclIntArray* strides, bool use_locking,
aclOpExecutor* executor) {
L0_DFX(ScatterNdUpdateV2AiCore, self, indices, updates, use_locking);
auto retAicore =
ADD_TO_LAUNCHER_LIST_AICORE(ScatterNdUpdateV2,
OP_INPUT(self, indices, updates), OP_OUTPUT(self), OP_ATTR(strides, use_locking));
CHECK_RET(retAicore == ACLNN_SUCCESS, nullptr);
return self;
}
// AiCPU的执行逻辑
inline static const aclTensor* ScatterNdUpdateV2AiCPU(const aclTensor* self, const aclTensor* indices,
const aclTensor* updates, bool use_locking,
aclOpExecutor* executor) {
L0_DFX(ScatterNdUpdateV2AiCPU, self, indices, updates, use_locking);
static internal::AicpuTaskSpace space("ScatterNdUpdateV2", ge::DEPEND_IN_SHAPE, true);
space.SetRef(0);
auto ret = ADD_TO_LAUNCHER_LIST_AICPU(ScatterNdUpdateV2, OP_ATTR_NAMES({"Tindices", "T", "use_locking"}),
OP_INPUT(self, indices, updates), OP_OUTPUT(self),
OP_ATTR(indices->GetDataType(), updates->GetDataType(), use_locking));
CHECK_RET(ret == ACLNN_SUCCESS, nullptr);
return self;
}
const aclTensor* ScatterNdUpdateV2(const aclTensor* self, const aclTensor* indices, const aclTensor* updates,
const aclIntArray* strides, bool use_locking, aclOpExecutor* executor) {
if (IsAiCoreSupport(self)) {
return ScatterNdUpdateV2AiCore(self, indices, updates, strides, use_locking, executor);
} else {
return ScatterNdUpdateV2AiCPU(self, indices, updates, use_locking, executor);
}
}
} // namespace l0op

24
csrc/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.h Normal file
View File

@@ -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.
*/
/*!
* \file scatter_nd_update_v2.h
* \brief
*/
#ifndef OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_
#define OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_
#include "opdev/op_executor.h"
namespace l0op {
const aclTensor* ScatterNdUpdateV2(const aclTensor* self, const aclTensor* indices, const aclTensor* updates,
const aclIntArray* strides, bool use_locking, aclOpExecutor* executor);
} // namespace l0op
#endif // OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_

50
csrc/moe/scatter_nd_update_v2/op_host/scatter_nd_update_infershape.cpp Normal file
View File

@@ -0,0 +1,50 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_v2_infershape.cpp
* \brief
*/
#include "register/op_impl_registry.h"
#include "log/log.h"
#include "util/shape_util.h"
using namespace ge;
namespace ops {
static graphStatus InferDataType4ScatterNdUpdateV2(gert::InferDataTypeContext* context)
{
OP_LOGD(context->GetNodeName(), "Begin to do ScatterNdUpdateV2InferDtype.");
auto var_dtype = context->GetInputDataType(0);
context->SetOutputDataType(0, var_dtype);
OP_LOGD(context->GetNodeName(), "End to do ScatterNdUpdateV2InferDtype.");
return GRAPH_SUCCESS;
}
static ge::graphStatus InferShape4ScatterNdUpdateV2(gert::InferShapeContext* context)
{
OP_LOGD(context->GetNodeName(), "Begin to do ScatterNdUpdateV2InferShape.");
const gert::Shape* var_in_shape = context->GetInputShape(0);
gert::Shape* var_out_shape = context->GetOutputShape(0);
OP_CHECK_NULL_WITH_CONTEXT(context, var_in_shape);
OP_CHECK_NULL_WITH_CONTEXT(context, var_out_shape);
if (Ops::Base::IsUnknownRank(*var_in_shape)) {
OP_LOGD(context->GetNodeName(), "input shape is UnknownRank, set output shape to (-2, )");
Ops::Base::SetUnknownRank(*var_out_shape);
return ge::GRAPH_SUCCESS;
}
*var_out_shape = *var_in_shape;
OP_LOGD(context->GetNodeName(), "End to do ScatterNdUpdateV2InferShape.");
return ge::GRAPH_SUCCESS;
}
IMPL_OP_INFERSHAPE(ScatterNdUpdateV2)
.InferShape(InferShape4ScatterNdUpdateV2)
.InferDataType(InferDataType4ScatterNdUpdateV2);
} // namespace ops

84
csrc/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_def.cpp Normal file
View File

@@ -0,0 +1,84 @@
/**
* Copyright (c) Huawei Technologies Co., Ltd. 2025. 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.
*/
/*!
* \file scatter_nd_update_v2.cpp
* \brief ScatterNdUpdateV2 ophost
*/
#include "register/op_def_registry.h"
namespace ops {
class ScatterNdUpdateV2 : public OpDef {
public:
explicit ScatterNdUpdateV2(const char* name) : OpDef(name) {
this->Input("var")
.ParamType(REQUIRED)
.DataType(
{ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8,
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8})
.Format(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Input("indices")
.ParamType(REQUIRED)
.DataType(
{ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
.Format(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Input("updates")
.ParamType(REQUIRED)
.DataType(
{ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8,
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8})
.Format(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Output("var")
.ParamType(REQUIRED)
.DataType(
{ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8,
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8})
.Format(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat(
{ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
this->Attr("strides").AttrType(REQUIRED).ListInt();
this->Attr("use_locking").AttrType(OPTIONAL).Bool(false);
OpAICoreConfig aicore_config;
aicore_config.DynamicCompileStaticFlag(true)
.DynamicFormatFlag(true)
.DynamicRankSupportFlag(true)
.DynamicShapeSupportFlag(true);
this->AICore().AddConfig("ascend910b", aicore_config);
this->AICore().AddConfig("ascend910_93", aicore_config);
}
};
OP_ADD(ScatterNdUpdateV2);
} // namespace ops

403
csrc/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.cpp Normal file
View File

@@ -0,0 +1,403 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_v2_tiling.cpp
* \brief
*/
#include "register/op_impl_registry.h"
#include "util/math_util.h"
#include "platform/platform_infos_def.h"
#include "log/log.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling_base/tiling_util.h"
#include "tiling_base/tiling_key.h"
#include "scatter_nd_update_v2_tiling.h"
namespace optiling {
// using namespace Ops::NN::Optiling;
constexpr uint64_t MAX_DIM_NUM = 8;
constexpr uint64_t MAX_LENGTH_INT32 = (1LL << 31) - 1;
constexpr uint64_t MAX_FLOAT_EXPRESS_INT32 = (1LL << 24) - 1;
constexpr uint64_t SORT_USE_GM_NUM = 2;
constexpr uint64_t SORT_BLOCK_LENGTH = 4096;
constexpr uint64_t GATHER_USE_NUM = 2;
constexpr uint64_t ALIGNED_NUM = 8;
constexpr uint64_t ALIGNED_SIZE = 32;
constexpr uint64_t ATTR_STRIDE = 0;
class ScatterNdUpdateV2Tiling {
public:
explicit ScatterNdUpdateV2Tiling(gert::TilingContext* context) : tilingContext_(context){}
ge::graphStatus Init();
ge::graphStatus SetKernelTiling();
void TilingDataPrint() const;
private:
inline bool IsSort(uint64_t totalLength, uint64_t indexRow);
inline bool IsLinearIndex(uint64_t totalLength);
inline size_t CalcWorkSpaceSize(uint64_t indexRow);
inline void SetTilingKeyMode();
inline void GetDtypeSize();
inline void Tiling4Scatter(uint64_t totalLength, uint64_t indexRow);
inline void Tiling4LinearIndex(uint64_t indexRow, uint64_t indexDim);
ScatterNdUpdateV2TilingData tilingData_;
gert::TilingContext* tilingContext_ = nullptr;
uint64_t coreNum_ = 0;
uint64_t tilingKey_ = 0;
uint64_t ubSize_ = 0;
uint64_t isLinearIndex_ = false;
uint64_t isSort_ = false;
uint64_t sortWorkspace_ = 0;
uint64_t dataTypeSize_ = 0;
uint64_t isInt64Indices_ = false;
uint64_t needLargeIndexKernel_ = false;
private:
// LinearIndex
uint64_t indexDim_ = 0;
uint64_t blockLength_ = 0;
uint64_t blockNum_ = 0;
uint64_t blockRemainLength_ = 0;
uint64_t tailBlockNum_ = 0;
uint64_t frontBlockNum_ = 0;
uint64_t frontCoreNum_ = 0;
uint64_t tailCoreNum_ = 0;
uint64_t indicesMask_[MAX_DIM_NUM] = {0};
// Scatter
uint64_t scatterLength_ = 1;
uint64_t tailRow_ = 0;
uint64_t frontRow_ = 0;
uint64_t frontNum_ = 0;
uint64_t tailNum_ = 0;
uint64_t ubLengthForUpdates_ = 0;
uint64_t scatterAlignLength_ = 0;
uint64_t formDim_ = 0;
uint64_t copyRow_ = 0;
uint64_t scatterTileNum_ = 1;
uint64_t scatterTileLength_ = 0;
uint64_t scatterTileTail_ = 0;
uint64_t scatterTileAlignLength_ = 0;
};
inline void ScatterNdUpdateV2Tiling::SetTilingKeyMode()
{
// tilingKey: indexType * 10 + sortFlag (indexType: 1=int32, 2=int64(cast), 3=int64(large))
uint64_t indexType;
if (!isInt64Indices_) {
indexType = 1;
} else if (needLargeIndexKernel_) {
indexType = 3;
} else {
indexType = 2;
}
uint64_t sortFlag = (indexType == 3) ? 0 : (isSort_ ? 1 : 0);
tilingKey_ = indexType * 10 + sortFlag;
tilingContext_->SetTilingKey(tilingKey_);
OP_LOGD(tilingContext_, "isLinearIndex=%lu, isSort=%lu, isInt64Indices=%lu, needLargeIndexKernel=%lu, tilingKey=%lu (indexType=%lu, sortFlag=%lu)",
isLinearIndex_, isSort_, isInt64Indices_, needLargeIndexKernel_, tilingKey_, indexType, sortFlag);
}
inline bool ScatterNdUpdateV2Tiling::IsLinearIndex(uint64_t totalLength)
{
return totalLength <= MAX_LENGTH_INT32;
}
inline bool ScatterNdUpdateV2Tiling::IsSort(uint64_t totalLength, uint64_t indexRow)
{
return totalLength <= MAX_FLOAT_EXPRESS_INT32;
}
inline void ScatterNdUpdateV2Tiling::Tiling4LinearIndex(uint64_t indexRow, uint64_t indexDim)
{
OP_LOGD(tilingContext_, "linearIndexTiling start");
auto attrs = tilingContext_->GetAttrs();
auto stridesPtr = attrs->GetListInt(ATTR_STRIDE);
for (uint64_t i = 0; i < indexDim; ++i) {
indicesMask_[i] = static_cast<uint64_t>(stridesPtr->GetData()[i]);
}
uint64_t coeff = isInt64Indices_ ? (2 * indexDim + 3) : (indexDim + 3);
uint64_t maxBlockLength = ubSize_ / coeff / sizeof(int);
blockLength_ = (maxBlockLength / ALIGNED_SIZE) * ALIGNED_SIZE;
blockLength_ = std::min(blockLength_, (uint64_t)SORT_BLOCK_LENGTH);
blockNum_ = indexRow / blockLength_;
blockRemainLength_ = indexRow % blockLength_;
if (blockNum_ == 0) {
tailBlockNum_ = 0;
frontBlockNum_ = 0;
frontCoreNum_ = 1;
tailCoreNum_ = 0;
} else {
tailBlockNum_ = blockNum_ / coreNum_;
frontBlockNum_ = tailBlockNum_ + 1;
frontCoreNum_ = blockNum_ % coreNum_;
tailCoreNum_ = tailBlockNum_ == 0 ? 0 : coreNum_ - frontCoreNum_;
}
OP_LOGD(tilingContext_, "linearIndexTiling finish");
}
inline void ScatterNdUpdateV2Tiling::Tiling4Scatter(uint64_t totalLength, uint64_t indexRow)
{
OP_LOGD(tilingContext_, "scatterTiling start new");
uint64_t scatterAlignNum = ALIGNED_SIZE / dataTypeSize_;
tailRow_ = totalLength / coreNum_;
frontRow_ = tailRow_ + 1;
frontNum_ = totalLength % coreNum_;
tailNum_ = tailRow_ == 0 ? 0 : coreNum_ - frontNum_;
ubLengthForUpdates_ = ((ubSize_ - SORT_BLOCK_LENGTH * SORT_USE_GM_NUM * sizeof(int)) / ALIGNED_SIZE * ALIGNED_SIZE) / dataTypeSize_;
scatterAlignLength_ = (scatterLength_ + scatterAlignNum - 1) & ~(scatterAlignNum - 1);
formDim_ = scatterAlignLength_ / ubLengthForUpdates_;
scatterTileLength_ = std::min(scatterLength_, ubLengthForUpdates_);
if (scatterTileLength_ == 0) {
scatterTileLength_ = 1;
}
scatterTileNum_ = (scatterLength_ + scatterTileLength_ - 1) / scatterTileLength_;
scatterTileTail_ = scatterLength_ - (scatterTileNum_ - 1) * scatterTileLength_;
scatterTileAlignLength_ = (scatterTileLength_ + scatterAlignNum - 1) & ~(scatterAlignNum - 1);
if (scatterTileNum_ > 1) {
copyRow_ = 1;
} else {
copyRow_ = formDim_ == 0 ? ubLengthForUpdates_ / scatterAlignLength_ : 1;
}
OP_LOGD(tilingContext_, "scatterTiling finish");
}
inline void ScatterNdUpdateV2Tiling::GetDtypeSize()
{
uint64_t varDtype = tilingContext_->GetInputDesc(0)->GetDataType();
switch (varDtype){
case ge::DT_FLOAT:
dataTypeSize_ = 4;
break;
case ge::DT_BF16:
dataTypeSize_ = 2;
break;
case ge::DT_FLOAT16:
dataTypeSize_ = 2;
break;
case ge::DT_BOOL:
dataTypeSize_ = 1;
break;
case ge::DT_INT64:
dataTypeSize_ = 8;
break;
case ge::DT_INT32:
dataTypeSize_ = 4;
break;
case ge::DT_INT16:
dataTypeSize_ = 2;
break;
case ge::DT_INT8:
dataTypeSize_ = 1;
break;
default:
break;
}
}
ge::graphStatus ScatterNdUpdateV2Tiling::SetKernelTiling()
{
tilingContext_->SetBlockDim(coreNum_);
tilingData_.linearIndexTiling.set_indexDim(indexDim_);
tilingData_.linearIndexTiling.set_ubSize(ubSize_);
tilingData_.linearIndexTiling.set_indicesMask(indicesMask_);
tilingData_.linearIndexTiling.set_coreNum(coreNum_);
tilingData_.linearIndexTiling.set_blockLength(blockLength_);
tilingData_.linearIndexTiling.set_blockNum(blockNum_);
tilingData_.linearIndexTiling.set_blockRemainLength(blockRemainLength_);
tilingData_.linearIndexTiling.set_tailBlockNum(tailBlockNum_);
tilingData_.linearIndexTiling.set_frontBlockNum(frontBlockNum_);
tilingData_.linearIndexTiling.set_frontCoreNum(frontCoreNum_);
tilingData_.linearIndexTiling.set_tailCoreNum(tailCoreNum_);
tilingData_.linearIndexTiling.set_sortWorkspace(sortWorkspace_);
tilingData_.linearIndexTiling.set_isInt64Indices(isInt64Indices_);
tilingData_.linearIndexTiling.set_needLargeIndexKernel(needLargeIndexKernel_);
tilingData_.scatterTiling.set_scatterLength(scatterLength_);
tilingData_.scatterTiling.set_tailRow(tailRow_);
tilingData_.scatterTiling.set_frontRow(frontRow_);
tilingData_.scatterTiling.set_frontNum(frontNum_);
tilingData_.scatterTiling.set_tailNum(tailNum_);
tilingData_.scatterTiling.set_ubLengthForUpdates(ubLengthForUpdates_);
tilingData_.scatterTiling.set_scatterAlignLength(scatterAlignLength_);
tilingData_.scatterTiling.set_formDim(formDim_);
tilingData_.scatterTiling.set_copyRow(copyRow_);
tilingData_.scatterTiling.set_scatterTileNum(scatterTileNum_);
tilingData_.scatterTiling.set_scatterTileLength(scatterTileLength_);
tilingData_.scatterTiling.set_scatterTileTail(scatterTileTail_);
tilingData_.scatterTiling.set_scatterTileAlignLength(scatterTileAlignLength_);
tilingData_.SaveToBuffer(
tilingContext_->GetRawTilingData()->GetData(), tilingContext_->GetRawTilingData()->GetCapacity());
tilingContext_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
TilingDataPrint();
return ge::GRAPH_SUCCESS;
}
inline size_t ScatterNdUpdateV2Tiling::CalcWorkSpaceSize(uint64_t indexRow)
{
auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingContext_->GetPlatformInfo());
size_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
size_t indexRowAligned = (indexRow + ALIGNED_NUM - 1) & ~(ALIGNED_NUM - 1);
sortWorkspace_ = indexRowAligned;
size_t totalWorkspace = sysWorkspaceSize;
if (isLinearIndex_) {
totalWorkspace += sortWorkspace_ * SORT_USE_GM_NUM * sizeof(int);
}
if (isSort_) {
totalWorkspace += sortWorkspace_ * SORT_USE_GM_NUM * sizeof(int);
}
return totalWorkspace;
}
void ScatterNdUpdateV2Tiling::TilingDataPrint() const
{
OP_LOGD(tilingContext_, "coreNum: %lu", coreNum_);
OP_LOGD(tilingContext_, "tilingKey: %lu", tilingKey_);
OP_LOGD(tilingContext_, "isInt64Indices: %lu", isInt64Indices_);
OP_LOGD(tilingContext_, "needLargeIndexKernel: %lu", needLargeIndexKernel_);
OP_LOGD(tilingContext_, "tiling for LinearIndex--------");
OP_LOGD(tilingContext_, "indexDim: %lu", indexDim_);
OP_LOGD(tilingContext_, "ubSize: %lu", ubSize_);
OP_LOGD(tilingContext_, "blockLength: %lu", blockLength_);
OP_LOGD(tilingContext_, "blockNum: %lu", blockNum_);
OP_LOGD(tilingContext_, "blockRemainLength: %lu", blockRemainLength_);
OP_LOGD(tilingContext_, "tailBlockNum: %lu", tailBlockNum_);
OP_LOGD(tilingContext_, "frontBlockNum: %lu", frontBlockNum_);
OP_LOGD(tilingContext_, "frontCoreNum: %lu", frontCoreNum_);
OP_LOGD(tilingContext_, "tailCoreNum: %lu", tailCoreNum_);
OP_LOGD(tilingContext_, "sortWorkspace: %lu", sortWorkspace_);
for (size_t i = 0; i < indexDim_; i++) {
OP_LOGD(tilingContext_, "indicesMask[%lu]: %lu", i, indicesMask_[i]);
}
OP_LOGD(tilingContext_, "tiling for Scatter------------");
OP_LOGD(tilingContext_, "scatterLength: %lu", scatterLength_);
OP_LOGD(tilingContext_, "tailRow: %lu", tailRow_);
OP_LOGD(tilingContext_, "frontRow: %lu", frontRow_);
OP_LOGD(tilingContext_, "frontNum: %lu", frontNum_);
OP_LOGD(tilingContext_, "tailNum: %lu", tailNum_);
OP_LOGD(tilingContext_, "ubLengthForUpdates: %lu", ubLengthForUpdates_);
OP_LOGD(tilingContext_, "scatterAlignLength: %lu", scatterAlignLength_);
OP_LOGD(tilingContext_, "formDim: %lu", formDim_);
OP_LOGD(tilingContext_, "copyRow: %lu", copyRow_);
OP_LOGD(tilingContext_, "scatterTileNum: %lu", scatterTileNum_);
OP_LOGD(tilingContext_, "scatterTileLength: %lu", scatterTileLength_);
OP_LOGD(tilingContext_, "scatterTileTail: %lu", scatterTileTail_);
OP_LOGD(tilingContext_, "scatterTileAlignLength: %lu", scatterTileAlignLength_);
}
ge::graphStatus ScatterNdUpdateV2Tiling::Init()
{
OP_LOGD(tilingContext_, "Tiling initing");
auto compileInfo = static_cast<const ScatterNdUpdateV2CompileInfo*>(tilingContext_->GetCompileInfo());
auto varRefShape = tilingContext_->GetInputShape(0)->GetStorageShape();
auto indicesShape = tilingContext_->GetInputShape(1)->GetStorageShape();
auto updatesShape = tilingContext_->GetInputShape(2)->GetStorageShape();
uint64_t varDimNum = varRefShape.GetDimNum();
indexDim_ = indicesShape.GetDim(indicesShape.GetDimNum() - 1);
auto indicesDtype = tilingContext_->GetInputDesc(1)->GetDataType();
isInt64Indices_ = (indicesDtype == ge::DT_INT64);
OP_LOGD(tilingContext_, "indicesDtype=%d, isInt64Indices=%lu", indicesDtype, isInt64Indices_);
uint64_t totalLength = 1;
for (uint64_t i = 0; i < indexDim_; ++i) {
totalLength *= varRefShape.GetDim(i);
}
if (isInt64Indices_) {
needLargeIndexKernel_ = !IsLinearIndex(totalLength);
}
if (varDimNum > indexDim_) {
for (uint64_t i = indexDim_; i < varDimNum; i++) {
scatterLength_ *= varRefShape.GetDim(i);
}
}
uint64_t indexRow = 1;
for (uint64_t i = 0; i < indicesShape.GetDimNum() - 1; i++) {
indexRow *= indicesShape.GetDim(i);
}
if (needLargeIndexKernel_) {
isSort_ = false;
isLinearIndex_ = false;
} else {
isSort_ = false;
isLinearIndex_ = IsLinearIndex(totalLength);
}
coreNum_ = std::min(compileInfo->totalCoreNum,
std::min(static_cast<uint64_t>(totalLength), static_cast<uint64_t>(indexRow)));
coreNum_ = coreNum_ == 0 ? 1 : coreNum_;
ubSize_ = compileInfo->ubSizePlatForm;
GetDtypeSize();
Tiling4LinearIndex(indexRow, indexDim_);
uint64_t maxPhysicalOffset = 0;
for (uint64_t i = 0; i < indexDim_; ++i) {
maxPhysicalOffset += (varRefShape.GetDim(i) - 1) * indicesMask_[i];
}
uint64_t totalPhysicalRange = maxPhysicalOffset + scatterLength_;
if (!needLargeIndexKernel_) {
isSort_ = IsSort(totalPhysicalRange, indexRow);
}
SetTilingKeyMode();
tilingContext_->SetScheduleMode(1);
Tiling4Scatter(totalPhysicalRange, indexRow);
size_t* currentWorkSpace = tilingContext_->GetWorkspaceSizes(1);
currentWorkSpace[0] = CalcWorkSpaceSize(indexRow);
OP_LOGD(tilingContext_, "Tiling inited");
return ge::GRAPH_SUCCESS;
}
ge::graphStatus Tiling4ScatterNdUpdateV2(gert::TilingContext* context)
{
if (context == nullptr) {
OP_LOGE("ScatterNdUpdateV2", "The context is nullptr.");
return ge::GRAPH_FAILED;
}
OP_LOGD(context, "Tiling for ScatterNdUpdateV2 start.");
ScatterNdUpdateV2Tiling tilingOp(context);
if (tilingOp.Init() != ge::GRAPH_SUCCESS) {
OP_LOGE(context, "Tiling init fail");
return ge::GRAPH_FAILED;
}
OP_LOGD(context, "Tiling for ScatterNdUpdateV2 end.");
return tilingOp.SetKernelTiling();
}
ge::graphStatus TilingPrepare4ScatterNdUpdateV2(gert::TilingParseContext* context)
{
OP_LOGD(context, "Tiling Prepare For ScatterNdUpdateV2 start.");
auto compileInfo = context->GetCompiledInfo<ScatterNdUpdateV2CompileInfo>();
OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
auto platformInfo = context->GetPlatformInfo();
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
compileInfo->totalCoreNum = ascendcPlatform.GetCoreNumAiv();
if (compileInfo->totalCoreNum == 0) {
OP_LOGE(context, "coreNum %lu", compileInfo->totalCoreNum);
return ge::GRAPH_FAILED;
}
uint64_t ubSizePlatForm;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
compileInfo->ubSizePlatForm = ubSizePlatForm;
OP_LOGD(context, "ubSizePlatForm is %lu.", compileInfo->ubSizePlatForm);
OP_LOGD(context, "Tiling Prepare For ScatterNdUpdateV2 end.");
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(ScatterNdUpdateV2).Tiling(Tiling4ScatterNdUpdateV2).TilingParse<ScatterNdUpdateV2CompileInfo>(TilingPrepare4ScatterNdUpdateV2);
}

71
csrc/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.h Normal file
View File

@@ -0,0 +1,71 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_v2_tiling.h
* \brief
*/
#ifndef SCATTER_ND_UPDATE_V2_TILING_H
#define SCATTER_ND_UPDATE_V2_TILING_H
#include "register/tilingdata_base.h"
constexpr uint64_t MAX_DIM_NUM = 8;
namespace optiling {
BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2ScatterTiling)
TILING_DATA_FIELD_DEF(uint64_t, scatterLength)
TILING_DATA_FIELD_DEF(uint64_t, tailRow)
TILING_DATA_FIELD_DEF(uint64_t, frontRow)
TILING_DATA_FIELD_DEF(uint64_t, frontNum)
TILING_DATA_FIELD_DEF(uint64_t, tailNum)
TILING_DATA_FIELD_DEF(uint64_t, ubLengthForUpdates)
TILING_DATA_FIELD_DEF(uint64_t, scatterAlignLength)
TILING_DATA_FIELD_DEF(uint64_t, formDim)
TILING_DATA_FIELD_DEF(uint64_t, copyRow)
TILING_DATA_FIELD_DEF(uint64_t, scatterTileNum)
TILING_DATA_FIELD_DEF(uint64_t, scatterTileLength)
TILING_DATA_FIELD_DEF(uint64_t, scatterTileTail)
TILING_DATA_FIELD_DEF(uint64_t, scatterTileAlignLength)
END_TILING_DATA_DEF
REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2ScatterTilingOp, ScatterNdUpdateV2ScatterTiling)
BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2LinearIndexTiling)
TILING_DATA_FIELD_DEF(uint64_t, coreNum)
TILING_DATA_FIELD_DEF(uint64_t, ubSize)
TILING_DATA_FIELD_DEF(uint64_t, indexDim)
TILING_DATA_FIELD_DEF(uint64_t, blockLength)
TILING_DATA_FIELD_DEF(uint64_t, blockNum)
TILING_DATA_FIELD_DEF(uint64_t, blockRemainLength)
TILING_DATA_FIELD_DEF(uint64_t, tailBlockNum)
TILING_DATA_FIELD_DEF(uint64_t, frontBlockNum)
TILING_DATA_FIELD_DEF(uint64_t, frontCoreNum)
TILING_DATA_FIELD_DEF(uint64_t, tailCoreNum)
TILING_DATA_FIELD_DEF(uint64_t, sortWorkspace)
TILING_DATA_FIELD_DEF_ARR(uint64_t, MAX_DIM_NUM, indicesMask)
TILING_DATA_FIELD_DEF(uint64_t, isInt64Indices)
TILING_DATA_FIELD_DEF(uint64_t, needLargeIndexKernel)
END_TILING_DATA_DEF
REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2LinearIndexTilingOp, ScatterNdUpdateV2LinearIndexTiling)
BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2TilingData)
TILING_DATA_FIELD_DEF_STRUCT(ScatterNdUpdateV2ScatterTiling, scatterTiling)
TILING_DATA_FIELD_DEF_STRUCT(ScatterNdUpdateV2LinearIndexTiling, linearIndexTiling)
END_TILING_DATA_DEF
REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2, ScatterNdUpdateV2TilingData)
REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2TilingDataOp, ScatterNdUpdateV2TilingData)
struct ScatterNdUpdateV2CompileInfo {
uint64_t totalCoreNum = 0;
uint64_t ubSizePlatForm = 0;
};
} // namespace optiling
#endif // SCATTER_ND_UPDATE_V2_TILING_H

70
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_common.h Normal file
View File

@@ -0,0 +1,70 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_common.h
* \brief ScatterNdUpdateV2 公共定义和工具函数
*/
#ifndef SCATTER_ND_UPDATE_V2_COMMON_H
#define SCATTER_ND_UPDATE_V2_COMMON_H
#include "kernel_operator.h"
namespace ScatterNdUpdateV2 {
using namespace AscendC;
// 公共常量定义
constexpr uint64_t DOUBLE_BUFFER = 1;
constexpr uint64_t SORT_RES_NUM = 2;
constexpr uint64_t SORT_TMP_NUM = 3;
constexpr uint64_t ALIGNED_BLOCK_NUM = 32;
constexpr uint64_t ALIGN_NUM = 8; // 32 字节对齐 = 8 个 int32
constexpr uint64_t ALIGNED_SIZE = 512;
// 公共同步函数
__aicore__ inline void PipeMte2ToS()
{
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
SetFlag<HardEvent::MTE2_S>(eventID);
WaitFlag<HardEvent::MTE2_S>(eventID);
}
__aicore__ inline void PipeMte3ToS()
{
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
SetFlag<HardEvent::MTE3_S>(eventID);
WaitFlag<HardEvent::MTE3_S>(eventID);
}
__aicore__ inline void PipeVToMte3()
{
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
SetFlag<HardEvent::V_MTE3>(eventID);
WaitFlag<HardEvent::V_MTE3>(eventID);
}
// 计算 block 分布参数
__aicore__ inline void CalcBlockDistribution(
uint64_t blockIdx, uint64_t frontNum, uint64_t frontRow, uint64_t tailRow,
uint64_t& computeRow, uint64_t& start)
{
if (blockIdx >= frontNum) {
computeRow = tailRow;
start = frontNum * frontRow + (blockIdx - frontNum) * computeRow;
} else {
computeRow = frontRow;
start = blockIdx * computeRow;
}
}
} // namespace ScatterNdUpdateV2
#endif // SCATTER_ND_UPDATE_V2_COMMON_H

175
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_large_index.h Normal file
View File

@@ -0,0 +1,175 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_large_index.h
* \brief LargeIndex Kernel (index > 2^31-1)
*/
#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "scatter_nd_update_common.h"
namespace ScatterNdUpdateV2 {
template<typename T>
class LargeIndexKernel {
public:
__aicore__ inline LargeIndexKernel() = delete;
__aicore__ inline LargeIndexKernel(
GM_ADDR indices, GM_ADDR updates, GM_ADDR output,
const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe)
{
InitParams(tiling);
InitBuffers(pipe);
SetGmAddr(indices, updates, output, tiling);
}
__aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling)
{
blockIdx_ = GetBlockIdx();
CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow,
tiling.scatterTiling.tailRow, computeRow_, start_);
end_ = start_ + computeRow_;
startInt64_ = static_cast<int64_t>(start_);
endInt64_ = static_cast<int64_t>(end_);
indexDim_ = tiling.linearIndexTiling.indexDim;
blockLength_ = tiling.linearIndexTiling.blockLength;
blockNum_ = tiling.linearIndexTiling.blockNum;
blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength;
scatterLength_ = tiling.scatterTiling.scatterLength;
ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates;
scatterTileNum_ = tiling.scatterTiling.scatterTileNum;
scatterTileLength_ = tiling.scatterTiling.scatterTileLength;
scatterTileTail_ = tiling.scatterTiling.scatterTileTail;
for (uint64_t i = 0; i < indexDim_; ++i) {
indicesMask_[i] = tiling.linearIndexTiling.indicesMask[i];
}
}
__aicore__ inline void InitBuffers(TPipe& pipe)
{
uint64_t indicesInt64Size = ((blockLength_ * indexDim_ * 2) + ALIGN_NUM - 1) & ~(ALIGN_NUM - 1);
uint64_t updateBufBytes = (ubLengthForUpdates_ * sizeof(T) + 31) & ~31ULL;
pipe.InitBuffer(indicesBuf, indicesInt64Size * sizeof(int));
pipe.InitBuffer(updateBuf, updateBufBytes);
indicesInt64Local = indicesBuf.Get<int>().ReinterpretCast<int64_t>();
updateLocal = updateBuf.Get<T>();
}
__aicore__ inline void SetGmAddr(GM_ADDR indices, GM_ADDR updates, GM_ADDR output,
const ScatterNdUpdateV2TilingData& tiling)
{
indicesGmInt64_.SetGlobalBuffer((__gm__ int64_t*)indices);
updatesGm_.SetGlobalBuffer((__gm__ T*)updates);
outputGm_.SetGlobalBuffer((__gm__ T*)output);
}
__aicore__ inline void Process()
{
for (uint64_t blockIdx = 0; blockIdx < blockNum_; ++blockIdx) {
ProcessOneBlock(blockIdx, false);
}
if (blockRemainLength_ != 0) {
ProcessOneBlock(blockNum_, true);
}
}
__aicore__ inline void ProcessOneBlock(uint64_t blockIdx, bool isTail)
{
uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_;
CopyInInt64(blockIdx, isTail);
for (uint64_t i = 0; i < copyRow; ++i) {
int64_t linearIndex = ComputeLinearIndex(i);
if (linearIndex >= startInt64_ && linearIndex < endInt64_) {
ScatterUpdate(i, linearIndex);
}
}
}
__aicore__ inline void CopyInInt64(uint64_t blockIdx, bool isTail)
{
uint64_t indicesOffset = blockIdx * blockLength_ * indexDim_;
uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_;
DataCopyExtParams copyParams{1, static_cast<uint32_t>(copyRow * indexDim_ * sizeof(int64_t)), 0, 0, 0};
DataCopyPadExtParams<int64_t> padParams{true, 0, 0, 0};
DataCopyPad(indicesInt64Local, indicesGmInt64_[indicesOffset], copyParams, padParams);
PipeMte2ToS();
}
__aicore__ inline int64_t ComputeLinearIndex(uint64_t rowIdx)
{
int64_t linearIndex = 0;
for (uint64_t dim = 0; dim < indexDim_; ++dim) {
int64_t idxValue = indicesInt64Local.GetValue(rowIdx * indexDim_ + dim);
int64_t stride = static_cast<int64_t>(indicesMask_[dim]);
linearIndex += idxValue * stride;
}
return linearIndex;
}
__aicore__ inline void ScatterUpdate(uint64_t rowIdx, int64_t linearIndex)
{
for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) {
uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_;
uint64_t gmOffset = rowIdx * scatterLength_ + tileIdx * scatterTileLength_;
DataCopyExtParams updateCopyParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
DataCopyPad(updateLocal, updatesGm_[gmOffset], updateCopyParams, padParams);
PipeMte2ToS();
uint64_t outOffset = static_cast<uint64_t>(linearIndex) + tileIdx * scatterTileLength_;
DataCopyExtParams outParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
DataCopyPad(outputGm_[outOffset], updateLocal, outParams);
PipeMte3ToS();
}
}
private:
GlobalTensor<int64_t> indicesGmInt64_;
GlobalTensor<T> updatesGm_;
GlobalTensor<T> outputGm_;
TBuf<TPosition::VECCALC> indicesBuf;
TBuf<TPosition::VECCALC> updateBuf;
LocalTensor<int64_t> indicesInt64Local;
LocalTensor<T> updateLocal;
uint64_t blockIdx_;
uint64_t computeRow_;
uint64_t start_;
uint64_t end_;
int64_t startInt64_;
int64_t endInt64_;
uint64_t indexDim_;
uint64_t blockLength_;
uint64_t blockNum_;
uint64_t blockRemainLength_;
uint64_t indicesMask_[8];
uint64_t scatterLength_;
uint64_t ubLengthForUpdates_;
uint64_t scatterTileNum_;
uint64_t scatterTileLength_;
uint64_t scatterTileTail_;
};
} // namespace ScatterNdUpdateV2

300
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_linear_index.h Normal file
View File

@@ -0,0 +1,300 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_linear_index.h
* \brief LinearIndex Kernel
*/
#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "scatter_nd_update_common.h"
namespace ScatterNdUpdateV2 {
template<bool isSort, typename IndicesT = int>
class LinearIndexKernel {
public:
__aicore__ inline LinearIndexKernel() = delete;
__aicore__ inline LinearIndexKernel(
GM_ADDR indices, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe)
{
InitParams(tiling);
InitBuffers(pipe);
SetGmAddr(indices, workSpace, tiling);
}
__aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling)
{
blockIdx_ = GetBlockIdx();
frontCoreNum_ = tiling.linearIndexTiling.frontCoreNum;
tailCoreNum_ = tiling.linearIndexTiling.tailCoreNum;
frontBlockNum_ = tiling.linearIndexTiling.frontBlockNum;
tailBlockNum_ = tiling.linearIndexTiling.tailBlockNum;
if (blockIdx_ >= frontCoreNum_) {
computeNum_ = tailBlockNum_;
} else {
computeNum_ = frontBlockNum_;
}
ubSize_ = tiling.linearIndexTiling.ubSize;
coreNum_ = tiling.linearIndexTiling.coreNum;
blockNum_ = tiling.linearIndexTiling.blockNum;
blockLength_ = tiling.linearIndexTiling.blockLength;
blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength;
indexDim_ = tiling.linearIndexTiling.indexDim;
indicesMask_ = tiling.linearIndexTiling.indicesMask;
}
template<bool isInt64, bool needSort>
__aicore__ inline void InitBuffersUnified()
{
uint64_t offset = 0;
indicesLocal = allUbLocal[offset];
offset += blockLength_;
uint64_t indicesOffset = offset;
if constexpr (isInt64) {
indicesInt64Local = allUbLocal[offset].ReinterpretCast<int64_t>();
indicesOriginLocal = allUbLocal[offset];
offset += blockLength_ * indexDim_ * 2;
} else {
indicesOriginLocal = allUbLocal[offset];
offset += blockLength_ * indexDim_;
}
addTmpLocal = allUbLocal[offset];
offset += blockLength_;
rangeLocal = allUbLocal[offset];
offset += blockLength_;
if constexpr (isSort) {
resLocal = allUbLocal[indicesOffset].ReinterpretCast<float>();
indicesOffset += blockLength_ * 2;
posIdxLocal = allUbLocal[indicesOffset];
indicesOffset += blockLength_;
sortTmpLocal = allUbLocal[indicesOffset].ReinterpretCast<float>();
}
}
__aicore__ inline void InitBuffers(TPipe& pipe)
{
pipe.InitBuffer(allUbBuf, ubSize_);
allUbLocal = allUbBuf.Get<int>();
if constexpr (isSort) {
if constexpr (std::is_same_v<IndicesT, int64_t>) {
InitBuffersUnified<true, true>();
} else {
InitBuffersUnified<false, true>();
}
} else {
if constexpr (std::is_same_v<IndicesT, int64_t>) {
InitBuffersUnified<true, false>();
} else {
InitBuffersUnified<false, false>();
}
}
}
__aicore__ inline void SetGmAddr(GM_ADDR indices, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling)
{
indiceAddrOffset_ =
blockIdx_ < tiling.linearIndexTiling.frontCoreNum ?
tiling.linearIndexTiling.frontBlockNum * blockLength_ * blockIdx_ :
tiling.linearIndexTiling.frontCoreNum * tiling.linearIndexTiling.frontBlockNum * blockLength_ +
(blockIdx_ - tiling.linearIndexTiling.frontCoreNum) * tiling.linearIndexTiling.tailBlockNum *
blockLength_;
sortedIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + indiceAddrOffset_);
if constexpr (isSort) {
posIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + tiling.linearIndexTiling.sortWorkspace + indiceAddrOffset_);
}
if constexpr (std::is_same_v<IndicesT, int64_t>) {
indicesGmInt64_.SetGlobalBuffer((__gm__ int64_t*)indices + indiceAddrOffset_ * indexDim_);
} else {
indicesGm_.SetGlobalBuffer((__gm__ int*)indices + indiceAddrOffset_ * indexDim_);
}
}
__aicore__ inline void Process()
{
if constexpr (isSort) {
for (uint64_t i = 0; i < computeNum_; i++) {
ProcessOneWithSort(i, false);
}
uint64_t lastActiveCore = (blockNum_ == 0) ? 0 :
(tailCoreNum_ == 0 ? frontCoreNum_ - 1 : frontCoreNum_ + tailCoreNum_ - 1);
if (blockIdx_ == lastActiveCore && blockRemainLength_ != 0) {
ProcessOneWithSort(computeNum_, true);
}
} else {
for (uint64_t i = 0; i < computeNum_; i++) {
ProcessOne(i, false);
}
uint64_t lastActiveCore = (blockNum_ == 0) ? 0 :
(tailCoreNum_ == 0 ? frontCoreNum_ - 1 : frontCoreNum_ + tailCoreNum_ - 1);
if (blockIdx_ == lastActiveCore && blockRemainLength_ != 0) {
ProcessOne(computeNum_, true);
}
}
}
__aicore__ inline void ProcessOne(uint64_t idx, bool isTail)
{
CopyIn(idx, isTail);
if constexpr (std::is_same_v<IndicesT, int64_t>) {
CastToInt32(idx, isTail);
}
Compute4LinearIndex(idx, isTail);
CopyOut(idx, isTail);
}
__aicore__ inline void ProcessOneWithSort(uint64_t idx, bool isTail)
{
CopyIn(idx, isTail);
if constexpr (std::is_same_v<IndicesT, int64_t>) {
CastToInt32(idx, isTail);
}
Compute4LinearIndex(idx, isTail);
ComputeForSort(idx, isTail);
CopyOut(idx, isTail);
}
__aicore__ inline void CopyIn(uint64_t process, bool isTail)
{
uint64_t indicesOffset = process * blockLength_ * indexDim_;
uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_;
if constexpr (std::is_same_v<IndicesT, int64_t>) {
DataCopyExtParams copyParams{1, static_cast<uint32_t>(copyRow * indexDim_ * sizeof(int64_t)), 0, 0, 0};
DataCopyPadExtParams<int64_t> padParams{true, 0, 0, 0};
DataCopyPad(indicesInt64Local, indicesGmInt64_[indicesOffset], copyParams, padParams);
} else {
DataCopyExtParams copyParams{1, static_cast<uint32_t>(copyRow * indexDim_ * sizeof(int)), 0, 0, 0};
DataCopyPadExtParams<int> padParams{true, 0, 0, 0};
DataCopyPad(indicesOriginLocal, indicesGm_[indicesOffset], copyParams, padParams);
}
PipeMte2ToS();
}
__aicore__ inline void CastToInt32(uint64_t process, bool isTail)
{
uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_;
uint64_t totalElements = computeRow * indexDim_;
Cast(indicesOriginLocal, indicesInt64Local, RoundMode::CAST_NONE, totalElements);
PipeBarrier<PIPE_V>();
}
__aicore__ inline void Compute4LinearIndex(uint64_t process, bool isTail)
{
uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_;
int32_t malValue = indexDim_ * sizeof(int);
Duplicate<int>(indicesLocal, 0, computeRow);
CreateVecIndex(rangeLocal, (int)0, computeRow);
PipeBarrier<PIPE_V>();
Muls(rangeLocal, rangeLocal, malValue, computeRow);
PipeBarrier<PIPE_V>();
for (int i = 0; i < indexDim_; ++i) {
if (i != 0) {
Adds(rangeLocal, rangeLocal, (int)(sizeof(int)), computeRow);
PipeBarrier<PIPE_V>();
}
LocalTensor<uint32_t> rangeLocalCasted = rangeLocal.ReinterpretCast<uint32_t>();
Gather(addTmpLocal, indicesOriginLocal, rangeLocalCasted, (uint32_t)0, (uint32_t)computeRow);
PipeBarrier<PIPE_V>();
Muls(addTmpLocal, addTmpLocal, (int)indicesMask_[i], computeRow);
PipeBarrier<PIPE_V>();
Add(indicesLocal, indicesLocal, addTmpLocal, computeRow);
PipeBarrier<PIPE_V>();
}
if constexpr (!isSort) {
PipeVToMte3();
}
}
__aicore__ inline void ComputeForSort(uint64_t process, bool isTail)
{
LocalTensor<float> indicesLocalFp32 = indicesLocal.ReinterpretCast<float>();
uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_;
uint64_t computeRowAligned = (computeRow + ALIGNED_BLOCK_NUM - 1) & ~(ALIGNED_BLOCK_NUM - 1);
uint64_t repeatTimes = computeRowAligned / ALIGNED_BLOCK_NUM;
uint64_t repeatId = computeRow / ALIGNED_BLOCK_NUM;
uint64_t repeatRemain = computeRow % ALIGNED_BLOCK_NUM;
int addValue = indiceAddrOffset_ + process * blockLength_;
Cast(indicesLocalFp32, indicesLocal, RoundMode::CAST_ROUND, computeRowAligned);
if (repeatRemain != 0) {
// 对齐处理：不足32的部分设为-1
Duplicate<int>(rangeLocal, -1, (uint32_t)ALIGNED_BLOCK_NUM);
PipeBarrier<PIPE_V>();
Cast(rangeLocal, indicesLocalFp32[ALIGNED_BLOCK_NUM * repeatId], RoundMode::CAST_ROUND, (uint32_t)repeatRemain);
PipeBarrier<PIPE_V>();
Cast(indicesLocalFp32[ALIGNED_BLOCK_NUM * repeatId], rangeLocal, RoundMode::CAST_ROUND, (uint32_t)ALIGNED_BLOCK_NUM);
PipeBarrier<PIPE_V>();
}
Duplicate<int>(posIdxLocal, -1, computeRowAligned);
PipeBarrier<PIPE_V>();
CreateVecIndex<int>(posIdxLocal, 0U, computeRow);
LocalTensor<uint32_t> posIdxULocal = posIdxLocal.ReinterpretCast<uint32_t>();
PipeBarrier<PIPE_V>();
Sort<float, true>(resLocal, indicesLocalFp32, posIdxULocal, sortTmpLocal, repeatTimes);
PipeBarrier<PIPE_V>();
Extract(indicesLocalFp32, posIdxULocal, resLocal, repeatTimes);
PipeBarrier<PIPE_V>();
Cast(indicesLocal, indicesLocalFp32, RoundMode::CAST_ROUND, computeRowAligned);
PipeBarrier<PIPE_V>();
Adds(posIdxLocal, posIdxLocal, addValue, computeRow);
PipeBarrier<PIPE_V>();
PipeVToMte3();
}
__aicore__ inline void CopyOut(uint64_t process, bool isTail)
{
uint64_t outOffset = process * blockLength_;
uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_;
DataCopyExtParams copyParams{1, static_cast<uint32_t>(copyRow * sizeof(int)), 0, 0, 0};
DataCopyPad(sortedIndicesGm_[outOffset], indicesLocal, copyParams);
if constexpr (isSort) {
DataCopyPad(posIndicesGm_[outOffset], posIdxLocal, copyParams);
}
PipeMte3ToS();
}
private:
GlobalTensor<int> indicesGm_;
GlobalTensor<int64_t> indicesGmInt64_;
GlobalTensor<int> sortedIndicesGm_;
GlobalTensor<int> posIndicesGm_;
TBuf<TPosition::VECCALC> allUbBuf;
LocalTensor<int> allUbLocal;
LocalTensor<int> indicesLocal;
LocalTensor<int> indicesOriginLocal;
LocalTensor<int64_t> indicesInt64Local;
LocalTensor<int> addTmpLocal;
LocalTensor<int> rangeLocal;
LocalTensor<int> posIdxLocal;
LocalTensor<float> sortTmpLocal;
LocalTensor<float> resLocal;
uint64_t ubSize_;
uint64_t coreNum_;
uint64_t blockIdx_;
uint64_t indexDim_;
uint64_t computeNum_;
uint64_t blockNum_;
uint64_t blockLength_;
uint64_t blockRemainLength_;
uint64_t frontCoreNum_;
uint64_t tailCoreNum_;
uint64_t frontBlockNum_;
uint64_t tailBlockNum_;
const uint64_t* indicesMask_;
uint64_t indiceAddrOffset_;
};
} // namespace ScatterNdUpdateV2

152
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_no_sort.h Normal file
View File

@@ -0,0 +1,152 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_no_sort.h
* \brief Scatter Kernel (NoSort)
*/
#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "scatter_nd_update_common.h"
namespace ScatterNdUpdateV2 {
constexpr uint64_t ALIGNED_SIZE_INDEX = 8;
template<typename T>
class ScatterNdUpdateV2KernelNoSort {
public:
__aicore__ inline ScatterNdUpdateV2KernelNoSort() = delete;
__aicore__ inline ScatterNdUpdateV2KernelNoSort(
GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe)
{
InitParam(tiling);
InitBuffers(pipe);
SetGmAddr(updates, output, workSpace, tiling);
}
__aicore__ inline void InitParam(const ScatterNdUpdateV2TilingData& tiling)
{
blockIdx_ = GetBlockIdx();
CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow,
tiling.scatterTiling.tailRow, computeRow_, start_);
end_ = start_ + computeRow_;
totalIndexRow_ = tiling.linearIndexTiling.blockNum * tiling.linearIndexTiling.blockLength
+ tiling.linearIndexTiling.blockRemainLength;
scatterLength_ = tiling.scatterTiling.scatterLength;
scatterAlignLength_ = tiling.scatterTiling.scatterAlignLength;
ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates;
scatterTileNum_ = tiling.scatterTiling.scatterTileNum;
scatterTileLength_ = tiling.scatterTiling.scatterTileLength;
scatterTileTail_ = tiling.scatterTiling.scatterTileTail;
scatterTileAlignLength_ = tiling.scatterTiling.scatterTileAlignLength;
CalcIndexTileParams();
}
__aicore__ inline void CalcIndexTileParams()
{
uint64_t ubSizeBytes = ubLengthForUpdates_ * sizeof(T);
uint64_t updateSizeBytes = scatterTileLength_ * sizeof(T);
uint64_t remainBytes = ubSizeBytes - updateSizeBytes;
indexTileLength_ = (remainBytes / sizeof(int) / ALIGNED_SIZE_INDEX) * ALIGNED_SIZE_INDEX;
if (indexTileLength_ == 0) {
indexTileLength_ = ALIGNED_SIZE_INDEX;
}
indexTileNum_ = (totalIndexRow_ + indexTileLength_ - 1) / indexTileLength_;
indexTileTail_ = totalIndexRow_ - (indexTileNum_ - 1) * indexTileLength_;
}
__aicore__ inline void InitBuffers(TPipe& pipe)
{
pipe.InitBuffer(indexQue_, DOUBLE_BUFFER, indexTileLength_ * sizeof(int));
pipe.InitBuffer(updateQue_, DOUBLE_BUFFER, scatterTileLength_ * sizeof(T));
}
__aicore__ inline void SetGmAddr(GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling)
{
linearIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace);
updatesGm_.SetGlobalBuffer((__gm__ T*)updates);
outputGm_.SetGlobalBuffer((__gm__ T*)output);
}
__aicore__ inline void Process()
{
for (uint64_t tileIdx = 0; tileIdx < indexTileNum_; ++tileIdx) {
uint64_t curTileLen = (tileIdx == indexTileNum_ - 1) ? indexTileTail_ : indexTileLength_;
uint64_t gmOffset = tileIdx * indexTileLength_;
LocalTensor<int> indexLocal = indexQue_.AllocTensor<int>();
DataCopyExtParams indexCopyParams{1, static_cast<uint32_t>(curTileLen * sizeof(int)), 0, 0, 0};
DataCopyPadExtParams<int> padParams{true, 0, 0, 0};
DataCopyPad(indexLocal, linearIndicesGm_[gmOffset], indexCopyParams, padParams);
indexQue_.EnQue(indexLocal);
PipeMte2ToS();
LocalTensor<int> indexData = indexQue_.DeQue<int>();
for (uint64_t i = 0; i < curTileLen; ++i) {
int64_t linearIndex = static_cast<int64_t>(indexData.GetValue(i));
if (linearIndex >= (int64_t)start_ && linearIndex < (int64_t)end_) {
uint64_t idx = gmOffset + i;
ProcessOneIndex(idx, linearIndex);
}
}
indexQue_.FreeTensor<int>(indexLocal);
}
}
__aicore__ inline void ProcessOneIndex(uint64_t idx, int64_t linearIndex)
{
for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) {
uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_;
LocalTensor<T> updateLocal = updateQue_.AllocTensor<T>();
uint64_t gmOffset = idx * scatterLength_ + tileIdx * scatterTileLength_;
DataCopyExtParams updateCopyParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
DataCopyPad(updateLocal, updatesGm_[gmOffset], updateCopyParams, padParams);
PipeMte2ToS();
uint64_t outOffset = linearIndex + tileIdx * scatterTileLength_;
DataCopyExtParams outParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
DataCopyPad(outputGm_[outOffset], updateLocal, outParams);
PipeMte3ToS();
updateQue_.FreeTensor<T>(updateLocal);
}
}
private:
GlobalTensor<int> linearIndicesGm_;
GlobalTensor<T> updatesGm_;
GlobalTensor<T> outputGm_;
TQue<TPosition::VECOUT, DOUBLE_BUFFER> indexQue_;
TQue<TPosition::VECOUT, DOUBLE_BUFFER> updateQue_;
uint64_t blockIdx_;
uint64_t computeRow_;
uint64_t start_;
uint64_t end_;
uint64_t totalIndexRow_;
uint64_t scatterLength_;
uint64_t scatterAlignLength_;
uint64_t ubLengthForUpdates_;
uint64_t scatterTileNum_;
uint64_t scatterTileLength_;
uint64_t scatterTileTail_;
uint64_t scatterTileAlignLength_;
uint64_t indexTileLength_;
uint64_t indexTileNum_;
uint64_t indexTileTail_;
};
} // namespace ScatterNdUpdateV2

71
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.cpp Normal file
View File

@@ -0,0 +1,71 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_v2.cpp
* \brief ScatterNdUpdateV2 算子入口
*/
#include "scatter_nd_update_v2.h"
#include "scatter_nd_update_linear_index.h"
#include "scatter_nd_update_no_sort.h"
#include "scatter_nd_update_large_index.h"
extern "C" __global__ __aicore__ void scatter_nd_update_v2(GM_ADDR varRef, GM_ADDR indices,
GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, GM_ADDR tiling) {
if (workSpace == nullptr) {
return;
}
GM_ADDR user = AscendC::GetUserWorkspace(workSpace);
if (user == nullptr) {
return;
}
GET_TILING_DATA(tilingData, tiling);
AscendC::TPipe tpipe;
#if (defined(DTYPE_VAR))
// tilingKey: indexType * 10 + sortFlag
// indexType: 1=int32, 2=int64(cast), 3=int64(large); sortFlag: 0=非排序, 1=排序
if (TILING_KEY_IS(11)) {
ScatterNdUpdateV2::LinearIndexKernel<true, int> op1(indices, workSpace, tilingData, tpipe);
op1.Process();
AscendC::SyncAll();
tpipe.Destroy();
AscendC::TPipe pipe;
ScatterNdUpdateV2::ScatterNdUpdateV2Kernel<DTYPE_VAR> op2(updates, output, workSpace, tilingData, pipe);
op2.Process();
} else if (TILING_KEY_IS(10)) {
ScatterNdUpdateV2::LinearIndexKernel<false, int> op1(indices, workSpace, tilingData, tpipe);
op1.Process();
AscendC::SyncAll();
tpipe.Destroy();
AscendC::TPipe pipe;
ScatterNdUpdateV2::ScatterNdUpdateV2KernelNoSort<DTYPE_VAR> op2(updates, output, workSpace, tilingData, pipe);
op2.Process();
} else if (TILING_KEY_IS(21)) {
ScatterNdUpdateV2::LinearIndexKernel<true, int64_t> op1(indices, workSpace, tilingData, tpipe);
op1.Process();
AscendC::SyncAll();
tpipe.Destroy();
AscendC::TPipe pipe;
ScatterNdUpdateV2::ScatterNdUpdateV2Kernel<DTYPE_VAR> op2(updates, output, workSpace, tilingData, pipe);
op2.Process();
} else if (TILING_KEY_IS(20)) {
ScatterNdUpdateV2::LinearIndexKernel<false, int64_t> op1(indices, workSpace, tilingData, tpipe);
op1.Process();
AscendC::SyncAll();
tpipe.Destroy();
AscendC::TPipe pipe;
ScatterNdUpdateV2::ScatterNdUpdateV2KernelNoSort<DTYPE_VAR> op2(updates, output, workSpace, tilingData, pipe);
op2.Process();
} else if (TILING_KEY_IS(30)) {
ScatterNdUpdateV2::LargeIndexKernel<DTYPE_VAR> op(indices, updates, output, tilingData, tpipe);
op.Process();
}
#endif
}

249
csrc/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.h Normal file
View File

@@ -0,0 +1,249 @@
/**
* Copyright (c) 2025-2026 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 scatter_nd_update_v2.h
* \brief Scatter Kernel (Sort)
*/
#include "kernel_operator.h"
#include "kernel_tiling/kernel_tiling.h"
#include "scatter_nd_update_common.h"
namespace ScatterNdUpdateV2 {
template<typename T>
class ScatterNdUpdateV2Kernel {
public:
__aicore__ inline ScatterNdUpdateV2Kernel() = delete;
__aicore__ inline ScatterNdUpdateV2Kernel(
GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe)
{
InitParams(tiling);
InitBuffers(pipe);
SetGmAddr(updates, output, workSpace, tiling);
}
__aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling)
{
blockIdx_ = GetBlockIdx();
CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow,
tiling.scatterTiling.tailRow, computeRow_, start_);
end_ = start_ + computeRow_;
blockNum_ = tiling.linearIndexTiling.blockNum;
blockLength_ = tiling.linearIndexTiling.blockLength;
blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength;
coreNum_ = tiling.linearIndexTiling.coreNum;
scatterLength_ = tiling.scatterTiling.scatterLength;
scatterAlignLength_ = tiling.scatterTiling.scatterAlignLength;
ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates;
formDim_ = tiling.scatterTiling.formDim;
copyRow_ = tiling.scatterTiling.copyRow;
scatterTileNum_ = tiling.scatterTiling.scatterTileNum;
scatterTileLength_ = tiling.scatterTiling.scatterTileLength;
scatterTileTail_ = tiling.scatterTiling.scatterTileTail;
scatterTileAlignLength_ = tiling.scatterTiling.scatterTileAlignLength;
}
__aicore__ inline void InitBuffers(TPipe& pipe)
{
pipe.InitBuffer(indiceQue_, DOUBLE_BUFFER, blockLength_ * sizeof(int));
pipe.InitBuffer(posIdxQue_, DOUBLE_BUFFER, blockLength_ * sizeof(int));
pipe.InitBuffer(updateQue_, DOUBLE_BUFFER, ubLengthForUpdates_ * sizeof(T));
}
__aicore__ inline void SetGmAddr(GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling)
{
sortedIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace);
posIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + tiling.linearIndexTiling.sortWorkspace);
updatesGm_.SetGlobalBuffer((__gm__ T*)updates);
outputGm_.SetGlobalBuffer((__gm__ T*)output);
}
__aicore__ inline void Process()
{
for (uint64_t i = 0; i < blockNum_; ++i) {
CopyIndicesIn(i, false);
Compute(i, false);
PipeMte3ToS();
}
if (blockRemainLength_ != 0) {
CopyIndicesIn(blockNum_, true);
Compute(blockNum_, true);
PipeMte3ToS();
}
}
__aicore__ inline void CopyIndicesIn(uint64_t process, bool isTail)
{
uint64_t copyNum = isTail ? blockRemainLength_ : blockLength_;
LocalTensor<int> indiceLocal = indiceQue_.AllocTensor<int>();
LocalTensor<int> posIdxLocal = posIdxQue_.AllocTensor<int>();
uint64_t indicesOffset = isTail ? (blockNum_ * blockLength_) : (process * blockLength_);
DataCopyParams indiceCopyParams{1, static_cast<uint16_t>(copyNum * sizeof(int)), 0, 0};
DataCopyPadParams padParams{true, 0, 0, 0};
DataCopyPad(indiceLocal, sortedIndicesGm_[indicesOffset], indiceCopyParams, padParams);
DataCopyPad(posIdxLocal, posIndicesGm_[indicesOffset], indiceCopyParams, padParams);
PipeMte2ToS();
PipeBarrier<PIPE_V>();
UpdateSearchParam(indiceLocal, isTail);
indiceQue_.EnQue<int>(indiceLocal);
posIdxQue_.EnQue<int>(posIdxLocal);
}
__aicore__ inline void CopyUpdateIn(LocalTensor<T> &updateLocal, uint64_t gmIdx, uint64_t ubIdx, uint64_t tileIdx, uint64_t tileLength)
{
uint64_t gmOffset = gmIdx * scatterLength_ + tileIdx * scatterTileLength_;
uint64_t ubOffset = ubIdx * scatterTileAlignLength_;
DataCopyExtParams updateCopyParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
DataCopyPad(updateLocal[ubOffset], updatesGm_[gmOffset], updateCopyParams, padParams);
PipeMte2ToS();
}
// 降序数组：二分查找边界
__aicore__ inline int64_t findFirstLt(LocalTensor<int> &indiceLocal, int64_t target, bool isTail)
{
int64_t left = 0;
int64_t right = (isTail ? blockRemainLength_ : blockLength_) - 1;
int64_t res = isTail ? blockRemainLength_ : blockLength_;
while (left <= right) {
int64_t mid = left + (right - left) / 2;
int64_t value = indiceLocal.GetValue(mid);
if (value < target) {
res = mid;
right = mid - 1;
} else {
left = mid + 1;
}
}
return res;
}
__aicore__ inline int64_t findLastGe(LocalTensor<int> &indiceLocal, int64_t target, bool isTail)
{
int64_t left = 0;
int64_t right = (isTail ? blockRemainLength_ : blockLength_) - 1;
int64_t res = -1;
while (left <= right) {
int64_t mid = left + (right - left) / 2;
int64_t value = indiceLocal.GetValue(mid);
if (value >= target) {
res = mid;
left = mid + 1;
} else {
right = mid - 1;
}
}
return res;
}
__aicore__ inline void UpdateSearchParam(LocalTensor<int> &indiceLocal, bool isTail)
{
int64_t searchNum = isTail ? blockRemainLength_ : blockLength_;
leftBound_ = findFirstLt(indiceLocal, end_, isTail);
rightBound_ = findLastGe(indiceLocal, start_, isTail);
isValidBound_ = (leftBound_ < searchNum && rightBound_ != -1 && leftBound_ <= rightBound_);
}
__aicore__ inline void Compute(uint64_t process, bool isTail)
{
LocalTensor<int> indiceLocal = indiceQue_.DeQue<int>();
LocalTensor<int> posIdxLocal = posIdxQue_.DeQue<int>();
if (!isValidBound_) {
indiceQue_.FreeTensor<int>(indiceLocal);
posIdxQue_.FreeTensor<int>(posIdxLocal);
return;
}
for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) {
uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_;
uint64_t inUbNum = 0;
LocalTensor<T> updateLocal;
lastProcessedIdx_ = -1;
for (int64_t i = rightBound_; i >= leftBound_; --i) {
if (inUbNum == 0) {
updateLocal = updateQue_.AllocTensor<T>();
}
int64_t posIdx = posIdxLocal.GetValue(i);
CopyUpdateIn(updateLocal, posIdx, inUbNum, tileIdx, tileLength);
inUbNum++;
if (inUbNum == copyRow_) {
updateQue_.EnQue<T>(updateLocal);
CopyOut(inUbNum, i, indiceLocal, posIdxLocal, tileIdx, tileLength);
inUbNum = 0;
}
if (i == leftBound_ && inUbNum != 0) {
updateQue_.EnQue<T>(updateLocal);
CopyOut(inUbNum, i, indiceLocal, posIdxLocal, tileIdx, tileLength);
}
}
}
indiceQue_.FreeTensor<int>(indiceLocal);
posIdxQue_.FreeTensor<int>(posIdxLocal);
}
__aicore__ inline void CopyOut(uint64_t inUbNum, int64_t curIdx, LocalTensor<int> &indiceLocal,
LocalTensor<int> &posIdxLocal, uint64_t tileIdx, uint64_t tileLength)
{
LocalTensor<T> updateLocal = updateQue_.DeQue<T>();
DataCopyExtParams outParams{1, static_cast<uint32_t>(tileLength * sizeof(T)), 0, 0, 0};
for (int64_t i = curIdx + inUbNum - 1; i >= curIdx; --i) {
int64_t curIdxValue = indiceLocal.GetValue(i);
if (curIdxValue == lastProcessedIdx_) continue;
lastProcessedIdx_ = curIdxValue;
uint64_t outOffset = curIdxValue + tileIdx * scatterTileLength_;
uint64_t updateOffset = (curIdx + inUbNum - 1 - i) * scatterTileAlignLength_;
DataCopyPad(outputGm_[outOffset], updateLocal[updateOffset], outParams);
}
PipeMte3ToS();
updateQue_.FreeTensor<T>(updateLocal);
}
private:
GlobalTensor<int> sortedIndicesGm_;
GlobalTensor<int> posIndicesGm_;
GlobalTensor<T> updatesGm_;
GlobalTensor<T> outputGm_;
TQue<TPosition::VECIN, DOUBLE_BUFFER> indiceQue_;
TQue<TPosition::VECIN, DOUBLE_BUFFER> posIdxQue_;
TQue<TPosition::VECOUT, DOUBLE_BUFFER> updateQue_;
uint64_t blockIdx_;
uint64_t computeRow_;
uint64_t start_;
uint64_t end_;
uint64_t blockNum_;
uint64_t blockLength_;
uint64_t blockRemainLength_;
uint64_t scatterLength_;
uint64_t scatterAlignLength_;
uint64_t ubLengthForUpdates_;
uint64_t formDim_;
uint64_t copyRow_;
uint64_t coreNum_;
uint64_t scatterTileNum_;
uint64_t scatterTileLength_;
uint64_t scatterTileTail_;
uint64_t scatterTileAlignLength_;
int64_t leftBound_;
int64_t rightBound_;
bool isValidBound_;
int64_t lastProcessedIdx_;
};
} // namespace ScatterNdUpdateV2