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

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