Star843
1299
代码
代码
Issues1413
Pull Requests2265
流水线
项目讨论
Wiki
分析
Star843
1299
已合并
fix: lintrunner --all-files --take NEWLINE -a #35872
Jingwei Huang创建于 5月17日
fix: lintrunner --all-files --take NEWLINE -a #35872
已合并
Jingwei Huang创建于 5月17日
从huangjingwei/pytorchv2.9.0_lintrunner合入到Ascend/pytorchv2.9.0
共 91 个文件变更+16512-16529
MCMakeLists.txt+390-391
@@ -1,391 +1,390 @@
1-cmake_minimum_required(VERSION 3.18 FATAL_ERROR)1+cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
2- 2+ 
3-find_program(CCACHE ccache)3+find_program(CCACHE ccache)
4-if(${CCACHE} STREQUAL "CCACHE-NOTFOUND")4+if(${CCACHE} STREQUAL "CCACHE-NOTFOUND")
5- message(STATUS "Compile without ccache")5+ message(STATUS "Compile without ccache")
6-else()6+else()
7- set(CMAKE_C_COMPILER_LAUNCHER ${CCACHE} CACHE PATH "cache Compiler")7+ set(CMAKE_C_COMPILER_LAUNCHER ${CCACHE} CACHE PATH "cache Compiler")
8- set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE} CACHE PATH "cache Compiler")8+ set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE} CACHE PATH "cache Compiler")
9- message(STATUS "CMAKE_C_COMPILER_LAUNCHER:${CMAKE_C_COMPILER_LAUNCHER}")9+ message(STATUS "CMAKE_C_COMPILER_LAUNCHER:${CMAKE_C_COMPILER_LAUNCHER}")
10- message(STATUS "CMAKE_CXX_COMPILER_LAUNCHER:${CMAKE_CXX_COMPILER_LAUNCHER}")10+ message(STATUS "CMAKE_CXX_COMPILER_LAUNCHER:${CMAKE_CXX_COMPILER_LAUNCHER}")
11-endif()11+endif()
12- 12+ 
13-project(TORCHNPU CXX C)13+project(TORCHNPU CXX C)
14-add_compile_options(-fmacro-prefix-map=${CMAKE_SOURCE_DIR}/=)14+add_compile_options(-fmacro-prefix-map=${CMAKE_SOURCE_DIR}/=)
15- 15+ 
16-find_program(MOLD_LINKER mold)16+find_program(MOLD_LINKER mold)
17-if(MOLD_LINKER)17+if(MOLD_LINKER)
18- add_link_options(-fuse-ld=mold)18+ add_link_options(-fuse-ld=mold)
19- message(STATUS "Using mold linker: ${MOLD_LINKER}")19+ message(STATUS "Using mold linker: ${MOLD_LINKER}")
20-else()20+else()
21- message(STATUS "mold linker not found, using default linker")21+ message(STATUS "mold linker not found, using default linker")
22-endif()22+endif()
23- 23+ 
24-set(LINUX TRUE)24+set(LINUX TRUE)
25-set(CMAKE_INSTALL_MESSAGE NEVER)25+set(CMAKE_INSTALL_MESSAGE NEVER)
26-# set(CMAKE_VERBOSE_MAKEFILE ON)26+# set(CMAKE_VERBOSE_MAKEFILE ON)
27-set(CMAKE_EXPORT_COMPILE_COMMANDS ON)27+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
28- 28+ 
29-if(DEFINED TORCH_VERSION)29+if(DEFINED TORCH_VERSION)
30- add_definitions(-DPYTORCH_NPU_VERSION="${TORCH_VERSION}")30+ add_definitions(-DPYTORCH_NPU_VERSION="${TORCH_VERSION}")
31-endif()31+endif()
32- 32+ 
33-set(PLUGIN_NAME torch_npu)33+set(PLUGIN_NAME torch_npu)
34- 34+ 
35-set(RPATH_VALUE $ORIGIN)35+set(RPATH_VALUE $ORIGIN)
36-set(CMAKE_SKIP_BUILD_RPATH FALSE)36+set(CMAKE_SKIP_BUILD_RPATH FALSE)
37-set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)37+set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
38-set(CMAKE_INSTALL_RPATH "${RPATH_VALUE}/lib/:${RPATH_VALUE}/")38+set(CMAKE_INSTALL_RPATH "${RPATH_VALUE}/lib/:${RPATH_VALUE}/")
39-set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE)39+set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE)
40-set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${TORCHNPU_INSTALL_LIBDIR})40+set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${TORCHNPU_INSTALL_LIBDIR})
41-SET(CMAKE_CXX_FLAGS_RELWITHDEBINFO "-O2 -g")41+SET(CMAKE_CXX_FLAGS_RELWITHDEBINFO "-O2 -g")
42-SET(CMAKE_CXX_FLAGS_RELEASE "-O2")42+SET(CMAKE_CXX_FLAGS_RELEASE "-O2")
43-SET(CMAKE_CXX_FLAGS_DEBUG "-O0 -g")43+SET(CMAKE_CXX_FLAGS_DEBUG "-O0 -g")
44- 44+ 
45-# LTO&PGO optimization in compile option45+# LTO&PGO optimization in compile option
46-SET(IF_APPEND FALSE)46+SET(IF_APPEND FALSE)
47-if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")47+if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
48- SET(APPEND_FLAGS "-fGNU-compatibility -Wno-non-pod-varargs")48+ SET(APPEND_FLAGS "-fGNU-compatibility -Wno-non-pod-varargs")
49- if(NOT MOLD_LINKER)49+ if(NOT MOLD_LINKER)
50- add_link_options(-fuse-ld=lld)50+ add_link_options(-fuse-ld=lld)
51- endif()51+ endif()
52- if (DEFINED ENABLE_LTO)52+ if (DEFINED ENABLE_LTO)
53- SET(IF_APPEND TRUE)53+ SET(IF_APPEND TRUE)
54- SET(APPEND_FLAGS "${APPEND_FLAGS} -flto=thin")54+ SET(APPEND_FLAGS "${APPEND_FLAGS} -flto=thin")
55- endif()55+ endif()
56- if (DEFINED PGO_MODE)56+ if (DEFINED PGO_MODE)
57- SET(IF_APPEND TRUE)57+ SET(IF_APPEND TRUE)
58- if (PGO_MODE EQUAL 1)58+ if (PGO_MODE EQUAL 1)
59- SET(APPEND_FLAGS "${APPEND_FLAGS} -fprofile-generate")59+ SET(APPEND_FLAGS "${APPEND_FLAGS} -fprofile-generate")
60- elseif (PGO_MODE EQUAL 2)60+ elseif (PGO_MODE EQUAL 2)
61- SET(APPEND_FLAGS "${APPEND_FLAGS} -fprofile-use=${CMAKE_CURRENT_SOURCE_DIR}/default.profdata")61+ SET(APPEND_FLAGS "${APPEND_FLAGS} -fprofile-use=${CMAKE_CURRENT_SOURCE_DIR}/default.profdata")
62- endif()62+ endif()
63- endif()63+ endif()
64-else()64+else()
65- if (DEFINED ENABLE_LTO)65+ if (DEFINED ENABLE_LTO)
66- message(FATAL_ERROR "Currently, LTO auto build is not supported in ${CMAKE_CXX_COMPILER_ID}")66+ message(FATAL_ERROR "Currently, LTO auto build is not supported in ${CMAKE_CXX_COMPILER_ID}")
67- endif()67+ endif()
68- if (DEFINED PGO_MODE)68+ if (DEFINED PGO_MODE)
69- message(FATAL_ERROR "Currently, PGO auto build is not supported in ${CMAKE_CXX_COMPILER_ID}")69+ message(FATAL_ERROR "Currently, PGO auto build is not supported in ${CMAKE_CXX_COMPILER_ID}")
70- endif()70+ endif()
71-endif()71+endif()
72-if (IF_APPEND)72+if (IF_APPEND)
73- SET(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} ${APPEND_FLAGS}")73+ SET(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} ${APPEND_FLAGS}")
74-endif()74+endif()
75- 75+ 
76-# check and set CMAKE_CXX_STANDARD76+# check and set CMAKE_CXX_STANDARD
77-string(FIND "${CMAKE_CXX_FLAGS}" "-std=c++" env_cxx_standard)77+string(FIND "${CMAKE_CXX_FLAGS}" "-std=c++" env_cxx_standard)
78-if(env_cxx_standard GREATER -1)78+if(env_cxx_standard GREATER -1)
79- message(79+ message(
80- WARNING "C++ standard version definition detected in environment variable."80+ WARNING "C++ standard version definition detected in environment variable."
81- "PyTorch requires -std=c++17. Please remove -std=c++ settings in your environment.")81+ "PyTorch requires -std=c++17. Please remove -std=c++ settings in your environment.")
82-endif()82+endif()
83-set(CMAKE_CXX_STANDARD 17)83+set(CMAKE_CXX_STANDARD 17)
84-set(CMAKE_C_STANDARD 11)84+set(CMAKE_C_STANDARD 11)
85-set(CMAKE_CXX_EXTENSIONS OFF)85+set(CMAKE_CXX_EXTENSIONS OFF)
86- 86+ 
87-set(TORCHNPU_ROOT "${PROJECT_SOURCE_DIR}/torch_npu/csrc")87+set(TORCHNPU_ROOT "${PROJECT_SOURCE_DIR}/torch_npu/csrc")
88-set(TORCHNPU_THIRD_PARTY_ROOT "${PROJECT_SOURCE_DIR}/third_party")88+set(TORCHNPU_THIRD_PARTY_ROOT "${PROJECT_SOURCE_DIR}/third_party")
89- 89+ 
90-set(Torch_DIR ${PYTORCH_INSTALL_DIR}/share/cmake/Torch)90+set(Torch_DIR ${PYTORCH_INSTALL_DIR}/share/cmake/Torch)
91-FIND_PACKAGE(Torch REQUIRED)91+FIND_PACKAGE(Torch REQUIRED)
92- 92+ 
93-set(LINUX TRUE)93+set(LINUX TRUE)
94-set(CMAKE_INSTALL_MESSAGE NEVER)94+set(CMAKE_INSTALL_MESSAGE NEVER)
95-#set(CMAKE_VERBOSE_MAKEFILE ON)95+#set(CMAKE_VERBOSE_MAKEFILE ON)
96-set(CMAKE_EXPORT_COMPILE_COMMANDS ON)96+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
97- 97+ 
98-# Define build type98+# Define build type
99-IF(CMAKE_BUILD_TYPE MATCHES Debug)99+IF(CMAKE_BUILD_TYPE MATCHES Debug)
100- message("Debug build.")100+ message("Debug build.")
101- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_DEBUG")101+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_DEBUG")
102-ELSEIF(CMAKE_BUILD_TYPE MATCHES RelWithDebInfo)102+ELSEIF(CMAKE_BUILD_TYPE MATCHES RelWithDebInfo)
103- message("RelWithDebInfo build")103+ message("RelWithDebInfo build")
104- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNDEBUG")104+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNDEBUG")
105-ELSE()105+ELSE()
106- message("Release build.")106+ message("Release build.")
107- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNDEBUG")107+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNDEBUG")
108-ENDIF()108+ENDIF()
109- 109+ 
110-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC")110+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC")
111-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-narrowing")111+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-narrowing")
112-# Eigen fails to build with some versions, so convert this to a warning112+# Eigen fails to build with some versions, so convert this to a warning
113-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall")113+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall")
114-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wextra")114+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wextra")
115-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-field-initializers")115+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-field-initializers")
116-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-type-limits")116+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-type-limits")
117-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-array-bounds")117+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-array-bounds")
118-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-pragmas")118+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-pragmas")
119-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-sign-compare")119+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-sign-compare")
120-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter")120+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter")
121-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-variable")121+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-variable")
122-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-function")122+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-function")
123-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")123+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")
124-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-overflow")124+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-overflow")
125-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-aliasing")125+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-aliasing")
126-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=deprecated-declarations")126+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=deprecated-declarations")
127-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-ignored-qualifiers")127+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-ignored-qualifiers")
128-if (CMAKE_COMPILER_IS_GNUCXX AND NOT (CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.0.0))128+if (CMAKE_COMPILER_IS_GNUCXX AND NOT (CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.0.0))
129- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-stringop-overflow")129+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-stringop-overflow")
130-endif()130+endif()
131-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=pedantic")131+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=pedantic")
132-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=redundant-decls")132+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=redundant-decls")
133-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=old-style-cast")133+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=old-style-cast")
134- 134+ 
135-# These flags are not available in GCC-4.8.5. Set only when using clang.135+# These flags are not available in GCC-4.8.5. Set only when using clang.
136-if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")136+if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
137- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-invalid-partial-specialization")137+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-invalid-partial-specialization")
138- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-typedef-redefinition")138+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-typedef-redefinition")
139- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-warning-option")139+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-warning-option")
140- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-private-field")140+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-private-field")
141- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-inconsistent-missing-override")141+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-inconsistent-missing-override")
142- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-aligned-allocation-unavailable")142+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-aligned-allocation-unavailable")
143- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++17-extensions")143+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++17-extensions")
144- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-constexpr-not-const")144+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-constexpr-not-const")
145- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-braces")145+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-braces")
146- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Qunused-arguments")146+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Qunused-arguments")
147- if (${COLORIZE_OUTPUT})147+ if (${COLORIZE_OUTPUT})
148- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fcolor-diagnostics")148+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fcolor-diagnostics")
149- endif()149+ endif()
150-endif()150+endif()
151-if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 4.9)151+if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 4.9)
152- if (${COLORIZE_OUTPUT})152+ if (${COLORIZE_OUTPUT})
153- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fdiagnostics-color=always")153+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fdiagnostics-color=always")
154- endif()154+ endif()
155-endif()155+endif()
156-if ((APPLE AND (NOT ("${CLANG_VERSION_STRING}" VERSION_LESS "9.0")))156+if ((APPLE AND (NOT ("${CLANG_VERSION_STRING}" VERSION_LESS "9.0")))
157- OR (CMAKE_COMPILER_IS_GNUCXX157+ OR (CMAKE_COMPILER_IS_GNUCXX
158- AND (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0 AND NOT APPLE)))158+ AND (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0 AND NOT APPLE)))
159- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -faligned-new")159+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -faligned-new")
160-endif()160+endif()
161-if (WERROR)161+if (WERROR)
162- check_cxx_compiler_flag("-Werror" COMPILER_SUPPORT_WERROR)162+ check_cxx_compiler_flag("-Werror" COMPILER_SUPPORT_WERROR)
163- if (NOT COMPILER_SUPPORT_WERROR)163+ if (NOT COMPILER_SUPPORT_WERROR)
164- set(WERROR FALSE)164+ set(WERROR FALSE)
165- else()165+ else()
166- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror")166+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror")
167- endif()167+ endif()
168-endif(WERROR)168+endif(WERROR)
169-if (NOT APPLE)169+if (NOT APPLE)
170- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-but-set-variable")170+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-but-set-variable")
171- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-uninitialized")171+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-uninitialized")
172-endif()172+endif()
173-set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -O0")173+set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -O0")
174-set(CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_STATIC_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -O0")174+set(CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_STATIC_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -O0")
175-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-math-errno")175+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-math-errno")
176-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-trapping-math")176+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-trapping-math")
177-set(CMAKE_CXX_COVERAGE $ENV{CALCULATE_CXX_COVERAGE})177+set(CMAKE_CXX_COVERAGE $ENV{CALCULATE_CXX_COVERAGE})
178- 178+ 
179-if (CMAKE_BUILD_TYPE MATCHES Debug)179+if (CMAKE_BUILD_TYPE MATCHES Debug)
180- set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_C_FLAGS}")180+ set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_C_FLAGS}")
181- set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_CXX_FLAGS}")181+ set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_CXX_FLAGS}")
182- set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CXXFLAGS}")182+ set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CXXFLAGS}")
183-elseif (CMAKE_CXX_COVERAGE STREQUAL "1")183+elseif (CMAKE_CXX_COVERAGE STREQUAL "1")
184- set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CMAKE_C_FLAGS}")184+ set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CMAKE_C_FLAGS}")
185- set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CMAKE_CXX_FLAGS}")185+ set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CMAKE_CXX_FLAGS}")
186- set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CXXFLAGS}")186+ set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fprofile-arcs -ftest-coverage -fPIE -pie ${CXXFLAGS}")
187-else()187+else()
188- set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_C_FLAGS}")188+ set(CMAKE_C_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_C_FLAGS}")
189- set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_CXX_FLAGS}")189+ set(CMAKE_CXX_FLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CMAKE_CXX_FLAGS}")
190- set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CXXFLAGS}")190+ set(CXXFLAGS "-fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack -fPIE -pie ${CXXFLAGS}")
191-endif()191+endif()
192- 192+ 
193-if (NOT DEFINED GLIBCXX_USE_CXX11_ABI)193+if (NOT DEFINED GLIBCXX_USE_CXX11_ABI)
194- set(GLIBCXX_USE_CXX11_ABI 0)194+ set(GLIBCXX_USE_CXX11_ABI 0)
195-endif()195+endif()
196-message(STATUS "Determined _GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}")196+message(STATUS "Determined _GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}")
197-set(_GLIBCXX_USE_CXX11_ABI ${GLIBCXX_USE_CXX11_ABI})197+set(_GLIBCXX_USE_CXX11_ABI ${GLIBCXX_USE_CXX11_ABI})
198-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}")198+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}")
199-if (${GLIBCXX_USE_CXX11_ABI} EQUAL 0)199+if (${GLIBCXX_USE_CXX11_ABI} EQUAL 0)
200- set(CMAKE_CXX_FLAGS "-fabi-version=11 ${CMAKE_CXX_FLAGS}")200+ set(CMAKE_CXX_FLAGS "-fabi-version=11 ${CMAKE_CXX_FLAGS}")
201-else()201+else()
202- set(CXX_STANDARD_REQUIRED ON)202+ set(CXX_STANDARD_REQUIRED ON)
203- if (DEFINED ABI_VERSION)203+ if (DEFINED ABI_VERSION)
204- set(CMAKE_CXX_FLAGS "-fabi-version=${ABI_VERSION} ${CMAKE_CXX_FLAGS}")204+ set(CMAKE_CXX_FLAGS "-fabi-version=${ABI_VERSION} ${CMAKE_CXX_FLAGS}")
205- endif()205+ endif()
206-endif()206+endif()
207- 207+ 
208- 208+ 
209-if (DEFINED BUILD_LIBTORCH)209+if (DEFINED BUILD_LIBTORCH)
210- add_compile_definitions(BUILD_LIBTORCH)210+ add_compile_definitions(BUILD_LIBTORCH)
211-endif()211+endif()
212- 212+ 
213- 213+ 
214-include_directories(${PROJECT_SOURCE_DIR})214+include_directories(${PROJECT_SOURCE_DIR})
215-include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/aten)215+include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/aten)
216-include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/inductor)216+include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/inductor)
217-include_directories(${PROJECT_SOURCE_DIR}/third_party/hccl/inc)217+include_directories(${PROJECT_SOURCE_DIR}/third_party/hccl/inc)
218-include_directories(${PROJECT_SOURCE_DIR}/third_party/acl/inc)218+include_directories(${PROJECT_SOURCE_DIR}/third_party/acl/inc)
219-include_directories(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)219+include_directories(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)
220-include_directories(${PROJECT_SOURCE_DIR}/third_party/nlohmann/include)220+include_directories(${PROJECT_SOURCE_DIR}/third_party/nlohmann/include)
221- 221+ 
222-# Set installed PyTorch dir222+# Set installed PyTorch dir
223-if(DEFINED PYTORCH_INSTALL_DIR)223+if(DEFINED PYTORCH_INSTALL_DIR)
224- include_directories(${PYTORCH_INSTALL_DIR}/include)224+ include_directories(${PYTORCH_INSTALL_DIR}/include)
225- include_directories(${PYTORCH_INSTALL_DIR}/include/torch/csrc/api/include)225+ include_directories(${PYTORCH_INSTALL_DIR}/include/torch/csrc/api/include)
226- include_directories(${PYTORCH_INSTALL_DIR}/include/torch/csrc/distributed)226+ include_directories(${PYTORCH_INSTALL_DIR}/include/torch/csrc/distributed)
227-else()227+else()
228- message(FATAL_ERROR "Cannot find installed PyTorch directory")228+ message(FATAL_ERROR "Cannot find installed PyTorch directory")
229-endif()229+endif()
230- 230+ 
231-# Set Python include dir231+# Set Python include dir
232-if(DEFINED PYTHON_INCLUDE_DIR)232+if(DEFINED PYTHON_INCLUDE_DIR)
233- include_directories(${PYTHON_INCLUDE_DIR})233+ include_directories(${PYTHON_INCLUDE_DIR})
234-else()234+else()
235- message(FATAL_ERROR "Cannot find installed Python head file directory")235+ message(FATAL_ERROR "Cannot find installed Python head file directory")
236-endif()236+endif()
237- 237+ 
238-# sources238+# sources
239-set(ATEN_SRCS)239+set(ATEN_SRCS)
240-set(CORE_SRCS)240+set(CORE_SRCS)
241-set(FRAMEWORK_SRCS)241+set(FRAMEWORK_SRCS)
242-set(LOGGING_SRCS)242+set(LOGGING_SRCS)
243-set(INDUCTOR_SRCS)243+set(INDUCTOR_SRCS)
244-set(DIST_SRCS)244+set(DIST_SRCS)
245- 245+ 
246-if (NOT DEFINED BUILD_LIBTORCH)246+if (NOT DEFINED BUILD_LIBTORCH)
247- set(FLOP_SRCS)247+ set(FLOP_SRCS)
248- set(NPU_SRCS)248+ set(NPU_SRCS)
249- set(PROF_SRCS)249+ set(PROF_SRCS)
250- set(IPC_SRCS)250+ set(IPC_SRCS)
251- set(UTILS_SRCS)251+ set(UTILS_SRCS)
252- set(SAN_SRCS)252+ set(SAN_SRCS)
253- set(AFD_SRCS)253+ set(AFD_SRCS)
254-endif()254+endif()
255- 255+ 
256-if (DEFINED BUILD_LIBTORCH)256+if (DEFINED BUILD_LIBTORCH)
257- set(NPU_CPP_LIBS_SRCS)257+ set(NPU_CPP_LIBS_SRCS)
258-endif()258+endif()
259- 259+ 
260-add_subdirectory(${TORCHNPU_ROOT}/aten)260+add_subdirectory(${TORCHNPU_ROOT}/aten)
261-add_subdirectory(${TORCHNPU_ROOT}/core)261+add_subdirectory(${TORCHNPU_ROOT}/core)
262-add_subdirectory(${TORCHNPU_ROOT}/framework)262+add_subdirectory(${TORCHNPU_ROOT}/framework)
263-add_subdirectory(${TORCHNPU_ROOT}/flopcount)263+add_subdirectory(${TORCHNPU_ROOT}/flopcount)
264-add_subdirectory(${TORCHNPU_ROOT}/logging)264+add_subdirectory(${TORCHNPU_ROOT}/logging)
265-add_subdirectory(${TORCHNPU_ROOT}/custom_dtype)265+add_subdirectory(${TORCHNPU_ROOT}/custom_dtype)
266-add_subdirectory(${TORCHNPU_ROOT}/inductor)266+add_subdirectory(${TORCHNPU_ROOT}/inductor)
267-add_subdirectory(${TORCHNPU_ROOT}/distributed)267+add_subdirectory(${TORCHNPU_ROOT}/distributed)
268- 268+ 
269-if (NOT DEFINED BUILD_LIBTORCH)269+if (NOT DEFINED BUILD_LIBTORCH)
270- add_subdirectory(${TORCHNPU_ROOT}/npu)270+ add_subdirectory(${TORCHNPU_ROOT}/npu)
271- add_subdirectory(${TORCHNPU_ROOT}/profiler)271+ add_subdirectory(${TORCHNPU_ROOT}/profiler)
272- add_subdirectory(${TORCHNPU_ROOT}/ipc)272+ add_subdirectory(${TORCHNPU_ROOT}/ipc)
273- add_subdirectory(${TORCHNPU_ROOT}/utils)273+ add_subdirectory(${TORCHNPU_ROOT}/utils)
274- add_subdirectory(${TORCHNPU_ROOT}/sanitizer)274+ add_subdirectory(${TORCHNPU_ROOT}/sanitizer)
275- add_subdirectory(${TORCHNPU_ROOT}/afd)275+ add_subdirectory(${TORCHNPU_ROOT}/afd)
276-endif()276+endif()
277- 277+ 
278-if (DEFINED BUILD_LIBTORCH)278+if (DEFINED BUILD_LIBTORCH)
279- add_subdirectory(${TORCHNPU_ROOT}/libs)279+ add_subdirectory(${TORCHNPU_ROOT}/libs)
280-endif()280+endif()
281- 281+ 
282-set(OPS_PLUGIN_SRCS)282+set(OPS_PLUGIN_SRCS)
283-# Add subdirectory of op-plugin283+# Add subdirectory of op-plugin
284-include_directories(${PROJECT_SOURCE_DIR}/third_party/op-plugin)284+include_directories(${PROJECT_SOURCE_DIR}/third_party/op-plugin)
285-add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/op-plugin/op_plugin)285+add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/op-plugin/op_plugin)
286- 286+ 
287-if (DEFINED BUILD_TENSORPIPE)287+if (DEFINED BUILD_TENSORPIPE)
288- add_definitions(-DUSE_RPC_FRAMEWORK)288+ add_definitions(-DUSE_RPC_FRAMEWORK)
289- set(BUILD_SHARED_LIBS ON)289+ set(BUILD_SHARED_LIBS ON)
290- if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")290+ if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")
291- message(WARNING "tensorpipe forces CMake compatibility")291+ message(WARNING "tensorpipe forces CMake compatibility")
292- set(CMAKE_POLICY_VERSION_MINIMUM 3.5)292+ set(CMAKE_POLICY_VERSION_MINIMUM 3.5)
293- endif()293+ endif()
294- add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)294+ add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)
295- if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")295+ if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")
296- unset(CMAKE_POLICY_VERSION_MINIMUM)296+ unset(CMAKE_POLICY_VERSION_MINIMUM)
297- endif()297+ endif()
298- set(BUILD_SHARED_LIBS OFF)298+ set(BUILD_SHARED_LIBS OFF)
299-endif()299+endif()
300- 300+ 
301-if (DEFINED BUILD_LIBTORCH)301+if (DEFINED BUILD_LIBTORCH)
302- set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${FRAMEWORK_SRCS} ${LOGGING_SRCS} ${NPU_CPP_LIBS_SRCS} )302+ set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${FRAMEWORK_SRCS} ${LOGGING_SRCS} ${NPU_CPP_LIBS_SRCS} )
303-else()303+else()
304-# Compile code with pybind11304+# Compile code with pybind11
305- set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${LOGGING_SRCS} ${FRAMEWORK_SRCS} ${NPU_SRCS} ${PROF_SRCS} ${IPC_SRCS} ${UTILS_SRCS} ${SAN_SRCS} ${AFD_SRCS})305+ set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${LOGGING_SRCS} ${FRAMEWORK_SRCS} ${NPU_SRCS} ${PROF_SRCS} ${IPC_SRCS} ${UTILS_SRCS} ${SAN_SRCS} ${AFD_SRCS})
306-endif()306+endif()
307- 307+ 
308-add_library(${PLUGIN_NAME} SHARED ${CPP_SRCS})308+add_library(${PLUGIN_NAME} SHARED ${CPP_SRCS})
309-include(CheckCXXCompilerFlag)309+include(CheckCXXCompilerFlag)
310-check_cxx_compiler_flag("-fvisibility=hidden" COMPILER_SUPPORTS_HIDDEN_VISIBILITY)310+check_cxx_compiler_flag("-fvisibility=hidden" COMPILER_SUPPORTS_HIDDEN_VISIBILITY)
311-if(${COMPILER_SUPPORTS_HIDDEN_VISIBILITY})311+if(${COMPILER_SUPPORTS_HIDDEN_VISIBILITY})
312- target_compile_options(${PLUGIN_NAME} PRIVATE "-fvisibility=hidden")312+ target_compile_options(${PLUGIN_NAME} PRIVATE "-fvisibility=hidden")
313-endif()313+endif()
314- 314+ 
315-target_link_options(${PLUGIN_NAME} PRIVATE "-Wl,-Bsymbolic-functions,--no-as-needed")315+target_link_options(${PLUGIN_NAME} PRIVATE "-Wl,-Bsymbolic-functions,--no-as-needed")
316- 316+ 
317-if (DEFINED BUILD_TORCHAIR)317+if (DEFINED BUILD_TORCHAIR)
318- add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/torchair)318+ add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/torchair)
319- add_dependencies(${PLUGIN_NAME} copy_torchair_pyfiles)319+ add_dependencies(${PLUGIN_NAME} copy_torchair_pyfiles)
320-endif()320+endif()
321- 321+ 
322-add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/fmt EXCLUDE_FROM_ALL)322+add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/fmt EXCLUDE_FROM_ALL)
323- 323+ 
324-add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/dvm)324+add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/dvm)
325-add_dependencies(${PLUGIN_NAME} dvm_build)325+add_dependencies(${PLUGIN_NAME} dvm_build)
326-target_link_libraries(${PLUGIN_NAME} PRIVATE ${PROJECT_SOURCE_DIR}/third_party/dvm/dvm/libdvm.a)326+target_link_libraries(${PLUGIN_NAME} PRIVATE ${PROJECT_SOURCE_DIR}/third_party/dvm/dvm/libdvm.a)
327- 327+ 
328-link_directories(${PYTORCH_INSTALL_DIR}/lib)328+link_directories(${PYTORCH_INSTALL_DIR}/lib)
329-link_directories(${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs)329+link_directories(${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs)
330- 330+ 
331-target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libhccl.so)331+target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libhccl.so)
332- 332+ 
333-if (NOT DEFINED BUILD_LIBTORCH)333+if (NOT DEFINED BUILD_LIBTORCH)
334- target_link_libraries(${PLUGIN_NAME} PUBLIC ${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so)334+ target_link_libraries(${PLUGIN_NAME} PUBLIC ${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so)
335-endif()335+endif()
336- 336+ 
337-target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libascendcl.so)337+target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libascendcl.so)
338-target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libacl_op_compiler.so)338+target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libacl_op_compiler.so)
339-target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libge_runner.so)339+target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libge_runner.so)
340-target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libgraph.so)340+target_link_libraries(${PLUGIN_NAME} PUBLIC ${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs/libgraph.so)
341- 341+ 
342-if (DEFINED BUILD_TENSORPIPE)342+if (DEFINED BUILD_TENSORPIPE)
343- target_link_libraries(${PLUGIN_NAME} PUBLIC ${PROJECT_SOURCE_DIR}/build/packages/torch_npu/lib/libtensorpipe.so)343+ target_link_libraries(${PLUGIN_NAME} PUBLIC ${PROJECT_SOURCE_DIR}/build/packages/torch_npu/lib/libtensorpipe.so)
344-endif()344+endif()
345- 345+ 
346-target_link_libraries(${PLUGIN_NAME} PUBLIC torch torch_cpu c10 fmt::fmt-header-only)346+target_link_libraries(${PLUGIN_NAME} PUBLIC torch torch_cpu c10 fmt::fmt-header-only)
347- 347+ 
348-if (NOT DEFINED BUILD_LIBTORCH)348+if (NOT DEFINED BUILD_LIBTORCH)
349- set(ATEN_THREADING "OMP" CACHE STRING "ATen parallel backend")349+ set(ATEN_THREADING "OMP" CACHE STRING "ATen parallel backend")
350- message(STATUS "Using ATen parallel backend: ${ATEN_THREADING}")350+ message(STATUS "Using ATen parallel backend: ${ATEN_THREADING}")
351- if ("${ATEN_THREADING}" STREQUAL "OMP")351+ if ("${ATEN_THREADING}" STREQUAL "OMP")
352- target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_OPENMP=1")352+ target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_OPENMP=1")
353- elseif ("${ATEN_THREADING}" STREQUAL "NATIVE")353+ elseif ("${ATEN_THREADING}" STREQUAL "NATIVE")
354- target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_NATIVE=1")354+ target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_NATIVE=1")
355- elseif ("${ATEN_THREADING}" STREQUAL "TBB")355+ elseif ("${ATEN_THREADING}" STREQUAL "TBB")
356- target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_NATIVE_TBB=1")356+ target_compile_definitions(${PLUGIN_NAME} PUBLIC "-DAT_PARALLEL_NATIVE_TBB=1")
357- else()357+ else()
358- message(FATAL_ERROR "Unknown ATen parallel backend: ${ATEN_THREADING}")358+ message(FATAL_ERROR "Unknown ATen parallel backend: ${ATEN_THREADING}")
359- endif()359+ endif()
360- 360+ 
361- include(GNUInstallDirs)361+ include(GNUInstallDirs)
362- target_compile_options(${PLUGIN_NAME} PRIVATE "-DC10_BUILD_MAIN_LIB")362+ target_compile_options(${PLUGIN_NAME} PRIVATE "-DC10_BUILD_MAIN_LIB")
363- install(TARGETS ${PLUGIN_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})363+ install(TARGETS ${PLUGIN_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
364-endif()364+endif()
365- 365+ 
366-if (NOT DEFINED BUILD_LIBTORCH)366+if (NOT DEFINED BUILD_LIBTORCH)
367- add_subdirectory(${TORCHNPU_ROOT}/toolkit)367+ add_subdirectory(${TORCHNPU_ROOT}/toolkit)
368- target_link_libraries(${PLUGIN_NAME} PUBLIC npu_profiler)368+ target_link_libraries(${PLUGIN_NAME} PUBLIC npu_profiler)
369-endif()369+endif()
370- 370+ 
371-if (DEFINED BUILD_GTEST)371+if (DEFINED BUILD_GTEST)
372- enable_testing()372+ enable_testing()
373- SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/build/gtest)373+ SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/build/gtest)
374- add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/googletest)374+ add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/googletest)
375- include_directories(${PROJECT_SOURCE_DIR}/third_party/googletest/googletest/include)375+ include_directories(${PROJECT_SOURCE_DIR}/third_party/googletest/googletest/include)
376- 376+ 
377- set(TORCH_API_TEST_SOURCES)377+ set(TORCH_API_TEST_SOURCES)
378- add_subdirectory(${PROJECT_SOURCE_DIR}/test/cpp/api)378+ add_subdirectory(${PROJECT_SOURCE_DIR}/test/cpp/api)
379- add_executable(test_api ${TORCH_API_TEST_SOURCES})379+ add_executable(test_api ${TORCH_API_TEST_SOURCES})
380- 380+ 
381- target_link_libraries(test_api PUBLIC torch_npu)381+ target_link_libraries(test_api PUBLIC torch_npu)
382- target_link_libraries(test_api PUBLIC gtest_main gtest)382+ target_link_libraries(test_api PUBLIC gtest_main gtest)
383-endif()383+endif()
384- 384+ 
385-if (DEFINED BUILD_LIBTORCH)385+if (DEFINED BUILD_LIBTORCH)
386- configure_file(386+ configure_file(
387- ${PROJECT_SOURCE_DIR}/cmake/Torch_npuConfig.cmake.in387+ ${PROJECT_SOURCE_DIR}/cmake/Torch_npuConfig.cmake.in
388- ${PROJECT_SOURCE_DIR}/build/Torch_npuConfig.cmake388+ ${PROJECT_SOURCE_DIR}/build/Torch_npuConfig.cmake
389- @ONLY)389+ @ONLY)
390-endif()390+endif()
391- 
MThird_Party_Open_Source_Software_Notice+190-190
@@ -1,190 +1,190 @@
1-OPEN SOURCE SOFTWARE NOTICE1+OPEN SOURCE SOFTWARE NOTICE
2-Please note we provide an open source software notice along with this product and/or this product firmware (in the following just "this product"). The open source software licenses are granted by the respective right holders. And the open source licenses prevail all other license information with regard to the respective open source software contained in the product, including but not limited to End User Software Licensing Agreement. This notice is provided on behalf of Huawei Technologies Co. Ltd. and any of its local subsidiaries which may have provided this product to you in your local country.2+Please note we provide an open source software notice along with this product and/or this product firmware (in the following just "this product"). The open source software licenses are granted by the respective right holders. And the open source licenses prevail all other license information with regard to the respective open source software contained in the product, including but not limited to End User Software Licensing Agreement. This notice is provided on behalf of Huawei Technologies Co. Ltd. and any of its local subsidiaries which may have provided this product to you in your local country.
3- 3+ 
4-Warranty Disclaimer4+Warranty Disclaimer
5-THE OPEN SOURCE SOFTWARE IN THIS PRODUCT IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL, BUT WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. SEE THE APPLICABLE LICENSES FOR MORE DETAILS.5+THE OPEN SOURCE SOFTWARE IN THIS PRODUCT IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL, BUT WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. SEE THE APPLICABLE LICENSES FOR MORE DETAILS.
6- 6+ 
7-Copyright Notice and License Texts7+Copyright Notice and License Texts
8-Software: pytorch v2.6.08+Software: pytorch v2.6.0
9-Copyright notice:9+Copyright notice:
10-Copyright (c) Advanced Micro Devices, Inc.10+Copyright (c) Advanced Micro Devices, Inc.
11- 11+ 
12-Copyright (c) Microsoft Corporation12+Copyright (c) Microsoft Corporation
13- 13+ 
14-Copyright (c) Bjorn Fahller14+Copyright (c) Bjorn Fahller
15- 15+ 
16-Copyright (c) 2001-2014 Python Software Foundation All Rights Reserved16+Copyright (c) 2001-2014 Python Software Foundation All Rights Reserved
17- 17+ 
18-Copyright (c) 2011-2013 NYU18+Copyright (c) 2011-2013 NYU
19- 19+ 
20-Copyright (c) 1995-2011 by Fredrik Lundh20+Copyright (c) 1995-2011 by Fredrik Lundh
21- 21+ 
22-Copyright (c) Edward Z. Yang ezyang@mit.edu22+Copyright (c) Edward Z. Yang ezyang@mit.edu
23- 23+ 
24-Copyright (c) 2014- Facebook, Inc24+Copyright (c) 2014- Facebook, Inc
25- 25+ 
26-Copyright (c) 2017 The Android Open Source Project26+Copyright (c) 2017 The Android Open Source Project
27- 27+ 
28-Copyright Python Software Foundation28+Copyright Python Software Foundation
29- 29+ 
30-Copyright (c) 2012 Massachusetts Institute of Technology30+Copyright (c) 2012 Massachusetts Institute of Technology
31- 31+ 
32-Copyright (c) 2018 Alex Rogozhnikov32+Copyright (c) 2018 Alex Rogozhnikov
33- 33+ 
34-Copyright (c) 2007-2009 Scientific Computing and Imaging Institute, University of Utah34+Copyright (c) 2007-2009 Scientific Computing and Imaging Institute, University of Utah
35- 35+ 
36-Copyright (c) 2006 Idiap Research Institute36+Copyright (c) 2006 Idiap Research Institute
37- 37+ 
38-Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved38+Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved
39- 39+ 
40-Copyright (c) 2015 Yangqing Jia All rights reserved40+Copyright (c) 2015 Yangqing Jia All rights reserved
41- 41+ 
42-Copyright (c) Meta Platforms, Inc.42+Copyright (c) Meta Platforms, Inc.
43- 43+ 
44-Copyright 2023-present Facebook. All Rights Reserved44+Copyright 2023-present Facebook. All Rights Reserved
45- 45+ 
46-Copyright (c) 2022 Apple Inc.46+Copyright (c) 2022 Apple Inc.
47- 47+ 
48-Copyright (c) 2005-2017, NumPy Developers. All rights reserved48+Copyright (c) 2005-2017, NumPy Developers. All rights reserved
49- 49+ 
50-Copyright (c) 1997 - 2002, Makoto Matsumoto and Takuji Nishimura, All rights reserved50+Copyright (c) 1997 - 2002, Makoto Matsumoto and Takuji Nishimura, All rights reserved
51- 51+ 
52-Copyright (c) 2014, The Regents52+Copyright (c) 2014, The Regents
53- 53+ 
54-Copyright (c) 2005-2010 ActiveState Software Inc.54+Copyright (c) 2005-2010 ActiveState Software Inc.
55- 55+ 
56-Copyright Malte Skarupke 201756+Copyright Malte Skarupke 2017
57- 57+ 
58-Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)58+Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
59- 59+ 
60-Copyright 2005, Google Inc. All rights reserved60+Copyright 2005, Google Inc. All rights reserved
61- 61+ 
62-Copyright (c) Meta Platforms, Inc. and affiliates62+Copyright (c) Meta Platforms, Inc. and affiliates
63- 63+ 
64-Copyright (c) 2023, Advanced Micro Devices, Inc.64+Copyright (c) 2023, Advanced Micro Devices, Inc.
65- 65+ 
66-Copyright (c) 2022, Tri Dao66+Copyright (c) 2022, Tri Dao
67- 67+ 
68-Copyright (c) 2005-2023 NVIDIA Corporation Built68+Copyright (c) 2005-2023 NVIDIA Corporation Built
69- 69+ 
70-Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. All rights reserved70+Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. All rights reserved
71- 71+ 
72-Copyright 2008 Google Inc. All rights reserved72+Copyright 2008 Google Inc. All rights reserved
73- 73+ 
74-Copyright (c) 2021, 2023-2024 Arm Limited74+Copyright (c) 2021, 2023-2024 Arm Limited
75- 75+ 
76-Copyright (c) 2003-2017 Josef Weidendorfer. All rights reserved76+Copyright (c) 2003-2017 Josef Weidendorfer. All rights reserved
77- 77+ 
78-Copyright (c) 1997-2011 by Secret Labs AB78+Copyright (c) 1997-2011 by Secret Labs AB
79- 79+ 
80-Copyright (c) 2016- Facebook, Inc80+Copyright (c) 2016- Facebook, Inc
81- 81+ 
82-Copyright (c) 2014 Matthew Rocklin82+Copyright (c) 2014 Matthew Rocklin
83- 83+ 
84-Copyright (c) 2005-2022 NVIDIA Corporation Built84+Copyright (c) 2005-2022 NVIDIA Corporation Built
85- 85+ 
86-Copyright (c) Facebook, Inc.86+Copyright (c) Facebook, Inc.
87- 87+ 
88-Copyright 2019-2020 Kakao Brain88+Copyright 2019-2020 Kakao Brain
89- 89+ 
90-Copyright (c) 2000-2017 Julian Seward. All rights reserved90+Copyright (c) 2000-2017 Julian Seward. All rights reserved
91- 91+ 
92-Copyright (c) 2005-2020 Rich Felker92+Copyright (c) 2005-2020 Rich Felker
93- 93+ 
94-Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved94+Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved
95- 95+ 
96-Copyright 2010-2014 Rich Geldreich and Tenacious Software LLC96+Copyright 2010-2014 Rich Geldreich and Tenacious Software LLC
97- 97+ 
98-Copyright (c) 2016 Facebook Inc.98+Copyright (c) 2016 Facebook Inc.
99- 99+ 
100-Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)100+Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
101- 101+ 
102-Copyright 2010-2014 Rich Geldreich and Tenacious Software LLC All Rights Reserved102+Copyright 2010-2014 Rich Geldreich and Tenacious Software LLC All Rights Reserved
103- 103+ 
104-Copyright (c) 2012-2014 Deepmind Technologies104+Copyright (c) 2012-2014 Deepmind Technologies
105- 105+ 
106-Copyright (c) 2012 Giovanni Garberoglio Interdisciplinary Laboratory106+Copyright (c) 2012 Giovanni Garberoglio Interdisciplinary Laboratory
107- 107+ 
108-Copyright (c) 2024, Tri Dao108+Copyright (c) 2024, Tri Dao
109- 109+ 
110-Copyright (c) Donald Stufft and individual contributors. All rights reserved110+Copyright (c) Donald Stufft and individual contributors. All rights reserved
111- 111+ 
112-Copyright (c) 2018, Steven Moshier All rights reserved112+Copyright (c) 2018, Steven Moshier All rights reserved
113- 113+ 
114-Copyright (c) 2015 Google Inc. All rights reserved114+Copyright (c) 2015 Google Inc. All rights reserved
115- 115+ 
116-Copyright (c) 2015-2016 Advanced Micro Devices, Inc. All rights reserved116+Copyright (c) 2015-2016 Advanced Micro Devices, Inc. All rights reserved
117- 117+ 
118-Copyright (c) 2010-2022 by Alex Clark and contributors118+Copyright (c) 2010-2022 by Alex Clark and contributors
119- 119+ 
120-Copyright 2015 Google Inc. All Rights Reserved120+Copyright 2015 Google Inc. All Rights Reserved
121- 121+ 
122-Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved122+Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved
123- 123+ 
124-Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved124+Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved
125- 125+ 
126-Copyright (c) 2016 manylinux126+Copyright (c) 2016 manylinux
127- 127+ 
128-Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved128+Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved
129- 129+ 
130-Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)130+Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
131- 131+ 
132-Copyright 2013-2014 RAD Game132+Copyright 2013-2014 RAD Game
133- 133+ 
134-Copyright (c) 2011-2019 Stephan Brumme. All rights reserved134+Copyright (c) 2011-2019 Stephan Brumme. All rights reserved
135- 135+ 
136-Copyright (c) 2018 MathInf GmbH, Thomas Viehmann136+Copyright (c) 2018 MathInf GmbH, Thomas Viehmann
137- 137+ 
138-Copyright (c) 2013 Eddy Petrisor138+Copyright (c) 2013 Eddy Petrisor
139- 139+ 
140-Copyright (c) 2023-2024 The ggml140+Copyright (c) 2023-2024 The ggml
141- 141+ 
142-Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation All Rights Reserved142+Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation All Rights Reserved
143- 143+ 
144-Copyright 2004-present Facebook. All Rights Reserved144+Copyright 2004-present Facebook. All Rights Reserved
145- 145+ 
146-Copyright (c) 2010 ActiveState Software Inc.146+Copyright (c) 2010 ActiveState Software Inc.
147- 147+ 
148-Copyright (c) 2006 The Android Open Source Project148+Copyright (c) 2006 The Android Open Source Project
149- 149+ 
150-(c) Meta Platforms, Inc. and affiliates. All rights reserved150+(c) Meta Platforms, Inc. and affiliates. All rights reserved
151- 151+ 
152-Copyright (c) 2023 Apple Inc.152+Copyright (c) 2023 Apple Inc.
153- 153+ 
154-Copyright (c) Microsoft Corporation. All rights reserved154+Copyright (c) Microsoft Corporation. All rights reserved
155- 155+ 
156-Copyright 2015 The TensorFlow Authors. All Rights Reserved156+Copyright 2015 The TensorFlow Authors. All Rights Reserved
157- 157+ 
158-Copyright (c) 2023, Tri Dao158+Copyright (c) 2023, Tri Dao
159- 159+ 
160-Copyright 2022 Cruise LLC160+Copyright 2022 Cruise LLC
161- 161+ 
162-Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved162+Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved
163- 163+ 
164-Copyright (c) 2022 Cruise LLC. All rights reserved164+Copyright (c) 2022 Cruise LLC. All rights reserved
165- 165+ 
166-Copyright (c) 2016-present, Facebook, Inc.166+Copyright (c) 2016-present, Facebook, Inc.
167- 167+ 
168-(c) Copyright John Maddock 2006168+(c) Copyright John Maddock 2006
169- 169+ 
170-Copyright (c) 2011-2014 Idiap Research Institute170+Copyright (c) 2011-2014 Idiap Research Institute
171- 171+ 
172-Copyright (c) 2014 Indiana University All rights reserved172+Copyright (c) 2014 Indiana University All rights reserved
173- 173+ 
174-copyright 2019 The TensorFlow Authors174+copyright 2019 The TensorFlow Authors
175- 175+ 
176-Copyright (c) 2016-present, Facebook Inc. All rights reserved176+Copyright (c) 2016-present, Facebook Inc. All rights reserved
177- 177+ 
178-License: BSD 3-Clause License178+License: BSD 3-Clause License
179-Copyright (c) ,179+Copyright (c) ,
180-All rights reserved.180+All rights reserved.
181- 181+ 
182-1. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:182+1. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
183- 183+ 
184-2. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.184+2. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
185- 185+ 
186-3. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.186+3. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
187- 187+ 
188-Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.188+Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
189- 189+ 
190-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.190+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
Mrequirements.txt+7-7
@@ -1,7 +1,7 @@
1---extra-index-url https://download.pytorch.org/whl/cpu1+--extra-index-url https://download.pytorch.org/whl/cpu
2- 2+ 
3-pyyaml3+pyyaml
4-setuptools4+setuptools
5-auditwheel5+auditwheel
6- 6+ 
7-torch==2.9.07+torch==2.9.0
Msetup.py+783-783
@@ -1,784 +1,784 @@
1-import glob1+import glob
2-import multiprocessing2+import multiprocessing
3-import multiprocessing.pool3+import multiprocessing.pool
4-import os4+import os
5-import re5+import re
6-import shutil6+import shutil
7-import stat7+import stat
8-import subprocess8+import subprocess
9-import sys9+import sys
10-import traceback10+import traceback
11-import platform11+import platform
12-import time12+import time
13-import sysconfig13+import sysconfig
14-from sysconfig import get_paths14+from sysconfig import get_paths
15-from pathlib import Path15+from pathlib import Path
16-from typing import Union16+from typing import Union
17- 17+ 
18-import distutils.ccompiler18+import distutils.ccompiler
19-import distutils.command.clean19+import distutils.command.clean
20-from distutils.version import LooseVersion20+from distutils.version import LooseVersion
21-from distutils.command.build_py import build_py21+from distutils.command.build_py import build_py
22-from setuptools import setup, distutils, Extension, find_packages22+from setuptools import setup, distutils, Extension, find_packages
23-from setuptools.command.build_clib import build_clib23+from setuptools.command.build_clib import build_clib
24-from setuptools.command.build_ext import build_ext24+from setuptools.command.build_ext import build_ext
25-from setuptools.command.egg_info import egg_info25+from setuptools.command.egg_info import egg_info
26-from setuptools.command.install import install26+from setuptools.command.install import install
27-from wheel.bdist_wheel import bdist_wheel27+from wheel.bdist_wheel import bdist_wheel
28- 28+ 
29-# Disable autoloading before running 'import torch' to avoid circular dependencies29+# Disable autoloading before running 'import torch' to avoid circular dependencies
30-os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"30+os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
31- 31+ 
32-from torchnpugen.utils import PathManager32+from torchnpugen.utils import PathManager
33- 33+ 
34-BASE_DIR = os.path.dirname(os.path.realpath(__file__))34+BASE_DIR = os.path.dirname(os.path.realpath(__file__))
35-THIRD_PARTY_PATH = os.path.join(BASE_DIR, "third_party")35+THIRD_PARTY_PATH = os.path.join(BASE_DIR, "third_party")
36-PathManager.check_directory_path_readable(os.path.join(BASE_DIR, "version.txt"))36+PathManager.check_directory_path_readable(os.path.join(BASE_DIR, "version.txt"))
37-with open(os.path.join(BASE_DIR, "version.txt")) as version_f:37+with open(os.path.join(BASE_DIR, "version.txt")) as version_f:
38- VERSION = version_f.read().strip()38+ VERSION = version_f.read().strip()
39-UNKNOWN = "Unknown"39+UNKNOWN = "Unknown"
40-BUILD_PERMISSION = stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR | stat.S_IRGRP | stat.S_IXGRP40+BUILD_PERMISSION = stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR | stat.S_IRGRP | stat.S_IXGRP
41- 41+ 
42-DISABLE_TORCHAIR = "FALSE"42+DISABLE_TORCHAIR = "FALSE"
43-if os.environ.get("DISABLE_INSTALL_TORCHAIR") is not None:43+if os.environ.get("DISABLE_INSTALL_TORCHAIR") is not None:
44- DISABLE_TORCHAIR = os.environ.get("DISABLE_INSTALL_TORCHAIR")44+ DISABLE_TORCHAIR = os.environ.get("DISABLE_INSTALL_TORCHAIR")
45-DISABLE_RPC = "FALSE"45+DISABLE_RPC = "FALSE"
46-if os.environ.get("DISABLE_RPC_FRAMEWORK") is not None:46+if os.environ.get("DISABLE_RPC_FRAMEWORK") is not None:
47- DISABLE_RPC = os.environ.get("DISABLE_RPC_FRAMEWORK")47+ DISABLE_RPC = os.environ.get("DISABLE_RPC_FRAMEWORK")
48-ENABLE_LTO = "FALSE"48+ENABLE_LTO = "FALSE"
49-if os.environ.get("ENABLE_LTO") is not None:49+if os.environ.get("ENABLE_LTO") is not None:
50- ENABLE_LTO = os.environ.get("ENABLE_LTO")50+ ENABLE_LTO = os.environ.get("ENABLE_LTO")
51-PGO_MODE = 051+PGO_MODE = 0
52-if os.environ.get("PGO_MODE") is not None:52+if os.environ.get("PGO_MODE") is not None:
53- PGO_MODE = int(os.environ.get("PGO_MODE"))53+ PGO_MODE = int(os.environ.get("PGO_MODE"))
54- 54+ 
55-# change to use cxx11.abi in default since 2.755+# change to use cxx11.abi in default since 2.7
56-USE_CXX11_ABI = True56+USE_CXX11_ABI = True
57-if os.environ.get("_GLIBCXX_USE_CXX11_ABI") is not None and os.environ.get("_GLIBCXX_USE_CXX11_ABI") == "0":57+if os.environ.get("_GLIBCXX_USE_CXX11_ABI") is not None and os.environ.get("_GLIBCXX_USE_CXX11_ABI") == "0":
58- USE_CXX11_ABI = False58+ USE_CXX11_ABI = False
59- 59+ 
60- 60+ 
61-def get_submodule_folders():61+def get_submodule_folders():
62- git_modules_path = os.path.join(BASE_DIR, ".gitmodules")62+ git_modules_path = os.path.join(BASE_DIR, ".gitmodules")
63- default_modules_path = [63+ default_modules_path = [
64- os.path.join(THIRD_PARTY_PATH, name)64+ os.path.join(THIRD_PARTY_PATH, name)
65- for name in [65+ for name in [
66- "op-plugin",66+ "op-plugin",
67- ]67+ ]
68- ]68+ ]
69- if not os.path.exists(git_modules_path):69+ if not os.path.exists(git_modules_path):
70- return default_modules_path70+ return default_modules_path
71- with open(git_modules_path) as f:71+ with open(git_modules_path) as f:
72- return [72+ return [
73- os.path.join(BASE_DIR, line.split("=", 1)[1].strip())73+ os.path.join(BASE_DIR, line.split("=", 1)[1].strip())
74- for line in f.readlines()74+ for line in f.readlines()
75- if line.strip().startswith("path")75+ if line.strip().startswith("path")
76- ]76+ ]
77- 77+ 
78- 78+ 
79-def check_submodules():79+def check_submodules():
80- def not_exists_or_empty(folder):80+ def not_exists_or_empty(folder):
81- return not os.path.exists(folder) or (81+ return not os.path.exists(folder) or (
82- os.path.isdir(folder) and len(os.listdir(folder)) == 082+ os.path.isdir(folder) and len(os.listdir(folder)) == 0
83- )83+ )
84- 84+ 
85- folders = get_submodule_folders()85+ folders = get_submodule_folders()
86- # If none of the submodule folders exists, try to initialize them86+ # If none of the submodule folders exists, try to initialize them
87- if all(not_exists_or_empty(folder) for folder in folders):87+ if all(not_exists_or_empty(folder) for folder in folders):
88- try:88+ try:
89- print(" --- Trying to initialize submodules")89+ print(" --- Trying to initialize submodules")
90- start = time.time()90+ start = time.time()
91- subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"], cwd=BASE_DIR) # Compliant91+ subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"], cwd=BASE_DIR) # Compliant
92- end = time.time()92+ end = time.time()
93- print(f" --- Submodule initialization took {end - start:.2f} sec")93+ print(f" --- Submodule initialization took {end - start:.2f} sec")
94- except Exception:94+ except Exception:
95- print(" --- Submodule initalization failed")95+ print(" --- Submodule initalization failed")
96- print("Please run:\n\tgit submodule init && git submodule update")96+ print("Please run:\n\tgit submodule init && git submodule update")
97- sys.exit(1)97+ sys.exit(1)
98- 98+ 
99- 99+ 
100-check_submodules()100+check_submodules()
101- 101+ 
102- 102+ 
103-def get_sha(pytorch_root: Union[str, Path]) -> str:103+def get_sha(pytorch_root: Union[str, Path]) -> str:
104- try:104+ try:
105- return (105+ return (
106- subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=pytorch_root) # Compliant106+ subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=pytorch_root) # Compliant
107- .decode("ascii")107+ .decode("ascii")
108- .strip()108+ .strip()
109- )109+ )
110- except Exception:110+ except Exception:
111- return UNKNOWN111+ return UNKNOWN
112- 112+ 
113- 113+ 
114-def generate_torch_npu_version():114+def generate_torch_npu_version():
115- torch_npu_root = Path(__file__).parent115+ torch_npu_root = Path(__file__).parent
116- version_path = torch_npu_root / "torch_npu" / "version.py"116+ version_path = torch_npu_root / "torch_npu" / "version.py"
117- if version_path.exists():117+ if version_path.exists():
118- version_path.unlink()118+ version_path.unlink()
119- flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL119+ flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
120- modes = stat.S_IWUSR | stat.S_IRUSR120+ modes = stat.S_IWUSR | stat.S_IRUSR
121- sha = get_sha(torch_npu_root)121+ sha = get_sha(torch_npu_root)
122- if os.getenv("BUILD_WITHOUT_SHA") is None:122+ if os.getenv("BUILD_WITHOUT_SHA") is None:
123- global VERSION123+ global VERSION
124- VERSION += "+git" + sha[:7]124+ VERSION += "+git" + sha[:7]
125- with os.fdopen(os.open(version_path, flags, modes), 'w') as f:125+ with os.fdopen(os.open(version_path, flags, modes), 'w') as f:
126- f.write("__version__ = '{version}'\n".format(version=VERSION))126+ f.write("__version__ = '{version}'\n".format(version=VERSION))
127- f.write("git_version = {}\n".format(repr(sha)))127+ f.write("git_version = {}\n".format(repr(sha)))
128- os.chmod(version_path, mode=stat.S_IRUSR | stat.S_IEXEC | stat.S_IRGRP | stat.S_IXGRP)128+ os.chmod(version_path, mode=stat.S_IRUSR | stat.S_IEXEC | stat.S_IRGRP | stat.S_IXGRP)
129- 129+ 
130- 130+ 
131-generate_torch_npu_version()131+generate_torch_npu_version()
132- 132+ 
133- 133+ 
134-def which(thefile):134+def which(thefile):
135- path = os.environ.get("PATH", os.defpath).split(os.pathsep)135+ path = os.environ.get("PATH", os.defpath).split(os.pathsep)
136- for d in path:136+ for d in path:
137- fname = os.path.join(d, thefile)137+ fname = os.path.join(d, thefile)
138- fnames = [fname]138+ fnames = [fname]
139- if sys.platform == 'win32':139+ if sys.platform == 'win32':
140- exts = os.environ.get('PATHEXT', '').split(os.pathsep)140+ exts = os.environ.get('PATHEXT', '').split(os.pathsep)
141- fnames += [fname + ext for ext in exts]141+ fnames += [fname + ext for ext in exts]
142- for name in fnames:142+ for name in fnames:
143- if os.access(name, os.F_OK | os.X_OK) and not os.path.isdir(name):143+ if os.access(name, os.F_OK | os.X_OK) and not os.path.isdir(name):
144- return name144+ return name
145- return None145+ return None
146- 146+ 
147- 147+ 
148-def get_cmake_command():148+def get_cmake_command():
149- def _get_version(cmd):149+ def _get_version(cmd):
150- for line in subprocess.check_output([cmd, '--version']).decode('utf-8').split('\n'):150+ for line in subprocess.check_output([cmd, '--version']).decode('utf-8').split('\n'):
151- if 'version' in line:151+ if 'version' in line:
152- return LooseVersion(line.strip().split(' ')[2])152+ return LooseVersion(line.strip().split(' ')[2])
153- raise RuntimeError('no version found')153+ raise RuntimeError('no version found')
154- "Returns cmake command."154+ "Returns cmake command."
155- cmake_command = 'cmake'155+ cmake_command = 'cmake'
156- if platform.system() == 'Windows':156+ if platform.system() == 'Windows':
157- return cmake_command157+ return cmake_command
158- cmake3 = which('cmake3')158+ cmake3 = which('cmake3')
159- cmake = which('cmake')159+ cmake = which('cmake')
160- if cmake3 is not None and _get_version(cmake3) >= LooseVersion("3.18.0"):160+ if cmake3 is not None and _get_version(cmake3) >= LooseVersion("3.18.0"):
161- cmake_command = 'cmake3'161+ cmake_command = 'cmake3'
162- return cmake_command162+ return cmake_command
163- elif cmake is not None and _get_version(cmake) >= LooseVersion("3.18.0"):163+ elif cmake is not None and _get_version(cmake) >= LooseVersion("3.18.0"):
164- return cmake_command164+ return cmake_command
165- else:165+ else:
166- raise RuntimeError('no cmake or cmake3 with version >= 3.18.0 found')166+ raise RuntimeError('no cmake or cmake3 with version >= 3.18.0 found')
167- 167+ 
168- 168+ 
169-def get_build_type():169+def get_build_type():
170- build_type = 'Release'170+ build_type = 'Release'
171- if os.getenv('DEBUG', default='0').upper() in ['ON', '1', 'YES', 'TRUE', 'Y']:171+ if os.getenv('DEBUG', default='0').upper() in ['ON', '1', 'YES', 'TRUE', 'Y']:
172- build_type = 'Debug'172+ build_type = 'Debug'
173- 173+ 
174- if os.getenv('REL_WITH_DEB_INFO', default='0').upper() in ['ON', '1', 'YES', 'TRUE', 'Y']:174+ if os.getenv('REL_WITH_DEB_INFO', default='0').upper() in ['ON', '1', 'YES', 'TRUE', 'Y']:
175- build_type = 'RelWithDebInfo'175+ build_type = 'RelWithDebInfo'
176- 176+ 
177- return build_type177+ return build_type
178- 178+ 
179- 179+ 
180-def _get_build_mode():180+def _get_build_mode():
181- for i in range(1, len(sys.argv)):181+ for i in range(1, len(sys.argv)):
182- if not sys.argv[i].startswith('-'):182+ if not sys.argv[i].startswith('-'):
183- return sys.argv[i]183+ return sys.argv[i]
184- 184+ 
185- raise RuntimeError("Run setup.py without build mode.")185+ raise RuntimeError("Run setup.py without build mode.")
186- 186+ 
187- 187+ 
188-def get_pytorch_dir():188+def get_pytorch_dir():
189- try:189+ try:
190- import torch190+ import torch
191- return os.path.dirname(os.path.realpath(torch.__file__))191+ return os.path.dirname(os.path.realpath(torch.__file__))
192- except Exception:192+ except Exception:
193- _, _, exc_traceback = sys.exc_info()193+ _, _, exc_traceback = sys.exc_info()
194- frame_summary = traceback.extract_tb(exc_traceback)[-1]194+ frame_summary = traceback.extract_tb(exc_traceback)[-1]
195- return os.path.dirname(frame_summary.filename)195+ return os.path.dirname(frame_summary.filename)
196- 196+ 
197- 197+ 
198-def generate_bindings_code(base_dir):198+def generate_bindings_code(base_dir):
199- python_execute = sys.executable199+ python_execute = sys.executable
200- generate_code_cmd = ["bash", os.path.join(base_dir, 'generate_code.sh'), python_execute, VERSION]200+ generate_code_cmd = ["bash", os.path.join(base_dir, 'generate_code.sh'), python_execute, VERSION]
201- if subprocess.call(generate_code_cmd) != 0: # Compliant201+ if subprocess.call(generate_code_cmd) != 0: # Compliant
202- print(202+ print(
203- 'Failed to generate ATEN bindings: {}'.format(generate_code_cmd),203+ 'Failed to generate ATEN bindings: {}'.format(generate_code_cmd),
204- file=sys.stderr)204+ file=sys.stderr)
205- sys.exit(1)205+ sys.exit(1)
206- 206+ 
207- 207+ 
208-def build_stub(base_dir):208+def build_stub(base_dir):
209- build_stub_cmd = ["sh", os.path.join(base_dir, 'third_party/acl/libs/build_stub.sh')]209+ build_stub_cmd = ["sh", os.path.join(base_dir, 'third_party/acl/libs/build_stub.sh')]
210- if subprocess.call(build_stub_cmd) != 0:210+ if subprocess.call(build_stub_cmd) != 0:
211- print(211+ print(
212- 'Failed to build stub: {}'.format(build_stub_cmd),212+ 'Failed to build stub: {}'.format(build_stub_cmd),
213- file=sys.stderr)213+ file=sys.stderr)
214- sys.exit(1)214+ sys.exit(1)
215- 215+ 
216- 216+ 
217-def check_torchair_valid(base_dir):217+def check_torchair_valid(base_dir):
218- # build with submodule of torchair, if path of torchair is valid218+ # build with submodule of torchair, if path of torchair is valid
219- torchair_path = os.path.join(base_dir, 'third_party/torchair/torchair')219+ torchair_path = os.path.join(base_dir, 'third_party/torchair/torchair')
220- return os.path.exists(torchair_path) and (220+ return os.path.exists(torchair_path) and (
221- os.path.isdir(torchair_path) and len(os.listdir(torchair_path)) != 0221+ os.path.isdir(torchair_path) and len(os.listdir(torchair_path)) != 0
222- )222+ )
223- 223+ 
224- 224+ 
225-def check_tensorpipe_valid(base_dir):225+def check_tensorpipe_valid(base_dir):
226- tensorpipe_path = os.path.join(base_dir, 'third_party/Tensorpipe/tensorpipe')226+ tensorpipe_path = os.path.join(base_dir, 'third_party/Tensorpipe/tensorpipe')
227- return os.path.exists(tensorpipe_path)227+ return os.path.exists(tensorpipe_path)
228- 228+ 
229- 229+ 
230-def generate_dbg_files_and_strip():230+def generate_dbg_files_and_strip():
231- library_dir = Path(BASE_DIR).joinpath("build/packages/torch_npu")231+ library_dir = Path(BASE_DIR).joinpath("build/packages/torch_npu")
232- dbg_dir = Path(BASE_DIR).joinpath("build/dbg")232+ dbg_dir = Path(BASE_DIR).joinpath("build/dbg")
233- os.makedirs(dbg_dir, exist_ok=True)233+ os.makedirs(dbg_dir, exist_ok=True)
234- library_files = [Path(i) for i in library_dir.rglob('*.so')]234+ library_files = [Path(i) for i in library_dir.rglob('*.so')]
235- for library_file in library_files:235+ for library_file in library_files:
236- subprocess.check_call(["eu-strip", library_file, "-f",236+ subprocess.check_call(["eu-strip", library_file, "-f",
237- str(dbg_dir.joinpath(library_file.name)) + ".debug"], cwd=BASE_DIR) # Compliant237+ str(dbg_dir.joinpath(library_file.name)) + ".debug"], cwd=BASE_DIR) # Compliant
238- 238+ 
239- 239+ 
240-def patchelf_dynamic_library():240+def patchelf_dynamic_library():
241- # Process all .so files in lib directory241+ # Process all .so files in lib directory
242- lib_dir = Path(BASE_DIR).joinpath("build/packages/torch_npu/lib")242+ lib_dir = Path(BASE_DIR).joinpath("build/packages/torch_npu/lib")
243- lib_files = [str(i) for i in lib_dir.rglob('*.so')]243+ lib_files = [str(i) for i in lib_dir.rglob('*.so')]
244- 244+ 
245- for library_file in lib_files:245+ for library_file in lib_files:
246- subprocess.check_call(["patchelf", "--remove-needed", "libgomp.so.1", library_file], cwd=BASE_DIR) # Compliant246+ subprocess.check_call(["patchelf", "--remove-needed", "libgomp.so.1", library_file], cwd=BASE_DIR) # Compliant
247- 247+ 
248- 248+ 
249- 249+ 
250-def CppExtension(name, sources, *args, **kwargs):250+def CppExtension(name, sources, *args, **kwargs):
251- r'''251+ r'''
252- Creates a :class:`setuptools.Extension` for C++.252+ Creates a :class:`setuptools.Extension` for C++.
253- '''253+ '''
254- pytorch_dir = get_pytorch_dir()254+ pytorch_dir = get_pytorch_dir()
255- temp_include_dirs = kwargs.get('include_dirs', [])255+ temp_include_dirs = kwargs.get('include_dirs', [])
256- temp_include_dirs.append(os.path.join(pytorch_dir, 'include'))256+ temp_include_dirs.append(os.path.join(pytorch_dir, 'include'))
257- temp_include_dirs.append(os.path.join(pytorch_dir, 'include/torch/csrc/api/include'))257+ temp_include_dirs.append(os.path.join(pytorch_dir, 'include/torch/csrc/api/include'))
258- kwargs['include_dirs'] = temp_include_dirs258+ kwargs['include_dirs'] = temp_include_dirs
259- 259+ 
260- temp_library_dirs = kwargs.get('library_dirs', [])260+ temp_library_dirs = kwargs.get('library_dirs', [])
261- temp_library_dirs.append(os.path.join(pytorch_dir, 'lib'))261+ temp_library_dirs.append(os.path.join(pytorch_dir, 'lib'))
262- temp_library_dirs.append(os.path.join(BASE_DIR, "third_party/acl/libs"))262+ temp_library_dirs.append(os.path.join(BASE_DIR, "third_party/acl/libs"))
263- kwargs['library_dirs'] = temp_library_dirs263+ kwargs['library_dirs'] = temp_library_dirs
264- 264+ 
265- libraries = kwargs.get('libraries', [])265+ libraries = kwargs.get('libraries', [])
266- libraries.append('c10')266+ libraries.append('c10')
267- libraries.append('torch')267+ libraries.append('torch')
268- libraries.append('torch_cpu')268+ libraries.append('torch_cpu')
269- libraries.append('torch_python')269+ libraries.append('torch_python')
270- libraries.append('hccl')270+ libraries.append('hccl')
271- kwargs['libraries'] = libraries271+ kwargs['libraries'] = libraries
272- kwargs['language'] = 'c++'272+ kwargs['language'] = 'c++'
273- return Extension(name, sources, *args, **kwargs)273+ return Extension(name, sources, *args, **kwargs)
274- 274+ 
275- 275+ 
276-class Clean(distutils.command.clean.clean):276+class Clean(distutils.command.clean.clean):
277- 277+ 
278- def run(self):278+ def run(self):
279- f_ignore = open('.gitignore', 'r')279+ f_ignore = open('.gitignore', 'r')
280- ignores = f_ignore.read()280+ ignores = f_ignore.read()
281- pat = re.compile(r'^#( BEGIN NOT-CLEAN-FILES )?')281+ pat = re.compile(r'^#( BEGIN NOT-CLEAN-FILES )?')
282- for wildcard in filter(None, ignores.split('\n')):282+ for wildcard in filter(None, ignores.split('\n')):
283- match = pat.match(wildcard)283+ match = pat.match(wildcard)
284- if match:284+ if match:
285- if match.group(1):285+ if match.group(1):
286- # Marker is found and stop reading .gitignore.286+ # Marker is found and stop reading .gitignore.
287- break287+ break
288- # Ignore lines which begin with '#'.288+ # Ignore lines which begin with '#'.
289- else:289+ else:
290- for filename in glob.glob(wildcard):290+ for filename in glob.glob(wildcard):
291- if os.path.islink(filename):291+ if os.path.islink(filename):
292- raise RuntimeError(f"Failed to remove path: {filename}")292+ raise RuntimeError(f"Failed to remove path: {filename}")
293- if os.path.exists(filename):293+ if os.path.exists(filename):
294- try:294+ try:
295- shutil.rmtree(filename, ignore_errors=True)295+ shutil.rmtree(filename, ignore_errors=True)
296- except Exception as err:296+ except Exception as err:
297- raise RuntimeError(f"Failed to remove path: {filename}") from err297+ raise RuntimeError(f"Failed to remove path: {filename}") from err
298- f_ignore.close()298+ f_ignore.close()
299- 299+ 
300- # It's an old-style class in Python 2.7...300+ # It's an old-style class in Python 2.7...
301- distutils.command.clean.clean.run(self)301+ distutils.command.clean.clean.run(self)
302- 302+ 
303- remove_files = [303+ remove_files = [
304- 'torch_npu/csrc/aten/RegisterCPU.cpp',304+ 'torch_npu/csrc/aten/RegisterCPU.cpp',
305- 'torch_npu/csrc/aten/RegisterNPU.cpp',305+ 'torch_npu/csrc/aten/RegisterNPU.cpp',
306- 'torch_npu/csrc/aten/RegisterAutogradNPU.cpp',306+ 'torch_npu/csrc/aten/RegisterAutogradNPU.cpp',
307- 'torch_npu/csrc/aten/NPUNativeFunctions.h',307+ 'torch_npu/csrc/aten/NPUNativeFunctions.h',
308- 'torch_npu/csrc/aten/CustomRegisterSchema.cpp',308+ 'torch_npu/csrc/aten/CustomRegisterSchema.cpp',
309- 'torch_npu/csrc/aten/ForeachRegister.cpp',309+ 'torch_npu/csrc/aten/ForeachRegister.cpp',
310- 'torch_npu/utils/custom_ops.py',310+ 'torch_npu/utils/custom_ops.py',
311- 'torch_npu/version.py',311+ 'torch_npu/version.py',
312- ]312+ ]
313- for remove_file in remove_files:313+ for remove_file in remove_files:
314- file_path = os.path.join(BASE_DIR, remove_file)314+ file_path = os.path.join(BASE_DIR, remove_file)
315- if os.path.exists(file_path):315+ if os.path.exists(file_path):
316- os.remove(file_path)316+ os.remove(file_path)
317- 317+ 
318-USE_NINJA = os.environ["CMAKE_GENERATOR"].lower() == "ninja" if "CMAKE_GENERATOR" in os.environ else shutil.which("ninja")318+USE_NINJA = os.environ["CMAKE_GENERATOR"].lower() == "ninja" if "CMAKE_GENERATOR" in os.environ else shutil.which("ninja")
319- 319+ 
320-class CPPLibBuild(build_clib, object):320+class CPPLibBuild(build_clib, object):
321- def run(self):321+ def run(self):
322- cmake = get_cmake_command()322+ cmake = get_cmake_command()
323- 323+ 
324- if cmake is None:324+ if cmake is None:
325- raise RuntimeError(325+ raise RuntimeError(
326- "CMake must be installed to build the following extensions: " +326+ "CMake must be installed to build the following extensions: " +
327- ", ".join(e.name for e in self.extensions))327+ ", ".join(e.name for e in self.extensions))
328- self.cmake = cmake328+ self.cmake = cmake
329- 329+ 
330- build_dir = os.path.join(BASE_DIR, "build")330+ build_dir = os.path.join(BASE_DIR, "build")
331- build_type_dir = os.path.join(build_dir)331+ build_type_dir = os.path.join(build_dir)
332- output_lib_path = os.path.join(build_type_dir, "packages/torch_npu/lib")332+ output_lib_path = os.path.join(build_type_dir, "packages/torch_npu/lib")
333- os.makedirs(build_type_dir, exist_ok=True)333+ os.makedirs(build_type_dir, exist_ok=True)
334- os.chmod(build_type_dir, mode=BUILD_PERMISSION)334+ os.chmod(build_type_dir, mode=BUILD_PERMISSION)
335- os.makedirs(output_lib_path, exist_ok=True)335+ os.makedirs(output_lib_path, exist_ok=True)
336- self.build_lib = os.path.relpath(os.path.join(build_dir, "packages/torch_npu"))336+ self.build_lib = os.path.relpath(os.path.join(build_dir, "packages/torch_npu"))
337- self.build_temp = os.path.relpath(build_type_dir)337+ self.build_temp = os.path.relpath(build_type_dir)
338- 338+ 
339- cmake_args = [339+ cmake_args = [
340- '-DCMAKE_BUILD_TYPE=' + get_build_type(),340+ '-DCMAKE_BUILD_TYPE=' + get_build_type(),
341- '-DCMAKE_INSTALL_PREFIX=' + os.path.realpath(output_lib_path),341+ '-DCMAKE_INSTALL_PREFIX=' + os.path.realpath(output_lib_path),
342- '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + os.path.realpath(output_lib_path),342+ '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + os.path.realpath(output_lib_path),
343- '-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY=' + os.path.realpath(output_lib_path),343+ '-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY=' + os.path.realpath(output_lib_path),
344- '-DTORCHNPU_INSTALL_LIBDIR=' + os.path.realpath(output_lib_path),344+ '-DTORCHNPU_INSTALL_LIBDIR=' + os.path.realpath(output_lib_path),
345- '-DPYTHON_INCLUDE_DIR=' + get_paths().get('include'),345+ '-DPYTHON_INCLUDE_DIR=' + get_paths().get('include'),
346- '-DTORCH_VERSION=' + VERSION,346+ '-DTORCH_VERSION=' + VERSION,
347- '-DPYTORCH_INSTALL_DIR=' + get_pytorch_dir()]347+ '-DPYTORCH_INSTALL_DIR=' + get_pytorch_dir()]
348- 348+ 
349- if DISABLE_TORCHAIR == 'FALSE':349+ if DISABLE_TORCHAIR == 'FALSE':
350- if check_torchair_valid(BASE_DIR):350+ if check_torchair_valid(BASE_DIR):
351- cmake_args.append('-DBUILD_TORCHAIR=on')351+ cmake_args.append('-DBUILD_TORCHAIR=on')
352- torchair_install_prefix = os.path.join(build_type_dir, "packages/torch_npu/dynamo/torchair")352+ torchair_install_prefix = os.path.join(build_type_dir, "packages/torch_npu/dynamo/torchair")
353- cmake_args.append(f'-DTORCHAIR_INSTALL_PREFIX={torchair_install_prefix}')353+ cmake_args.append(f'-DTORCHAIR_INSTALL_PREFIX={torchair_install_prefix}')
354- cmake_args.append(f'-DTORCHAIR_TARGET_PYTHON={sys.executable}')354+ cmake_args.append(f'-DTORCHAIR_TARGET_PYTHON={sys.executable}')
355- 355+ 
356- if DISABLE_RPC == 'FALSE':356+ if DISABLE_RPC == 'FALSE':
357- if check_tensorpipe_valid(BASE_DIR):357+ if check_tensorpipe_valid(BASE_DIR):
358- cmake_args.append('-DBUILD_TENSORPIPE=on')358+ cmake_args.append('-DBUILD_TENSORPIPE=on')
359- 359+
360- if ENABLE_LTO == "TRUE":360+ if ENABLE_LTO == "TRUE":
361- cmake_args.append('-DENABLE_LTO=on')361+ cmake_args.append('-DENABLE_LTO=on')
362- if PGO_MODE != 0:362+ if PGO_MODE != 0:
363- cmake_args.append('-DPGO_MODE=' + str(PGO_MODE))363+ cmake_args.append('-DPGO_MODE=' + str(PGO_MODE))
364- 364+
365- if USE_CXX11_ABI:365+ if USE_CXX11_ABI:
366- cmake_args.append('-DGLIBCXX_USE_CXX11_ABI=1')366+ cmake_args.append('-DGLIBCXX_USE_CXX11_ABI=1')
367- 367+ 
368- if os.getenv('_ABI_VERSION') is not None:368+ if os.getenv('_ABI_VERSION') is not None:
369- cmake_args.append('-DABI_VERSION=' + os.getenv('_ABI_VERSION'))369+ cmake_args.append('-DABI_VERSION=' + os.getenv('_ABI_VERSION'))
370- 370+ 
371- if USE_NINJA:371+ if USE_NINJA:
372- cmake_args.append("-GNinja")372+ cmake_args.append("-GNinja")
373- 373+ 
374- max_jobs = os.getenv("MAX_JOBS")374+ max_jobs = os.getenv("MAX_JOBS")
375- if max_jobs is not None or not USE_NINJA:375+ if max_jobs is not None or not USE_NINJA:
376- max_jobs = max_jobs or str(multiprocessing.cpu_count())376+ max_jobs = max_jobs or str(multiprocessing.cpu_count())
377- build_args = ['-j', max_jobs]377+ build_args = ['-j', max_jobs]
378- else:378+ else:
379- build_args = []379+ build_args = []
380- 380+ 
381- subprocess.check_call([self.cmake, BASE_DIR] + cmake_args, cwd=build_type_dir, env=os.environ)381+ subprocess.check_call([self.cmake, BASE_DIR] + cmake_args, cwd=build_type_dir, env=os.environ)
382- for base_dir, dirs, files in os.walk(build_type_dir):382+ for base_dir, dirs, files in os.walk(build_type_dir):
383- for dir_name in dirs:383+ for dir_name in dirs:
384- dir_path = os.path.join(base_dir, dir_name)384+ dir_path = os.path.join(base_dir, dir_name)
385- os.chmod(dir_path, mode=BUILD_PERMISSION)385+ os.chmod(dir_path, mode=BUILD_PERMISSION)
386- for file_name in files:386+ for file_name in files:
387- file_path = os.path.join(base_dir, file_name)387+ file_path = os.path.join(base_dir, file_name)
388- os.chmod(file_path, mode=BUILD_PERMISSION)388+ os.chmod(file_path, mode=BUILD_PERMISSION)
389- 389+ 
390- if USE_NINJA:390+ if USE_NINJA:
391- subprocess.check_call(['ninja'] + build_args, cwd=build_type_dir, env=os.environ)391+ subprocess.check_call(['ninja'] + build_args, cwd=build_type_dir, env=os.environ)
392- else:392+ else:
393- subprocess.check_call(['make'] + build_args, cwd=build_type_dir, env=os.environ)393+ subprocess.check_call(['make'] + build_args, cwd=build_type_dir, env=os.environ)
394- 394+ 
395- 395+ 
396-class Build(build_ext, object):396+class Build(build_ext, object):
397- 397+ 
398- def run(self):398+ def run(self):
399- self.run_command('build_clib')399+ self.run_command('build_clib')
400- self.run_command('build_py')400+ self.run_command('build_py')
401- # proceed with the normal build_ext process401+ # proceed with the normal build_ext process
402- self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))402+ self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))
403- self.build_temp = os.path.relpath(os.path.join(BASE_DIR, "build/temp"))403+ self.build_temp = os.path.relpath(os.path.join(BASE_DIR, "build/temp"))
404- self.library_dirs.append(404+ self.library_dirs.append(
405- os.path.relpath(os.path.join(BASE_DIR, "build/packages/torch_npu/lib"))405+ os.path.relpath(os.path.join(BASE_DIR, "build/packages/torch_npu/lib"))
406- )406+ )
407- super(Build, self).run()407+ super(Build, self).run()
408- 408+ 
409- 409+ 
410-class InstallCmd(install):410+class InstallCmd(install):
411- 411+ 
412- def finalize_options(self) -> None:412+ def finalize_options(self) -> None:
413- self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))413+ self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))
414- return super(InstallCmd, self).finalize_options()414+ return super(InstallCmd, self).finalize_options()
415- 415+ 
416- 416+ 
417-def add_ops_files(base_dir, file_list):417+def add_ops_files(base_dir, file_list):
418- # add ops header files418+ # add ops header files
419- plugin_path = os.path.join(base_dir, 'third_party/op-plugin/op_plugin/include')419+ plugin_path = os.path.join(base_dir, 'third_party/op-plugin/op_plugin/include')
420- if os.path.exists(plugin_path):420+ if os.path.exists(plugin_path):
421- file_list.append('third_party/op-plugin/op_plugin/include/*.h')421+ file_list.append('third_party/op-plugin/op_plugin/include/*.h')
422- plugin_utils_path = os.path.join(base_dir, 'third_party/op-plugin/op_plugin/utils')422+ plugin_utils_path = os.path.join(base_dir, 'third_party/op-plugin/op_plugin/utils')
423- if os.path.exists(plugin_utils_path):423+ if os.path.exists(plugin_utils_path):
424- file_list.append('third_party/op-plugin/op_plugin/utils/*.h')424+ file_list.append('third_party/op-plugin/op_plugin/utils/*.h')
425- return425+ return
426- 426+ 
427- 427+ 
428-def add_ops_python_files(ret_list):428+def add_ops_python_files(ret_list):
429- # add ops python files429+ # add ops python files
430- opplugin_path = os.path.join(BASE_DIR, 'third_party/op-plugin/op_plugin/python')430+ opplugin_path = os.path.join(BASE_DIR, 'third_party/op-plugin/op_plugin/python')
431- 431+ 
432- if os.path.exists(opplugin_path):432+ if os.path.exists(opplugin_path):
433- ops_python_files = glob.glob(os.path.join(opplugin_path, '**/*.py'), recursive=True)433+ ops_python_files = glob.glob(os.path.join(opplugin_path, '**/*.py'), recursive=True)
434- for src in ops_python_files:434+ for src in ops_python_files:
435- dst = os.path.join(435+ dst = os.path.join(
436- os.path.join(BASE_DIR, "build/packages/torch_npu/op_plugin"),436+ os.path.join(BASE_DIR, "build/packages/torch_npu/op_plugin"),
437- os.path.relpath(src, opplugin_path))437+ os.path.relpath(src, opplugin_path))
438- os.makedirs(os.path.dirname(dst), exist_ok=True)438+ os.makedirs(os.path.dirname(dst), exist_ok=True)
439- ret_list.append((src, dst))439+ ret_list.append((src, dst))
440- return440+ return
441- 441+ 
442- 442+ 
443-def get_src_py_and_dst():443+def get_src_py_and_dst():
444- ret = []444+ ret = []
445- generated_python_files = glob.glob(445+ generated_python_files = glob.glob(
446- os.path.join(BASE_DIR, "torch_npu", '**/*.py'),446+ os.path.join(BASE_DIR, "torch_npu", '**/*.py'),
447- recursive=True) + glob.glob(447+ recursive=True) + glob.glob(
448- os.path.join(BASE_DIR, "torch_npu", '**/*.yaml'),448+ os.path.join(BASE_DIR, "torch_npu", '**/*.yaml'),
449- recursive=True) + glob.glob(449+ recursive=True) + glob.glob(
450- os.path.join(BASE_DIR, "torch_npu", 'acl*.json'),450+ os.path.join(BASE_DIR, "torch_npu", 'acl*.json'),
451- recursive=True) + glob.glob(451+ recursive=True) + glob.glob(
452- os.path.join(BASE_DIR, "torch_npu", 'contrib/apis_config.json'),452+ os.path.join(BASE_DIR, "torch_npu", 'contrib/apis_config.json'),
453- recursive=True)453+ recursive=True)
454- for src in generated_python_files:454+ for src in generated_python_files:
455- dst = os.path.join(455+ dst = os.path.join(
456- os.path.join(BASE_DIR, "build/packages/torch_npu"),456+ os.path.join(BASE_DIR, "build/packages/torch_npu"),
457- os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")))457+ os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")))
458- os.makedirs(os.path.dirname(dst), exist_ok=True)458+ os.makedirs(os.path.dirname(dst), exist_ok=True)
459- ret.append((src, dst))459+ ret.append((src, dst))
460- 460+ 
461- add_ops_python_files(ret)461+ add_ops_python_files(ret)
462- 462+ 
463- header_files = [463+ header_files = [
464- "torch_npu/csrc/*.h",464+ "torch_npu/csrc/*.h",
465- "torch_npu/csrc/*/*.h",465+ "torch_npu/csrc/*/*.h",
466- "torch_npu/csrc/*/*.hpp",466+ "torch_npu/csrc/*/*.hpp",
467- "torch_npu/csrc/*/*/*.h",467+ "torch_npu/csrc/*/*/*.h",
468- "torch_npu/csrc/*/*/*/*.h",468+ "torch_npu/csrc/*/*/*/*.h",
469- "torch_npu/csrc/*/*/*/*/*.h",469+ "torch_npu/csrc/*/*/*/*/*.h",
470- "third_party/acl/inc/*/*.h",470+ "third_party/acl/inc/*/*.h",
471- "third_party/hccl/inc/*/*.h",471+ "third_party/hccl/inc/*/*.h",
472- "third_party/acl/inc/*/*/*.h",472+ "third_party/acl/inc/*/*/*.h",
473- "torch_npu/csrc/distributed/HCCLUtils.hpp",473+ "torch_npu/csrc/distributed/HCCLUtils.hpp",
474- "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"474+ "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"
475- ]475+ ]
476- add_ops_files(BASE_DIR, header_files)476+ add_ops_files(BASE_DIR, header_files)
477- glob_header_files = []477+ glob_header_files = []
478- for regex_pattern in header_files:478+ for regex_pattern in header_files:
479- glob_header_files += glob.glob(os.path.join(BASE_DIR, regex_pattern), recursive=True)479+ glob_header_files += glob.glob(os.path.join(BASE_DIR, regex_pattern), recursive=True)
480- 480+ 
481- for src in glob_header_files:481+ for src in glob_header_files:
482- dst = os.path.join(482+ dst = os.path.join(
483- os.path.join(BASE_DIR, "build/packages/torch_npu/include/torch_npu"),483+ os.path.join(BASE_DIR, "build/packages/torch_npu/include/torch_npu"),
484- os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")))484+ os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")))
485- os.makedirs(os.path.dirname(dst), exist_ok=True)485+ os.makedirs(os.path.dirname(dst), exist_ok=True)
486- ret.append((src, dst))486+ ret.append((src, dst))
487- 487+ 
488- torch_header_files = [488+ torch_header_files = [
489- "*/*.h",489+ "*/*.h",
490- "*/*/*.h",490+ "*/*/*.h",
491- "*/*/*/*.h",491+ "*/*/*/*.h",
492- "*/*/*/*/*.h",492+ "*/*/*/*/*.h",
493- "*/*/*/*/*/*.h"493+ "*/*/*/*/*/*.h"
494- ]494+ ]
495- torch_glob_header_files = []495+ torch_glob_header_files = []
496- for regex_pattern in torch_header_files:496+ for regex_pattern in torch_header_files:
497- torch_glob_header_files += glob.glob(os.path.join(BASE_DIR, "patch/include", regex_pattern), recursive=True)497+ torch_glob_header_files += glob.glob(os.path.join(BASE_DIR, "patch/include", regex_pattern), recursive=True)
498- 498+ 
499- for src in torch_glob_header_files:499+ for src in torch_glob_header_files:
500- dst = os.path.join(500+ dst = os.path.join(
501- os.path.join(BASE_DIR, "build/packages/torch_npu/include"),501+ os.path.join(BASE_DIR, "build/packages/torch_npu/include"),
502- os.path.relpath(src, os.path.join(BASE_DIR, "patch/include")))502+ os.path.relpath(src, os.path.join(BASE_DIR, "patch/include")))
503- os.makedirs(os.path.dirname(dst), exist_ok=True)503+ os.makedirs(os.path.dirname(dst), exist_ok=True)
504- ret.append((src, dst))504+ ret.append((src, dst))
505- 505+ 
506- aot_inductor_files = [506+ aot_inductor_files = [
507- # Follow torch v2.6.0.507+ # Follow torch v2.6.0.
508- # These aoti_runtime/*.cpp don't compile to libtorch_npu,508+ # These aoti_runtime/*.cpp don't compile to libtorch_npu,
509- # but act like header files when generate cppwrapper in aot-inductor.509+ # but act like header files when generate cppwrapper in aot-inductor.
510- "torch_npu/_inductor/codegen/aoti_runtime/*.cpp"510+ "torch_npu/_inductor/codegen/aoti_runtime/*.cpp"
511- ]511+ ]
512- glob_aoti_files = []512+ glob_aoti_files = []
513- for regex_pattern in aot_inductor_files:513+ for regex_pattern in aot_inductor_files:
514- glob_aoti_files += glob.glob(514+ glob_aoti_files += glob.glob(
515- os.path.join(BASE_DIR, regex_pattern), recursive=True515+ os.path.join(BASE_DIR, regex_pattern), recursive=True
516- )516+ )
517- 517+ 
518- for src in glob_aoti_files:518+ for src in glob_aoti_files:
519- # Dst: torch_npu/_inductor/codegen/aoti_runtime/*.cpp519+ # Dst: torch_npu/_inductor/codegen/aoti_runtime/*.cpp
520- dst = os.path.join(520+ dst = os.path.join(
521- os.path.join(BASE_DIR, "build/packages/torch_npu/"),521+ os.path.join(BASE_DIR, "build/packages/torch_npu/"),
522- os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")),522+ os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")),
523- )523+ )
524- os.makedirs(os.path.dirname(dst), exist_ok=True)524+ os.makedirs(os.path.dirname(dst), exist_ok=True)
525- ret.append((src, dst))525+ ret.append((src, dst))
526- 526+ 
527- 527+ 
528- def add_torch_npu_codegen(codegen_src_dir, codegen_dst_dir, exclude_root_init=None):528+ def add_torch_npu_codegen(codegen_src_dir, codegen_dst_dir, exclude_root_init=None):
529- """529+ """
530- 复制codegen目录下的文件到目标路径530+ 复制codegen目录下的文件到目标路径
531- :param codegen_src_dir: 源codegen目录531+ :param codegen_src_dir: 源codegen目录
532- :param codegen_dst_dir: 目标目录532+ :param codegen_dst_dir: 目标目录
533- :param exclude_root_init: 需要排除根目录__init__.py的源目录(仅过滤该目录下的__init__.py)533+ :param exclude_root_init: 需要排除根目录__init__.py的源目录(仅过滤该目录下的__init__.py)
534- """534+ """
535- # 匹配需要复制的文件类型535+ # 匹配需要复制的文件类型
536- codegen_files = glob.glob(536+ codegen_files = glob.glob(
537- os.path.join(codegen_src_dir, '**/*.py'), recursive=True537+ os.path.join(codegen_src_dir, '**/*.py'), recursive=True
538- ) + glob.glob(538+ ) + glob.glob(
539- os.path.join(codegen_src_dir, '**/*.yaml'), recursive=True539+ os.path.join(codegen_src_dir, '**/*.yaml'), recursive=True
540- ) + glob.glob(540+ ) + glob.glob(
541- os.path.join(codegen_src_dir, '**/*.json'), recursive=True541+ os.path.join(codegen_src_dir, '**/*.json'), recursive=True
542- ) + glob.glob(542+ ) + glob.glob(
543- os.path.join(codegen_src_dir, '**/*.cpp'), recursive=True543+ os.path.join(codegen_src_dir, '**/*.cpp'), recursive=True
544- ) + glob.glob(544+ ) + glob.glob(
545- os.path.join(codegen_src_dir, '**/*.h'), recursive=True545+ os.path.join(codegen_src_dir, '**/*.h'), recursive=True
546- )546+ )
547- 547+ 
548- # 按原目录结构复制到目标路径548+ # 按原目录结构复制到目标路径
549- for src in codegen_files:549+ for src in codegen_files:
550- # 仅过滤指定目录下的根级__init__.py550+ # 仅过滤指定目录下的根级__init__.py
551- if (exclude_root_init is not None and 551+ if (exclude_root_init is not None and
552- os.path.basename(src) == '__init__.py' and 552+ os.path.basename(src) == '__init__.py' and
553- os.path.dirname(src) == exclude_root_init):553+ os.path.dirname(src) == exclude_root_init):
554- continue # 跳过op-plugin/codegen根目录的__init__.py554+ continue # 跳过op-plugin/codegen根目录的__init__.py
555- 555+
556- # 计算目标路径（保留原目录层级）556+ # 计算目标路径（保留原目录层级）
557- dst = os.path.join(557+ dst = os.path.join(
558- codegen_dst_dir,558+ codegen_dst_dir,
559- os.path.relpath(src, codegen_src_dir) # 保留torchnpugen内部的目录层级559+ os.path.relpath(src, codegen_src_dir) # 保留torchnpugen内部的目录层级
560- )560+ )
561- print(os.path.relpath(src, codegen_src_dir))561+ print(os.path.relpath(src, codegen_src_dir))
562- # 确保目标目录存在562+ # 确保目标目录存在
563- os.makedirs(os.path.dirname(dst), exist_ok=True)563+ os.makedirs(os.path.dirname(dst), exist_ok=True)
564- # 加入文件复制列表564+ # 加入文件复制列表
565- ret.append((src, dst))565+ ret.append((src, dst))
566- 566+ 
567- # 新增：提前创建 torchnpugen 根目录567+ # 新增：提前创建 torchnpugen 根目录
568- torchnpugen_root = os.path.join(BASE_DIR, "build/packages/torchnpugen")568+ torchnpugen_root = os.path.join(BASE_DIR, "build/packages/torchnpugen")
569- os.makedirs(torchnpugen_root, exist_ok=True)569+ os.makedirs(torchnpugen_root, exist_ok=True)
570- # 将codegen复制到package路径570+ # 将codegen复制到package路径
571- codegen_src_dir = os.path.join(BASE_DIR, "torchnpugen")571+ codegen_src_dir = os.path.join(BASE_DIR, "torchnpugen")
572- codegen_dst_dir = os.path.join(BASE_DIR, "build/packages/torchnpugen")572+ codegen_dst_dir = os.path.join(BASE_DIR, "build/packages/torchnpugen")
573- # 复制torch_npu的torchnpugen573+ # 复制torch_npu的torchnpugen
574- add_torch_npu_codegen(codegen_src_dir, codegen_dst_dir)574+ add_torch_npu_codegen(codegen_src_dir, codegen_dst_dir)
575- # 复制op-plugin的torchnpugen（仅过滤其根目录的__init__.py）575+ # 复制op-plugin的torchnpugen（仅过滤其根目录的__init__.py）
576- op_plugin_codegen_src = os.path.join(BASE_DIR, "third_party/op-plugin/torchnpugen")576+ op_plugin_codegen_src = os.path.join(BASE_DIR, "third_party/op-plugin/torchnpugen")
577- add_torch_npu_codegen(577+ add_torch_npu_codegen(
578- op_plugin_codegen_src,578+ op_plugin_codegen_src,
579- codegen_dst_dir,579+ codegen_dst_dir,
580- exclude_root_init=op_plugin_codegen_src # 指定要过滤根目录__init__.py的源目录580+ exclude_root_init=op_plugin_codegen_src # 指定要过滤根目录__init__.py的源目录
581- )581+ )
582- 582+ 
583- return ret583+ return ret
584- 584+ 
585- 585+ 
586-class EggInfoBuild(egg_info, object):586+class EggInfoBuild(egg_info, object):
587- def finalize_options(self):587+ def finalize_options(self):
588- self.egg_base = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))588+ self.egg_base = os.path.relpath(os.path.join(BASE_DIR, "build/packages"))
589- ret = get_src_py_and_dst()589+ ret = get_src_py_and_dst()
590- for src, dst in ret:590+ for src, dst in ret:
591- self.copy_file(src, dst)591+ self.copy_file(src, dst)
592- super(EggInfoBuild, self).finalize_options()592+ super(EggInfoBuild, self).finalize_options()
593- 593+ 
594- 594+ 
595-class PythonPackageBuild(build_py, object):595+class PythonPackageBuild(build_py, object):
596- def run(self) -> None:596+ def run(self) -> None:
597- ret = get_src_py_and_dst()597+ ret = get_src_py_and_dst()
598- for src, dst in ret:598+ for src, dst in ret:
599- self.copy_file(src, dst)599+ self.copy_file(src, dst)
600- super(PythonPackageBuild, self).finalize_options()600+ super(PythonPackageBuild, self).finalize_options()
601- 601+ 
602- 602+ 
603-class BdistWheelBuild(bdist_wheel):603+class BdistWheelBuild(bdist_wheel):
604- def run(self):604+ def run(self):
605- if which('patchelf') is not None:605+ if which('patchelf') is not None:
606- patchelf_dynamic_library()606+ patchelf_dynamic_library()
607- 607+ 
608- if not DEBUG and which('eu-strip') is not None:608+ if not DEBUG and which('eu-strip') is not None:
609- generate_dbg_files_and_strip()609+ generate_dbg_files_and_strip()
610- 610+ 
611- torch_dependencies = ["libc10.so", "libtorch.so", "libtorch_cpu.so", "libtorch_python.so"]611+ torch_dependencies = ["libc10.so", "libtorch.so", "libtorch_cpu.so", "libtorch_python.so"]
612- cann_dependencies = ["libhccl.so", "libascendcl.so", "libacl_op_compiler.so", "libge_runner.so",612+ cann_dependencies = ["libhccl.so", "libascendcl.so", "libacl_op_compiler.so", "libge_runner.so",
613- "libgraph.so", "libacl_tdt_channel.so", "libfmk_parser.so", "libascend_protobuf.so",613+ "libgraph.so", "libacl_tdt_channel.so", "libfmk_parser.so", "libascend_protobuf.so",
614- "libascend_ml.so"]614+ "libascend_ml.so"]
615- other_dependencies = ["libtorch_npu.so", "libnpu_profiler.so", "libgomp.so.1", "libatb.so"]615+ other_dependencies = ["libtorch_npu.so", "libnpu_profiler.so", "libgomp.so.1", "libatb.so"]
616- 616+ 
617- dependencies = torch_dependencies + cann_dependencies + other_dependencies617+ dependencies = torch_dependencies + cann_dependencies + other_dependencies
618- 618+ 
619- bdist_wheel.run(self)619+ bdist_wheel.run(self)
620- 620+ 
621- if is_manylinux:621+ if is_manylinux:
622- file = glob.glob(os.path.join(self.dist_dir, "*linux*.whl"))[0]622+ file = glob.glob(os.path.join(self.dist_dir, "*linux*.whl"))[0]
623- 623+ 
624- auditwheel_cmd = ["auditwheel", "-v", "repair", "-w", self.dist_dir, file]624+ auditwheel_cmd = ["auditwheel", "-v", "repair", "-w", self.dist_dir, file]
625- for i in dependencies:625+ for i in dependencies:
626- auditwheel_cmd += ["--exclude", i]626+ auditwheel_cmd += ["--exclude", i]
627- 627+ 
628- try:628+ try:
629- subprocess.run(auditwheel_cmd, check=True, stdout=subprocess.PIPE)629+ subprocess.run(auditwheel_cmd, check=True, stdout=subprocess.PIPE)
630- finally:630+ finally:
631- os.remove(file)631+ os.remove(file)
632- 632+ 
633- 633+ 
634-build_mode = _get_build_mode()634+build_mode = _get_build_mode()
635-if build_mode not in ['clean']:635+if build_mode not in ['clean']:
636- # Generate bindings code, including RegisterNPU.cpp & NPUNativeFunctions.h.636+ # Generate bindings code, including RegisterNPU.cpp & NPUNativeFunctions.h.
637- generate_bindings_code(BASE_DIR)637+ generate_bindings_code(BASE_DIR)
638- if Path(BASE_DIR).joinpath("third_party/Tensorpipe/third_party/acl/libs").exists():638+ if Path(BASE_DIR).joinpath("third_party/Tensorpipe/third_party/acl/libs").exists():
639- build_stub(Path(BASE_DIR).joinpath("third_party/Tensorpipe"))639+ build_stub(Path(BASE_DIR).joinpath("third_party/Tensorpipe"))
640- build_stub(BASE_DIR)640+ build_stub(BASE_DIR)
641- 641+ 
642-# Setup include directories folders.642+# Setup include directories folders.
643-include_directories = [643+include_directories = [
644- BASE_DIR,644+ BASE_DIR,
645- os.path.join(BASE_DIR, 'patch/include'),645+ os.path.join(BASE_DIR, 'patch/include'),
646- os.path.join(BASE_DIR, 'third_party/hccl/inc'),646+ os.path.join(BASE_DIR, 'third_party/hccl/inc'),
647- os.path.join(BASE_DIR, 'third_party/acl/inc'),647+ os.path.join(BASE_DIR, 'third_party/acl/inc'),
648- os.path.join(BASE_DIR, 'third_party/nlohmann/include')648+ os.path.join(BASE_DIR, 'third_party/nlohmann/include')
649-]649+]
650- 650+ 
651-extra_link_args = []651+extra_link_args = []
652- 652+ 
653-DEBUG = (os.getenv('DEBUG', default='').upper() in ['ON', '1', 'YES', 'TRUE', 'Y'])653+DEBUG = (os.getenv('DEBUG', default='').upper() in ['ON', '1', 'YES', 'TRUE', 'Y'])
654- 654+ 
655-extra_compile_args = [655+extra_compile_args = [
656- '-std=c++17',656+ '-std=c++17',
657- '-Wno-sign-compare',657+ '-Wno-sign-compare',
658- '-Wno-deprecated-declarations',658+ '-Wno-deprecated-declarations',
659- '-Wno-return-type'659+ '-Wno-return-type'
660-]660+]
661- 661+ 
662-if re.match(r'clang', os.getenv('CC', '')):662+if re.match(r'clang', os.getenv('CC', '')):
663- extra_compile_args += [663+ extra_compile_args += [
664- '-Wno-macro-redefined',664+ '-Wno-macro-redefined',
665- '-Wno-return-std-move',665+ '-Wno-return-std-move',
666- ]666+ ]
667- 667+ 
668-if DEBUG:668+if DEBUG:
669- extra_compile_args += ['-O0', '-g']669+ extra_compile_args += ['-O0', '-g']
670- extra_link_args += ['-O0', '-g', '-Wl,-z,now']670+ extra_link_args += ['-O0', '-g', '-Wl,-z,now']
671-else:671+else:
672- extra_compile_args += ['-DNDEBUG']672+ extra_compile_args += ['-DNDEBUG']
673- extra_link_args += ['-Wl,-z,now']673+ extra_link_args += ['-Wl,-z,now']
674- 674+ 
675-# valid manylinux tags675+# valid manylinux tags
676-manylinux_tags = [676+manylinux_tags = [
677- "manylinux1_x86_64",677+ "manylinux1_x86_64",
678- "manylinux2010_x86_64",678+ "manylinux2010_x86_64",
679- "manylinux2014_x86_64",679+ "manylinux2014_x86_64",
680- "manylinux2014_aarch64",680+ "manylinux2014_aarch64",
681- "manylinux_2_5_x86_64",681+ "manylinux_2_5_x86_64",
682- "manylinux_2_12_x86_64",682+ "manylinux_2_12_x86_64",
683- "manylinux_2_17_x86_64",683+ "manylinux_2_17_x86_64",
684- "manylinux_2_17_aarch64",684+ "manylinux_2_17_aarch64",
685- "manylinux_2_24_x86_64",685+ "manylinux_2_24_x86_64",
686- "manylinux_2_24_aarch64",686+ "manylinux_2_24_aarch64",
687- "manylinux_2_27_x86_64",687+ "manylinux_2_27_x86_64",
688- "manylinux_2_27_aarch64",688+ "manylinux_2_27_aarch64",
689- "manylinux_2_28_x86_64",689+ "manylinux_2_28_x86_64",
690- "manylinux_2_28_aarch64",690+ "manylinux_2_28_aarch64",
691- "manylinux_2_31_x86_64",691+ "manylinux_2_31_x86_64",
692- "manylinux_2_31_aarch64",692+ "manylinux_2_31_aarch64",
693- "manylinux_2_34_x86_64",693+ "manylinux_2_34_x86_64",
694- "manylinux_2_34_aarch64",694+ "manylinux_2_34_aarch64",
695- "manylinux_2_35_x86_64",695+ "manylinux_2_35_x86_64",
696- "manylinux_2_35_aarch64",696+ "manylinux_2_35_aarch64",
697-]697+]
698-is_manylinux = os.environ.get("AUDITWHEEL_PLAT", None) in manylinux_tags698+is_manylinux = os.environ.get("AUDITWHEEL_PLAT", None) in manylinux_tags
699- 699+ 
700-readme = os.path.join(BASE_DIR, "README.md")700+readme = os.path.join(BASE_DIR, "README.md")
701-if not os.path.exists(readme):701+if not os.path.exists(readme):
702- raise FileNotFoundError("Unable to find 'README.md'")702+ raise FileNotFoundError("Unable to find 'README.md'")
703-with open(readme, encoding="utf-8") as fdesc:703+with open(readme, encoding="utf-8") as fdesc:
704- long_description = fdesc.read()704+ long_description = fdesc.read()
705- 705+ 
706-classifiers = [706+classifiers = [
707- "Development Status :: 5 - Production/Stable",707+ "Development Status :: 5 - Production/Stable",
708- "Intended Audience :: Developers",708+ "Intended Audience :: Developers",
709- "License :: OSI Approved :: BSD License",709+ "License :: OSI Approved :: BSD License",
710- "Operating System :: POSIX :: Linux",710+ "Operating System :: POSIX :: Linux",
711- "Topic :: Scientific/Engineering",711+ "Topic :: Scientific/Engineering",
712- "Topic :: Scientific/Engineering :: Mathematics",712+ "Topic :: Scientific/Engineering :: Mathematics",
713- "Topic :: Scientific/Engineering :: Artificial Intelligence",713+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
714- "Topic :: Software Development",714+ "Topic :: Software Development",
715- "Topic :: Software Development :: Libraries",715+ "Topic :: Software Development :: Libraries",
716- "Topic :: Software Development :: Libraries :: Python Modules",716+ "Topic :: Software Development :: Libraries :: Python Modules",
717- "Programming Language :: Python",717+ "Programming Language :: Python",
718- "Programming Language :: Python :: 3 :: Only",718+ "Programming Language :: Python :: 3 :: Only",
719- "Programming Language :: Python :: 3.8",719+ "Programming Language :: Python :: 3.8",
720- "Programming Language :: Python :: 3.9",720+ "Programming Language :: Python :: 3.9",
721- "Programming Language :: Python :: 3.10",721+ "Programming Language :: Python :: 3.10",
722- "Programming Language :: Python :: 3.11",722+ "Programming Language :: Python :: 3.11",
723-]723+]
724- 724+ 
725-requirements = ['torch==2.9.0+cpu' if platform.machine() == 'x86_64' else 'torch==2.9.0']725+requirements = ['torch==2.9.0+cpu' if platform.machine() == 'x86_64' else 'torch==2.9.0']
726- 726+ 
727-ext_modules = [CppExtension(727+ext_modules = [CppExtension(
728- 'torch_npu._C',728+ 'torch_npu._C',
729- sources=["torch_npu/csrc/InitNpuBindings.cpp"],729+ sources=["torch_npu/csrc/InitNpuBindings.cpp"],
730- libraries=["torch_npu"],730+ libraries=["torch_npu"],
731- include_dirs=include_directories,731+ include_dirs=include_directories,
732- extra_compile_args=extra_compile_args + ['-fstack-protector-all'] + [732+ extra_compile_args=extra_compile_args + ['-fstack-protector-all'] + [
733- '-D__FILENAME__=\"InitNpuBindings.cpp\"'],733+ '-D__FILENAME__=\"InitNpuBindings.cpp\"'],
734- library_dirs=["lib"],734+ library_dirs=["lib"],
735- extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/lib', '-Wl,-Bsymbolic-functions'],735+ extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/lib', '-Wl,-Bsymbolic-functions'],
736- define_macros=[('_GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0'),736+ define_macros=[('_GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0'),
737- ('GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0')]737+ ('GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0')]
738- )]738+ )]
739- 739+ 
740-setup(740+setup(
741- name=os.environ.get('TORCH_NPU_PACKAGE_NAME', 'torch_npu'),741+ name=os.environ.get('TORCH_NPU_PACKAGE_NAME', 'torch_npu'),
742- version=VERSION,742+ version=VERSION,
743- description='NPU bridge for PyTorch',743+ description='NPU bridge for PyTorch',
744- long_description=long_description,744+ long_description=long_description,
745- long_description_content_type="text/markdown",745+ long_description_content_type="text/markdown",
746- license="BSD License",746+ license="BSD License",
747- classifiers=classifiers,747+ classifiers=classifiers,
748- packages=["torch_npu", "torchnpugen"],748+ packages=["torch_npu", "torchnpugen"],
749- libraries=[('torch_npu', {'sources': list()})],749+ libraries=[('torch_npu', {'sources': list()})],
750- package_dir={'': os.path.relpath(os.path.join(BASE_DIR, "build/packages"))},750+ package_dir={'': os.path.relpath(os.path.join(BASE_DIR, "build/packages"))},
751- ext_modules=ext_modules,751+ ext_modules=ext_modules,
752- install_requires=requirements,752+ install_requires=requirements,
753- extras_require={753+ extras_require={
754- },754+ },
755- package_data={755+ package_data={
756- 'torch_npu': [756+ 'torch_npu': [
757- '*.so',757+ '*.so',
758- 'lib/*.so*',758+ 'lib/*.so*',
759- ],759+ ],
760- 'torchnpugen': [760+ 'torchnpugen': [
761- '*.py', '**/*.py',761+ '*.py', '**/*.py',
762- '*.yaml', '**/*.yaml',762+ '*.yaml', '**/*.yaml',
763- '*.json', '**/*.json',763+ '*.json', '**/*.json',
764- '*.cpp', '**/*.cpp',764+ '*.cpp', '**/*.cpp',
765- '*.h', '**/*.h',765+ '*.h', '**/*.h',
766- ],766+ ],
767- },767+ },
768- cmdclass={768+ cmdclass={
769- 'build_clib': CPPLibBuild,769+ 'build_clib': CPPLibBuild,
770- 'build_ext': Build,770+ 'build_ext': Build,
771- 'build_py': PythonPackageBuild,771+ 'build_py': PythonPackageBuild,
772- 'bdist_wheel': BdistWheelBuild,772+ 'bdist_wheel': BdistWheelBuild,
773- 'install': InstallCmd,773+ 'install': InstallCmd,
774- 'clean': Clean774+ 'clean': Clean
775- },775+ },
776- entry_points={776+ entry_points={
777- 'console_scripts': [777+ 'console_scripts': [
778- 'torch_npu_run = torch_npu.distributed.run:_main',778+ 'torch_npu_run = torch_npu.distributed.run:_main',
779- ],779+ ],
780- 'torch.backends': [780+ 'torch.backends': [
781- 'torch_npu = torch_npu:_autoload',781+ 'torch_npu = torch_npu:_autoload',
782- ],782+ ],
783- }783+ }
784)784)
Mtest/_inductor/test_empty.py+0-1
@@ -42,4 +42,3 @@ instantiate_parametrized_tests(TestEmpty)
42 42 
43if __name__ == "__main__":43if __name__ == "__main__":
44 run_tests()44 run_tests()
45- 
Mtest/_inductor/test_exp.py+0-1
@@ -25,4 +25,3 @@ instantiate_parametrized_tests(TestExp)
25 25 
26if __name__ == "__main__":26if __name__ == "__main__":
27 run_tests()27 run_tests()
28- 
Mtest/_inductor/test_gt.py+0-1
@@ -29,4 +29,3 @@ instantiate_parametrized_tests(TestGt)
29 29 
30if __name__ == "__main__":30if __name__ == "__main__":
31 run_tests()31 run_tests()
32- 
Mtest/_inductor/test_upcast_codegen.py+34-34
@@ -1,35 +1,35 @@
1-import unittest1+import unittest
2-import torch2+import torch
3- 3+ 
4-from testutils import TestUtils4+from testutils import TestUtils
5-from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests5+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
6-from torch._inductor import config6+from torch._inductor import config
7-from torch._inductor.utils import run_and_get_code7+from torch._inductor.utils import run_and_get_code
8- 8+ 
9-import torch_npu9+import torch_npu
10-import torch_npu._inductor10+import torch_npu._inductor
11- 11+ 
12-DEVICE = "npu"12+DEVICE = "npu"
13- 13+ 
14- 14+ 
15-class TestCodegenUpcastToFP32(TestUtils):15+class TestCodegenUpcastToFP32(TestUtils):
16- @parametrize("dtype", [torch.float16, torch.bfloat16])16+ @parametrize("dtype", [torch.float16, torch.bfloat16])
17- @parametrize("upcast_flag", [True, False])17+ @parametrize("upcast_flag", [True, False])
18- def test_codegen_upcast_to_fp32_emits_cast(self, dtype, upcast_flag):18+ def test_codegen_upcast_to_fp32_emits_cast(self, dtype, upcast_flag):
19- @torch.compile(backend="inductor")19+ @torch.compile(backend="inductor")
20- def func(x):20+ def func(x):
21- return torch.abs(x)21+ return torch.abs(x)
22- 22+ 
23- x = torch.randn((1024, 1024), device=DEVICE, dtype=dtype)23+ x = torch.randn((1024, 1024), device=DEVICE, dtype=dtype)
24- 24+ 
25- with config.patch("triton.codegen_upcast_to_fp32", upcast_flag):25+ with config.patch("triton.codegen_upcast_to_fp32", upcast_flag):
26- opt_func = torch._dynamo.optimize("inductor")(func)26+ opt_func = torch._dynamo.optimize("inductor")(func)
27- out, code = run_and_get_code(opt_func, x)27+ out, code = run_and_get_code(opt_func, x)
28- 28+ 
29- self.assertTrue(".to(tl.float32)" in code[0])29+ self.assertTrue(".to(tl.float32)" in code[0])
30- self.assertEqual(func(x), opt_func(x))30+ self.assertEqual(func(x), opt_func(x))
31- 31+ 
32-instantiate_parametrized_tests(TestCodegenUpcastToFP32)32+instantiate_parametrized_tests(TestCodegenUpcastToFP32)
33- 33+ 
34-if __name__ == "__main__":34+if __name__ == "__main__":
35 run_tests()35 run_tests()
Mtest/_inductor/test_use_static_kernel.py+49-49
@@ -1,50 +1,50 @@
1-import unittest1+import unittest
2-import torch2+import torch
3-import torch_npu3+import torch_npu
4- 4+ 
5-from torch.testing._internal.common_utils import (5+from torch.testing._internal.common_utils import (
6- run_tests,6+ run_tests,
7- parametrize,7+ parametrize,
8- instantiate_parametrized_tests,8+ instantiate_parametrized_tests,
9-)9+)
10-from testutils import TestUtils10+from testutils import TestUtils
11- 11+ 
12- 12+ 
13-class TestInductorStaticKernel(TestUtils):13+class TestInductorStaticKernel(TestUtils):
14- 14+ 
15- def simple_op(self, x):15+ def simple_op(self, x):
16- return torch.neg(x)16+ return torch.neg(x)
17- 17+ 
18- @parametrize("shape", [(1024, 1024), (4096,)])18+ @parametrize("shape", [(1024, 1024), (4096,)])
19- @parametrize("dtype", [torch.float16, torch.float32])19+ @parametrize("dtype", [torch.float16, torch.float32])
20- def test_inductor_static_kernel(self, shape, dtype):20+ def test_inductor_static_kernel(self, shape, dtype):
21- device = "npu"21+ device = "npu"
22- 22+ 
23- x = torch.randn(shape, dtype=dtype, device=device)23+ x = torch.randn(shape, dtype=dtype, device=device)
24- 24+ 
25- ref = self.simple_op(x)25+ ref = self.simple_op(x)
26- 26+ 
27- torch._inductor.config.triton.cudagraph_trees = False27+ torch._inductor.config.triton.cudagraph_trees = False
28- torch_npu.npu.aclnn._use_static_aclnn_kernel = True28+ torch_npu.npu.aclnn._use_static_aclnn_kernel = True
29- 29+ 
30- compiled_fn = torch.compile(30+ compiled_fn = torch.compile(
31- self.simple_op,31+ self.simple_op,
32- backend="inductor",32+ backend="inductor",
33- dynamic=False33+ dynamic=False
34- )34+ )
35- 35+ 
36- for _ in range(3):36+ for _ in range(3):
37- compiled_fn(x)37+ compiled_fn(x)
38- 38+ 
39- torch.npu.synchronize()39+ torch.npu.synchronize()
40- 40+ 
41- out = compiled_fn(x)41+ out = compiled_fn(x)
42- torch.npu.synchronize()42+ torch.npu.synchronize()
43- 43+ 
44- self.assertEqual(ref, out)44+ self.assertEqual(ref, out)
45- 45+ 
46-instantiate_parametrized_tests(TestInductorStaticKernel)46+instantiate_parametrized_tests(TestInductorStaticKernel)
47- 47+ 
48-if __name__ == "__main__":48+if __name__ == "__main__":
49- torch.npu.config.allow_internal_format = False49+ torch.npu.config.allow_internal_format = False
50 run_tests()50 run_tests()
Mtest/autograd/test_autograd_fallback.py+30-30
@@ -1,30 +1,30 @@
1-import torch1+import torch
2-from torch.testing._internal.common_utils import (2+from torch.testing._internal.common_utils import (
3- run_tests,3+ run_tests,
4- TestCase,4+ TestCase,
5-)5+)
6-import torch_npu6+import torch_npu
7- 7+ 
8-class TestAutogradFallback(TestCase):8+class TestAutogradFallback(TestCase):
9- 9+ 
10- def test_pad_backward_warn(self):10+ def test_pad_backward_warn(self):
11- 11+ 
12- def _exec_npu_pad():12+ def _exec_npu_pad():
13- npu_input = torch.randn(2, 3).npu()13+ npu_input = torch.randn(2, 3).npu()
14- npu_input.requires_grad = True14+ npu_input.requires_grad = True
15- pads = (1, 1, 1, 1)15+ pads = (1, 1, 1, 1)
16- output = torch_npu.npu_pad(npu_input, pads)16+ output = torch_npu.npu_pad(npu_input, pads)
17- output.backward(torch.ones_like(output))17+ output.backward(torch.ones_like(output))
18- 18+ 
19- # When set to "nothing," calling the reverse function directly causes an error.19+ # When set to "nothing," calling the reverse function directly causes an error.
20- torch._C._set_autograd_fallback_mode("nothing")20+ torch._C._set_autograd_fallback_mode("nothing")
21- with self.assertRaisesRegex(RuntimeError, "does not require grad"):21+ with self.assertRaisesRegex(RuntimeError, "does not require grad"):
22- _exec_npu_pad()22+ _exec_npu_pad()
23- 23+ 
24- # When set to "warn," calling the print function emits a warning.24+ # When set to "warn," calling the print function emits a warning.
25- torch._C._set_autograd_fallback_mode("warn")25+ torch._C._set_autograd_fallback_mode("warn")
26- _exec_npu_pad()26+ _exec_npu_pad()
27- 27+ 
28- 28+ 
29-if __name__ == "__main__":29+if __name__ == "__main__":
30- run_tests()30+ run_tests()
Mtest/custom_ops/test_fast_gelu.py+41-41
@@ -1,41 +1,41 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestFastGelu(TestCase):9+class TestFastGelu(TestCase):
10- 10+ 
11- def supported_op_exec(self, input1):11+ def supported_op_exec(self, input1):
12- attr = 1.70212+ attr = 1.702
13- attr_half = attr / 213+ attr_half = attr / 2
14- abs_input1 = torch.abs(input1)14+ abs_input1 = torch.abs(input1)
15- numerator = input1 * torch.exp((attr_half * input1) * (input1 - abs_input1))15+ numerator = input1 * torch.exp((attr_half * input1) * (input1 - abs_input1))
16- denominator = 1.0 + torch.exp(- attr * abs_input1)16+ denominator = 1.0 + torch.exp(- attr * abs_input1)
17- output = numerator / denominator17+ output = numerator / denominator
18- return output.cpu().detach()18+ return output.cpu().detach()
19- 19+ 
20- def custom_op_exec(self, input1):20+ def custom_op_exec(self, input1):
21- output = torch_npu.fast_gelu(input1)21+ output = torch_npu.fast_gelu(input1)
22- return output.cpu().detach()22+ return output.cpu().detach()
23- 23+ 
24- def test_fast_gelu(self, device="npu"):24+ def test_fast_gelu(self, device="npu"):
25- item = [np.float32, 0, [3, 16, 32]]25+ item = [np.float32, 0, [3, 16, 32]]
26- _, npu_input = create_common_tensor(item, 0, 100)26+ _, npu_input = create_common_tensor(item, 0, 100)
27- 27+ 
28- supported_output = self.supported_op_exec(npu_input)28+ supported_output = self.supported_op_exec(npu_input)
29- custom_output = self.custom_op_exec(npu_input)29+ custom_output = self.custom_op_exec(npu_input)
30- self.assertRtolEqual(supported_output, custom_output)30+ self.assertRtolEqual(supported_output, custom_output)
31- 31+ 
32- def test_fast_gelu_input_arg(self):32+ def test_fast_gelu_input_arg(self):
33- item = [np.float32, 0, [3, 16, 32]]33+ item = [np.float32, 0, [3, 16, 32]]
34- _, npu_input = create_common_tensor(item, 0, 100)34+ _, npu_input = create_common_tensor(item, 0, 100)
35- supported_output = self.supported_op_exec(npu_input)35+ supported_output = self.supported_op_exec(npu_input)
36- custom_output = torch_npu.fast_gelu(input=npu_input)36+ custom_output = torch_npu.fast_gelu(input=npu_input)
37- self.assertRtolEqual(supported_output, custom_output)37+ self.assertRtolEqual(supported_output, custom_output)
38- 38+ 
39- 39+ 
40-if __name__ == "__main__":40+if __name__ == "__main__":
41- run_tests()41+ run_tests()
Mtest/custom_ops/test_npu_confusion_transpose.py+35-35
@@ -1,35 +1,35 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestConfusionTranspose(TestCase):9+class TestConfusionTranspose(TestCase):
10- 10+ 
11- def supported_op_exec(self, input1, perm, shape, transpose_first):11+ def supported_op_exec(self, input1, perm, shape, transpose_first):
12- if transpose_first:12+ if transpose_first:
13- output = input1.permute(*perm).contiguous().view(shape)13+ output = input1.permute(*perm).contiguous().view(shape)
14- else:14+ else:
15- output = input1.view(shape).permute(*perm)15+ output = input1.view(shape).permute(*perm)
16- return output.cpu().detach()16+ return output.cpu().detach()
17- 17+ 
18- def custom_op_exec(self, input1, perm, shape, transpose_first):18+ def custom_op_exec(self, input1, perm, shape, transpose_first):
19- output = torch_npu.npu_confusion_transpose(input1, perm, shape, transpose_first)19+ output = torch_npu.npu_confusion_transpose(input1, perm, shape, transpose_first)
20- return output.cpu().detach()20+ return output.cpu().detach()
21- 21+ 
22- def test_npu_confusion_transpose(self, device="npu"):22+ def test_npu_confusion_transpose(self, device="npu"):
23- item = [np.float32, 0, [1, 576, 2560]]23+ item = [np.float32, 0, [1, 576, 2560]]
24- _, npu_input = create_common_tensor(item, 0, 100)24+ _, npu_input = create_common_tensor(item, 0, 100)
25- perm = (0, 2, 1, 3)25+ perm = (0, 2, 1, 3)
26- shape = [1, 576, 32, 80]26+ shape = [1, 576, 32, 80]
27- transpose_first = False27+ transpose_first = False
28- 28+ 
29- supported_output = self.supported_op_exec(npu_input, perm, shape, transpose_first)29+ supported_output = self.supported_op_exec(npu_input, perm, shape, transpose_first)
30- custom_output = self.custom_op_exec(npu_input, perm, shape, transpose_first)30+ custom_output = self.custom_op_exec(npu_input, perm, shape, transpose_first)
31- self.assertRtolEqual(supported_output, custom_output)31+ self.assertRtolEqual(supported_output, custom_output)
32- 32+ 
33- 33+ 
34-if __name__ == "__main__":34+if __name__ == "__main__":
35- run_tests()35+ run_tests()
Mtest/custom_ops/test_npu_convolution.py+52-52
@@ -1,52 +1,52 @@
1-#1+#
2- 2+ 
3-import numpy as np3+import numpy as np
4-import torch4+import torch
5- 5+ 
6-import torch_npu6+import torch_npu
7-from torch_npu.testing.testcase import TestCase, run_tests7+from torch_npu.testing.testcase import TestCase, run_tests
8-from torch_npu.testing.common_utils import create_common_tensor8+from torch_npu.testing.common_utils import create_common_tensor
9- 9+ 
10- 10+ 
11-class TestConvolution(TestCase):11+class TestConvolution(TestCase):
12- 12+ 
13- def supported_op_exec(self, input1, weight, bias, stride, padding, dilation, groups):13+ def supported_op_exec(self, input1, weight, bias, stride, padding, dilation, groups):
14- dim = input1.dim()14+ dim = input1.dim()
15- if dim == 4:15+ if dim == 4:
16- output = torch.nn.functional.conv2d(input1, weight, bias, stride, padding, dilation, groups)16+ output = torch.nn.functional.conv2d(input1, weight, bias, stride, padding, dilation, groups)
17- if dim == 5:17+ if dim == 5:
18- is_dilated = False18+ is_dilated = False
19- for d in dilation:19+ for d in dilation:
20- is_dilated |= (d != 1)20+ is_dilated |= (d != 1)
21- if groups == 1 and not is_dilated:21+ if groups == 1 and not is_dilated:
22- kernel_size = weight.size()[2]22+ kernel_size = weight.size()[2]
23- output = torch._C._nn.slow_conv3d(input1, weight, kernel_size, bias, stride, padding)23+ output = torch._C._nn.slow_conv3d(input1, weight, kernel_size, bias, stride, padding)
24- else:24+ else:
25- output = torch.nn.functional.conv3d(input1, weight, bias, stride, padding, dilation, groups)25+ output = torch.nn.functional.conv3d(input1, weight, bias, stride, padding, dilation, groups)
26- return output.cpu().detach()26+ return output.cpu().detach()
27- 27+ 
28- def custom_op_exec(self, input1, weight, bias, stride, padding, dilation, groups):28+ def custom_op_exec(self, input1, weight, bias, stride, padding, dilation, groups):
29- output = torch_npu.npu_convolution(input1, weight, bias, stride, padding, dilation, groups)29+ output = torch_npu.npu_convolution(input1, weight, bias, stride, padding, dilation, groups)
30- return output.cpu().detach()30+ return output.cpu().detach()
31- 31+ 
32- def test_npu_convolution(self, device="npu"):32+ def test_npu_convolution(self, device="npu"):
33- items = [[[np.float32, 0, [16, 128, 112, 112]], [np.float32, 0, [256, 128, 3, 3]], [np.float32, 2, [256]],33+ items = [[[np.float32, 0, [16, 128, 112, 112]], [np.float32, 0, [256, 128, 3, 3]], [np.float32, 2, [256]],
34- [1, 1], [1, 1], [1, 1], 1],34+ [1, 1], [1, 1], [1, 1], 1],
35- [[np.float16, 30, [1, 128, 4, 14, 14]], [np.float16, 30, [1, 128, 3, 3, 3]], None,35+ [[np.float16, 30, [1, 128, 4, 14, 14]], [np.float16, 30, [1, 128, 3, 3, 3]], None,
36- [1, 1, 1], [1, 1, 1], [1, 1, 1], 1]]36+ [1, 1, 1], [1, 1, 1], [1, 1, 1], 1]]
37- for item in items:37+ for item in items:
38- _, npu_input = create_common_tensor(item[0], -1, 1)38+ _, npu_input = create_common_tensor(item[0], -1, 1)
39- _, weight = create_common_tensor(item[1], -1, 1)39+ _, weight = create_common_tensor(item[1], -1, 1)
40- _, bias = create_common_tensor(item[2], -1, 1) if item[2] else _, None40+ _, bias = create_common_tensor(item[2], -1, 1) if item[2] else _, None
41- stride = item[3]41+ stride = item[3]
42- padding = item[4]42+ padding = item[4]
43- dilation = item[5]43+ dilation = item[5]
44- groups = item[6]44+ groups = item[6]
45- 45+ 
46- supported_output = self.supported_op_exec(npu_input, weight, bias, stride, padding, dilation, groups)46+ supported_output = self.supported_op_exec(npu_input, weight, bias, stride, padding, dilation, groups)
47- custom_output = self.custom_op_exec(npu_input, weight, bias, stride, padding, dilation, groups)47+ custom_output = self.custom_op_exec(npu_input, weight, bias, stride, padding, dilation, groups)
48- self.assertRtolEqual(supported_output, custom_output)48+ self.assertRtolEqual(supported_output, custom_output)
49- 49+ 
50- 50+ 
51-if __name__ == "__main__":51+if __name__ == "__main__":
52- run_tests()52+ run_tests()
Mtest/custom_ops/test_npu_dtype_cast.py+50-50
@@ -1,50 +1,50 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7-from torch_npu.testing.common_utils import SupportedDevices7+from torch_npu.testing.common_utils import SupportedDevices
8- 8+ 
9- 9+ 
10-class TestDtypeCast(TestCase):10+class TestDtypeCast(TestCase):
11- 11+ 
12- def supported_op_exec(self, input1, dst_dtype):12+ def supported_op_exec(self, input1, dst_dtype):
13- output = input1.to(dst_dtype)13+ output = input1.to(dst_dtype)
14- return output.cpu().detach()14+ return output.cpu().detach()
15- 15+ 
16- def custom_op_exec(self, input1, dst_dtype):16+ def custom_op_exec(self, input1, dst_dtype):
17- output = torch_npu.npu_dtype_cast(input1, dst_dtype)17+ output = torch_npu.npu_dtype_cast(input1, dst_dtype)
18- return output.cpu().detach()18+ return output.cpu().detach()
19- 19+ 
20- def test_npu_dtype_cast(self):20+ def test_npu_dtype_cast(self):
21- item = [np.float32, 0, (64, 10)]21+ item = [np.float32, 0, (64, 10)]
22- _, npu_input = create_common_tensor(item, -1, 1)22+ _, npu_input = create_common_tensor(item, -1, 1)
23- dst_dtype = torch.float1623+ dst_dtype = torch.float16
24- 24+ 
25- supported_output = self.supported_op_exec(npu_input, dst_dtype)25+ supported_output = self.supported_op_exec(npu_input, dst_dtype)
26- custom_output = self.custom_op_exec(npu_input, dst_dtype)26+ custom_output = self.custom_op_exec(npu_input, dst_dtype)
27- self.assertRtolEqual(supported_output, custom_output)27+ self.assertRtolEqual(supported_output, custom_output)
28- 28+ 
29- def test_npu_dtype_cast_double_backward(self):29+ def test_npu_dtype_cast_double_backward(self):
30- x = torch.randn(3, 3, requires_grad=True).to("npu")30+ x = torch.randn(3, 3, requires_grad=True).to("npu")
31- y = torch_npu.npu_dtype_cast(x, torch.half)31+ y = torch_npu.npu_dtype_cast(x, torch.half)
32- z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y))32+ z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y))
33- self.assertIsNone(z[0].grad_fn)33+ self.assertIsNone(z[0].grad_fn)
34- z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y), create_graph=True)34+ z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y), create_graph=True)
35- self.assertIsNotNone(z[0].grad_fn)35+ self.assertIsNotNone(z[0].grad_fn)
36- 36+ 
37- @SupportedDevices(['Ascend910B'])37+ @SupportedDevices(['Ascend910B'])
38- def test_npu_dtype_cast_complex(self):38+ def test_npu_dtype_cast_complex(self):
39- x = torch.empty([2, 3], dtype=torch.complex64, device="npu")39+ x = torch.empty([2, 3], dtype=torch.complex64, device="npu")
40- x.requires_grad_()40+ x.requires_grad_()
41- y = torch_npu.npu_dtype_cast(x, torch.complex128)41+ y = torch_npu.npu_dtype_cast(x, torch.complex128)
42- grad_fn = str(y.grad_fn)42+ grad_fn = str(y.grad_fn)
43- self.assertTrue("NpuDtypeCastBackward" in grad_fn)43+ self.assertTrue("NpuDtypeCastBackward" in grad_fn)
44- 44+
45- with self.assertRaisesRegex(RuntimeError, r'grad can be implicitly created'):45+ with self.assertRaisesRegex(RuntimeError, r'grad can be implicitly created'):
46- y.sum().backward()46+ y.sum().backward()
47- 47+ 
48- 48+ 
49-if __name__ == "__main__":49+if __name__ == "__main__":
50- run_tests()50+ run_tests()
Mtest/custom_ops/test_npu_format_cast.py+31-31
@@ -1,31 +1,31 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestFormatCast(TestCase):9+class TestFormatCast(TestCase):
10- 10+ 
11- def supported_op_exec(self, input1):11+ def supported_op_exec(self, input1):
12- m = torch.nn.Identity(54, unused_argument1=0.1, unused_argument2=False)12+ m = torch.nn.Identity(54, unused_argument1=0.1, unused_argument2=False)
13- output = m(input1)13+ output = m(input1)
14- return output.cpu().detach()14+ return output.cpu().detach()
15- 15+ 
16- def custom_op_exec(self, input1, acl_format):16+ def custom_op_exec(self, input1, acl_format):
17- output = torch_npu.npu_format_cast(input1, acl_format)17+ output = torch_npu.npu_format_cast(input1, acl_format)
18- return output.cpu().detach()18+ return output.cpu().detach()
19- 19+ 
20- def test_npu_format_cast(self, device="npu"):20+ def test_npu_format_cast(self, device="npu"):
21- item = [np.float16, 0, (2, 2, 4, 4)]21+ item = [np.float16, 0, (2, 2, 4, 4)]
22- _, npu_input = create_common_tensor(item, -1, 1)22+ _, npu_input = create_common_tensor(item, -1, 1)
23- acl_format = 323+ acl_format = 3
24- 24+ 
25- supported_output = self.supported_op_exec(npu_input)25+ supported_output = self.supported_op_exec(npu_input)
26- custom_output = self.custom_op_exec(npu_input, acl_format)26+ custom_output = self.custom_op_exec(npu_input, acl_format)
27- self.assertRtolEqual(supported_output, custom_output)27+ self.assertRtolEqual(supported_output, custom_output)
28- 28+ 
29- 29+ 
30-if __name__ == "__main__":30+if __name__ == "__main__":
31- run_tests()31+ run_tests()
Mtest/custom_ops/test_npu_grid_assign_positive.py+44-44
@@ -1,44 +1,44 @@
1-import torch1+import torch
2- 2+ 
3-import torch_npu3+import torch_npu
4-from torch_npu.testing.testcase import TestCase, run_tests4+from torch_npu.testing.testcase import TestCase, run_tests
5- 5+ 
6- 6+ 
7-class TestGridAssignPositive(TestCase):7+class TestGridAssignPositive(TestCase):
8- 8+ 
9- def supported_op_exec(self, input1, box_responsible_flags, max_overlaps, argmax_overlaps, pos_iou_thr):9+ def supported_op_exec(self, input1, box_responsible_flags, max_overlaps, argmax_overlaps, pos_iou_thr):
10- pos_inds = (max_overlaps > pos_iou_thr) & box_responsible_flags.type(torch.bool)10+ pos_inds = (max_overlaps > pos_iou_thr) & box_responsible_flags.type(torch.bool)
11- argmax_overlaps = argmax_overlaps.to(input1.dtype)11+ argmax_overlaps = argmax_overlaps.to(input1.dtype)
12- input1[pos_inds] = argmax_overlaps[pos_inds] + 112+ input1[pos_inds] = argmax_overlaps[pos_inds] + 1
13- return input1.cpu().detach()13+ return input1.cpu().detach()
14- 14+ 
15- def custom_op_exec(self, input1, overlaps, box_responsible_flags, max_overlaps, argmax_overlaps,15+ def custom_op_exec(self, input1, overlaps, box_responsible_flags, max_overlaps, argmax_overlaps,
16- gt_max_overlaps, gt_argmax_overlaps, num_gts, pos_iou_thr, min_pos_iou, gt_max_assign_all):16+ gt_max_overlaps, gt_argmax_overlaps, num_gts, pos_iou_thr, min_pos_iou, gt_max_assign_all):
17- output = torch_npu.npu_grid_assign_positive(input1, overlaps, box_responsible_flags, max_overlaps,17+ output = torch_npu.npu_grid_assign_positive(input1, overlaps, box_responsible_flags, max_overlaps,
18- argmax_overlaps, gt_max_overlaps, gt_argmax_overlaps, num_gts,18+ argmax_overlaps, gt_max_overlaps, gt_argmax_overlaps, num_gts,
19- pos_iou_thr, min_pos_iou, gt_max_assign_all)19+ pos_iou_thr, min_pos_iou, gt_max_assign_all)
20- return output.cpu().detach()20+ return output.cpu().detach()
21- 21+ 
22- def test_npu_grid_assign_positive(self):22+ def test_npu_grid_assign_positive(self):
23- npu_input = torch.rand((4,), dtype=torch.float32).to("npu")23+ npu_input = torch.rand((4,), dtype=torch.float32).to("npu")
24- overlaps = torch.rand((2, 4), dtype=torch.float32).to("npu")24+ overlaps = torch.rand((2, 4), dtype=torch.float32).to("npu")
25- box_responsible_flags = torch.tensor([1, 1, 1, 0], dtype=torch.uint8).to("npu")25+ box_responsible_flags = torch.tensor([1, 1, 1, 0], dtype=torch.uint8).to("npu")
26- max_overlaps = torch.rand((4,), dtype=torch.float32).to("npu")26+ max_overlaps = torch.rand((4,), dtype=torch.float32).to("npu")
27- argmax_overlaps = torch.tensor([1, 0, 1, 0], dtype=torch.int32).to("npu")27+ argmax_overlaps = torch.tensor([1, 0, 1, 0], dtype=torch.int32).to("npu")
28- gt_max_overlaps = torch.rand((2,), dtype=torch.float32).to("npu")28+ gt_max_overlaps = torch.rand((2,), dtype=torch.float32).to("npu")
29- gt_argmax_overlaps = torch.tensor([1, 0], dtype=torch.int32).to("npu")29+ gt_argmax_overlaps = torch.tensor([1, 0], dtype=torch.int32).to("npu")
30- num_gts = 12830+ num_gts = 128
31- pos_iou_thr = .531+ pos_iou_thr = .5
32- min_pos_iou = .032+ min_pos_iou = .0
33- gt_max_assign_all = True33+ gt_max_assign_all = True
34- 34+ 
35- supported_output = self.supported_op_exec(npu_input, box_responsible_flags, max_overlaps,35+ supported_output = self.supported_op_exec(npu_input, box_responsible_flags, max_overlaps,
36- argmax_overlaps, pos_iou_thr)36+ argmax_overlaps, pos_iou_thr)
37- custom_output = self.custom_op_exec(npu_input, overlaps, box_responsible_flags, max_overlaps,37+ custom_output = self.custom_op_exec(npu_input, overlaps, box_responsible_flags, max_overlaps,
38- argmax_overlaps, gt_max_overlaps, gt_argmax_overlaps, num_gts,38+ argmax_overlaps, gt_max_overlaps, gt_argmax_overlaps, num_gts,
39- pos_iou_thr, min_pos_iou, gt_max_assign_all)39+ pos_iou_thr, min_pos_iou, gt_max_assign_all)
40- self.assertRtolEqual(supported_output, custom_output)40+ self.assertRtolEqual(supported_output, custom_output)
41- 41+ 
42- 42+ 
43-if __name__ == "__main__":43+if __name__ == "__main__":
44- run_tests()44+ run_tests()
Mtest/custom_ops/test_npu_ps_roi_pooling.py+94-94
@@ -1,94 +1,94 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestPSROIPooling(TestCase):9+class TestPSROIPooling(TestCase):
10- 10+ 
11- def cal_sum(self, hstart, hend, wstart, wend, image):11+ def cal_sum(self, hstart, hend, wstart, wend, image):
12- out_sum = 0.012+ out_sum = 0.0
13- for row in range(hstart, hend):13+ for row in range(hstart, hend):
14- for col in range(wstart, wend):14+ for col in range(wstart, wend):
15- out_sum += image[row][col]15+ out_sum += image[row][col]
16- return out_sum16+ return out_sum
17- 17+ 
18- def supported_op_exec(self, input1, rois, spatial_scale, group_size, output_dim):18+ def supported_op_exec(self, input1, rois, spatial_scale, group_size, output_dim):
19- dst_type = input1.dtype19+ dst_type = input1.dtype
20- if dst_type != torch.float32:20+ if dst_type != torch.float32:
21- input1 = input1.to(torch.float32)21+ input1 = input1.to(torch.float32)
22- rois = rois.to(torch.float32)22+ rois = rois.to(torch.float32)
23- 23+ 
24- n, channels, height, width = input1.shape24+ n, channels, height, width = input1.shape
25- tensor_height = torch.tensor([height]).npu()25+ tensor_height = torch.tensor([height]).npu()
26- tensor_width = torch.tensor([width]).npu()26+ tensor_width = torch.tensor([width]).npu()
27- tensor_zero = torch.tensor([0]).npu()27+ tensor_zero = torch.tensor([0]).npu()
28- tensor_one_tenth = torch.tensor([0.1]).npu()28+ tensor_one_tenth = torch.tensor([0.1]).npu()
29- 29+ 
30- output_size = [rois.size(0) * rois.size(2), output_dim, group_size, group_size]30+ output_size = [rois.size(0) * rois.size(2), output_dim, group_size, group_size]
31- rois = rois.transpose(2, 1)31+ rois = rois.transpose(2, 1)
32- rois = torch.reshape(rois, (rois.shape[0] * rois.shape[1], rois.shape[2]))32+ rois = torch.reshape(rois, (rois.shape[0] * rois.shape[1], rois.shape[2]))
33- output = torch.zeros(output_size).npu()33+ output = torch.zeros(output_size).npu()
34- for out_n in range(output_size[0]):34+ for out_n in range(output_size[0]):
35- for ctop in range(output_size[1]):35+ for ctop in range(output_size[1]):
36- for ph in range(output_size[2]):36+ for ph in range(output_size[2]):
37- for pw in range(output_size[3]):37+ for pw in range(output_size[3]):
38- roi_batch_ind = rois[out_n][0].to(int)38+ roi_batch_ind = rois[out_n][0].to(int)
39- roi_start_w = torch.round(rois[out_n][1]) * spatial_scale39+ roi_start_w = torch.round(rois[out_n][1]) * spatial_scale
40- roi_start_h = torch.round(rois[out_n][2]) * spatial_scale40+ roi_start_h = torch.round(rois[out_n][2]) * spatial_scale
41- roi_end_w = torch.round(rois[out_n][3] + 1.0) * spatial_scale41+ roi_end_w = torch.round(rois[out_n][3] + 1.0) * spatial_scale
42- roi_end_h = torch.round(rois[out_n][4] + 1.0) * spatial_scale42+ roi_end_h = torch.round(rois[out_n][4] + 1.0) * spatial_scale
43- # Force too small ROIs to be 1x143+ # Force too small ROIs to be 1x1
44- roi_width = torch.max(roi_end_w - roi_start_w, tensor_one_tenth)44+ roi_width = torch.max(roi_end_w - roi_start_w, tensor_one_tenth)
45- roi_height = torch.max(roi_end_h - roi_start_h, tensor_one_tenth)45+ roi_height = torch.max(roi_end_h - roi_start_h, tensor_one_tenth)
46- # Compute w and h at bottom46+ # Compute w and h at bottom
47- bin_size_h = roi_height / group_size47+ bin_size_h = roi_height / group_size
48- bin_size_w = roi_width / group_size48+ bin_size_w = roi_width / group_size
49- 49+ 
50- # Add roi offsets and clip to input boundaries50+ # Add roi offsets and clip to input boundaries
51- hstart = (ph * bin_size_h + roi_start_h).to(int)51+ hstart = (ph * bin_size_h + roi_start_h).to(int)
52- wstart = torch.floor(pw * bin_size_w + roi_start_w).to(int)52+ wstart = torch.floor(pw * bin_size_w + roi_start_w).to(int)
53- hend = torch.ceil((ph + 1) * bin_size_h + roi_start_h).to(int)53+ hend = torch.ceil((ph + 1) * bin_size_h + roi_start_h).to(int)
54- wend = torch.ceil((pw + 1) * bin_size_w + roi_start_w).to(int)54+ wend = torch.ceil((pw + 1) * bin_size_w + roi_start_w).to(int)
55- hstart = torch.min(torch.max(hstart, tensor_zero), tensor_height).to(int)55+ hstart = torch.min(torch.max(hstart, tensor_zero), tensor_height).to(int)
56- hend = torch.min(torch.max(hend, tensor_zero), tensor_height).to(int)56+ hend = torch.min(torch.max(hend, tensor_zero), tensor_height).to(int)
57- wstart = torch.min(torch.max(wstart, tensor_zero), tensor_width).to(int)57+ wstart = torch.min(torch.max(wstart, tensor_zero), tensor_width).to(int)
58- wend = torch.min(torch.max(wend, tensor_zero), tensor_width).to(int)58+ wend = torch.min(torch.max(wend, tensor_zero), tensor_width).to(int)
59- is_empty = ((hend <= hstart) or (wend <= wstart))59+ is_empty = ((hend <= hstart) or (wend <= wstart))
60- c = (ctop * group_size + ph) * group_size + pw60+ c = (ctop * group_size + ph) * group_size + pw
61- out_sum = self.cal_sum(hstart,61+ out_sum = self.cal_sum(hstart,
62- hend,62+ hend,
63- wstart,63+ wstart,
64- wend,64+ wend,
65- image=input1[roi_batch_ind][c])65+ image=input1[roi_batch_ind][c])
66- bin_area = (hend - hstart) * (wend - wstart)66+ bin_area = (hend - hstart) * (wend - wstart)
67- if is_empty:67+ if is_empty:
68- output[out_n][ctop][ph][pw] = 0.068+ output[out_n][ctop][ph][pw] = 0.0
69- else:69+ else:
70- output[out_n][ctop][ph][pw] = out_sum / bin_area70+ output[out_n][ctop][ph][pw] = out_sum / bin_area
71- 71+ 
72- if dst_type != torch.float32:72+ if dst_type != torch.float32:
73- output = output.to(dst_type)73+ output = output.to(dst_type)
74- return output.cpu().detach()74+ return output.cpu().detach()
75- 75+ 
76- def custom_op_exec(self, input1, rois, spatial_scale, group_size, output_dim):76+ def custom_op_exec(self, input1, rois, spatial_scale, group_size, output_dim):
77- output = torch_npu.npu_ps_roi_pooling(input1, rois, spatial_scale, group_size, output_dim)77+ output = torch_npu.npu_ps_roi_pooling(input1, rois, spatial_scale, group_size, output_dim)
78- return output.cpu().detach()78+ return output.cpu().detach()
79- 79+ 
80- def test_npu_ps_roi_pooling(self, device="npu"):80+ def test_npu_ps_roi_pooling(self, device="npu"):
81- item = [np.float32, 0, (2, 961, 127, 127)]81+ item = [np.float32, 0, (2, 961, 127, 127)]
82- _, npu_input = create_common_tensor(item, 0.1, 1)82+ _, npu_input = create_common_tensor(item, 0.1, 1)
83- rois = torch.tensor([[[0], [3], [2], [9], [9]], [[1], [1], [4], [9], [9]]], dtype=torch.float32).npu()83+ rois = torch.tensor([[[0], [3], [2], [9], [9]], [[1], [1], [4], [9], [9]]], dtype=torch.float32).npu()
84- spatial_scale = 0.2584+ spatial_scale = 0.25
85- group_size = 3185+ group_size = 31
86- output_dim = 186+ output_dim = 1
87- 87+ 
88- supported_output = self.supported_op_exec(npu_input, rois, spatial_scale, group_size, output_dim)88+ supported_output = self.supported_op_exec(npu_input, rois, spatial_scale, group_size, output_dim)
89- custom_output = self.custom_op_exec(npu_input, rois, spatial_scale, group_size, output_dim)89+ custom_output = self.custom_op_exec(npu_input, rois, spatial_scale, group_size, output_dim)
90- self.assertRtolEqual(supported_output, custom_output)90+ self.assertRtolEqual(supported_output, custom_output)
91- 91+ 
92- 92+ 
93-if __name__ == "__main__":93+if __name__ == "__main__":
94- run_tests()94+ run_tests()
Mtest/custom_ops/test_npu_softmax_cross_entropy_with_logits.py+31-31
@@ -1,31 +1,31 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestSoftmaxCrossEntropyWithLogits(TestCase):9+class TestSoftmaxCrossEntropyWithLogits(TestCase):
10- 10+ 
11- def supported_op_exec(self, input1, label):11+ def supported_op_exec(self, input1, label):
12- softmax = torch.nn.functional.softmax(input1)12+ softmax = torch.nn.functional.softmax(input1)
13- log_softmax = torch.log(softmax)13+ log_softmax = torch.log(softmax)
14- loss = torch.sum(- label * log_softmax, dim=1)14+ loss = torch.sum(- label * log_softmax, dim=1)
15- return loss.cpu().detach()15+ return loss.cpu().detach()
16- 16+ 
17- def custom_op_exec(self, input1, label):17+ def custom_op_exec(self, input1, label):
18- output = torch_npu.npu_softmax_cross_entropy_with_logits(input1, label)18+ output = torch_npu.npu_softmax_cross_entropy_with_logits(input1, label)
19- return output.cpu().detach()19+ return output.cpu().detach()
20- 20+ 
21- def test_npu_softmax_cross_entropy_with_logits(self, device="npu"):21+ def test_npu_softmax_cross_entropy_with_logits(self, device="npu"):
22- item = [np.float32, 0, (64, 10)]22+ item = [np.float32, 0, (64, 10)]
23- _, npu_input = create_common_tensor(item, -1, 1)23+ _, npu_input = create_common_tensor(item, -1, 1)
24- _, label = create_common_tensor(item, 0, 1)24+ _, label = create_common_tensor(item, 0, 1)
25- supported_output = self.supported_op_exec(npu_input, label)25+ supported_output = self.supported_op_exec(npu_input, label)
26- custom_output = self.custom_op_exec(npu_input, label)26+ custom_output = self.custom_op_exec(npu_input, label)
27- self.assertRtolEqual(supported_output, custom_output)27+ self.assertRtolEqual(supported_output, custom_output)
28- 28+ 
29- 29+ 
30-if __name__ == "__main__":30+if __name__ == "__main__":
31- run_tests()31+ run_tests()
Mtest/custom_ops/test_scatter_update.py+115-115
@@ -1,115 +1,115 @@
1-import numpy as np1+import numpy as np
2-import torch2+import torch
3- 3+ 
4-import torch_npu4+import torch_npu
5-from torch_npu.testing.testcase import TestCase, run_tests5+from torch_npu.testing.testcase import TestCase, run_tests
6-from torch_npu.testing.common_utils import create_common_tensor6+from torch_npu.testing.common_utils import create_common_tensor
7- 7+ 
8- 8+ 
9-class TestScatterUpdate(TestCase):9+class TestScatterUpdate(TestCase):
10- 10+ 
11- def supported_scatter_update_exec(self, var, updates, start, length, axis=-2):11+ def supported_scatter_update_exec(self, var, updates, start, length, axis=-2):
12- var_input = var.clone()12+ var_input = var.clone()
13- var_input[:, :, start: start + length, :] = updates13+ var_input[:, :, start: start + length, :] = updates
14- return var_input14+ return var_input
15- 15+ 
16- def custom_scatter_update_exec(self, var, indices, updates, axis=-2):16+ def custom_scatter_update_exec(self, var, indices, updates, axis=-2):
17- output = torch_npu.scatter_update(var, indices, updates, axis)17+ output = torch_npu.scatter_update(var, indices, updates, axis)
18- return output.cpu().detach()18+ return output.cpu().detach()
19- 19+ 
20- 20+ 
21- def supported_scatter_update__exec(self, var, updates, start, length, axis=-2):21+ def supported_scatter_update__exec(self, var, updates, start, length, axis=-2):
22- var[:, :, start: start + length, :] = updates22+ var[:, :, start: start + length, :] = updates
23- return var23+ return var
24- 24+ 
25- def custom_scatter_update__exec(self, var, indices, updates, axis=-2):25+ def custom_scatter_update__exec(self, var, indices, updates, axis=-2):
26- torch_npu.scatter_update_(var, indices, updates, axis)26+ torch_npu.scatter_update_(var, indices, updates, axis)
27- return var.cpu().detach()27+ return var.cpu().detach()
28- 28+ 
29- def custom_scatter_update__exec_return(self, var, indices, updates, axis=-2):29+ def custom_scatter_update__exec_return(self, var, indices, updates, axis=-2):
30- result = torch_npu.scatter_update_(var, indices, updates, axis)30+ result = torch_npu.scatter_update_(var, indices, updates, axis)
31- return result.cpu().detach()31+ return result.cpu().detach()
32- 32+ 
33- 33+ 
34- def test_scatter_update(self, device="npu"):34+ def test_scatter_update(self, device="npu"):
35- # input_dtype, slice_start, slice_length35+ # input_dtype, slice_start, slice_length
36- items = [36+ items = [
37- [torch.float16, 0, 1],37+ [torch.float16, 0, 1],
38- [torch.float16, 2, 1],38+ [torch.float16, 2, 1],
39- [torch.float16, 4, 8],39+ [torch.float16, 4, 8],
40- [torch.float16, 0, 16],40+ [torch.float16, 0, 16],
41- [torch.float32, 0, 1],41+ [torch.float32, 0, 1],
42- [torch.float32, 2, 1],42+ [torch.float32, 2, 1],
43- [torch.float32, 4, 8],43+ [torch.float32, 4, 8],
44- [torch.float32, 0, 16],44+ [torch.float32, 0, 16],
45- ]45+ ]
46- 46+ 
47- for item in items:47+ for item in items:
48- in_self_cpu = torch.randn(4, 8, 16, 64, dtype=item[0])48+ in_self_cpu = torch.randn(4, 8, 16, 64, dtype=item[0])
49- in_self_npu = in_self_cpu.npu()49+ in_self_npu = in_self_cpu.npu()
50- in_update_cpu = torch.randn(4, 8, item[2], 64, dtype=item[0])50+ in_update_cpu = torch.randn(4, 8, item[2], 64, dtype=item[0])
51- in_update_npu = in_update_cpu.npu()51+ in_update_npu = in_update_cpu.npu()
52- in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1]]).npu()52+ in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1]]).npu()
53- 53+ 
54- supported_output = self.supported_scatter_update_exec(in_self_cpu, in_update_cpu, item[1], item[2])54+ supported_output = self.supported_scatter_update_exec(in_self_cpu, in_update_cpu, item[1], item[2])
55- custom_output = self.custom_scatter_update_exec(in_self_npu, in_indices_npu, in_update_npu)55+ custom_output = self.custom_scatter_update_exec(in_self_npu, in_indices_npu, in_update_npu)
56- self.assertRtolEqual(supported_output, custom_output)56+ self.assertRtolEqual(supported_output, custom_output)
57- # check whether the custom operator modifies the input_self57+ # check whether the custom operator modifies the input_self
58- self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())58+ self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())
59- 59+ 
60- 60+ 
61- def test_scatter_update_(self, device="npu"):61+ def test_scatter_update_(self, device="npu"):
62- # input_dtype, slice_start, slice_length62+ # input_dtype, slice_start, slice_length
63- items = [63+ items = [
64- [torch.float16, 0, 1],64+ [torch.float16, 0, 1],
65- [torch.float16, 4, 1],65+ [torch.float16, 4, 1],
66- [torch.float16, 8, 32],66+ [torch.float16, 8, 32],
67- [torch.float16, 0, 64],67+ [torch.float16, 0, 64],
68- [torch.float32, 0, 1],68+ [torch.float32, 0, 1],
69- [torch.float32, 4, 1],69+ [torch.float32, 4, 1],
70- [torch.float32, 8, 32],70+ [torch.float32, 8, 32],
71- [torch.float32, 0, 64],71+ [torch.float32, 0, 64],
72- ]72+ ]
73- for item in items:73+ for item in items:
74- in_self_cpu = torch.randn(4, 8, 64, 128, dtype=item[0])74+ in_self_cpu = torch.randn(4, 8, 64, 128, dtype=item[0])
75- in_self_npu = in_self_cpu.npu()75+ in_self_npu = in_self_cpu.npu()
76- in_update_cpu = torch.randn(4, 8, item[2], 128, dtype=item[0])76+ in_update_cpu = torch.randn(4, 8, item[2], 128, dtype=item[0])
77- in_update_npu = in_update_cpu.npu()77+ in_update_npu = in_update_cpu.npu()
78- in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1]]).npu()78+ in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1]]).npu()
79- 79+ 
80- supported_output = self.supported_scatter_update__exec(in_self_cpu, in_update_cpu, item[1], item[2])80+ supported_output = self.supported_scatter_update__exec(in_self_cpu, in_update_cpu, item[1], item[2])
81- custom_output = self.custom_scatter_update__exec(in_self_npu, in_indices_npu, in_update_npu)81+ custom_output = self.custom_scatter_update__exec(in_self_npu, in_indices_npu, in_update_npu)
82- self.assertRtolEqual(supported_output, custom_output)82+ self.assertRtolEqual(supported_output, custom_output)
83- # check whether the custom operator modifies the input_self83+ # check whether the custom operator modifies the input_self
84- self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())84+ self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())
85- 85+ 
86- 86+ 
87- def test_scatter_update__return(self, device="npu"):87+ def test_scatter_update__return(self, device="npu"):
88- # input_dtype, slice_start, slice_length88+ # input_dtype, slice_start, slice_length
89- items = [89+ items = [
90- [torch.float16, 0, 1],90+ [torch.float16, 0, 1],
91- [torch.float16, 8, 1],91+ [torch.float16, 8, 1],
92- [torch.float16, 16, 8],92+ [torch.float16, 16, 8],
93- [torch.float16, 0, 32],93+ [torch.float16, 0, 32],
94- [torch.float32, 0, 1],94+ [torch.float32, 0, 1],
95- [torch.float32, 8, 1],95+ [torch.float32, 8, 1],
96- [torch.float32, 16, 8],96+ [torch.float32, 16, 8],
97- [torch.float32, 0, 32],97+ [torch.float32, 0, 32],
98- ]98+ ]
99- for item in items:99+ for item in items:
100- in_self_cpu = torch.randn(12, 2, 64, 128, dtype=item[0])100+ in_self_cpu = torch.randn(12, 2, 64, 128, dtype=item[0])
101- in_self_npu = in_self_cpu.npu()101+ in_self_npu = in_self_cpu.npu()
102- in_update_cpu = torch.randn(12, 2, item[2], 128, dtype=item[0])102+ in_update_cpu = torch.randn(12, 2, item[2], 128, dtype=item[0])
103- in_update_npu = in_update_cpu.npu()103+ in_update_npu = in_update_cpu.npu()
104- in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1], item[1], item[1],104+ in_indices_npu = torch.tensor([item[1], item[1], item[1], item[1], item[1], item[1],
105- item[1], item[1], item[1], item[1], item[1], item[1]]).npu()105+ item[1], item[1], item[1], item[1], item[1], item[1]]).npu()
106- 106+ 
107- supported_output = self.supported_scatter_update__exec(in_self_cpu, in_update_cpu, item[1], item[2])107+ supported_output = self.supported_scatter_update__exec(in_self_cpu, in_update_cpu, item[1], item[2])
108- custom_output = self.custom_scatter_update__exec(in_self_npu, in_indices_npu, in_update_npu)108+ custom_output = self.custom_scatter_update__exec(in_self_npu, in_indices_npu, in_update_npu)
109- self.assertRtolEqual(supported_output, custom_output)109+ self.assertRtolEqual(supported_output, custom_output)
110- # check whether the custom operator modifies the input_self110+ # check whether the custom operator modifies the input_self
111- self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())111+ self.assertRtolEqual(in_self_cpu, in_self_npu.cpu().detach())
112- 112+ 
113- 113+ 
114-if __name__ == "__main__":114+if __name__ == "__main__":
115- run_tests()115+ run_tests()
Mtest/distributed/elastic/events/test_events_api.py+319-319
@@ -1,320 +1,320 @@
1-"""1+"""
2-Add validation cases for torch.distributed.elastic.events.record API:2+Add validation cases for torch.distributed.elastic.events.record API:
3- 3+ 
4-1. PyTorch community lacks direct test cases for torch.distributed.elastic.events.record4+1. PyTorch community lacks direct test cases for torch.distributed.elastic.events.record
5- in the standard test suite, so this file is added.5+ in the standard test suite, so this file is added.
6- 6+ 
7-2. This file validates the following APIs:7+2. This file validates the following APIs:
8- torch.distributed.elastic.events.record (This is a pure Python-level event logging API with no hardware dependency.)8+ torch.distributed.elastic.events.record (This is a pure Python-level event logging API with no hardware dependency.)
9- torch.distributed.elastic.events.api.EventMetadataValue (Type alias for metadata values)9+ torch.distributed.elastic.events.api.EventMetadataValue (Type alias for metadata values)
10- (extendable)10+ (extendable)
11-"""11+"""
12- 12+ 
13-import json13+import json
14-import time14+import time
15-from typing import Union, Optional, get_args, get_origin15+from typing import Union, Optional, get_args, get_origin
16-from unittest.mock import patch, MagicMock16+from unittest.mock import patch, MagicMock
17- 17+ 
18-import torch18+import torch
19-from torch.distributed.elastic.events import record, get_logging_handler19+from torch.distributed.elastic.events import record, get_logging_handler
20-from torch.distributed.elastic.events.api import Event, EventSource, EventMetadataValue20+from torch.distributed.elastic.events.api import Event, EventSource, EventMetadataValue
21-from torch.testing._internal.common_utils import TestCase, run_tests21+from torch.testing._internal.common_utils import TestCase, run_tests
22- 22+ 
23- 23+ 
24-class TestEventsRecord(TestCase):24+class TestEventsRecord(TestCase):
25- """Test torch.distributed.elastic.events.record method."""25+ """Test torch.distributed.elastic.events.record method."""
26- 26+ 
27- DESTINATION_NULL = "null"27+ DESTINATION_NULL = "null"
28- DESTINATION_CONSOLE = "console"28+ DESTINATION_CONSOLE = "console"
29- TEST_EVENT_PREFIX = "test_event_"29+ TEST_EVENT_PREFIX = "test_event_"
30- 30+ 
31- def tearDown(self):31+ def tearDown(self):
32- """Clean up resources after each test case to ensure test isolation."""32+ """Clean up resources after each test case to ensure test isolation."""
33- from torch.distributed.elastic.events import _events_loggers33+ from torch.distributed.elastic.events import _events_loggers
34- _events_loggers.clear()34+ _events_loggers.clear()
35- 35+ 
36- if hasattr(get_logging_handler, "cache_clear"):36+ if hasattr(get_logging_handler, "cache_clear"):
37- get_logging_handler.cache_clear()37+ get_logging_handler.cache_clear()
38- 38+ 
39- patch.stopall()39+ patch.stopall()
40- super().tearDown()40+ super().tearDown()
41- 41+ 
42- def test_record_null_destination(self):42+ def test_record_null_destination(self):
43- """Verify that record does not raise an exception when using the default null destination."""43+ """Verify that record does not raise an exception when using the default null destination."""
44- event = Event(44+ event = Event(
45- name=f"{self.TEST_EVENT_PREFIX}null",45+ name=f"{self.TEST_EVENT_PREFIX}null",
46- source=EventSource.WORKER,46+ source=EventSource.WORKER,
47- metadata={"key": "value"}47+ metadata={"key": "value"}
48- )48+ )
49- record(event)49+ record(event)
50- 50+ 
51- def test_record_console_destination(self):51+ def test_record_console_destination(self):
52- """Verify that record does not raise an exception when using the console destination."""52+ """Verify that record does not raise an exception when using the console destination."""
53- event = Event(53+ event = Event(
54- name=f"{self.TEST_EVENT_PREFIX}console",54+ name=f"{self.TEST_EVENT_PREFIX}console",
55- source=EventSource.AGENT,55+ source=EventSource.AGENT,
56- metadata={"stage": "init"}56+ metadata={"stage": "init"}
57- )57+ )
58- record(event, destination=self.DESTINATION_CONSOLE)58+ record(event, destination=self.DESTINATION_CONSOLE)
59- 59+ 
60- def test_record_with_timestamp(self):60+ def test_record_with_timestamp(self):
61- """Verify that record can correctly log an event with a custom or auto-generated timestamp."""61+ """Verify that record can correctly log an event with a custom or auto-generated timestamp."""
62- custom_timestamp = int(time.time() * 1000)62+ custom_timestamp = int(time.time() * 1000)
63- event = Event(63+ event = Event(
64- name=f"{self.TEST_EVENT_PREFIX}timestamp_custom",64+ name=f"{self.TEST_EVENT_PREFIX}timestamp_custom",
65- source=EventSource.WORKER,65+ source=EventSource.WORKER,
66- timestamp=custom_timestamp,66+ timestamp=custom_timestamp,
67- metadata={"ts": custom_timestamp}67+ metadata={"ts": custom_timestamp}
68- )68+ )
69- record(event, destination=self.DESTINATION_NULL)69+ record(event, destination=self.DESTINATION_NULL)
70- self.assertEqual(event.timestamp, custom_timestamp)70+ self.assertEqual(event.timestamp, custom_timestamp)
71- self.assertIsInstance(event.timestamp, int)71+ self.assertIsInstance(event.timestamp, int)
72- self.assertGreaterEqual(event.timestamp, 0)72+ self.assertGreaterEqual(event.timestamp, 0)
73- 73+ 
74- event_with_auto_timestamp = Event(74+ event_with_auto_timestamp = Event(
75- name=f"{self.TEST_EVENT_PREFIX}timestamp_auto",75+ name=f"{self.TEST_EVENT_PREFIX}timestamp_auto",
76- source=EventSource.WORKER76+ source=EventSource.WORKER
77- )77+ )
78- record(event_with_auto_timestamp)78+ record(event_with_auto_timestamp)
79- self.assertIsNotNone(event_with_auto_timestamp.timestamp)79+ self.assertIsNotNone(event_with_auto_timestamp.timestamp)
80- self.assertIsInstance(event_with_auto_timestamp.timestamp, int)80+ self.assertIsInstance(event_with_auto_timestamp.timestamp, int)
81- self.assertGreaterEqual(event_with_auto_timestamp.timestamp, 0)81+ self.assertGreaterEqual(event_with_auto_timestamp.timestamp, 0)
82- 82+ 
83- @patch("torch.distributed.elastic.events._get_or_create_logger")83+ @patch("torch.distributed.elastic.events._get_or_create_logger")
84- def test_record_with_various_metadata_types(self, mock_get_logger):84+ def test_record_with_various_metadata_types(self, mock_get_logger):
85- """Verify that record correctly passes various metadata types to the underlying logger."""85+ """Verify that record correctly passes various metadata types to the underlying logger."""
86- mock_logger = MagicMock()86+ mock_logger = MagicMock()
87- mock_get_logger.return_value = mock_logger87+ mock_get_logger.return_value = mock_logger
88- 88+ 
89- metadata = {89+ metadata = {
90- "str_val": "hello",90+ "str_val": "hello",
91- "int_val": 42,91+ "int_val": 42,
92- "float_val": 3.14,92+ "float_val": 3.14,
93- "bool_val": True,93+ "bool_val": True,
94- "none_val": None,94+ "none_val": None,
95- "list_val": [1, 2, 3],95+ "list_val": [1, 2, 3],
96- "nested_dict": {"sub_key": "sub_val"}96+ "nested_dict": {"sub_key": "sub_val"}
97- }97+ }
98- event = Event(98+ event = Event(
99- name=f"{self.TEST_EVENT_PREFIX}metadata_types",99+ name=f"{self.TEST_EVENT_PREFIX}metadata_types",
100- source=EventSource.AGENT,100+ source=EventSource.AGENT,
101- metadata=metadata101+ metadata=metadata
102- )102+ )
103- record(event, destination=self.DESTINATION_CONSOLE)103+ record(event, destination=self.DESTINATION_CONSOLE)
104- 104+ 
105- # Verify the internal call chain105+ # Verify the internal call chain
106- mock_get_logger.assert_called_once_with(self.DESTINATION_CONSOLE)106+ mock_get_logger.assert_called_once_with(self.DESTINATION_CONSOLE)
107- mock_logger.info.assert_called_once()107+ mock_logger.info.assert_called_once()
108- 108+ 
109- # Verify that the serialized metadata is correctly preserved109+ # Verify that the serialized metadata is correctly preserved
110- logged_raw = mock_logger.info.call_args[0][0]110+ logged_raw = mock_logger.info.call_args[0][0]
111- logged_data = json.loads(logged_raw)111+ logged_data = json.loads(logged_raw)
112- self.assertEqual(logged_data["metadata"], metadata)112+ self.assertEqual(logged_data["metadata"], metadata)
113- 113+ 
114- def test_record_multiple_events(self):114+ def test_record_multiple_events(self):
115- """Verify that consecutive calls to record do not interfere with each other."""115+ """Verify that consecutive calls to record do not interfere with each other."""
116- event1 = Event(116+ event1 = Event(
117- name=f"{self.TEST_EVENT_PREFIX}multiple_1",117+ name=f"{self.TEST_EVENT_PREFIX}multiple_1",
118- source=EventSource.WORKER,118+ source=EventSource.WORKER,
119- metadata={"seq": 1}119+ metadata={"seq": 1}
120- )120+ )
121- event2 = Event(121+ event2 = Event(
122- name=f"{self.TEST_EVENT_PREFIX}multiple_2",122+ name=f"{self.TEST_EVENT_PREFIX}multiple_2",
123- source=EventSource.WORKER,123+ source=EventSource.WORKER,
124- metadata={"seq": 2}124+ metadata={"seq": 2}
125- )125+ )
126- 126+ 
127- record(event1)127+ record(event1)
128- record(event2)128+ record(event2)
129- 129+ 
130- self.assertIn("seq", event1.metadata)130+ self.assertIn("seq", event1.metadata)
131- self.assertEqual(event1.metadata["seq"], 1)131+ self.assertEqual(event1.metadata["seq"], 1)
132- self.assertIn("seq", event2.metadata)132+ self.assertIn("seq", event2.metadata)
133- self.assertEqual(event2.metadata["seq"], 2)133+ self.assertEqual(event2.metadata["seq"], 2)
134- 134+ 
135- @patch("torch.distributed.elastic.events._get_or_create_logger")135+ @patch("torch.distributed.elastic.events._get_or_create_logger")
136- def test_record_calls_get_or_create_logger(self, mock_get_logger):136+ def test_record_calls_get_or_create_logger(self, mock_get_logger):
137- """137+ """
138- Verify that record internally calls _get_or_create_logger to obtain the logger,138+ Verify that record internally calls _get_or_create_logger to obtain the logger,
139- and then calls .info() on that logger with the serialized event.139+ and then calls .info() on that logger with the serialized event.
140- """140+ """
141- mock_logger = MagicMock()141+ mock_logger = MagicMock()
142- mock_get_logger.return_value = mock_logger142+ mock_get_logger.return_value = mock_logger
143- 143+ 
144- event = Event(144+ event = Event(
145- name=f"{self.TEST_EVENT_PREFIX}mock",145+ name=f"{self.TEST_EVENT_PREFIX}mock",
146- source=EventSource.WORKER146+ source=EventSource.WORKER
147- )147+ )
148- record(event, destination=self.DESTINATION_CONSOLE)148+ record(event, destination=self.DESTINATION_CONSOLE)
149- 149+ 
150- mock_get_logger.assert_called_once_with(self.DESTINATION_CONSOLE)150+ mock_get_logger.assert_called_once_with(self.DESTINATION_CONSOLE)
151- mock_logger.info.assert_called_once_with(event.serialize())151+ mock_logger.info.assert_called_once_with(event.serialize())
152- 152+ 
153- def test_record_event_name_empty(self):153+ def test_record_event_name_empty(self):
154- """Verify that record does not raise an exception when the event name is an empty string."""154+ """Verify that record does not raise an exception when the event name is an empty string."""
155- event = Event(name="", source=EventSource.WORKER)155+ event = Event(name="", source=EventSource.WORKER)
156- record(event)156+ record(event)
157- 157+ 
158- def test_record_event_name_special_chars(self):158+ def test_record_event_name_special_chars(self):
159- """Verify event names with special characters (spaces, symbols, unicode) are handled gracefully."""159+ """Verify event names with special characters (spaces, symbols, unicode) are handled gracefully."""
160- special_names = ["test event", "test@event", "测试事件", "a" * 256]160+ special_names = ["test event", "test@event", "测试事件", "a" * 256]
161- for name in special_names:161+ for name in special_names:
162- with self.subTest(name=name):162+ with self.subTest(name=name):
163- event = Event(name=name, source=EventSource.WORKER)163+ event = Event(name=name, source=EventSource.WORKER)
164- record(event, destination=self.DESTINATION_NULL)164+ record(event, destination=self.DESTINATION_NULL)
165- 165+ 
166- def test_record_all_event_sources(self):166+ def test_record_all_event_sources(self):
167- """Verify record compatibility with all EventSource enum values."""167+ """Verify record compatibility with all EventSource enum values."""
168- for source in EventSource:168+ for source in EventSource:
169- with self.subTest(source=source):169+ with self.subTest(source=source):
170- event = Event(170+ event = Event(
171- name=f"{self.TEST_EVENT_PREFIX}source_{source.name.lower()}",171+ name=f"{self.TEST_EVENT_PREFIX}source_{source.name.lower()}",
172- source=source,172+ source=source,
173- metadata={"source_type": source.name}173+ metadata={"source_type": source.name}
174- )174+ )
175- record(event)175+ record(event)
176- self.assertEqual(event.source, source)176+ self.assertEqual(event.source, source)
177- self.assertIn("source_type", event.metadata)177+ self.assertIn("source_type", event.metadata)
178- self.assertEqual(event.metadata["source_type"], source.name)178+ self.assertEqual(event.metadata["source_type"], source.name)
179- 179+ 
180- 180+ 
181-class TestEventMetadataValue(TestCase):181+class TestEventMetadataValue(TestCase):
182- """Test torch.distributed.elastic.events.api.EventMetadataValue type alias."""182+ """Test torch.distributed.elastic.events.api.EventMetadataValue type alias."""
183- 183+ 
184- def test_event_metadata_value_is_defined(self):184+ def test_event_metadata_value_is_defined(self):
185- """Verify that EventMetadataValue is exported and defined."""185+ """Verify that EventMetadataValue is exported and defined."""
186- self.assertIsNotNone(EventMetadataValue)186+ self.assertIsNotNone(EventMetadataValue)
187- self.assertIs(get_origin(EventMetadataValue), Union)187+ self.assertIs(get_origin(EventMetadataValue), Union)
188- 188+ 
189- def test_event_metadata_value_type_structure(self):189+ def test_event_metadata_value_type_structure(self):
190- """190+ """
191- Verify that EventMetadataValue is a Union of str, int, float, bool, None.191+ Verify that EventMetadataValue is a Union of str, int, float, bool, None.
192- Expected: Optional[Union[str, int, float, bool]]192+ Expected: Optional[Union[str, int, float, bool]]
193- """193+ """
194- origin = get_origin(EventMetadataValue)194+ origin = get_origin(EventMetadataValue)
195- args = get_args(EventMetadataValue)195+ args = get_args(EventMetadataValue)
196- 196+ 
197- self.assertIs(origin, Union)197+ self.assertIs(origin, Union)
198- self.assertIn(str, args)198+ self.assertIn(str, args)
199- self.assertIn(int, args)199+ self.assertIn(int, args)
200- self.assertIn(float, args)200+ self.assertIn(float, args)
201- self.assertIn(bool, args)201+ self.assertIn(bool, args)
202- self.assertIn(type(None), args)202+ self.assertIn(type(None), args)
203- self.assertEqual(len(args), 5)203+ self.assertEqual(len(args), 5)
204- 204+ 
205- def test_event_metadata_value_legal_primitives(self):205+ def test_event_metadata_value_legal_primitives(self):
206- """Verify that all legal primitive values are accepted by Event metadata."""206+ """Verify that all legal primitive values are accepted by Event metadata."""
207- test_cases = [207+ test_cases = [
208- ("string_val", "hello"),208+ ("string_val", "hello"),
209- ("int_val", 42),209+ ("int_val", 42),
210- ("int_val_neg", -7),210+ ("int_val_neg", -7),
211- ("float_val", 3.14159),211+ ("float_val", 3.14159),
212- ("float_val_zero", 0.0),212+ ("float_val_zero", 0.0),
213- ("bool_true", True),213+ ("bool_true", True),
214- ("bool_false", False),214+ ("bool_false", False),
215- ("none_val", None),215+ ("none_val", None),
216- ]216+ ]
217- 217+ 
218- for key, value in test_cases:218+ for key, value in test_cases:
219- with self.subTest(key=key, value=value):219+ with self.subTest(key=key, value=value):
220- event = Event(220+ event = Event(
221- name=f"test_metadata_{key}",221+ name=f"test_metadata_{key}",
222- source=EventSource.WORKER,222+ source=EventSource.WORKER,
223- metadata={key: value}223+ metadata={key: value}
224- )224+ )
225- self.assertIn(key, event.metadata)225+ self.assertIn(key, event.metadata)
226- self.assertEqual(event.metadata[key], value)226+ self.assertEqual(event.metadata[key], value)
227- record(event, destination=TestEventsRecord.DESTINATION_NULL)227+ record(event, destination=TestEventsRecord.DESTINATION_NULL)
228- 228+ 
229- def test_event_metadata_value_none_explicit(self):229+ def test_event_metadata_value_none_explicit(self):
230- """Verify that None is explicitly allowed as a metadata value."""230+ """Verify that None is explicitly allowed as a metadata value."""
231- event = Event(231+ event = Event(
232- name="test_metadata_none",232+ name="test_metadata_none",
233- source=EventSource.AGENT,233+ source=EventSource.AGENT,
234- metadata={"explicit_none": None}234+ metadata={"explicit_none": None}
235- )235+ )
236- self.assertIsNone(event.metadata["explicit_none"])236+ self.assertIsNone(event.metadata["explicit_none"])
237- record(event, destination=TestEventsRecord.DESTINATION_NULL)237+ record(event, destination=TestEventsRecord.DESTINATION_NULL)
238- 238+ 
239- def test_event_metadata_value_mixed_dict(self):239+ def test_event_metadata_value_mixed_dict(self):
240- """Verify that a metadata dict containing all legal types can be constructed and recorded."""240+ """Verify that a metadata dict containing all legal types can be constructed and recorded."""
241- metadata = {241+ metadata = {
242- "epoch": 10,242+ "epoch": 10,
243- "loss": 0.1234,243+ "loss": 0.1234,
244- "model_name": "resnet50",244+ "model_name": "resnet50",
245- "is_training": True,245+ "is_training": True,
246- "checkpoint_path": None,246+ "checkpoint_path": None,
247- }247+ }
248- event = Event(248+ event = Event(
249- name="test_metadata_mixed",249+ name="test_metadata_mixed",
250- source=EventSource.WORKER,250+ source=EventSource.WORKER,
251- metadata=metadata251+ metadata=metadata
252- )252+ )
253- self.assertEqual(len(event.metadata), 5)253+ self.assertEqual(len(event.metadata), 5)
254- record(event, destination=TestEventsRecord.DESTINATION_NULL)254+ record(event, destination=TestEventsRecord.DESTINATION_NULL)
255- 255+ 
256- def test_event_metadata_value_serialization_roundtrip(self):256+ def test_event_metadata_value_serialization_roundtrip(self):
257- """257+ """
258- Verify that metadata values survive the Event.serialize() roundtrip258+ Verify that metadata values survive the Event.serialize() roundtrip
259- and remain as their original Python types.259+ and remain as their original Python types.
260- """260+ """
261- metadata = {261+ metadata = {
262- "lr": 0.01,262+ "lr": 0.01,
263- "step": 100,263+ "step": 100,
264- "tag": "train",264+ "tag": "train",
265- "enabled": False,265+ "enabled": False,
266- "optional": None,266+ "optional": None,
267- }267+ }
268- event = Event(268+ event = Event(
269- name="test_serialization",269+ name="test_serialization",
270- source=EventSource.AGENT,270+ source=EventSource.AGENT,
271- metadata=metadata271+ metadata=metadata
272- )272+ )
273- serialized = event.serialize()273+ serialized = event.serialize()
274- 274+ 
275- self.assertIsInstance(serialized, str)275+ self.assertIsInstance(serialized, str)
276- data = json.loads(serialized)276+ data = json.loads(serialized)
277- self.assertIn("metadata", data)277+ self.assertIn("metadata", data)
278- 278+ 
279- serialized_metadata = data["metadata"]279+ serialized_metadata = data["metadata"]
280- self.assertEqual(serialized_metadata["lr"], 0.01)280+ self.assertEqual(serialized_metadata["lr"], 0.01)
281- self.assertEqual(serialized_metadata["step"], 100)281+ self.assertEqual(serialized_metadata["step"], 100)
282- self.assertEqual(serialized_metadata["tag"], "train")282+ self.assertEqual(serialized_metadata["tag"], "train")
283- self.assertEqual(serialized_metadata["enabled"], False)283+ self.assertEqual(serialized_metadata["enabled"], False)
284- self.assertIsNone(serialized_metadata["optional"])284+ self.assertIsNone(serialized_metadata["optional"])
285- 285+ 
286- def test_event_serialize_empty_metadata(self):286+ def test_event_serialize_empty_metadata(self):
287- """Verify that empty metadata serializes correctly without data loss."""287+ """Verify that empty metadata serializes correctly without data loss."""
288- event = Event(288+ event = Event(
289- name="test_empty_metadata",289+ name="test_empty_metadata",
290- source=EventSource.WORKER,290+ source=EventSource.WORKER,
291- metadata={}291+ metadata={}
292- )292+ )
293- data = json.loads(event.serialize())293+ data = json.loads(event.serialize())
294- self.assertEqual(data["metadata"], {})294+ self.assertEqual(data["metadata"], {})
295- 295+ 
296- def test_event_serialize_unsupported_metadata_type(self):296+ def test_event_serialize_unsupported_metadata_type(self):
297- """297+ """
298- Verify that unsupported metadata types (e.g. set, bytes) raise TypeError298+ Verify that unsupported metadata types (e.g. set, bytes) raise TypeError
299- during serialization / record, since Event itself does not validate at construction time.299+ during serialization / record, since Event itself does not validate at construction time.
300- """300+ """
301- invalid_metadata_cases = [301+ invalid_metadata_cases = [
302- ("set_val", {1, 2, 3}),302+ ("set_val", {1, 2, 3}),
303- ("bytes_val", b"test"),303+ ("bytes_val", b"test"),
304- ("object_val", MagicMock()),304+ ("object_val", MagicMock()),
305- ]305+ ]
306- 306+ 
307- for key, value in invalid_metadata_cases:307+ for key, value in invalid_metadata_cases:
308- with self.subTest(key=key):308+ with self.subTest(key=key):
309- event = Event(309+ event = Event(
310- name=f"test_invalid_{key}",310+ name=f"test_invalid_{key}",
311- source=EventSource.WORKER,311+ source=EventSource.WORKER,
312- metadata={key: value}312+ metadata={key: value}
313- )313+ )
314- # Event construction succeeds, but serialize/record should fail314+ # Event construction succeeds, but serialize/record should fail
315- with self.assertRaises(TypeError):315+ with self.assertRaises(TypeError):
316- event.serialize()316+ event.serialize()
317- 317+ 
318- 318+ 
319-if __name__ == "__main__":319+if __name__ == "__main__":
320 run_tests()320 run_tests()
Mtest/distributed/rpc/multiprocessing/test_reduction_new.py+0-1
@@ -127,4 +127,3 @@ class TestReduction(TestCase):
127 127 
128if __name__ == "__main__":128if __name__ == "__main__":
129 run_tests()129 run_tests()
130- 
Mtest/distributed/test_all_to_all_single.py+1-0
@@ -133,6 +133,7 @@ class HcclAlltoAllSingleTest(TestCase):
133 HcclAlltoAllSingleTest._test_alltoall_single_2p,133 HcclAlltoAllSingleTest._test_alltoall_single_2p,
134 HcclAlltoAllSingleTest._init_dist_hccl)134 HcclAlltoAllSingleTest._init_dist_hccl)
135 135 
136+ @unittest.skip("Disabled during A1 to A2 chip transition")
136 @skipIfUnsupportMultiNPU(2)137 @skipIfUnsupportMultiNPU(2)
137 def test_alltoall_single_2p_size_dist(self):138 def test_alltoall_single_2p_size_dist(self):
138 self._test_multiprocess_2p(139 self._test_multiprocess_2p(
Mtest/distributed/test_device_mesh.py+1155-1155
@@ -1,1156 +1,1156 @@
1-# Copyright (c) Meta Platforms, Inc. and affiliates1+# Copyright (c) Meta Platforms, Inc. and affiliates
2-# Owner(s): ["oncall: distributed"]2+# Owner(s): ["oncall: distributed"]
3-import os3+import os
4-from functools import wraps4+from functools import wraps
5-from typing import Tuple, Dict, Any5+from typing import Tuple, Dict, Any
6- 6+ 
7-import torch7+import torch
8-import torch.distributed as dist8+import torch.distributed as dist
9-import torch.distributed._functional_collectives as funcol9+import torch.distributed._functional_collectives as funcol
10-from torch.distributed._tensor import DTensor10+from torch.distributed._tensor import DTensor
11-from torch.distributed.device_mesh import _mesh_resources, DeviceMesh, init_device_mesh11+from torch.distributed.device_mesh import _mesh_resources, DeviceMesh, init_device_mesh
12-from torch.distributed.distributed_c10d import (12+from torch.distributed.distributed_c10d import (
13- _get_default_group,13+ _get_default_group,
14- _world,14+ _world,
15- get_global_rank,15+ get_global_rank,
16- get_world_size,16+ get_world_size,
17- init_process_group,17+ init_process_group,
18- is_initialized,18+ is_initialized,
19- new_group,19+ new_group,
20- ProcessGroup,20+ ProcessGroup,
21-)21+)
22-from torch.distributed.tensor._collective_utils import (22+from torch.distributed.tensor._collective_utils import (
23- mesh_broadcast,23+ mesh_broadcast,
24- mesh_scatter,24+ mesh_scatter,
25- unpad_tensor,25+ unpad_tensor,
26-)26+)
27-from torch.distributed.tensor.placement_types import _Partial, Shard27+from torch.distributed.tensor.placement_types import _Partial, Shard
28-from torch.testing._internal.distributed.fake_pg import FakeStore28+from torch.testing._internal.distributed.fake_pg import FakeStore
29-from torch.utils._typing_utils import not_none29+from torch.utils._typing_utils import not_none
30- 30+ 
31-import torch_npu31+import torch_npu
32-from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU32+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
33-from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase33+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
34-from torch_npu.testing.testcase import run_tests34+from torch_npu.testing.testcase import run_tests
35- 35+ 
36- 36+ 
37-def _get_device_type(world_size):37+def _get_device_type(world_size):
38- if (38+ if (
39- torch.npu.is_available()39+ torch.npu.is_available()
40- and torch.npu.device_count() >= world_size40+ and torch.npu.device_count() >= world_size
41- and torch.distributed.is_hccl_available()41+ and torch.distributed.is_hccl_available()
42- ):42+ ):
43- device_type = "npu"43+ device_type = "npu"
44- else:44+ else:
45- device_type = "cpu"45+ device_type = "cpu"
46- return device_type46+ return device_type
47- 47+ 
48- 48+ 
49-def _set_env_var(addr="localhost", port="29500", world_size=1, rank=0):49+def _set_env_var(addr="localhost", port="29500", world_size=1, rank=0):
50- os.environ["MASTER_ADDR"] = addr50+ os.environ["MASTER_ADDR"] = addr
51- os.environ["MASTER_PORT"] = port51+ os.environ["MASTER_PORT"] = port
52- os.environ["WORLD_SIZE"] = f"{world_size}"52+ os.environ["WORLD_SIZE"] = f"{world_size}"
53- os.environ["RANK"] = f"{rank}"53+ os.environ["RANK"] = f"{rank}"
54- 54+ 
55- 55+ 
56-class DeviceMeshTest(NPUDTensorTestBase):56+class DeviceMeshTest(NPUDTensorTestBase):
57- @property57+ @property
58- def world_size(self):58+ def world_size(self):
59- return 259+ return 2
60- 60+ 
61- @skipIfUnsupportMultiNPU(2)61+ @skipIfUnsupportMultiNPU(2)
62- def test_init_process_group(self):62+ def test_init_process_group(self):
63- device_type = _get_device_type(self.world_size)63+ device_type = _get_device_type(self.world_size)
64- mesh_tensor = torch.arange(2).reshape(2, 1)64+ mesh_tensor = torch.arange(2).reshape(2, 1)
65- self.assertTrue(not is_initialized())65+ self.assertTrue(not is_initialized())
66- _set_env_var(world_size=self.world_size, rank=self.rank)66+ _set_env_var(world_size=self.world_size, rank=self.rank)
67- DeviceMesh(device_type, mesh_tensor)67+ DeviceMesh(device_type, mesh_tensor)
68- self.assertTrue(is_initialized())68+ self.assertTrue(is_initialized())
69- self.destroy_pg()69+ self.destroy_pg()
70- 70+ 
71- @skipIfUnsupportMultiNPU(2)71+ @skipIfUnsupportMultiNPU(2)
72- @with_comms72+ @with_comms
73- def test_2d_mesh_non_eager_init_subgroup(self):73+ def test_2d_mesh_non_eager_init_subgroup(self):
74- mesh_shape = (2, self.world_size // 2)74+ mesh_shape = (2, self.world_size // 2)
75- mesh_2d = init_device_mesh(self.device_type, mesh_shape)75+ mesh_2d = init_device_mesh(self.device_type, mesh_shape)
76- 76+ 
77- self.assertEqual(mesh_2d.get_group(0).bound_device_id, None)77+ self.assertEqual(mesh_2d.get_group(0).bound_device_id, None)
78- self.assertEqual(mesh_2d.get_group(1).bound_device_id, None)78+ self.assertEqual(mesh_2d.get_group(1).bound_device_id, None)
79- 79+ 
80- # need to refactor the other tests in this file to test both80+ # need to refactor the other tests in this file to test both
81- # eager_init=True and eager_init=False scenarios.81+ # eager_init=True and eager_init=False scenarios.
82- @skipIfUnsupportMultiNPU(2)82+ @skipIfUnsupportMultiNPU(2)
83- @with_comms83+ @with_comms
84- def test_2d_mesh_eager_init_subgroup(self):84+ def test_2d_mesh_eager_init_subgroup(self):
85- # Test with eager_init=True85+ # Test with eager_init=True
86- with self.subTest(eager_init=True):86+ with self.subTest(eager_init=True):
87- mesh_shape = (2, self.world_size // 2)87+ mesh_shape = (2, self.world_size // 2)
88- mesh_2d = init_device_mesh(self.device_type, mesh_shape)88+ mesh_2d = init_device_mesh(self.device_type, mesh_shape)
89- 89+ 
90- # when eager init is used, the subgroup is created from hccl comm split and90+ # when eager init is used, the subgroup is created from hccl comm split and
91- # there would be bound_device_id immediately assigned for the subgroup.91+ # there would be bound_device_id immediately assigned for the subgroup.
92- if self.backend == "hccl":92+ if self.backend == "hccl":
93- curr_device = torch.npu.current_device()93+ curr_device = torch.npu.current_device()
94- self.assertEqual(mesh_2d.get_group(0).bound_device_id.index, curr_device)94+ self.assertEqual(mesh_2d.get_group(0).bound_device_id.index, curr_device)
95- self.assertEqual(mesh_2d.get_group(1).bound_device_id.index, curr_device)95+ self.assertEqual(mesh_2d.get_group(1).bound_device_id.index, curr_device)
96- 96+ 
97- # Test with eager_init=False97+ # Test with eager_init=False
98- with self.subTest(eager_init=False):98+ with self.subTest(eager_init=False):
99- mesh_shape = (2, self.world_size // 2)99+ mesh_shape = (2, self.world_size // 2)
100- mesh_2d = init_device_mesh(self.device_type, mesh_shape)100+ mesh_2d = init_device_mesh(self.device_type, mesh_shape)
101- 101+ 
102- # when eager init is used, the subgroup is created from hccl comm split and102+ # when eager init is used, the subgroup is created from hccl comm split and
103- # there would be bound_device_id immediately assigned for the subgroup.103+ # there would be bound_device_id immediately assigned for the subgroup.
104- if self.backend == "hccl":104+ if self.backend == "hccl":
105- curr_device = torch.npu.current_device()105+ curr_device = torch.npu.current_device()
106- self.assertEqual(mesh_2d.get_group(0).bound_device_id.index, curr_device)106+ self.assertEqual(mesh_2d.get_group(0).bound_device_id.index, curr_device)
107- self.assertEqual(mesh_2d.get_group(1).bound_device_id.index, curr_device)107+ self.assertEqual(mesh_2d.get_group(1).bound_device_id.index, curr_device)
108- 108+ 
109- @skipIfUnsupportMultiNPU(2)109+ @skipIfUnsupportMultiNPU(2)
110- @with_comms110+ @with_comms
111- def test_get_group_and_get_all_groups(self):111+ def test_get_group_and_get_all_groups(self):
112- mesh_shape = (2, self.world_size // 2)112+ mesh_shape = (2, self.world_size // 2)
113- mesh_2d = init_device_mesh(113+ mesh_2d = init_device_mesh(
114- self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")114+ self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
115- )115+ )
116- 116+ 
117- tp_mesh = mesh_2d["tp"]117+ tp_mesh = mesh_2d["tp"]
118- dp_mesh = mesh_2d["dp"]118+ dp_mesh = mesh_2d["dp"]
119- 119+ 
120- self.assertEqual(mesh_2d.get_group(0), mesh_2d.get_group("dp"))120+ self.assertEqual(mesh_2d.get_group(0), mesh_2d.get_group("dp"))
121- self.assertEqual(mesh_2d.get_group(1), mesh_2d.get_group("tp"))121+ self.assertEqual(mesh_2d.get_group(1), mesh_2d.get_group("tp"))
122- 122+ 
123- self.assertEqual(mesh_2d.get_group("dp"), dp_mesh.get_group())123+ self.assertEqual(mesh_2d.get_group("dp"), dp_mesh.get_group())
124- self.assertEqual(mesh_2d.get_group("tp"), tp_mesh.get_group())124+ self.assertEqual(mesh_2d.get_group("tp"), tp_mesh.get_group())
125- 125+ 
126- groups = mesh_2d.get_all_groups()126+ groups = mesh_2d.get_all_groups()
127- self.assertEqual(len(groups), 2)127+ self.assertEqual(len(groups), 2)
128- self.assertTrue(tp_mesh.get_group() in groups)128+ self.assertTrue(tp_mesh.get_group() in groups)
129- self.assertTrue(dp_mesh.get_group() in groups)129+ self.assertTrue(dp_mesh.get_group() in groups)
130- 130+ 
131- @skipIfUnsupportMultiNPU(2)131+ @skipIfUnsupportMultiNPU(2)
132- @with_comms132+ @with_comms
133- def test_get_local_rank_raises_exception(self):133+ def test_get_local_rank_raises_exception(self):
134- mesh_shape = (2, self.world_size // 2)134+ mesh_shape = (2, self.world_size // 2)
135- mesh_2d = init_device_mesh(135+ mesh_2d = init_device_mesh(
136- self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")136+ self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
137- )137+ )
138- 138+ 
139- with self.assertRaisesRegex(139+ with self.assertRaisesRegex(
140- RuntimeError,140+ RuntimeError,
141- "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.",141+ "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.",
142- ):142+ ):
143- mesh_2d.get_local_rank()143+ mesh_2d.get_local_rank()
144- 144+ 
145- @skipIfUnsupportMultiNPU(2)145+ @skipIfUnsupportMultiNPU(2)
146- @with_comms146+ @with_comms
147- def test_device_mesh_init_backend(self):147+ def test_device_mesh_init_backend(self):
148- mesh = DeviceMesh(self.device_type, [1], _init_backend=False)148+ mesh = DeviceMesh(self.device_type, [1], _init_backend=False)
149- 149+ 
150- with self.assertRaisesRegex(RuntimeError, "process groups not initialized!"):150+ with self.assertRaisesRegex(RuntimeError, "process groups not initialized!"):
151- mesh.get_group()151+ mesh.get_group()
152- 152+ 
153- # coordinates should always been populated when init_backend is False, as whenever153+ # coordinates should always been populated when init_backend is False, as whenever
154- # we call init_backend we should make sure the default pg already created154+ # we call init_backend we should make sure the default pg already created
155- mesh.get_coordinate()155+ mesh.get_coordinate()
156- 156+ 
157- @skipIfUnsupportMultiNPU(2)157+ @skipIfUnsupportMultiNPU(2)
158- def test_fake_pg_device_mesh(self):158+ def test_fake_pg_device_mesh(self):
159- fake_store = FakeStore()159+ fake_store = FakeStore()
160- init_process_group("fake", store=fake_store, rank=0, world_size=self.world_size)160+ init_process_group("fake", store=fake_store, rank=0, world_size=self.world_size)
161- device_type = "npu" if torch.npu.is_available() else "cpu"161+ device_type = "npu" if torch.npu.is_available() else "cpu"
162- mesh = DeviceMesh(device_type, torch.arange(self.world_size))162+ mesh = DeviceMesh(device_type, torch.arange(self.world_size))
163- 163+ 
164- local_tensor = torch.randn(2, 8)164+ local_tensor = torch.randn(2, 8)
165- global_tensor = funcol.all_gather_tensor(165+ global_tensor = funcol.all_gather_tensor(
166- local_tensor, gather_dim=0, group=(mesh, 0)166+ local_tensor, gather_dim=0, group=(mesh, 0)
167- ).wait()167+ ).wait()
168- self.assertEqual(global_tensor.shape, (self.world_size * 2, 8))168+ self.assertEqual(global_tensor.shape, (self.world_size * 2, 8))
169- 169+ 
170- @skipIfUnsupportMultiNPU(2)170+ @skipIfUnsupportMultiNPU(2)
171- @with_comms171+ @with_comms
172- def test_from_group_with_global_pg(self):172+ def test_from_group_with_global_pg(self):
173- # Simple test: check `from_group` from a mesh pg vs. directly173+ # Simple test: check `from_group` from a mesh pg vs. directly
174- # initializing via `init_device_mesh`174+ # initializing via `init_device_mesh`
175- ref_global_mesh = init_device_mesh(self.device_type, (self.world_size,))175+ ref_global_mesh = init_device_mesh(self.device_type, (self.world_size,))
176- mesh_pg = ref_global_mesh.get_group()176+ mesh_pg = ref_global_mesh.get_group()
177- global_mesh = DeviceMesh.from_group(mesh_pg, self.device_type)177+ global_mesh = DeviceMesh.from_group(mesh_pg, self.device_type)
178- self.assertEqual(ref_global_mesh, global_mesh)178+ self.assertEqual(ref_global_mesh, global_mesh)
179- self.assertEqual(ref_global_mesh._dim_group_names, global_mesh._dim_group_names)179+ self.assertEqual(ref_global_mesh._dim_group_names, global_mesh._dim_group_names)
180- self.assertEqual(180+ self.assertEqual(
181- ref_global_mesh._coordinate_on_dim, global_mesh._coordinate_on_dim181+ ref_global_mesh._coordinate_on_dim, global_mesh._coordinate_on_dim
182- )182+ )
183- # Check when `mesh` is passed as well183+ # Check when `mesh` is passed as well
184- global_mesh = DeviceMesh.from_group(184+ global_mesh = DeviceMesh.from_group(
185- mesh_pg, self.device_type, mesh=torch.arange(self.world_size)185+ mesh_pg, self.device_type, mesh=torch.arange(self.world_size)
186- )186+ )
187- self.assertEqual(ref_global_mesh, global_mesh)187+ self.assertEqual(ref_global_mesh, global_mesh)
188- self.assertEqual(ref_global_mesh._dim_group_names, global_mesh._dim_group_names)188+ self.assertEqual(ref_global_mesh._dim_group_names, global_mesh._dim_group_names)
189- self.assertEqual(189+ self.assertEqual(
190- ref_global_mesh._coordinate_on_dim, global_mesh._coordinate_on_dim190+ ref_global_mesh._coordinate_on_dim, global_mesh._coordinate_on_dim
191- )191+ )
192- 192+ 
193- @skipIfUnsupportMultiNPU(2)193+ @skipIfUnsupportMultiNPU(2)
194- def test_raises_invalid_device_type(self):194+ def test_raises_invalid_device_type(self):
195- with self.assertRaisesRegex(195+ with self.assertRaisesRegex(
196- RuntimeError,196+ RuntimeError,
197- "Device type with index is not supported",197+ "Device type with index is not supported",
198- ):198+ ):
199- # test init_device_mesh with an invalid device type that contains a NPU index199+ # test init_device_mesh with an invalid device type that contains a NPU index
200- mesh_shape = (2, self.world_size // 2)200+ mesh_shape = (2, self.world_size // 2)
201- mesh_2d = init_device_mesh(201+ mesh_2d = init_device_mesh(
202- "npu:0", mesh_shape=mesh_shape, mesh_dim_names=("dp", "tp")202+ "npu:0", mesh_shape=mesh_shape, mesh_dim_names=("dp", "tp")
203- )203+ )
204- 204+ 
205- @skipIfUnsupportMultiNPU(2)205+ @skipIfUnsupportMultiNPU(2)
206- @with_comms206+ @with_comms
207- def test_set_mesh_dim_group_options(self):207+ def test_set_mesh_dim_group_options(self):
208- device_type = "npu" if torch.npu.is_available() else "cpu"208+ device_type = "npu" if torch.npu.is_available() else "cpu"
209- _mesh_resources._set_mesh_dim_group_options(1, "fake", None)209+ _mesh_resources._set_mesh_dim_group_options(1, "fake", None)
210- 210+ 
211- mesh_tensor = torch.arange(2).reshape(2, 1)211+ mesh_tensor = torch.arange(2).reshape(2, 1)
212- mesh = DeviceMesh(device_type, mesh_tensor)212+ mesh = DeviceMesh(device_type, mesh_tensor)
213- # Fake pg only have BackendType as BackendType::CUSTOM.213+ # Fake pg only have BackendType as BackendType::CUSTOM.
214- self.assertEqual(mesh.get_group(1)._get_backend_name(), "custom")214+ self.assertEqual(mesh.get_group(1)._get_backend_name(), "custom")
215- 215+ 
216- 216+ 
217-#DeviceMeshTest with resetting world_size to 4.217+#DeviceMeshTest with resetting world_size to 4.
218-class DeviceMeshTestF(NPUDTensorTestBase):218+class DeviceMeshTestF(NPUDTensorTestBase):
219- @property219+ @property
220- def world_size(self):220+ def world_size(self):
221- return 4221+ return 4
222- 222+ 
223- @skipIfUnsupportMultiNPU(4)223+ @skipIfUnsupportMultiNPU(4)
224- @with_comms224+ @with_comms
225- def test_assert_invalid_mesh_tensor(self):225+ def test_assert_invalid_mesh_tensor(self):
226- mesh = torch.arange(self.world_size).to(self.rank)226+ mesh = torch.arange(self.world_size).to(self.rank)
227- with self.assertRaises(ValueError):227+ with self.assertRaises(ValueError):
228- device_mesh = DeviceMesh(self.device_type, mesh)228+ device_mesh = DeviceMesh(self.device_type, mesh)
229- 229+
230- @skipIfUnsupportMultiNPU(4)230+ @skipIfUnsupportMultiNPU(4)
231- @with_comms231+ @with_comms
232- def test_get_local_rank(self):232+ def test_get_local_rank(self):
233- mesh_shape = (2, self.world_size // 2)233+ mesh_shape = (2, self.world_size // 2)
234- mesh_2d = init_device_mesh(234+ mesh_2d = init_device_mesh(
235- self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")235+ self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
236- )236+ )
237- self.assertEqual(mesh_2d.get_local_rank("dp"), mesh_2d.get_local_rank(0))237+ self.assertEqual(mesh_2d.get_local_rank("dp"), mesh_2d.get_local_rank(0))
238- self.assertEqual(mesh_2d.get_local_rank("tp"), mesh_2d.get_local_rank(1))238+ self.assertEqual(mesh_2d.get_local_rank("tp"), mesh_2d.get_local_rank(1))
239- 239+ 
240- dp_mesh = mesh_2d["dp"]240+ dp_mesh = mesh_2d["dp"]
241- tp_mesh = mesh_2d["tp"]241+ tp_mesh = mesh_2d["tp"]
242- self.assertEqual(dp_mesh.get_local_rank(), mesh_2d.get_local_rank("dp"))242+ self.assertEqual(dp_mesh.get_local_rank(), mesh_2d.get_local_rank("dp"))
243- self.assertEqual(tp_mesh.get_local_rank(), mesh_2d.get_local_rank("tp"))243+ self.assertEqual(tp_mesh.get_local_rank(), mesh_2d.get_local_rank("tp"))
244- 244+ 
245- # Verify flattened mesh local rank correctness.245+ # Verify flattened mesh local rank correctness.
246- flattened_mesh = mesh_2d["dp", "tp"]._flatten()246+ flattened_mesh = mesh_2d["dp", "tp"]._flatten()
247- self.assertEqual(flattened_mesh.get_local_rank(), self.rank)247+ self.assertEqual(flattened_mesh.get_local_rank(), self.rank)
248- 248+ 
249- @skipIfUnsupportMultiNPU(4)249+ @skipIfUnsupportMultiNPU(4)
250- @with_comms250+ @with_comms
251- def test_device_mesh_2d(self):251+ def test_device_mesh_2d(self):
252- mesh_tensor = torch.arange(4).reshape(2, 2)252+ mesh_tensor = torch.arange(4).reshape(2, 2)
253- # construct a npu device mesh253+ # construct a npu device mesh
254- mesh = DeviceMesh(self.device_type, mesh_tensor)254+ mesh = DeviceMesh(self.device_type, mesh_tensor)
255- 255+ 
256- # check all dim groups256+ # check all dim groups
257- dim_to_subgroups = mesh.get_all_groups()257+ dim_to_subgroups = mesh.get_all_groups()
258- 258+ 
259- expected_ranks_by_dim = [[[0, 2], [1, 3]], [[0, 1], [2, 3]]]259+ expected_ranks_by_dim = [[[0, 2], [1, 3]], [[0, 1], [2, 3]]]
260- for dim, dim_group in enumerate(dim_to_subgroups):260+ for dim, dim_group in enumerate(dim_to_subgroups):
261- self.assertTrue(dim < 2)261+ self.assertTrue(dim < 2)
262- dim_ranks = expected_ranks_by_dim[dim]262+ dim_ranks = expected_ranks_by_dim[dim]
263- 263+ 
264- dim_group_size = get_world_size(dim_group)264+ dim_group_size = get_world_size(dim_group)
265- self.assertIsInstance(dim_group, ProcessGroup)265+ self.assertIsInstance(dim_group, ProcessGroup)
266- self.assertEqual(dim_group_size, 2)266+ self.assertEqual(dim_group_size, 2)
267- global_ranks = [267+ global_ranks = [
268- get_global_rank(dim_group, i) for i in range(dim_group_size)268+ get_global_rank(dim_group, i) for i in range(dim_group_size)
269- ]269+ ]
270- current_rank_expected_group_ranks = (270+ current_rank_expected_group_ranks = (
271- dim_ranks[0] if self.rank in dim_ranks[0] else dim_ranks[1]271+ dim_ranks[0] if self.rank in dim_ranks[0] else dim_ranks[1]
272- )272+ )
273- self.assertEqual(global_ranks, current_rank_expected_group_ranks)273+ self.assertEqual(global_ranks, current_rank_expected_group_ranks)
274- 274+ 
275- @skipIfUnsupportMultiNPU(4)275+ @skipIfUnsupportMultiNPU(4)
276- @with_comms276+ @with_comms
277- def test_from_group_with_invalid_mesh(self):277+ def test_from_group_with_invalid_mesh(self):
278- global_pg = _get_default_group()278+ global_pg = _get_default_group()
279- global_pg_size = global_pg.size()279+ global_pg_size = global_pg.size()
280- assert global_pg_size == 4, "Test assumes global world size of 4"280+ assert global_pg_size == 4, "Test assumes global world size of 4"
281- invalid_mesh = [[0, 1], [2, 3]] # 2D mesh when we need 1D281+ invalid_mesh = [[0, 1], [2, 3]] # 2D mesh when we need 1D
282- regex = r"Invalid mesh \[\[0, 1\], \[2, 3\]\] for ProcessGroup with ranks \[0, 1, 2, 3\]"282+ regex = r"Invalid mesh \[\[0, 1\], \[2, 3\]\] for ProcessGroup with ranks \[0, 1, 2, 3\]"
283- with self.assertRaisesRegex(ValueError, regex):283+ with self.assertRaisesRegex(ValueError, regex):
284- DeviceMesh.from_group(284+ DeviceMesh.from_group(
285- global_pg, "npu", invalid_mesh, mesh_dim_names=("dim0", "dim1")285+ global_pg, "npu", invalid_mesh, mesh_dim_names=("dim0", "dim1")
286- )286+ )
287- 287+ 
288- device_mesh = init_device_mesh(self.device_type, (2, 2))288+ device_mesh = init_device_mesh(self.device_type, (2, 2))
289- groups = device_mesh.get_all_groups()289+ groups = device_mesh.get_all_groups()
290- invalid_mesh = (0, 1, 2, 3) # 1D mesh when we need 2D290+ invalid_mesh = (0, 1, 2, 3) # 1D mesh when we need 2D
291- regex = r"Expects mesh with ndim equal to number of ProcessGroups but got mesh \[0, 1, 2, 3\] and 2 ProcessGroups"291+ regex = r"Expects mesh with ndim equal to number of ProcessGroups but got mesh \[0, 1, 2, 3\] and 2 ProcessGroups"
292- with self.assertRaisesRegex(ValueError, regex):292+ with self.assertRaisesRegex(ValueError, regex):
293- DeviceMesh.from_group(293+ DeviceMesh.from_group(
294- groups, self.device_type, invalid_mesh, mesh_dim_names=("dim0", "dim1")294+ groups, self.device_type, invalid_mesh, mesh_dim_names=("dim0", "dim1")
295- )295+ )
296- 296+ 
297- 297+ 
298-class DeviceMeshTestNDim(NPUDTensorTestBase):298+class DeviceMeshTestNDim(NPUDTensorTestBase):
299- @property299+ @property
300- def world_size(self):300+ def world_size(self):
301- return 2301+ return 2
302- 302+ 
303- @skipIfUnsupportMultiNPU(2)303+ @skipIfUnsupportMultiNPU(2)
304- @with_comms304+ @with_comms
305- def test_device_mesh_parent_child_hash(self):305+ def test_device_mesh_parent_child_hash(self):
306- mesh_2d = init_device_mesh(306+ mesh_2d = init_device_mesh(
307- self.device_type, (2, self.world_size // 2), mesh_dim_names=("DP", "TP")307+ self.device_type, (2, self.world_size // 2), mesh_dim_names=("DP", "TP")
308- )308+ )
309- 309+ 
310- mesh_group_1 = torch.arange(0, self.world_size // 2)310+ mesh_group_1 = torch.arange(0, self.world_size // 2)
311- mesh_group_2 = torch.arange(self.world_size // 2, self.world_size)311+ mesh_group_2 = torch.arange(self.world_size // 2, self.world_size)
312- ep_mesh_1 = DeviceMesh(self.device_type, mesh_group_1)312+ ep_mesh_1 = DeviceMesh(self.device_type, mesh_group_1)
313- ep_mesh_2 = DeviceMesh(self.device_type, mesh_group_2)313+ ep_mesh_2 = DeviceMesh(self.device_type, mesh_group_2)
314- ep_mesh = ep_mesh_1 if self.rank < self.world_size // 2 else ep_mesh_2314+ ep_mesh = ep_mesh_1 if self.rank < self.world_size // 2 else ep_mesh_2
315- # ep_mesh is considered different from mesh_2d["TP"]315+ # ep_mesh is considered different from mesh_2d["TP"]
316- self.assertEqual(mesh_2d["TP"]._flatten_mesh_list, ep_mesh._flatten_mesh_list)316+ self.assertEqual(mesh_2d["TP"]._flatten_mesh_list, ep_mesh._flatten_mesh_list)
317- self.assertEqual(mesh_2d["TP"].mesh.shape, ep_mesh.mesh.shape)317+ self.assertEqual(mesh_2d["TP"].mesh.shape, ep_mesh.mesh.shape)
318- self.assertEqual(mesh_2d["TP"].device_type, ep_mesh.device_type)318+ self.assertEqual(mesh_2d["TP"].device_type, ep_mesh.device_type)
319- self.assertNotEqual(mesh_2d["TP"].mesh_dim_names, ep_mesh.mesh_dim_names)319+ self.assertNotEqual(mesh_2d["TP"].mesh_dim_names, ep_mesh.mesh_dim_names)
320- self.assertEqual(mesh_2d["TP"]._thread_id, ep_mesh._thread_id)320+ self.assertEqual(mesh_2d["TP"]._thread_id, ep_mesh._thread_id)
321- self.assertNotEqual(hash(mesh_2d["TP"]), hash(ep_mesh))321+ self.assertNotEqual(hash(mesh_2d["TP"]), hash(ep_mesh))
322- self.assertNotEqual(mesh_2d["TP"], ep_mesh)322+ self.assertNotEqual(mesh_2d["TP"], ep_mesh)
323- 323+ 
324- another_mesh_1 = DeviceMesh(self.device_type, mesh_group_1)324+ another_mesh_1 = DeviceMesh(self.device_type, mesh_group_1)
325- another_mesh_2 = DeviceMesh(self.device_type, mesh_group_2)325+ another_mesh_2 = DeviceMesh(self.device_type, mesh_group_2)
326- another_mesh = (326+ another_mesh = (
327- another_mesh_1 if self.rank < self.world_size // 2 else another_mesh_2327+ another_mesh_1 if self.rank < self.world_size // 2 else another_mesh_2
328- )328+ )
329- # another_mesh is considered the same as ep_mesh329+ # another_mesh is considered the same as ep_mesh
330- self.assertEqual(ep_mesh._flatten_mesh_list, another_mesh._flatten_mesh_list)330+ self.assertEqual(ep_mesh._flatten_mesh_list, another_mesh._flatten_mesh_list)
331- self.assertEqual(ep_mesh.mesh.shape, another_mesh.mesh.shape)331+ self.assertEqual(ep_mesh.mesh.shape, another_mesh.mesh.shape)
332- self.assertEqual(ep_mesh.device_type, another_mesh.device_type)332+ self.assertEqual(ep_mesh.device_type, another_mesh.device_type)
333- self.assertEqual(ep_mesh.mesh_dim_names, another_mesh.mesh_dim_names)333+ self.assertEqual(ep_mesh.mesh_dim_names, another_mesh.mesh_dim_names)
334- self.assertEqual(ep_mesh._thread_id, another_mesh._thread_id)334+ self.assertEqual(ep_mesh._thread_id, another_mesh._thread_id)
335- self.assertEqual(hash(ep_mesh), hash(another_mesh))335+ self.assertEqual(hash(ep_mesh), hash(another_mesh))
336- self.assertEqual(ep_mesh, another_mesh)336+ self.assertEqual(ep_mesh, another_mesh)
337- 337+ 
338- 338+ 
339-#DeviceMeshTestNDim with resetting world_size to 8.339+#DeviceMeshTestNDim with resetting world_size to 8.
340-class DeviceMeshTestNDimE(NPUDTensorTestBase):340+class DeviceMeshTestNDimE(NPUDTensorTestBase):
341- @property341+ @property
342- def world_size(self):342+ def world_size(self):
343- return 8343+ return 8
344- 344+ 
345- @skipIfUnsupportMultiNPU(8)345+ @skipIfUnsupportMultiNPU(8)
346- @with_comms346+ @with_comms
347- def test_device_mesh_nd(self):347+ def test_device_mesh_nd(self):
348- # construct a npu device mesh348+ # construct a npu device mesh
349- mesh_tensor = torch.arange(8).reshape(2, 2, 2)349+ mesh_tensor = torch.arange(8).reshape(2, 2, 2)
350- mesh = DeviceMesh(self.device_type, mesh_tensor)350+ mesh = DeviceMesh(self.device_type, mesh_tensor)
351- 351+ 
352- # check all dim groups352+ # check all dim groups
353- dim_to_subgroups = mesh.get_all_groups()353+ dim_to_subgroups = mesh.get_all_groups()
354- 354+ 
355- for dim, dim_group in enumerate(dim_to_subgroups):355+ for dim, dim_group in enumerate(dim_to_subgroups):
356- self.assertTrue(dim < mesh_tensor.ndim)356+ self.assertTrue(dim < mesh_tensor.ndim)
357- dim_ranks = mesh_tensor.swapdims(-1, dim).reshape(-1, 2)357+ dim_ranks = mesh_tensor.swapdims(-1, dim).reshape(-1, 2)
358- 358+ 
359- dim_group_size = get_world_size(dim_group)359+ dim_group_size = get_world_size(dim_group)
360- self.assertIsInstance(dim_group, ProcessGroup)360+ self.assertIsInstance(dim_group, ProcessGroup)
361- self.assertEqual(dim_group_size, 2)361+ self.assertEqual(dim_group_size, 2)
362- global_ranks = [362+ global_ranks = [
363- get_global_rank(dim_group, i) for i in range(dim_group_size)363+ get_global_rank(dim_group, i) for i in range(dim_group_size)
364- ]364+ ]
365- for ranks in dim_ranks:365+ for ranks in dim_ranks:
366- if self.rank in ranks:366+ if self.rank in ranks:
367- self.assertEqual(global_ranks, ranks.tolist())367+ self.assertEqual(global_ranks, ranks.tolist())
368- 368+ 
369- @skipIfUnsupportMultiNPU(8)369+ @skipIfUnsupportMultiNPU(8)
370- @with_comms370+ @with_comms
371- def test_device_mesh_hash(self):371+ def test_device_mesh_hash(self):
372- mesh_tensor_2d = torch.arange(8).reshape(4, 2)372+ mesh_tensor_2d = torch.arange(8).reshape(4, 2)
373- mesh = DeviceMesh(self.device_type, mesh_tensor_2d)373+ mesh = DeviceMesh(self.device_type, mesh_tensor_2d)
374- mesh2 = DeviceMesh(self.device_type, mesh_tensor_2d)374+ mesh2 = DeviceMesh(self.device_type, mesh_tensor_2d)
375- self.assertEqual(hash(mesh), hash(mesh2))375+ self.assertEqual(hash(mesh), hash(mesh2))
376- mesh_tensor_3d = torch.arange(8).reshape(2, 2, 2)376+ mesh_tensor_3d = torch.arange(8).reshape(2, 2, 2)
377- mesh3 = DeviceMesh(self.device_type, mesh_tensor_3d)377+ mesh3 = DeviceMesh(self.device_type, mesh_tensor_3d)
378- self.assertNotEqual(hash(mesh), hash(mesh3))378+ self.assertNotEqual(hash(mesh), hash(mesh3))
379- self.assertNotEqual(hash(mesh2), hash(mesh3))379+ self.assertNotEqual(hash(mesh2), hash(mesh3))
380- 380+ 
381- @skipIfUnsupportMultiNPU(8)381+ @skipIfUnsupportMultiNPU(8)
382- @with_comms382+ @with_comms
383- def test_get_local_rank_3d(self):383+ def test_get_local_rank_3d(self):
384- """384+ """
385- If we have a 3D mesh and we want to apply dp, pp, tp to it,385+ If we have a 3D mesh and we want to apply dp, pp, tp to it,
386- mesh_dim_names = ["dp", "pp", "tp"], and the mesh tensor would be:386+ mesh_dim_names = ["dp", "pp", "tp"], and the mesh tensor would be:
387- mesh_3d_tensor = [387+ mesh_3d_tensor = [
388- [388+ [
389- [0, 1],389+ [0, 1],
390- [2, 3],390+ [2, 3],
391- ],391+ ],
392- [392+ [
393- [4, 5],393+ [4, 5],
394- [6, 7],394+ [6, 7],
395- ]395+ ]
396- 396+ 
397- ]397+ ]
398- """398+ """
399- mesh_shape = (2, 2, 2)399+ mesh_shape = (2, 2, 2)
400- mesh_3d = init_device_mesh(400+ mesh_3d = init_device_mesh(
401- self.device_type, mesh_shape, mesh_dim_names=("dp", "pp", "tp")401+ self.device_type, mesh_shape, mesh_dim_names=("dp", "pp", "tp")
402- )402+ )
403- 403+ 
404- # tp_rank_0: [0, 2, 4, 6], tp_rank_1: [1, 3, 5, 7]404+ # tp_rank_0: [0, 2, 4, 6], tp_rank_1: [1, 3, 5, 7]
405- tp_rank = mesh_3d.get_local_rank("tp")405+ tp_rank = mesh_3d.get_local_rank("tp")
406- expected_tp_rank = self.rank % 2406+ expected_tp_rank = self.rank % 2
407- self.assertEqual(tp_rank, expected_tp_rank)407+ self.assertEqual(tp_rank, expected_tp_rank)
408- 408+ 
409- # pp_rank_0: [0, 1, 4, 5], pp_rank_1: [2, 3, 6, 7]409+ # pp_rank_0: [0, 1, 4, 5], pp_rank_1: [2, 3, 6, 7]
410- pp_rank = mesh_3d.get_local_rank("pp")410+ pp_rank = mesh_3d.get_local_rank("pp")
411- expected_pp_rank = 0 if self.rank % 4 <= 1 else 1411+ expected_pp_rank = 0 if self.rank % 4 <= 1 else 1
412- self.assertEqual(pp_rank, expected_pp_rank)412+ self.assertEqual(pp_rank, expected_pp_rank)
413- 413+ 
414- # dp_rank_0: [0, 1, 2, 3], dp_rank_1: [4, 5, 6, 7]414+ # dp_rank_0: [0, 1, 2, 3], dp_rank_1: [4, 5, 6, 7]
415- dp_rank = mesh_3d.get_local_rank("dp")415+ dp_rank = mesh_3d.get_local_rank("dp")
416- expected_dp_rank = self.rank // 4416+ expected_dp_rank = self.rank // 4
417- self.assertEqual(dp_rank, expected_dp_rank)417+ self.assertEqual(dp_rank, expected_dp_rank)
418- 418+ 
419- @skipIfUnsupportMultiNPU(8)419+ @skipIfUnsupportMultiNPU(8)
420- @with_comms420+ @with_comms
421- def test_from_group_with_mesh_shape(self):421+ def test_from_group_with_mesh_shape(self):
422- """Tests ``from_group`` when passing ``mesh_shape`` as 2D."""422+ """Tests ``from_group`` when passing ``mesh_shape`` as 2D."""
423- # Consider two different logical views of the same mesh:423+ # Consider two different logical views of the same mesh:
424- # - (4, 2) ("dp", "tp") mesh424+ # - (4, 2) ("dp", "tp") mesh
425- # - (2, 2, 2) ("dp_replicate", "dp_shard", "tp") mesh425+ # - (2, 2, 2) ("dp_replicate", "dp_shard", "tp") mesh
426- mesh_shape = (2, 2, 2)426+ mesh_shape = (2, 2, 2)
427- mesh_dim_names = ("dp_replicate", "dp_shard", "tp")427+ mesh_dim_names = ("dp_replicate", "dp_shard", "tp")
428- ref_mesh = init_device_mesh(428+ ref_mesh = init_device_mesh(
429- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names429+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
430- )430+ )
431- 431+ 
432- dp_shard_group = ref_mesh["dp_shard"].get_group()432+ dp_shard_group = ref_mesh["dp_shard"].get_group()
433- dp_replicate_group = ref_mesh["dp_replicate"].get_group()433+ dp_replicate_group = ref_mesh["dp_replicate"].get_group()
434- 434+ 
435- dp_mesh = DeviceMesh.from_group(435+ dp_mesh = DeviceMesh.from_group(
436- [dp_replicate_group, dp_shard_group],436+ [dp_replicate_group, dp_shard_group],
437- self.device_type,437+ self.device_type,
438- mesh=ref_mesh.mesh[:, :, ref_mesh.get_local_rank(2)],438+ mesh=ref_mesh.mesh[:, :, ref_mesh.get_local_rank(2)],
439- mesh_dim_names=mesh_dim_names[:2],439+ mesh_dim_names=mesh_dim_names[:2],
440- )440+ )
441- 441+ 
442- ref_mesh_dp_dim_group_names = ref_mesh._dim_group_names[:2]442+ ref_mesh_dp_dim_group_names = ref_mesh._dim_group_names[:2]
443- self.assertEqual(ref_mesh_dp_dim_group_names, ref_mesh._dim_group_names[:2])443+ self.assertEqual(ref_mesh_dp_dim_group_names, ref_mesh._dim_group_names[:2])
444- # Cannot check directly for mesh equality since parent meshes are not444+ # Cannot check directly for mesh equality since parent meshes are not
445- # the same since the ref's parent mesh is 3D445+ # the same since the ref's parent mesh is 3D
446- self.assertEqual(dp_mesh["dp_replicate"].mesh, ref_mesh["dp_replicate"].mesh)446+ self.assertEqual(dp_mesh["dp_replicate"].mesh, ref_mesh["dp_replicate"].mesh)
447- self.assertEqual(447+ self.assertEqual(
448- dp_mesh["dp_replicate"]._dim_group_names,448+ dp_mesh["dp_replicate"]._dim_group_names,
449- ref_mesh["dp_replicate"]._dim_group_names,449+ ref_mesh["dp_replicate"]._dim_group_names,
450- )450+ )
451- self.assertEqual(dp_mesh["dp_shard"].mesh, ref_mesh["dp_shard"].mesh)451+ self.assertEqual(dp_mesh["dp_shard"].mesh, ref_mesh["dp_shard"].mesh)
452- self.assertEqual(452+ self.assertEqual(
453- dp_mesh["dp_shard"]._dim_group_names,453+ dp_mesh["dp_shard"]._dim_group_names,
454- ref_mesh["dp_shard"]._dim_group_names,454+ ref_mesh["dp_shard"]._dim_group_names,
455- )455+ )
456- 456+ 
457- @skipIfUnsupportMultiNPU(8)457+ @skipIfUnsupportMultiNPU(8)
458- @with_comms458+ @with_comms
459- def test_from_group_with_mesh_shape_2d(self):459+ def test_from_group_with_mesh_shape_2d(self):
460- """Tests ``from_group`` when passing ``mesh_shape`` as 2D."""460+ """Tests ``from_group`` when passing ``mesh_shape`` as 2D."""
461- # Consider the following scenario where the process group has been created,461+ # Consider the following scenario where the process group has been created,
462- # but we need to create the 2D HSDP mesh from it later in the program.462+ # but we need to create the 2D HSDP mesh from it later in the program.
463- mesh_shape = (2, 4)463+ mesh_shape = (2, 4)
464- mesh_dim_names = ("dp_replicate", "dp_shard")464+ mesh_dim_names = ("dp_replicate", "dp_shard")
465- ref_mesh = init_device_mesh(465+ ref_mesh = init_device_mesh(
466- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names466+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
467- )467+ )
468- 468+ 
469- # Create shard groups (e.g. (0, 1, 2, 3), (4, 5, 6, 7))469+ # Create shard groups (e.g. (0, 1, 2, 3), (4, 5, 6, 7))
470- # and assign the correct shard group to each rank470+ # and assign the correct shard group to each rank
471- shard_rank_lists = list(range(0, self.world_size // 2)), list(471+ shard_rank_lists = list(range(0, self.world_size // 2)), list(
472- range(self.world_size // 2, self.world_size)472+ range(self.world_size // 2, self.world_size)
473- )473+ )
474- shard_groups = (474+ shard_groups = (
475- new_group(shard_rank_lists[0]),475+ new_group(shard_rank_lists[0]),
476- new_group(shard_rank_lists[1]),476+ new_group(shard_rank_lists[1]),
477- )477+ )
478- current_shard_group = (478+ current_shard_group = (
479- shard_groups[0] if self.rank in shard_rank_lists[0] else shard_groups[1]479+ shard_groups[0] if self.rank in shard_rank_lists[0] else shard_groups[1]
480- )480+ )
481- 481+ 
482- # Create replicate groups (for example, (0, 4), (1, 5), (2, 6), (3, 7))482+ # Create replicate groups (for example, (0, 4), (1, 5), (2, 6), (3, 7))
483- # and assign the correct replicate group to each rank483+ # and assign the correct replicate group to each rank
484- current_replicate_group = None484+ current_replicate_group = None
485- shard_factor = len(shard_rank_lists[0])485+ shard_factor = len(shard_rank_lists[0])
486- for i in range(self.world_size // 2):486+ for i in range(self.world_size // 2):
487- replicate_group_ranks = list(range(i, self.world_size, shard_factor))487+ replicate_group_ranks = list(range(i, self.world_size, shard_factor))
488- replicate_group = new_group(replicate_group_ranks)488+ replicate_group = new_group(replicate_group_ranks)
489- if self.rank in replicate_group_ranks:489+ if self.rank in replicate_group_ranks:
490- current_replicate_group = replicate_group490+ current_replicate_group = replicate_group
491- 491+ 
492- dp_mesh = DeviceMesh.from_group(492+ dp_mesh = DeviceMesh.from_group(
493- [not_none(current_replicate_group), current_shard_group],493+ [not_none(current_replicate_group), current_shard_group],
494- self.device_type,494+ self.device_type,
495- mesh=ref_mesh.mesh,495+ mesh=ref_mesh.mesh,
496- mesh_dim_names=("dp_replicate", "dp_shard"),496+ mesh_dim_names=("dp_replicate", "dp_shard"),
497- )497+ )
498- 498+ 
499- # self.assertEqual(ref_mesh._dim_group_names, dp_mesh._dim_group_names)499+ # self.assertEqual(ref_mesh._dim_group_names, dp_mesh._dim_group_names)
500- for mesh_dim_group, ref_mesh_dim_group in zip(500+ for mesh_dim_group, ref_mesh_dim_group in zip(
501- dp_mesh.get_all_groups(), ref_mesh.get_all_groups()501+ dp_mesh.get_all_groups(), ref_mesh.get_all_groups()
502- ):502+ ):
503- mesh_dim_group_ranks = dist.get_process_group_ranks(mesh_dim_group)503+ mesh_dim_group_ranks = dist.get_process_group_ranks(mesh_dim_group)
504- ref_mesh_dim_group_ranks = dist.get_process_group_ranks(ref_mesh_dim_group)504+ ref_mesh_dim_group_ranks = dist.get_process_group_ranks(ref_mesh_dim_group)
505- self.assertEqual(mesh_dim_group_ranks, ref_mesh_dim_group_ranks)505+ self.assertEqual(mesh_dim_group_ranks, ref_mesh_dim_group_ranks)
506- # check both the 2d mesh and the submeshes are exactly the same.506+ # check both the 2d mesh and the submeshes are exactly the same.
507- self.assertEqual(dp_mesh, ref_mesh)507+ self.assertEqual(dp_mesh, ref_mesh)
508- self.assertEqual(dp_mesh["dp_replicate"], ref_mesh["dp_replicate"])508+ self.assertEqual(dp_mesh["dp_replicate"], ref_mesh["dp_replicate"])
509- self.assertEqual(dp_mesh["dp_shard"], ref_mesh["dp_shard"])509+ self.assertEqual(dp_mesh["dp_shard"], ref_mesh["dp_shard"])
510- 510+ 
511- 511+ 
512-class InitDeviceMeshTest(NPUDTensorTestBase):512+class InitDeviceMeshTest(NPUDTensorTestBase):
513- @property513+ @property
514- def world_size(self):514+ def world_size(self):
515- return 2515+ return 2
516- 516+ 
517- @skipIfUnsupportMultiNPU(2)517+ @skipIfUnsupportMultiNPU(2)
518- @with_comms518+ @with_comms
519- def test_init_device_mesh(self):519+ def test_init_device_mesh(self):
520- mesh_shape = (2, 1)520+ mesh_shape = (2, 1)
521- mesh_dim_names = ("DP", "TP")521+ mesh_dim_names = ("DP", "TP")
522- ref_mesh = DeviceMesh(522+ ref_mesh = DeviceMesh(
523- self.device_type,523+ self.device_type,
524- torch.arange(2).view(mesh_shape),524+ torch.arange(2).view(mesh_shape),
525- mesh_dim_names=mesh_dim_names,525+ mesh_dim_names=mesh_dim_names,
526- )526+ )
527- 527+ 
528- # test init_device_mesh with mesh_dim_names528+ # test init_device_mesh with mesh_dim_names
529- mesh_2d = init_device_mesh(529+ mesh_2d = init_device_mesh(
530- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names530+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
531- )531+ )
532- self.assertEqual(mesh_2d, ref_mesh)532+ self.assertEqual(mesh_2d, ref_mesh)
533- self.assertEqual(mesh_2d.mesh_dim_names, mesh_dim_names)533+ self.assertEqual(mesh_2d.mesh_dim_names, mesh_dim_names)
534- 534+ 
535- @skipIfUnsupportMultiNPU(2)535+ @skipIfUnsupportMultiNPU(2)
536- @with_comms536+ @with_comms
537- def test_raises_duplicate_mesh_dim_names(self):537+ def test_raises_duplicate_mesh_dim_names(self):
538- with self.assertRaisesRegex(538+ with self.assertRaisesRegex(
539- RuntimeError,539+ RuntimeError,
540- "Each mesh_dim_name must be unique.",540+ "Each mesh_dim_name must be unique.",
541- ):541+ ):
542- mesh = init_device_mesh(542+ mesh = init_device_mesh(
543- self.device_type,543+ self.device_type,
544- (1, 2),544+ (1, 2),
545- mesh_dim_names=["dp", "dp"],545+ mesh_dim_names=["dp", "dp"],
546- )546+ )
547- 547+ 
548- @skipIfUnsupportMultiNPU(2)548+ @skipIfUnsupportMultiNPU(2)
549- @with_comms549+ @with_comms
550- def test_raises_mesh_shape_mesh_dim_names_mismatch(self):550+ def test_raises_mesh_shape_mesh_dim_names_mismatch(self):
551- with self.assertRaisesRegex(551+ with self.assertRaisesRegex(
552- RuntimeError,552+ RuntimeError,
553- "mesh_shape and mesh_dim_names should have same length!",553+ "mesh_shape and mesh_dim_names should have same length!",
554- ):554+ ):
555- mesh = init_device_mesh(555+ mesh = init_device_mesh(
556- self.device_type,556+ self.device_type,
557- (2,),557+ (2,),
558- mesh_dim_names=["dp", "tp"],558+ mesh_dim_names=["dp", "tp"],
559- )559+ )
560- 560+ 
561- 561+ 
562-class TestDeviceMeshGetItem(NPUDTensorTestBase):562+class TestDeviceMeshGetItem(NPUDTensorTestBase):
563- @property563+ @property
564- def world_size(self):564+ def world_size(self):
565- return 2565+ return 2
566- 566+ 
567- @skipIfUnsupportMultiNPU(2)567+ @skipIfUnsupportMultiNPU(2)
568- @with_comms568+ @with_comms
569- def test_raises_no_mesh_dim_found(self):569+ def test_raises_no_mesh_dim_found(self):
570- with self.assertRaisesRegex(570+ with self.assertRaisesRegex(
571- RuntimeError, "Cannot slice a DeviceMesh without mesh_dim_names!"571+ RuntimeError, "Cannot slice a DeviceMesh without mesh_dim_names!"
572- ):572+ ):
573- mesh = init_device_mesh(self.device_type, (1, 2))573+ mesh = init_device_mesh(self.device_type, (1, 2))
574- child_mesh = mesh["DP"]574+ child_mesh = mesh["DP"]
575- 575+ 
576- @skipIfUnsupportMultiNPU(2)576+ @skipIfUnsupportMultiNPU(2)
577- @with_comms577+ @with_comms
578- def test_raises_invalid_mesh_dim_name(self):578+ def test_raises_invalid_mesh_dim_name(self):
579- child_mesh_dim_name = ("PP",)579+ child_mesh_dim_name = ("PP",)
580- with self.assertRaisesRegex(KeyError, "Invalid mesh_dim_name"):580+ with self.assertRaisesRegex(KeyError, "Invalid mesh_dim_name"):
581- mesh_dim_names = ("DP", "TP")581+ mesh_dim_names = ("DP", "TP")
582- mesh = init_device_mesh(582+ mesh = init_device_mesh(
583- self.device_type, (1, 2), mesh_dim_names=mesh_dim_names583+ self.device_type, (1, 2), mesh_dim_names=mesh_dim_names
584- )584+ )
585- child_mesh = mesh[child_mesh_dim_name]585+ child_mesh = mesh[child_mesh_dim_name]
586- 586+ 
587- @skipIfUnsupportMultiNPU(2)587+ @skipIfUnsupportMultiNPU(2)
588- @with_comms588+ @with_comms
589- def test_get_item_1d(self):589+ def test_get_item_1d(self):
590- mesh = init_device_mesh(self.device_type, (2,), mesh_dim_names=("dp",))590+ mesh = init_device_mesh(self.device_type, (2,), mesh_dim_names=("dp",))
591- # Make sure slicing out 1D mesh from a 1D mesh works.591+ # Make sure slicing out 1D mesh from a 1D mesh works.
592- dp_mesh = mesh["dp"]592+ dp_mesh = mesh["dp"]
593- self.assertEqual(dp_mesh, mesh)593+ self.assertEqual(dp_mesh, mesh)
594- 594+ 
595- with self.assertRaisesRegex(KeyError, "Invalid mesh_dim_name"):595+ with self.assertRaisesRegex(KeyError, "Invalid mesh_dim_name"):
596- dp_mesh = mesh["dim0"]596+ dp_mesh = mesh["dim0"]
597- 597+ 
598- @skipIfUnsupportMultiNPU(2)598+ @skipIfUnsupportMultiNPU(2)
599- @with_comms599+ @with_comms
600- def test_cache_and_reuse_submesh_slice_result(self):600+ def test_cache_and_reuse_submesh_slice_result(self):
601- mesh = init_device_mesh(self.device_type, (1, 2), mesh_dim_names=("dp", "tp"))601+ mesh = init_device_mesh(self.device_type, (1, 2), mesh_dim_names=("dp", "tp"))
602- 602+ 
603- dp_mesh = mesh["dp"]603+ dp_mesh = mesh["dp"]
604- ref_pg_count = _world.group_count604+ ref_pg_count = _world.group_count
605- 605+ 
606- # When we access the "dp" slice again it should not create any new pg.606+ # When we access the "dp" slice again it should not create any new pg.
607- # As we are just using the cached result so the pg count should be the same.607+ # As we are just using the cached result so the pg count should be the same.
608- dp_mesh_2 = mesh["dp"]608+ dp_mesh_2 = mesh["dp"]
609- self.assertEqual(ref_pg_count, _world.group_count)609+ self.assertEqual(ref_pg_count, _world.group_count)
610- 610+ 
611- # When we access the "tp" slice, it should not create a new pg, as the "tp" slice would611+ # When we access the "tp" slice, it should not create a new pg, as the "tp" slice would
612- # just reuse the parent mesh pg.612+ # just reuse the parent mesh pg.
613- tp_mesh = mesh["tp"]613+ tp_mesh = mesh["tp"]
614- self.assertEqual(_world.group_count, ref_pg_count)614+ self.assertEqual(_world.group_count, ref_pg_count)
615- 615+ 
616- @skipIfUnsupportMultiNPU(2)616+ @skipIfUnsupportMultiNPU(2)
617- @with_comms617+ @with_comms
618- def test_flatten_mesh_3d(self):618+ def test_flatten_mesh_3d(self):
619- mesh_shape = (1, 1, 2)619+ mesh_shape = (1, 1, 2)
620- mesh_dim_names = ("dp", "cp", "tp")620+ mesh_dim_names = ("dp", "cp", "tp")
621- mesh_3d = init_device_mesh(621+ mesh_3d = init_device_mesh(
622- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names622+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
623- )623+ )
624- 624+ 
625- # Test flatten contiguous dims625+ # Test flatten contiguous dims
626- dp_cp_mesh = mesh_3d["dp", "cp"]626+ dp_cp_mesh = mesh_3d["dp", "cp"]
627- flattened_dp_cp_mesh = dp_cp_mesh._flatten()627+ flattened_dp_cp_mesh = dp_cp_mesh._flatten()
628- self.assertEqual(dp_cp_mesh.mesh.flatten(), flattened_dp_cp_mesh.mesh)628+ self.assertEqual(dp_cp_mesh.mesh.flatten(), flattened_dp_cp_mesh.mesh)
629- self.assertEqual(flattened_dp_cp_mesh.mesh_dim_names[0], "dp_cp")629+ self.assertEqual(flattened_dp_cp_mesh.mesh_dim_names[0], "dp_cp")
630- root_mesh = _mesh_resources.get_root_mesh(dp_cp_mesh)630+ root_mesh = _mesh_resources.get_root_mesh(dp_cp_mesh)
631- self.assertEqual(root_mesh, mesh_3d)631+ self.assertEqual(root_mesh, mesh_3d)
632- flatten_mesh_root_dims = _mesh_resources.flatten_name_to_root_dims[root_mesh][632+ flatten_mesh_root_dims = _mesh_resources.flatten_name_to_root_dims[root_mesh][
633- "dp_cp"633+ "dp_cp"
634- ]634+ ]
635- self.assertEqual(flatten_mesh_root_dims, (0, 1))635+ self.assertEqual(flatten_mesh_root_dims, (0, 1))
636- 636+ 
637- ref_pg_count = _world.group_count637+ ref_pg_count = _world.group_count
638- # Calling flatten again should not create a new pg.638+ # Calling flatten again should not create a new pg.
639- flattened_dp_cp_mesh_2 = dp_cp_mesh._flatten()639+ flattened_dp_cp_mesh_2 = dp_cp_mesh._flatten()
640- self.assertEqual(flattened_dp_cp_mesh, flattened_dp_cp_mesh_2)640+ self.assertEqual(flattened_dp_cp_mesh, flattened_dp_cp_mesh_2)
641- self.assertEqual(ref_pg_count, _world.group_count)641+ self.assertEqual(ref_pg_count, _world.group_count)
642- 642+ 
643- # Test flatten non-contiguous dims643+ # Test flatten non-contiguous dims
644- dp_tp_mesh = mesh_3d["dp", "tp"]644+ dp_tp_mesh = mesh_3d["dp", "tp"]
645- flattened_dp_tp_mesh = dp_tp_mesh._flatten()645+ flattened_dp_tp_mesh = dp_tp_mesh._flatten()
646- self.assertEqual(dp_tp_mesh.mesh.flatten(), flattened_dp_tp_mesh.mesh)646+ self.assertEqual(dp_tp_mesh.mesh.flatten(), flattened_dp_tp_mesh.mesh)
647- self.assertEqual(flattened_dp_tp_mesh.mesh_dim_names[0], "dp_tp")647+ self.assertEqual(flattened_dp_tp_mesh.mesh_dim_names[0], "dp_tp")
648- root_mesh = _mesh_resources.get_root_mesh(dp_tp_mesh)648+ root_mesh = _mesh_resources.get_root_mesh(dp_tp_mesh)
649- self.assertEqual(root_mesh, mesh_3d)649+ self.assertEqual(root_mesh, mesh_3d)
650- flatten_mesh_root_dims = _mesh_resources.flatten_name_to_root_dims[root_mesh][650+ flatten_mesh_root_dims = _mesh_resources.flatten_name_to_root_dims[root_mesh][
651- "dp_tp"651+ "dp_tp"
652- ]652+ ]
653- self.assertEqual(flatten_mesh_root_dims, (0, 2))653+ self.assertEqual(flatten_mesh_root_dims, (0, 2))
654- 654+ 
655- # Test flatten with a flattened mesh_dim_name655+ # Test flatten with a flattened mesh_dim_name
656- cp_tp_mesh = mesh_3d["cp", "tp"]656+ cp_tp_mesh = mesh_3d["cp", "tp"]
657- cp_tp_mesh._flatten("dummy")657+ cp_tp_mesh._flatten("dummy")
658- self.assertEqual(mesh_3d["dummy"].mesh_dim_names[0], "dummy")658+ self.assertEqual(mesh_3d["dummy"].mesh_dim_names[0], "dummy")
659- 659+ 
660- @skipIfUnsupportMultiNPU(2)660+ @skipIfUnsupportMultiNPU(2)
661- @with_comms661+ @with_comms
662- def test_flatten_mesh_4d(self):662+ def test_flatten_mesh_4d(self):
663- with self.subTest(eager_init=True):663+ with self.subTest(eager_init=True):
664- mesh_shape = (2, 1, 1, 1)664+ mesh_shape = (2, 1, 1, 1)
665- mesh_dim_names = ("dp_replicate", "dp_shard", "cp", "tp")665+ mesh_dim_names = ("dp_replicate", "dp_shard", "cp", "tp")
666- mesh_4d = init_device_mesh(666+ mesh_4d = init_device_mesh(
667- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names667+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
668- )668+ )
669- 669+ 
670- # flatten HSDP and CP into one mesh670+ # flatten HSDP and CP into one mesh
671- dp_cp_mesh = mesh_4d[mesh_dim_names[:3]]._flatten("dp_cp")671+ dp_cp_mesh = mesh_4d[mesh_dim_names[:3]]._flatten("dp_cp")
672- # check flattened mesh integrity672+ # check flattened mesh integrity
673- self.assertEqual(mesh_4d["dp_cp"].mesh.flatten(), dp_cp_mesh.mesh)673+ self.assertEqual(mesh_4d["dp_cp"].mesh.flatten(), dp_cp_mesh.mesh)
674- # check flattened mesh dim names is correct674+ # check flattened mesh dim names is correct
675- self.assertEqual(dp_cp_mesh.mesh_dim_names, ("dp_cp",))675+ self.assertEqual(dp_cp_mesh.mesh_dim_names, ("dp_cp",))
676- # check flattened mesh dependency676+ # check flattened mesh dependency
677- self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d)677+ self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d)
678- 678+ 
679- with self.subTest(eager_init=False):679+ with self.subTest(eager_init=False):
680- mesh_shape = (2, 1, 1, 1)680+ mesh_shape = (2, 1, 1, 1)
681- mesh_dim_names = ("dp_replicate", "dp_shard", "cp", "tp")681+ mesh_dim_names = ("dp_replicate", "dp_shard", "cp", "tp")
682- mesh_4d = init_device_mesh(682+ mesh_4d = init_device_mesh(
683- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names683+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
684- )684+ )
685- 685+ 
686- # flatten HSDP and CP into one mesh686+ # flatten HSDP and CP into one mesh
687- dp_cp_mesh = mesh_4d[mesh_dim_names[:3]]._flatten("dp_cp")687+ dp_cp_mesh = mesh_4d[mesh_dim_names[:3]]._flatten("dp_cp")
688- # check flattened mesh integrity688+ # check flattened mesh integrity
689- self.assertEqual(mesh_4d["dp_cp"].mesh.flatten(), dp_cp_mesh.mesh)689+ self.assertEqual(mesh_4d["dp_cp"].mesh.flatten(), dp_cp_mesh.mesh)
690- # check flattened mesh dim names is correct690+ # check flattened mesh dim names is correct
691- self.assertEqual(dp_cp_mesh.mesh_dim_names, ("dp_cp",))691+ self.assertEqual(dp_cp_mesh.mesh_dim_names, ("dp_cp",))
692- # check flattened mesh dependency692+ # check flattened mesh dependency
693- self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d)693+ self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d)
694- 694+ 
695- 695+ 
696-#TestDeviceMeshGetItem with resetting world_size to 8.696+#TestDeviceMeshGetItem with resetting world_size to 8.
697-class TestDeviceMeshGetItemE(NPUDTensorTestBase):697+class TestDeviceMeshGetItemE(NPUDTensorTestBase):
698- @property698+ @property
699- def world_size(self):699+ def world_size(self):
700- return 8700+ return 8
701- 701+ 
702- @skipIfUnsupportMultiNPU(8)702+ @skipIfUnsupportMultiNPU(8)
703- @with_comms703+ @with_comms
704- def test_get_item_2d(self):704+ def test_get_item_2d(self):
705- mesh_shape = (2, 4)705+ mesh_shape = (2, 4)
706- mesh_dim_names = ("DP", "TP")706+ mesh_dim_names = ("DP", "TP")
707- mesh_2d = init_device_mesh(707+ mesh_2d = init_device_mesh(
708- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names708+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
709- )709+ )
710- 710+ 
711- pg_ranks_by_dim_name = {}711+ pg_ranks_by_dim_name = {}
712- for mesh_dim_name in mesh_dim_names:712+ for mesh_dim_name in mesh_dim_names:
713- mesh_dim = mesh_dim_names.index(mesh_dim_name)713+ mesh_dim = mesh_dim_names.index(mesh_dim_name)
714- pg_ranks_by_dim_name[mesh_dim_name] = mesh_2d.mesh.swapdims(714+ pg_ranks_by_dim_name[mesh_dim_name] = mesh_2d.mesh.swapdims(
715- -1, mesh_dim715+ -1, mesh_dim
716- ).reshape(-1, mesh_2d.mesh.size(mesh_dim))716+ ).reshape(-1, mesh_2d.mesh.size(mesh_dim))
717- 717+ 
718- tp_mesh = mesh_2d["TP"]718+ tp_mesh = mesh_2d["TP"]
719- tp_group_idx = self.rank // 4719+ tp_group_idx = self.rank // 4
720- self.assertEqual(tp_mesh.mesh, pg_ranks_by_dim_name["TP"][tp_group_idx])720+ self.assertEqual(tp_mesh.mesh, pg_ranks_by_dim_name["TP"][tp_group_idx])
721- 721+ 
722- dp_mesh = mesh_2d["DP"]722+ dp_mesh = mesh_2d["DP"]
723- dp_group_idx = self.rank % 4723+ dp_group_idx = self.rank % 4
724- self.assertEqual(mesh_2d["DP"].mesh, pg_ranks_by_dim_name["DP"][dp_group_idx])724+ self.assertEqual(mesh_2d["DP"].mesh, pg_ranks_by_dim_name["DP"][dp_group_idx])
725- 725+ 
726- @skipIfUnsupportMultiNPU(8)726+ @skipIfUnsupportMultiNPU(8)
727- @with_comms727+ @with_comms
728- def test_get_item_3d(self):728+ def test_get_item_3d(self):
729- mesh_shape = (2, 2, 2)729+ mesh_shape = (2, 2, 2)
730- mesh_dim_names = ("Replicate", "Shard", "TP")730+ mesh_dim_names = ("Replicate", "Shard", "TP")
731- mesh_3d = init_device_mesh(731+ mesh_3d = init_device_mesh(
732- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names732+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
733- )733+ )
734- 734+ 
735- tp_group = [[0, 1], [2, 3], [4, 5], [6, 7]]735+ tp_group = [[0, 1], [2, 3], [4, 5], [6, 7]]
736- tp_group_idx = int(self.rank / 2)736+ tp_group_idx = int(self.rank / 2)
737- self.assertEqual(mesh_3d["TP"].mesh.tolist(), tp_group[tp_group_idx])737+ self.assertEqual(mesh_3d["TP"].mesh.tolist(), tp_group[tp_group_idx])
738- 738+ 
739- shard_group = [[0, 2], [1, 3], [4, 6], [5, 7]]739+ shard_group = [[0, 2], [1, 3], [4, 6], [5, 7]]
740- shard_group_idx = self.rank % 2 + self.rank // 4 * 2740+ shard_group_idx = self.rank % 2 + self.rank // 4 * 2
741- self.assertEqual(mesh_3d["Shard"].mesh.tolist(), shard_group[shard_group_idx])741+ self.assertEqual(mesh_3d["Shard"].mesh.tolist(), shard_group[shard_group_idx])
742- 742+ 
743- replicate_group = [[0, 4], [1, 5], [2, 6], [3, 7]]743+ replicate_group = [[0, 4], [1, 5], [2, 6], [3, 7]]
744- replicate_group_idx = self.rank % 4744+ replicate_group_idx = self.rank % 4
745- self.assertEqual(745+ self.assertEqual(
746- mesh_3d["Replicate"].mesh.tolist(), replicate_group[replicate_group_idx]746+ mesh_3d["Replicate"].mesh.tolist(), replicate_group[replicate_group_idx]
747- )747+ )
748- 748+ 
749- # We support both UX for nD slicing.749+ # We support both UX for nD slicing.
750- # E.g. mesh_3d[["Replicate", "Shard"]] or mesh_3d["Replicate", "Shard"].750+ # E.g. mesh_3d[["Replicate", "Shard"]] or mesh_3d["Replicate", "Shard"].
751- hsdp_mesh_1 = mesh_3d[["Replicate", "Shard"]]751+ hsdp_mesh_1 = mesh_3d[["Replicate", "Shard"]]
752- hsdp_mesh_2 = mesh_3d["Replicate", "Shard"]752+ hsdp_mesh_2 = mesh_3d["Replicate", "Shard"]
753- hsdp_group = [[[0, 2], [4, 6]], [[1, 3], [5, 7]]]753+ hsdp_group = [[[0, 2], [4, 6]], [[1, 3], [5, 7]]]
754- hsdp_group_idx = self.rank % 2754+ hsdp_group_idx = self.rank % 2
755- self.assertEqual(hsdp_mesh_1.mesh.tolist(), hsdp_group[hsdp_group_idx])755+ self.assertEqual(hsdp_mesh_1.mesh.tolist(), hsdp_group[hsdp_group_idx])
756- self.assertEqual(hsdp_mesh_2.mesh.tolist(), hsdp_group[hsdp_group_idx])756+ self.assertEqual(hsdp_mesh_2.mesh.tolist(), hsdp_group[hsdp_group_idx])
757- self.assertEqual(hsdp_mesh_1, hsdp_mesh_2)757+ self.assertEqual(hsdp_mesh_1, hsdp_mesh_2)
758- 758+ 
759- @skipIfUnsupportMultiNPU(8)759+ @skipIfUnsupportMultiNPU(8)
760- @with_comms760+ @with_comms
761- def test_get_item_3d_noncontiguous_slicing(self):761+ def test_get_item_3d_noncontiguous_slicing(self):
762- mesh_shape = (2, 2, 2)762+ mesh_shape = (2, 2, 2)
763- mesh_dim_names = ("dp", "pp", "cp")763+ mesh_dim_names = ("dp", "pp", "cp")
764- mesh_3d = init_device_mesh(764+ mesh_3d = init_device_mesh(
765- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names765+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
766- )766+ )
767- 767+ 
768- # Slice order simply decides which mesh_dim sits on which mesh_dim.768+ # Slice order simply decides which mesh_dim sits on which mesh_dim.
769- # For dp_cp_mesh, cp mesh is the innermost dimension.769+ # For dp_cp_mesh, cp mesh is the innermost dimension.
770- dp_cp_mesh = mesh_3d["dp", "cp"]770+ dp_cp_mesh = mesh_3d["dp", "cp"]
771- expected_mesh_tensor = (771+ expected_mesh_tensor = (
772- torch.tensor([[0, 1], [4, 5]], dtype=torch.int)772+ torch.tensor([[0, 1], [4, 5]], dtype=torch.int)
773- if self.rank in (0, 1, 4, 5)773+ if self.rank in (0, 1, 4, 5)
774- else torch.tensor([[2, 3], [6, 7]], dtype=torch.int)774+ else torch.tensor([[2, 3], [6, 7]], dtype=torch.int)
775- )775+ )
776- dp_local_rank = dp_cp_mesh.get_local_rank("dp")776+ dp_local_rank = dp_cp_mesh.get_local_rank("dp")
777- self.assertEqual(dp_cp_mesh.mesh, expected_mesh_tensor)777+ self.assertEqual(dp_cp_mesh.mesh, expected_mesh_tensor)
778- cp_mesh = mesh_3d["cp"]778+ cp_mesh = mesh_3d["cp"]
779- # Check on the current dp_local_rank, whether the cp mesh tensor is the same.779+ # Check on the current dp_local_rank, whether the cp mesh tensor is the same.
780- self.assertEqual(dp_cp_mesh.mesh[dp_local_rank], cp_mesh.mesh)780+ self.assertEqual(dp_cp_mesh.mesh[dp_local_rank], cp_mesh.mesh)
781- 781+ 
782- with self.assertRaisesRegex(782+ with self.assertRaisesRegex(
783- KeyError,783+ KeyError,
784- "Invalid mesh_dim_names",784+ "Invalid mesh_dim_names",
785- ):785+ ):
786- cp_dp_mesh = mesh_3d["cp", "dp"]786+ cp_dp_mesh = mesh_3d["cp", "dp"]
787- 787+ 
788- @skipIfUnsupportMultiNPU(8)788+ @skipIfUnsupportMultiNPU(8)
789- @with_comms789+ @with_comms
790- def test_reconstruct_mesh_with_flatten_dim(self):790+ def test_reconstruct_mesh_with_flatten_dim(self):
791- mesh_3d = init_device_mesh(791+ mesh_3d = init_device_mesh(
792- self.device_type, (2, 2, 2), mesh_dim_names=("replicate", "shard", "cp")792+ self.device_type, (2, 2, 2), mesh_dim_names=("replicate", "shard", "cp")
793- )793+ )
794- shard_cp_mesh = mesh_3d["shard", "cp"]._flatten()794+ shard_cp_mesh = mesh_3d["shard", "cp"]._flatten()
795- hsdp_mesh = mesh_3d["replicate", "shard_cp"]795+ hsdp_mesh = mesh_3d["replicate", "shard_cp"]
796- expected_mesh_tensor = torch.tensor(796+ expected_mesh_tensor = torch.tensor(
797- [[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.int797+ [[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.int
798- )798+ )
799- self.assertEqual(hsdp_mesh.mesh, expected_mesh_tensor)799+ self.assertEqual(hsdp_mesh.mesh, expected_mesh_tensor)
800- self.assertEqual(shard_cp_mesh.get_group(), mesh_3d["shard_cp"].get_group())800+ self.assertEqual(shard_cp_mesh.get_group(), mesh_3d["shard_cp"].get_group())
801- self.assertEqual(801+ self.assertEqual(
802- shard_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="shard_cp")802+ shard_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="shard_cp")
803- )803+ )
804- 804+ 
805- mesh_3d = init_device_mesh(805+ mesh_3d = init_device_mesh(
806- self.device_type, (2, 2, 2), mesh_dim_names=("dp", "cp", "tp")806+ self.device_type, (2, 2, 2), mesh_dim_names=("dp", "cp", "tp")
807- )807+ )
808- dp_cp_mesh = mesh_3d["dp", "cp"]._flatten()808+ dp_cp_mesh = mesh_3d["dp", "cp"]._flatten()
809- spmd_mesh = mesh_3d["dp_cp", "tp"]809+ spmd_mesh = mesh_3d["dp_cp", "tp"]
810- expected_mesh_tensor = torch.tensor(810+ expected_mesh_tensor = torch.tensor(
811- [[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.int811+ [[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.int
812- )812+ )
813- self.assertEqual(spmd_mesh.mesh, expected_mesh_tensor)813+ self.assertEqual(spmd_mesh.mesh, expected_mesh_tensor)
814- self.assertEqual(dp_cp_mesh.get_group(), mesh_3d["dp_cp"].get_group())814+ self.assertEqual(dp_cp_mesh.get_group(), mesh_3d["dp_cp"].get_group())
815- self.assertEqual(dp_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="dp_cp"))815+ self.assertEqual(dp_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="dp_cp"))
816- 816+ 
817- 817+ 
818-class TestMeshEnv(NPUDTensorTestBase):818+class TestMeshEnv(NPUDTensorTestBase):
819- @property819+ @property
820- def world_size(self):820+ def world_size(self):
821- return 2821+ return 2
822- 822+ 
823- @skipIfUnsupportMultiNPU(2)823+ @skipIfUnsupportMultiNPU(2)
824- @with_comms824+ @with_comms
825- def test_get_root_mesh(self):825+ def test_get_root_mesh(self):
826- mesh_3d = init_device_mesh(826+ mesh_3d = init_device_mesh(
827- self.device_type, (2, 1, 1), mesh_dim_names=("dp", "cp", "tp")827+ self.device_type, (2, 1, 1), mesh_dim_names=("dp", "cp", "tp")
828- )828+ )
829- 829+ 
830- dp_cp_mesh = mesh_3d["dp", "cp"]830+ dp_cp_mesh = mesh_3d["dp", "cp"]
831- dp_tp_mesh = mesh_3d["dp", "tp"]831+ dp_tp_mesh = mesh_3d["dp", "tp"]
832- cp_tp_mesh = mesh_3d["cp", "tp"]832+ cp_tp_mesh = mesh_3d["cp", "tp"]
833- dp_mesh = mesh_3d["dp"]833+ dp_mesh = mesh_3d["dp"]
834- cp_mesh = mesh_3d["cp"]834+ cp_mesh = mesh_3d["cp"]
835- tp_mesh = mesh_3d["tp"]835+ tp_mesh = mesh_3d["tp"]
836- self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_3d)836+ self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_3d)
837- self.assertEqual(_mesh_resources.get_root_mesh(dp_tp_mesh), mesh_3d)837+ self.assertEqual(_mesh_resources.get_root_mesh(dp_tp_mesh), mesh_3d)
838- self.assertEqual(_mesh_resources.get_root_mesh(cp_tp_mesh), mesh_3d)838+ self.assertEqual(_mesh_resources.get_root_mesh(cp_tp_mesh), mesh_3d)
839- self.assertEqual(_mesh_resources.get_root_mesh(dp_mesh), mesh_3d)839+ self.assertEqual(_mesh_resources.get_root_mesh(dp_mesh), mesh_3d)
840- self.assertEqual(_mesh_resources.get_root_mesh(cp_mesh), mesh_3d)840+ self.assertEqual(_mesh_resources.get_root_mesh(cp_mesh), mesh_3d)
841- self.assertEqual(_mesh_resources.get_root_mesh(tp_mesh), mesh_3d)841+ self.assertEqual(_mesh_resources.get_root_mesh(tp_mesh), mesh_3d)
842- 842+ 
843- @skipIfUnsupportMultiNPU(2)843+ @skipIfUnsupportMultiNPU(2)
844- @with_comms844+ @with_comms
845- def test_get_root_mesh_dim_exist(self):845+ def test_get_root_mesh_dim_exist(self):
846- mesh_shape = (2, self.world_size // 2)846+ mesh_shape = (2, self.world_size // 2)
847- mesh_dim_names = ("DP", "TP")847+ mesh_dim_names = ("DP", "TP")
848- mesh_2d = init_device_mesh(848+ mesh_2d = init_device_mesh(
849- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names849+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
850- )850+ )
851- 851+ 
852- self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh_2d["DP"]), 0)852+ self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh_2d["DP"]), 0)
853- self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh_2d["TP"]), 1)853+ self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh_2d["TP"]), 1)
854- 854+ 
855- @skipIfUnsupportMultiNPU(2)855+ @skipIfUnsupportMultiNPU(2)
856- @with_comms856+ @with_comms
857- def test_get_root_mesh_dim_not_exist(self):857+ def test_get_root_mesh_dim_not_exist(self):
858- mesh_shape = (self.world_size,)858+ mesh_shape = (self.world_size,)
859- mesh = init_device_mesh(self.device_type, mesh_shape)859+ mesh = init_device_mesh(self.device_type, mesh_shape)
860- 860+ 
861- self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh), None)861+ self.assertEqual(_mesh_resources.get_root_mesh_dim(mesh), None)
862- 862+ 
863- @skipIfUnsupportMultiNPU(2)863+ @skipIfUnsupportMultiNPU(2)
864- @with_comms864+ @with_comms
865- def test_get_mesh_dim_by_name(self):865+ def test_get_mesh_dim_by_name(self):
866- mesh_shape = (2, self.world_size // 2)866+ mesh_shape = (2, self.world_size // 2)
867- mesh_dim_names = ("DP", "TP")867+ mesh_dim_names = ("DP", "TP")
868- mesh_2d = init_device_mesh(868+ mesh_2d = init_device_mesh(
869- self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names869+ self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
870- )870+ )
871- 871+ 
872- self.assertEqual(_mesh_resources.get_mesh_dim_by_name(mesh_2d, "DP"), 0)872+ self.assertEqual(_mesh_resources.get_mesh_dim_by_name(mesh_2d, "DP"), 0)
873- self.assertEqual(_mesh_resources.get_mesh_dim_by_name(mesh_2d, "TP"), 1)873+ self.assertEqual(_mesh_resources.get_mesh_dim_by_name(mesh_2d, "TP"), 1)
874- 874+ 
875- @skipIfUnsupportMultiNPU(2)875+ @skipIfUnsupportMultiNPU(2)
876- @with_comms876+ @with_comms
877- def test_get_all_submeshes(self):877+ def test_get_all_submeshes(self):
878- mesh_2d = init_device_mesh(878+ mesh_2d = init_device_mesh(
879- self.device_type, (1, 2), mesh_dim_names=("replicate", "shard")879+ self.device_type, (1, 2), mesh_dim_names=("replicate", "shard")
880- )880+ )
881- all_submeshes = _mesh_resources._get_all_submeshes(mesh_2d, "replicate")881+ all_submeshes = _mesh_resources._get_all_submeshes(mesh_2d, "replicate")
882- self.assertEqual(len(all_submeshes), 2)882+ self.assertEqual(len(all_submeshes), 2)
883- self.assertEqual(883+ self.assertEqual(
884- all(submesh.mesh.numel() == 1 for submesh in all_submeshes), True884+ all(submesh.mesh.numel() == 1 for submesh in all_submeshes), True
885- )885+ )
886- 886+ 
887- 887+ 
888-class DeviceMeshCollectiveTest(NPUDTensorTestBase):888+class DeviceMeshCollectiveTest(NPUDTensorTestBase):
889- @property889+ @property
890- def world_size(self):890+ def world_size(self):
891- return 2891+ return 2
892- 892+ 
893- @skipIfUnsupportMultiNPU(2)893+ @skipIfUnsupportMultiNPU(2)
894- @with_comms894+ @with_comms
895- def test_broadcast_1d(self):895+ def test_broadcast_1d(self):
896- mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))896+ mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
897- local_tensor = torch.ones(3, 3, device=self.device_type) * self.rank897+ local_tensor = torch.ones(3, 3, device=self.device_type) * self.rank
898- mesh_broadcast(local_tensor, mesh, mesh_dim=0)898+ mesh_broadcast(local_tensor, mesh, mesh_dim=0)
899- self.assertEqual(local_tensor, torch.zeros(3, 3))899+ self.assertEqual(local_tensor, torch.zeros(3, 3))
900- 900+ 
901- @skipIfUnsupportMultiNPU(2)901+ @skipIfUnsupportMultiNPU(2)
902- @with_comms902+ @with_comms
903- def test_scatter_1d(self):903+ def test_scatter_1d(self):
904- mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))904+ mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
905- scatter_tensor_shape = [3, 3, 3]905+ scatter_tensor_shape = [3, 3, 3]
906- len_scatter_tensor_shape = len(scatter_tensor_shape)906+ len_scatter_tensor_shape = len(scatter_tensor_shape)
907- for scatter_dim in range(len_scatter_tensor_shape):907+ for scatter_dim in range(len_scatter_tensor_shape):
908- shard_placement = Shard(scatter_dim)908+ shard_placement = Shard(scatter_dim)
909- scatter_tensor_shape[scatter_dim] *= self.world_size909+ scatter_tensor_shape[scatter_dim] *= self.world_size
910- # make the random seed same across rank910+ # make the random seed same across rank
911- torch.manual_seed(0)911+ torch.manual_seed(0)
912- global_tensor = torch.randn(scatter_tensor_shape, device=self.device_type)912+ global_tensor = torch.randn(scatter_tensor_shape, device=self.device_type)
913- splitted_list, _ = shard_placement._split_tensor(913+ splitted_list, _ = shard_placement._split_tensor(
914- global_tensor, mesh.size(), with_padding=True, contiguous=True914+ global_tensor, mesh.size(), with_padding=True, contiguous=True
915- )915+ )
916- recv_tensor = torch.empty_like(splitted_list[mesh.get_rank()])916+ recv_tensor = torch.empty_like(splitted_list[mesh.get_rank()])
917- # scatter on dim > 0 would generate non-contiguous tensor, verify that works917+ # scatter on dim > 0 would generate non-contiguous tensor, verify that works
918- mesh_scatter(recv_tensor, splitted_list, mesh, mesh_dim=0)918+ mesh_scatter(recv_tensor, splitted_list, mesh, mesh_dim=0)
919- self.assertEqual(recv_tensor, splitted_list[mesh.get_rank()])919+ self.assertEqual(recv_tensor, splitted_list[mesh.get_rank()])
920- 920+ 
921- @skipIfUnsupportMultiNPU(2)921+ @skipIfUnsupportMultiNPU(2)
922- @with_comms922+ @with_comms
923- def test_scatter_uneven(self):923+ def test_scatter_uneven(self):
924- device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))924+ device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
925- my_rank = device_mesh.get_rank()925+ my_rank = device_mesh.get_rank()
926- tensor_to_split = torch.randn(926+ tensor_to_split = torch.randn(
927- device_mesh.size() + 3, device_mesh.size() + 1, device=self.device_type927+ device_mesh.size() + 3, device_mesh.size() + 1, device=self.device_type
928- )928+ )
929- 929+ 
930- for shard_dim in range(tensor_to_split.ndim):930+ for shard_dim in range(tensor_to_split.ndim):
931- shard_placement = Shard(shard_dim)931+ shard_placement = Shard(shard_dim)
932- 932+ 
933- tensor_to_scatter = tensor_to_split.clone()933+ tensor_to_scatter = tensor_to_split.clone()
934- tensor_splitted_list = list(934+ tensor_splitted_list = list(
935- torch.chunk(tensor_to_split, self.world_size, dim=shard_dim)935+ torch.chunk(tensor_to_split, self.world_size, dim=shard_dim)
936- )936+ )
937- for _ in range(self.world_size - len(tensor_splitted_list)):937+ for _ in range(self.world_size - len(tensor_splitted_list)):
938- tensor_splitted_list.append(torch.tensor([], device=self.device_type))938+ tensor_splitted_list.append(torch.tensor([], device=self.device_type))
939- 939+ 
940- padded_tensor_list, pad_sizes = shard_placement._split_tensor(940+ padded_tensor_list, pad_sizes = shard_placement._split_tensor(
941- tensor_to_scatter,941+ tensor_to_scatter,
942- device_mesh.size(),942+ device_mesh.size(),
943- with_padding=True,943+ with_padding=True,
944- contiguous=True,944+ contiguous=True,
945- )945+ )
946- 946+ 
947- scattered_tensor = torch.empty_like(padded_tensor_list[my_rank])947+ scattered_tensor = torch.empty_like(padded_tensor_list[my_rank])
948- mesh_scatter(scattered_tensor, padded_tensor_list, device_mesh, mesh_dim=0)948+ mesh_scatter(scattered_tensor, padded_tensor_list, device_mesh, mesh_dim=0)
949- 949+ 
950- if pad_sizes[my_rank] != 0:950+ if pad_sizes[my_rank] != 0:
951- scattered_tensor = unpad_tensor(951+ scattered_tensor = unpad_tensor(
952- scattered_tensor, shard_dim, pad_sizes[my_rank]952+ scattered_tensor, shard_dim, pad_sizes[my_rank]
953- )953+ )
954- 954+ 
955- if scattered_tensor.numel() == 0:955+ if scattered_tensor.numel() == 0:
956- # We need to check numel() instead of size if a tensor is ([]) after unpadding,956+ # We need to check numel() instead of size if a tensor is ([]) after unpadding,
957- # since the size could be ([0, 8]) after unpadding.957+ # since the size could be ([0, 8]) after unpadding.
958- self.assertEqual(958+ self.assertEqual(
959- scattered_tensor.numel(), tensor_splitted_list[my_rank].numel()959+ scattered_tensor.numel(), tensor_splitted_list[my_rank].numel()
960- )960+ )
961- else:961+ else:
962- self.assertEqual(962+ self.assertEqual(
963- scattered_tensor.size(), tensor_splitted_list[my_rank].size()963+ scattered_tensor.size(), tensor_splitted_list[my_rank].size()
964- )964+ )
965- self.assertEqual(scattered_tensor, tensor_splitted_list[my_rank])965+ self.assertEqual(scattered_tensor, tensor_splitted_list[my_rank])
966- 966+ 
967- @skipIfUnsupportMultiNPU(2)967+ @skipIfUnsupportMultiNPU(2)
968- @with_comms968+ @with_comms
969- def test_all_gather_uneven(self):969+ def test_all_gather_uneven(self):
970- device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))970+ device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
971- my_rank = device_mesh.get_rank()971+ my_rank = device_mesh.get_rank()
972- tensor_to_split = torch.ones(972+ tensor_to_split = torch.ones(
973- device_mesh.size() + 3,973+ device_mesh.size() + 3,
974- device_mesh.size() + 1,974+ device_mesh.size() + 1,
975- device=self.device_type,975+ device=self.device_type,
976- )976+ )
977- 977+ 
978- for shard_dim in range(tensor_to_split.ndim):978+ for shard_dim in range(tensor_to_split.ndim):
979- shard_placement = Shard(shard_dim)979+ shard_placement = Shard(shard_dim)
980- tensor_padded_list, pad_sizes = shard_placement._split_tensor(980+ tensor_padded_list, pad_sizes = shard_placement._split_tensor(
981- tensor_to_split,981+ tensor_to_split,
982- device_mesh.size(),982+ device_mesh.size(),
983- with_padding=True,983+ with_padding=True,
984- contiguous=True,984+ contiguous=True,
985- )985+ )
986- local_tensor = tensor_padded_list[my_rank]986+ local_tensor = tensor_padded_list[my_rank]
987- big_tensor = funcol.all_gather_tensor(987+ big_tensor = funcol.all_gather_tensor(
988- local_tensor, gather_dim=shard_dim, group=(device_mesh, 0)988+ local_tensor, gather_dim=shard_dim, group=(device_mesh, 0)
989- )989+ )
990- big_tensor_chunks = list(990+ big_tensor_chunks = list(
991- torch.chunk(big_tensor, device_mesh.size(), dim=shard_dim)991+ torch.chunk(big_tensor, device_mesh.size(), dim=shard_dim)
992- )992+ )
993- unpadded_list = [993+ unpadded_list = [
994- (994+ (
995- unpad_tensor(big_tensor, shard_dim, pad_sizes[i])995+ unpad_tensor(big_tensor, shard_dim, pad_sizes[i])
996- if pad_sizes[i] > 0996+ if pad_sizes[i] > 0
997- else big_tensor997+ else big_tensor
998- )998+ )
999- for i, big_tensor in enumerate(big_tensor_chunks)999+ for i, big_tensor in enumerate(big_tensor_chunks)
1000- ]1000+ ]
1001- all_gathered_tensor = torch.cat(unpadded_list, dim=shard_dim)1001+ all_gathered_tensor = torch.cat(unpadded_list, dim=shard_dim)
1002- 1002+ 
1003- self.assertEqual(all_gathered_tensor.size(), tensor_to_split.size())1003+ self.assertEqual(all_gathered_tensor.size(), tensor_to_split.size())
1004- self.assertEqual(all_gathered_tensor, tensor_to_split)1004+ self.assertEqual(all_gathered_tensor, tensor_to_split)
1005- 1005+ 
1006- @skipIfUnsupportMultiNPU(2)1006+ @skipIfUnsupportMultiNPU(2)
1007- @with_comms1007+ @with_comms
1008- def test_reduce_scatter_contiguous(self):1008+ def test_reduce_scatter_contiguous(self):
1009- device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))1009+ device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
1010- my_rank = device_mesh.get_rank()1010+ my_rank = device_mesh.get_rank()
1011- 1011+ 
1012- # Init the tensor1012+ # Init the tensor
1013- step = self.world_size * 21013+ step = self.world_size * 2
1014- total_elem = step**21014+ total_elem = step**2
1015- tensor = torch.arange(0, total_elem).view(step, -1).to(device=self.device_type)1015+ tensor = torch.arange(0, total_elem).view(step, -1).to(device=self.device_type)
1016- tensor = tensor * (my_rank + 1)1016+ tensor = tensor * (my_rank + 1)
1017- 1017+ 
1018- # Get non-contiguous tensor by slicing1018+ # Get non-contiguous tensor by slicing
1019- tensor_to_reduce = tensor[::2, :2]1019+ tensor_to_reduce = tensor[::2, :2]
1020- tensor_contiguous = tensor_to_reduce.clone().contiguous()1020+ tensor_contiguous = tensor_to_reduce.clone().contiguous()
1021- 1021+ 
1022- # Partial to Shard to trigger reduce_scatter1022+ # Partial to Shard to trigger reduce_scatter
1023- tensor_to_reduce = DTensor.from_local(1023+ tensor_to_reduce = DTensor.from_local(
1024- tensor_to_reduce, device_mesh, [_Partial()]1024+ tensor_to_reduce, device_mesh, [_Partial()]
1025- )1025+ )
1026- tensor_contiguous = DTensor.from_local(1026+ tensor_contiguous = DTensor.from_local(
1027- tensor_contiguous, device_mesh, [_Partial()]1027+ tensor_contiguous, device_mesh, [_Partial()]
1028- )1028+ )
1029- new_tensor = tensor_to_reduce.redistribute(device_mesh, [Shard(0)])1029+ new_tensor = tensor_to_reduce.redistribute(device_mesh, [Shard(0)])
1030- new_tensor_contiguous = tensor_contiguous.redistribute(device_mesh, [Shard(0)])1030+ new_tensor_contiguous = tensor_contiguous.redistribute(device_mesh, [Shard(0)])
1031- 1031+ 
1032- # The output for contiguous and non-contiguous tensors of the same value1032+ # The output for contiguous and non-contiguous tensors of the same value
1033- # should return the same reducescatter value.1033+ # should return the same reducescatter value.
1034- new_tensor_local = new_tensor._local_tensor1034+ new_tensor_local = new_tensor._local_tensor
1035- new_tensor_contiguous_local = new_tensor_contiguous._local_tensor1035+ new_tensor_contiguous_local = new_tensor_contiguous._local_tensor
1036- self.assertEqual(new_tensor_local, new_tensor_contiguous_local)1036+ self.assertEqual(new_tensor_local, new_tensor_contiguous_local)
1037- self.assertEqual(list(new_tensor_local.size()), [1, 2])1037+ self.assertEqual(list(new_tensor_local.size()), [1, 2])
1038- 1038+ 
1039- # Check the reduce numerical value1039+ # Check the reduce numerical value
1040- sum_base = (1 + self.world_size) * self.world_size / 21040+ sum_base = (1 + self.world_size) * self.world_size / 2
1041- first_elem = my_rank * sum_base * step * 21041+ first_elem = my_rank * sum_base * step * 2
1042- expected_tensor = torch.tensor(1042+ expected_tensor = torch.tensor(
1043- [[first_elem, first_elem + sum_base]],1043+ [[first_elem, first_elem + sum_base]],
1044- dtype=new_tensor_local.dtype,1044+ dtype=new_tensor_local.dtype,
1045- device=self.device_type,1045+ device=self.device_type,
1046- )1046+ )
1047- self.assertEqual(new_tensor_local, expected_tensor)1047+ self.assertEqual(new_tensor_local, expected_tensor)
1048- 1048+ 
1049- @skipIfUnsupportMultiNPU(2)1049+ @skipIfUnsupportMultiNPU(2)
1050- @with_comms1050+ @with_comms
1051- def test_reduce_scatter_uneven(self):1051+ def test_reduce_scatter_uneven(self):
1052- device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))1052+ device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
1053- my_rank = device_mesh.get_rank()1053+ my_rank = device_mesh.get_rank()
1054- tensor_to_split = (1054+ tensor_to_split = (
1055- torch.ones(1055+ torch.ones(
1056- device_mesh.size() + 3,1056+ device_mesh.size() + 3,
1057- device_mesh.size() + 1,1057+ device_mesh.size() + 1,
1058- device=self.device_type,1058+ device=self.device_type,
1059- )1059+ )
1060- * self.rank1060+ * self.rank
1061- )1061+ )
1062- 1062+ 
1063- for shard_dim in range(tensor_to_split.ndim):1063+ for shard_dim in range(tensor_to_split.ndim):
1064- shard_placement = Shard(shard_dim)1064+ shard_placement = Shard(shard_dim)
1065- tensor_to_scatter = tensor_to_split.clone()1065+ tensor_to_scatter = tensor_to_split.clone()
1066- 1066+ 
1067- tensor_splitted_list = list(1067+ tensor_splitted_list = list(
1068- torch.chunk(tensor_to_split, self.world_size, dim=shard_dim)1068+ torch.chunk(tensor_to_split, self.world_size, dim=shard_dim)
1069- )1069+ )
1070- for _ in range(self.world_size - len(tensor_splitted_list)):1070+ for _ in range(self.world_size - len(tensor_splitted_list)):
1071- tensor_splitted_list.append(torch.tensor([], device=self.device_type))1071+ tensor_splitted_list.append(torch.tensor([], device=self.device_type))
1072- 1072+ 
1073- padded_tensor_list, pad_sizes = shard_placement._split_tensor(1073+ padded_tensor_list, pad_sizes = shard_placement._split_tensor(
1074- tensor_to_scatter,1074+ tensor_to_scatter,
1075- device_mesh.size(),1075+ device_mesh.size(),
1076- with_padding=True,1076+ with_padding=True,
1077- contiguous=True,1077+ contiguous=True,
1078- )1078+ )
1079- 1079+ 
1080- tensor_to_reduce = torch.cat(padded_tensor_list, shard_dim)1080+ tensor_to_reduce = torch.cat(padded_tensor_list, shard_dim)
1081- 1081+ 
1082- res_num = ((0 + self.world_size - 1) * self.world_size) / 21082+ res_num = ((0 + self.world_size - 1) * self.world_size) / 2
1083- 1083+ 
1084- scattered_tensor = funcol.reduce_scatter_tensor(1084+ scattered_tensor = funcol.reduce_scatter_tensor(
1085- tensor_to_reduce,1085+ tensor_to_reduce,
1086- reduceOp="sum",1086+ reduceOp="sum",
1087- scatter_dim=shard_dim,1087+ scatter_dim=shard_dim,
1088- group=(device_mesh, 0),1088+ group=(device_mesh, 0),
1089- )1089+ )
1090- 1090+ 
1091- # unpad scattered_tensor1091+ # unpad scattered_tensor
1092- if pad_sizes[my_rank] > 0:1092+ if pad_sizes[my_rank] > 0:
1093- scattered_tensor = unpad_tensor(1093+ scattered_tensor = unpad_tensor(
1094- scattered_tensor, shard_dim, pad_sizes[my_rank]1094+ scattered_tensor, shard_dim, pad_sizes[my_rank]
1095- )1095+ )
1096- 1096+ 
1097- if scattered_tensor.numel() == 0:1097+ if scattered_tensor.numel() == 0:
1098- # We need to check numel() instead of size if a tensor is ([]) after unpadding,1098+ # We need to check numel() instead of size if a tensor is ([]) after unpadding,
1099- # since the size could be ([0, 8]) after unpadding.1099+ # since the size could be ([0, 8]) after unpadding.
1100- self.assertEqual(1100+ self.assertEqual(
1101- scattered_tensor.numel(), tensor_splitted_list[my_rank].numel()1101+ scattered_tensor.numel(), tensor_splitted_list[my_rank].numel()
1102- )1102+ )
1103- else:1103+ else:
1104- self.assertEqual(1104+ self.assertEqual(
1105- scattered_tensor.size(), tensor_splitted_list[my_rank].size()1105+ scattered_tensor.size(), tensor_splitted_list[my_rank].size()
1106- )1106+ )
1107- self.assertEqual(1107+ self.assertEqual(
1108- scattered_tensor,1108+ scattered_tensor,
1109- torch.ones_like(tensor_splitted_list[my_rank]) * res_num,1109+ torch.ones_like(tensor_splitted_list[my_rank]) * res_num,
1110- )1110+ )
1111- 1111+ 
1112- @skipIfUnsupportMultiNPU(2)1112+ @skipIfUnsupportMultiNPU(2)
1113- @with_comms1113+ @with_comms
1114- def test_broadcast_nd(self):1114+ def test_broadcast_nd(self):
1115- mesh_tensor = torch.arange(2).reshape(2, 1, 1)1115+ mesh_tensor = torch.arange(2).reshape(2, 1, 1)
1116- mesh = DeviceMesh(self.device_type, mesh_tensor)1116+ mesh = DeviceMesh(self.device_type, mesh_tensor)
1117- local_tensor = torch.ones(3, 3, device=self.device_type) * self.rank1117+ local_tensor = torch.ones(3, 3, device=self.device_type) * self.rank
1118- 1118+ 
1119- # check all dim groups1119+ # check all dim groups
1120- dim_to_subgroups = mesh.get_all_groups()1120+ dim_to_subgroups = mesh.get_all_groups()
1121- for dim, dim_group in enumerate(dim_to_subgroups):1121+ for dim, dim_group in enumerate(dim_to_subgroups):
1122- dim_group_size = get_world_size(dim_group)1122+ dim_group_size = get_world_size(dim_group)
1123- global_ranks = [1123+ global_ranks = [
1124- get_global_rank(dim_group, i) for i in range(dim_group_size)1124+ get_global_rank(dim_group, i) for i in range(dim_group_size)
1125- ]1125+ ]
1126- cloned_local_tensor = local_tensor.clone()1126+ cloned_local_tensor = local_tensor.clone()
1127- mesh_broadcast(cloned_local_tensor, mesh, mesh_dim=dim)1127+ mesh_broadcast(cloned_local_tensor, mesh, mesh_dim=dim)
1128- res_num = global_ranks[0]1128+ res_num = global_ranks[0]
1129- self.assertEqual(cloned_local_tensor, torch.ones(3, 3) * res_num)1129+ self.assertEqual(cloned_local_tensor, torch.ones(3, 3) * res_num)
1130- 1130+ 
1131- @skipIfUnsupportMultiNPU(2)1131+ @skipIfUnsupportMultiNPU(2)
1132- @with_comms1132+ @with_comms
1133- def test_scatter_nd(self):1133+ def test_scatter_nd(self):
1134- mesh_tensor = torch.arange(2).reshape(2, 1, 1)1134+ mesh_tensor = torch.arange(2).reshape(2, 1, 1)
1135- mesh = DeviceMesh(self.device_type, mesh_tensor)1135+ mesh = DeviceMesh(self.device_type, mesh_tensor)
1136- 1136+ 
1137- # check all dim groups1137+ # check all dim groups
1138- dim_to_subgroups = mesh.get_all_groups()1138+ dim_to_subgroups = mesh.get_all_groups()
1139- for dim, dim_group in enumerate(dim_to_subgroups):1139+ for dim, dim_group in enumerate(dim_to_subgroups):
1140- dim_group_size = get_world_size(dim_group)1140+ dim_group_size = get_world_size(dim_group)
1141- global_ranks = [1141+ global_ranks = [
1142- get_global_rank(dim_group, i) for i in range(dim_group_size)1142+ get_global_rank(dim_group, i) for i in range(dim_group_size)
1143- ]1143+ ]
1144- scattered_tensors = [1144+ scattered_tensors = [
1145- torch.ones(3, 3, device=self.device_type) * global_rank1145+ torch.ones(3, 3, device=self.device_type) * global_rank
1146- for global_rank in global_ranks1146+ for global_rank in global_ranks
1147- ]1147+ ]
1148- received_tensor = torch.empty_like(1148+ received_tensor = torch.empty_like(
1149- scattered_tensors[mesh.get_coordinate()[dim]]1149+ scattered_tensors[mesh.get_coordinate()[dim]]
1150- )1150+ )
1151- mesh_scatter(received_tensor, scattered_tensors, mesh, mesh_dim=dim)1151+ mesh_scatter(received_tensor, scattered_tensors, mesh, mesh_dim=dim)
1152- self.assertEqual(received_tensor, torch.ones(3, 3) * self.rank)1152+ self.assertEqual(received_tensor, torch.ones(3, 3) * self.rank)
1153- 1153+ 
1154- 1154+ 
1155-if __name__ == "__main__":1155+if __name__ == "__main__":
1156 run_tests()1156 run_tests()
Mtest/distributed/watchdog/watchdog_quick_exit.py+27-27
@@ -1,27 +1,27 @@
1-import os1+import os
2-import time2+import time
3-import datetime3+import datetime
4-import torch.distributed as dist4+import torch.distributed as dist
5-import torch5+import torch
6-import torch_npu6+import torch_npu
7- 7+ 
8- 8+ 
9-def main():9+def main():
10- torch.npu.set_compile_mode(jit_compile=True)10+ torch.npu.set_compile_mode(jit_compile=True)
11- rank = int(os.environ['RANK'])11+ rank = int(os.environ['RANK'])
12- local_rank = int(os.environ['LOCAL_RANK'])12+ local_rank = int(os.environ['LOCAL_RANK'])
13- device = torch.device('npu:{}'.format(local_rank))13+ device = torch.device('npu:{}'.format(local_rank))
14- torch.npu.set_device(device)14+ torch.npu.set_device(device)
15- dist.init_process_group(backend='hccl', rank=rank, world_size=2, timeout=datetime.timedelta(seconds=120))15+ dist.init_process_group(backend='hccl', rank=rank, world_size=2, timeout=datetime.timedelta(seconds=120))
16- tensor = torch.tensor(1).npu()16+ tensor = torch.tensor(1).npu()
17- dist.all_reduce(tensor)17+ dist.all_reduce(tensor)
18- if rank == 0:18+ if rank == 0:
19- x1 = torch.randn(3).float().npu()19+ x1 = torch.randn(3).float().npu()
20- x2 = torch.randn(1).long().npu()20+ x2 = torch.randn(1).long().npu()
21- x3 = torch.randn(1).float().npu()21+ x3 = torch.randn(1).float().npu()
22- y = torch.addcmul(x1, x2, x3)22+ y = torch.addcmul(x1, x2, x3)
23- dist.all_reduce(tensor)23+ dist.all_reduce(tensor)
24- 24+ 
25- 25+ 
26-if __name__ == "__main__":26+if __name__ == "__main__":
27- main()27+ main()
Mtest/dynamo/test_stream.py+32-32
@@ -1,32 +1,32 @@
1-# Owner(s): ["module: dynamo"]1+# Owner(s): ["module: dynamo"]
2-import functools2+import functools
3-import unittest3+import unittest
4-import torch4+import torch
5-import torch._dynamo.test_case5+import torch._dynamo.test_case
6-import torch_npu6+import torch_npu
7- 7+ 
8-requires_npu = functools.partial(unittest.skipIf, not torch.npu.is_available(), "requires npu")8+requires_npu = functools.partial(unittest.skipIf, not torch.npu.is_available(), "requires npu")
9- 9+ 
10- 10+ 
11-class StreamintoDynamoTests(torch._dynamo.test_case.TestCase):11+class StreamintoDynamoTests(torch._dynamo.test_case.TestCase):
12- 12+ 
13- @requires_npu()13+ @requires_npu()
14- def test_stream(self):14+ def test_stream(self):
15- def model_1(x):15+ def model_1(x):
16- a = x * x16+ a = x * x
17- s = torch.npu.Stream()17+ s = torch.npu.Stream()
18- s.wait_stream(torch.npu.current_stream())18+ s.wait_stream(torch.npu.current_stream())
19- with torch.npu.stream(s):19+ with torch.npu.stream(s):
20- b = x + a20+ b = x + a
21- return b21+ return b
22- inp = torch.randn(2, 8).npu()22+ inp = torch.randn(2, 8).npu()
23- m = torch.compile(model_1, backend="aot_eager", fullgraph=True)23+ m = torch.compile(model_1, backend="aot_eager", fullgraph=True)
24- output = m(inp)24+ output = m(inp)
25- output1 = model_1(inp)25+ output1 = model_1(inp)
26- torch.allclose(output, output1)26+ torch.allclose(output, output1)
27- 27+ 
28- 28+ 
29-if __name__ == "__main__":29+if __name__ == "__main__":
30- from torch._dynamo.test_case import run_tests30+ from torch._dynamo.test_case import run_tests
31- 31+ 
32- run_tests()32+ run_tests()
Mtest/get_failed_ut_from_log.py+0-1
@@ -44,4 +44,3 @@ if __name__ == "__main__":
44 args = parser.parse_args()44 args = parser.parse_args()
45 failed_ut = get_error_or_fail_ut(args.file)45 failed_ut = get_error_or_fail_ut(args.file)
46 write_to_json(ut_list=failed_ut)46 write_to_json(ut_list=failed_ut)
47- 
Mtest/nn/test_uninitialized_parameter_cls_to_become.py+35-35
@@ -1,36 +1,36 @@
1-import torch1+import torch
2-from torch.testing._internal.common_utils import TestCase, run_tests2+from torch.testing._internal.common_utils import TestCase, run_tests
3-import torch_npu3+import torch_npu
4- 4+ 
5-# 关闭NPU JIT编译，减少CI耗时5+# 关闭NPU JIT编译，减少CI耗时
6-torch_npu.npu.set_compile_mode(jit_compile=False)6+torch_npu.npu.set_compile_mode(jit_compile=False)
7- 7+ 
8- 8+ 
9-# 修复：将自定义属性设为类属性（确保实例化后必存在）9+# 修复：将自定义属性设为类属性（确保实例化后必存在）
10-class CustomParameter(torch.nn.Parameter):10+class CustomParameter(torch.nn.Parameter):
11- custom_attr = "custom_param" # 类属性，所有实例共享，无需__init__赋值11+ custom_attr = "custom_param" # 类属性，所有实例共享，无需__init__赋值
12- 12+
13- def __init__(self, data=None, requires_grad=True):13+ def __init__(self, data=None, requires_grad=True):
14- super().__init__(data, requires_grad)14+ super().__init__(data, requires_grad)
15- 15+ 
16- 16+ 
17-class TestUninitializedParameterClsToBecome(TestCase):17+class TestUninitializedParameterClsToBecome(TestCase):
18- 18+
19- def test_core_functionality_npu(self):19+ def test_core_functionality_npu(self):
20- """极简验证NPU环境下cls_to_become+materialize核心功能"""20+ """极简验证NPU环境下cls_to_become+materialize核心功能"""
21- # 1. 创建NPU未初始化参数21+ # 1. 创建NPU未初始化参数
22- uninit_param = torch.nn.parameter.UninitializedParameter(device="npu")22+ uninit_param = torch.nn.parameter.UninitializedParameter(device="npu")
23- # 2. 绑定自定义类23+ # 2. 绑定自定义类
24- uninit_param.cls_to_become = CustomParameter24+ uninit_param.cls_to_become = CustomParameter
25- # 3. 实例化参数25+ # 3. 实例化参数
26- uninit_param.materialize(shape=(3, 3))26+ uninit_param.materialize(shape=(3, 3))
27- 27+ 
28- # 核心断言（全部通过，无报错）28+ # 核心断言（全部通过，无报错）
29- self.assertEqual(uninit_param.shape, torch.Size((3, 3)))29+ self.assertEqual(uninit_param.shape, torch.Size((3, 3)))
30- self.assertEqual(uninit_param.device.type, "npu")30+ self.assertEqual(uninit_param.device.type, "npu")
31- self.assertIsInstance(uninit_param, CustomParameter)31+ self.assertIsInstance(uninit_param, CustomParameter)
32- self.assertEqual(uninit_param.custom_attr, "custom_param") # 现在能正常访问32+ self.assertEqual(uninit_param.custom_attr, "custom_param") # 现在能正常访问
33- 33+ 
34- 34+ 
35-if __name__ == "__main__":35+if __name__ == "__main__":
36 run_tests()36 run_tests()
Mtest/npu/test_amp.py+526-526
@@ -1,526 +1,526 @@
1-import unittest1+import unittest
2-from itertools import chain2+from itertools import chain
3- 3+ 
4-import torch4+import torch
5- 5+ 
6-import torch_npu6+import torch_npu
7-from torch_npu.npu.amp import GradScaler, autocast7+from torch_npu.npu.amp import GradScaler, autocast
8-from torch_npu.testing.common_utils import SupportedDevices8+from torch_npu.testing.common_utils import SupportedDevices
9-from torch_npu.testing.testcase import TestCase, run_tests9+from torch_npu.testing.testcase import TestCase, run_tests
10- 10+ 
11- 11+ 
12-def make_device_overflow_1():12+def make_device_overflow_1():
13- float_tensor = torch.tensor([40000.0], dtype=torch.float16).npu()13+ float_tensor = torch.tensor([40000.0], dtype=torch.float16).npu()
14- float_tensor = float_tensor + float_tensor14+ float_tensor = float_tensor + float_tensor
15- 15+ 
16- 16+ 
17-def make_device_overflow_2(model):17+def make_device_overflow_2(model):
18- for param in model.parameters():18+ for param in model.parameters():
19- if param.grad is not None:19+ if param.grad is not None:
20- param.grad = torch.full_like(param.grad, float("inf"))20+ param.grad = torch.full_like(param.grad, float("inf"))
21- break21+ break
22- 22+ 
23- 23+ 
24-class TestAmp(TestCase):24+class TestAmp(TestCase):
25- 25+ 
26- def test_grad_scaling_scale(self):26+ def test_grad_scaling_scale(self):
27- scaler = GradScaler(init_scale=2.)27+ scaler = GradScaler(init_scale=2.)
28- t0 = torch.full((1,), 4.0, dtype=torch.float32, device="npu")28+ t0 = torch.full((1,), 4.0, dtype=torch.float32, device="npu")
29- t1 = torch.full((1,), 4.0, dtype=torch.float32, device="npu")29+ t1 = torch.full((1,), 4.0, dtype=torch.float32, device="npu")
30- # Create some nested iterables of tensors on different devices.30+ # Create some nested iterables of tensors on different devices.
31- outputs = (t1.clone(), (t0.clone(), t1.clone()), [t0.clone(), (t1.clone(), t0.clone())])31+ outputs = (t1.clone(), (t0.clone(), t1.clone()), [t0.clone(), (t1.clone(), t0.clone())])
32- outputs = scaler.scale(outputs)32+ outputs = scaler.scale(outputs)
33- self.assertTrue(outputs[0] == 8.0 and outputs[1][0] == 8.0 and outputs[1][1] == 8.0 and33+ self.assertTrue(outputs[0] == 8.0 and outputs[1][0] == 8.0 and outputs[1][1] == 8.0 and
34- outputs[2][0] == 8.0 and outputs[2][1][0] == 8.0 and outputs[2][1][1] == 8.0)34+ outputs[2][0] == 8.0 and outputs[2][1][0] == 8.0 and outputs[2][1][1] == 8.0)
35- self.assertTrue(scaler._scale.device == t1.device)35+ self.assertTrue(scaler._scale.device == t1.device)
36- 36+ 
37- def test_grad_scaling_state_dict(self):37+ def test_grad_scaling_state_dict(self):
38- for lazy_init_scale in True, False:38+ for lazy_init_scale in True, False:
39- s0 = GradScaler(init_scale=3., growth_factor=4., backoff_factor=.5, growth_interval=2)39+ s0 = GradScaler(init_scale=3., growth_factor=4., backoff_factor=.5, growth_interval=2)
40- s1 = GradScaler(init_scale=6., growth_factor=7., backoff_factor=.8, growth_interval=1)40+ s1 = GradScaler(init_scale=6., growth_factor=7., backoff_factor=.8, growth_interval=1)
41- 41+ 
42- # sets a random value for load_state_dict to overwrite42+ # sets a random value for load_state_dict to overwrite
43- s1._init_growth_tracker = 743+ s1._init_growth_tracker = 7
44- 44+ 
45- if lazy_init_scale:45+ if lazy_init_scale:
46- # Dummy scale() call to ensure the scale tensor is lazily initialized.46+ # Dummy scale() call to ensure the scale tensor is lazily initialized.
47- s1.scale(torch.full((1,), 4.0, dtype=torch.float32, device="npu"))47+ s1.scale(torch.full((1,), 4.0, dtype=torch.float32, device="npu"))
48- self.assertTrue(isinstance(s1._scale, torch.npu.FloatTensor))48+ self.assertTrue(isinstance(s1._scale, torch.npu.FloatTensor))
49- 49+ 
50- s1.load_state_dict(s0.state_dict())50+ s1.load_state_dict(s0.state_dict())
51- 51+ 
52- self.assertTrue(s1.get_scale() == 3.)52+ self.assertTrue(s1.get_scale() == 3.)
53- self.assertTrue(s1.get_growth_factor() == 4.)53+ self.assertTrue(s1.get_growth_factor() == 4.)
54- self.assertTrue(s1.get_backoff_factor() == .5)54+ self.assertTrue(s1.get_backoff_factor() == .5)
55- self.assertTrue(s1.get_growth_interval() == 2)55+ self.assertTrue(s1.get_growth_interval() == 2)
56- self.assertTrue(s1._init_growth_tracker == 0)56+ self.assertTrue(s1._init_growth_tracker == 0)
57- 57+ 
58- def _create_scaling_models_optimizers(self, device="npu"):58+ def _create_scaling_models_optimizers(self, device="npu"):
59- # Create a module+optimizer that will use scaling, and a control module+optimizer59+ # Create a module+optimizer that will use scaling, and a control module+optimizer
60- # that will not use scaling, against which the scaling-enabled module+optimizer can be compared.60+ # that will not use scaling, against which the scaling-enabled module+optimizer can be compared.
61- mod_control = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)61+ mod_control = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)
62- mod_scaling = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)62+ mod_scaling = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device)
63- for c, s in zip(mod_control.parameters(), mod_scaling.parameters()):63+ for c, s in zip(mod_control.parameters(), mod_scaling.parameters()):
64- s.data.copy_(c.data)64+ s.data.copy_(c.data)
65- 65+ 
66- opt_control = torch.optim.SGD(mod_control.parameters(), lr=1.0)66+ opt_control = torch.optim.SGD(mod_control.parameters(), lr=1.0)
67- opt_scaling = torch.optim.SGD(mod_scaling.parameters(), lr=1.0)67+ opt_scaling = torch.optim.SGD(mod_scaling.parameters(), lr=1.0)
68- 68+ 
69- ret = (mod_control, mod_scaling, opt_control, opt_scaling)69+ ret = (mod_control, mod_scaling, opt_control, opt_scaling)
70- return ret70+ return ret
71- 71+ 
72- def _create_scaling_case(self, device="npu", dtype=torch.float):72+ def _create_scaling_case(self, device="npu", dtype=torch.float):
73- data = [(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),73+ data = [(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
74- (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),74+ (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
75- (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),75+ (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)),
76- (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device))]76+ (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device))]
77- 77+ 
78- loss_fn = torch.nn.MSELoss().npu()78+ loss_fn = torch.nn.MSELoss().npu()
79- 79+ 
80- skip_iter = 280+ skip_iter = 2
81- 81+ 
82- return self._create_scaling_models_optimizers(device=device) + (data, loss_fn, skip_iter)82+ return self._create_scaling_models_optimizers(device=device) + (data, loss_fn, skip_iter)
83- 83+ 
84- # _run_scaling_case generalizes some single-optimizer test logic to avoid too much copy-pasting below.84+ # _run_scaling_case generalizes some single-optimizer test logic to avoid too much copy-pasting below.
85- def _run_scaling_case(self, run, unskipped, skipped, atol=1e-7):85+ def _run_scaling_case(self, run, unskipped, skipped, atol=1e-7):
86- # Ensure scaling can be disabled without changing user control flow.86+ # Ensure scaling can be disabled without changing user control flow.
87- for enabled in True, False:87+ for enabled in True, False:
88- mod_control, mod_scaling, opt_control, opt_scaling, data, loss_fn, skip_iter = self._create_scaling_case()88+ mod_control, mod_scaling, opt_control, opt_scaling, data, loss_fn, skip_iter = self._create_scaling_case()
89- 89+ 
90- # For functionality, test with a modest initial scale, and an unrealistically-large growth factor90+ # For functionality, test with a modest initial scale, and an unrealistically-large growth factor
91- # so any potential errors with the growth factor handling will be magnified.91+ # so any potential errors with the growth factor handling will be magnified.
92- scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)92+ scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)
93- 93+ 
94- _ = run(data, mod_control, opt_control, scaler, loss_fn, skip_iter, False)94+ _ = run(data, mod_control, opt_control, scaler, loss_fn, skip_iter, False)
95- ret = run(data, mod_scaling, opt_scaling, scaler, loss_fn, skip_iter, True)95+ ret = run(data, mod_scaling, opt_scaling, scaler, loss_fn, skip_iter, True)
96- 96+ 
97- # Allows run() to optionally return a different scaler instance.97+ # Allows run() to optionally return a different scaler instance.
98- scaler = ret if ret else scaler98+ scaler = ret if ret else scaler
99- 99+ 
100- # If scaling was enabled, the scale factor should have been multiplied by the growth factor100+ # If scaling was enabled, the scale factor should have been multiplied by the growth factor
101- # len(data) - skipped times and the backoff factor "skipped" times.101+ # len(data) - skipped times and the backoff factor "skipped" times.
102- if enabled:102+ if enabled:
103- net_growth = scaler.get_growth_factor()**unskipped if unskipped > 0 else 1.0103+ net_growth = scaler.get_growth_factor()**unskipped if unskipped > 0 else 1.0
104- net_backoff = scaler.get_backoff_factor()**skipped if skipped > 0 else 1.0104+ net_backoff = scaler.get_backoff_factor()**skipped if skipped > 0 else 1.0
105- self.assertTrue(scaler.get_scale() == (128. * net_growth * net_backoff))105+ self.assertTrue(scaler.get_scale() == (128. * net_growth * net_backoff))
106- else:106+ else:
107- self.assertTrue(scaler.get_scale() == 1.0)107+ self.assertTrue(scaler.get_scale() == 1.0)
108- 108+ 
109- for c, s in zip(mod_control.parameters(), mod_scaling.parameters()):109+ for c, s in zip(mod_control.parameters(), mod_scaling.parameters()):
110- c = c.cpu().to(torch.float).detach().numpy()110+ c = c.cpu().to(torch.float).detach().numpy()
111- s = s.cpu().to(torch.float).detach().numpy()111+ s = s.cpu().to(torch.float).detach().numpy()
112- self.assertRtolEqual(c, s, atol)112+ self.assertRtolEqual(c, s, atol)
113- 113+ 
114- # Compares no scaling + no autocasting against scaling + autocasting.114+ # Compares no scaling + no autocasting against scaling + autocasting.
115- @SupportedDevices(['Ascend910A', 'Ascend910P'])115+ @SupportedDevices(['Ascend910A', 'Ascend910P'])
116- def test_grad_scaling_autocast_1(self):116+ def test_grad_scaling_autocast_1(self):
117- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):117+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
118- for i, (input_data, target) in enumerate(data):118+ for i, (input_data, target) in enumerate(data):
119- optimizer.zero_grad()119+ optimizer.zero_grad()
120- with torch.autocast('npu', enabled=try_scaling_api):120+ with torch.autocast('npu', enabled=try_scaling_api):
121- output = model(input_data)121+ output = model(input_data)
122- loss = loss_fn(output, target)122+ loss = loss_fn(output, target)
123- if try_scaling_api:123+ if try_scaling_api:
124- scaler.scale(loss).backward()124+ scaler.scale(loss).backward()
125- if i == skip_iter and scaler.is_enabled():125+ if i == skip_iter and scaler.is_enabled():
126- make_device_overflow_1()126+ make_device_overflow_1()
127- scaler.step(optimizer)127+ scaler.step(optimizer)
128- scaler.update()128+ scaler.update()
129- else:129+ else:
130- loss.backward()130+ loss.backward()
131- if (not scaler.is_enabled()) or (i != skip_iter):131+ if (not scaler.is_enabled()) or (i != skip_iter):
132- optimizer.step()132+ optimizer.step()
133- return scaler133+ return scaler
134- 134+ 
135- # sets atol=1e-3 because we're comparing pure fp32 arithmetic vs a mixture of fp16 and fp32135+ # sets atol=1e-3 because we're comparing pure fp32 arithmetic vs a mixture of fp16 and fp32
136- self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-3)136+ self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-3)
137- 137+ 
138- @SupportedDevices(['Ascend910B'])138+ @SupportedDevices(['Ascend910B'])
139- def test_grad_scaling_autocast_2(self):139+ def test_grad_scaling_autocast_2(self):
140- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):140+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
141- for i, (input_data, target) in enumerate(data):141+ for i, (input_data, target) in enumerate(data):
142- optimizer.zero_grad()142+ optimizer.zero_grad()
143- with torch.autocast('npu', enabled=try_scaling_api):143+ with torch.autocast('npu', enabled=try_scaling_api):
144- output = model(input_data)144+ output = model(input_data)
145- loss = loss_fn(output, target)145+ loss = loss_fn(output, target)
146- if try_scaling_api:146+ if try_scaling_api:
147- scaler.scale(loss).backward()147+ scaler.scale(loss).backward()
148- if i == skip_iter and scaler.is_enabled():148+ if i == skip_iter and scaler.is_enabled():
149- make_device_overflow_2(model)149+ make_device_overflow_2(model)
150- scaler.step(optimizer)150+ scaler.step(optimizer)
151- scaler.update()151+ scaler.update()
152- else:152+ else:
153- loss.backward()153+ loss.backward()
154- if (not scaler.is_enabled()) or (i != skip_iter):154+ if (not scaler.is_enabled()) or (i != skip_iter):
155- optimizer.step()155+ optimizer.step()
156- return scaler156+ return scaler
157- 157+ 
158- # sets atol=1e-3 because we're comparing pure fp32 arithmetic vs a mixture of fp16 and fp32158+ # sets atol=1e-3 because we're comparing pure fp32 arithmetic vs a mixture of fp16 and fp32
159- self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-3)159+ self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-3)
160- 160+ 
161- @SupportedDevices(['Ascend910A', 'Ascend910P'])161+ @SupportedDevices(['Ascend910A', 'Ascend910P'])
162- def test_grad_scaling_clipping_1(self):162+ def test_grad_scaling_clipping_1(self):
163- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):163+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
164- max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads164+ max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads
165- for i, (input_data, target) in enumerate(data):165+ for i, (input_data, target) in enumerate(data):
166- optimizer.zero_grad()166+ optimizer.zero_grad()
167- output = model(input_data)167+ output = model(input_data)
168- loss = loss_fn(output, target)168+ loss = loss_fn(output, target)
169- if try_scaling_api:169+ if try_scaling_api:
170- scaler.scale(loss).backward()170+ scaler.scale(loss).backward()
171- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm * scaler.get_scale())171+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm * scaler.get_scale())
172- if i == skip_iter and scaler.is_enabled():172+ if i == skip_iter and scaler.is_enabled():
173- make_device_overflow_1()173+ make_device_overflow_1()
174- scaler.step(optimizer)174+ scaler.step(optimizer)
175- scaler.update()175+ scaler.update()
176- else:176+ else:
177- loss.backward()177+ loss.backward()
178- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)178+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
179- if (not scaler.is_enabled()) or (i != skip_iter):179+ if (not scaler.is_enabled()) or (i != skip_iter):
180- optimizer.step()180+ optimizer.step()
181- 181+ 
182- self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-6)182+ self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-6)
183- 183+ 
184- @SupportedDevices(['Ascend910B'])184+ @SupportedDevices(['Ascend910B'])
185- def test_grad_scaling_clipping_2(self):185+ def test_grad_scaling_clipping_2(self):
186- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):186+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
187- max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads187+ max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads
188- for i, (input_data, target) in enumerate(data):188+ for i, (input_data, target) in enumerate(data):
189- optimizer.zero_grad()189+ optimizer.zero_grad()
190- output = model(input_data)190+ output = model(input_data)
191- loss = loss_fn(output, target)191+ loss = loss_fn(output, target)
192- if try_scaling_api:192+ if try_scaling_api:
193- scaler.scale(loss).backward()193+ scaler.scale(loss).backward()
194- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm * scaler.get_scale())194+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm * scaler.get_scale())
195- if i == skip_iter and scaler.is_enabled():195+ if i == skip_iter and scaler.is_enabled():
196- make_device_overflow_2(model)196+ make_device_overflow_2(model)
197- scaler.step(optimizer)197+ scaler.step(optimizer)
198- scaler.update()198+ scaler.update()
199- else:199+ else:
200- loss.backward()200+ loss.backward()
201- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)201+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
202- if (not scaler.is_enabled()) or (i != skip_iter):202+ if (not scaler.is_enabled()) or (i != skip_iter):
203- optimizer.step()203+ optimizer.step()
204- 204+ 
205- self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-6)205+ self._run_scaling_case(run, unskipped=3, skipped=1, atol=1e-6)
206- 206+ 
207- @SupportedDevices(['Ascend910A', 'Ascend910P'])207+ @SupportedDevices(['Ascend910A', 'Ascend910P'])
208- def test_grad_scaling_clipping_separate_unscale_1(self):208+ def test_grad_scaling_clipping_separate_unscale_1(self):
209- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):209+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
210- max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads210+ max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads
211- for i, (input_data, target) in enumerate(data):211+ for i, (input_data, target) in enumerate(data):
212- optimizer.zero_grad()212+ optimizer.zero_grad()
213- output = model(input_data)213+ output = model(input_data)
214- loss = loss_fn(output, target)214+ loss = loss_fn(output, target)
215- if try_scaling_api:215+ if try_scaling_api:
216- scaler.scale(loss).backward()216+ scaler.scale(loss).backward()
217- if i == skip_iter and scaler.is_enabled():217+ if i == skip_iter and scaler.is_enabled():
218- make_device_overflow_1()218+ make_device_overflow_1()
219- scaler.unscale_(optimizer)219+ scaler.unscale_(optimizer)
220- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)220+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
221- scaler.step(optimizer)221+ scaler.step(optimizer)
222- scaler.update()222+ scaler.update()
223- else:223+ else:
224- loss.backward()224+ loss.backward()
225- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)225+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
226- if (not scaler.is_enabled()) or (i != skip_iter):226+ if (not scaler.is_enabled()) or (i != skip_iter):
227- optimizer.step()227+ optimizer.step()
228- 228+ 
229- self._run_scaling_case(run, unskipped=3, skipped=1)229+ self._run_scaling_case(run, unskipped=3, skipped=1)
230- 230+ 
231- @SupportedDevices(['Ascend910B'])231+ @SupportedDevices(['Ascend910B'])
232- def test_grad_scaling_clipping_separate_unscale_2(self):232+ def test_grad_scaling_clipping_separate_unscale_2(self):
233- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):233+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
234- max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads234+ max_norm = 0.2 # A reasonable value that actually has an effect, based on printouts of grads
235- for i, (input_data, target) in enumerate(data):235+ for i, (input_data, target) in enumerate(data):
236- optimizer.zero_grad()236+ optimizer.zero_grad()
237- output = model(input_data)237+ output = model(input_data)
238- loss = loss_fn(output, target)238+ loss = loss_fn(output, target)
239- if try_scaling_api:239+ if try_scaling_api:
240- scaler.scale(loss).backward()240+ scaler.scale(loss).backward()
241- if i == skip_iter and scaler.is_enabled():241+ if i == skip_iter and scaler.is_enabled():
242- make_device_overflow_2(model)242+ make_device_overflow_2(model)
243- scaler.unscale_(optimizer)243+ scaler.unscale_(optimizer)
244- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)244+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
245- scaler.step(optimizer)245+ scaler.step(optimizer)
246- scaler.update()246+ scaler.update()
247- else:247+ else:
248- loss.backward()248+ loss.backward()
249- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)249+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
250- if (not scaler.is_enabled()) or (i != skip_iter):250+ if (not scaler.is_enabled()) or (i != skip_iter):
251- optimizer.step()251+ optimizer.step()
252- 252+ 
253- self._run_scaling_case(run, unskipped=3, skipped=1)253+ self._run_scaling_case(run, unskipped=3, skipped=1)
254- 254+ 
255- @SupportedDevices(['Ascend910A', 'Ascend910P'])255+ @SupportedDevices(['Ascend910A', 'Ascend910P'])
256- def test_grad_scaling_penalty_1(self):256+ def test_grad_scaling_penalty_1(self):
257- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):257+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
258- for i, (input_data, target) in enumerate(data):258+ for i, (input_data, target) in enumerate(data):
259- optimizer.zero_grad()259+ optimizer.zero_grad()
260- output = model(input_data)260+ output = model(input_data)
261- loss = loss_fn(output, target)261+ loss = loss_fn(output, target)
262- 262+ 
263- if try_scaling_api:263+ if try_scaling_api:
264- grad_params = torch.autograd.grad(scaler.scale(loss),264+ grad_params = torch.autograd.grad(scaler.scale(loss),
265- model.parameters(), create_graph=True)265+ model.parameters(), create_graph=True)
266- inv_scale = 1. / scaler.get_scale()266+ inv_scale = 1. / scaler.get_scale()
267- grad_params = [p * inv_scale for p in grad_params]267+ grad_params = [p * inv_scale for p in grad_params]
268- else:268+ else:
269- grad_params = torch.autograd.grad(loss, model.parameters(), create_graph=True)269+ grad_params = torch.autograd.grad(loss, model.parameters(), create_graph=True)
270- 270+ 
271- grad_norm = 0271+ grad_norm = 0
272- for grad in grad_params:272+ for grad in grad_params:
273- grad_norm += grad.pow(2).sum()273+ grad_norm += grad.pow(2).sum()
274- grad_norm = grad_norm.sqrt()274+ grad_norm = grad_norm.sqrt()
275- loss = loss + grad_norm275+ loss = loss + grad_norm
276- 276+ 
277- if try_scaling_api:277+ if try_scaling_api:
278- scaler.scale(loss).backward()278+ scaler.scale(loss).backward()
279- if i == skip_iter and scaler.is_enabled():279+ if i == skip_iter and scaler.is_enabled():
280- make_device_overflow_1()280+ make_device_overflow_1()
281- scaler.step(optimizer)281+ scaler.step(optimizer)
282- scaler.update()282+ scaler.update()
283- else:283+ else:
284- loss.backward()284+ loss.backward()
285- if (not scaler.is_enabled()) or (i != skip_iter):285+ if (not scaler.is_enabled()) or (i != skip_iter):
286- optimizer.step()286+ optimizer.step()
287- 287+ 
288- self._run_scaling_case(run, unskipped=3, skipped=1)288+ self._run_scaling_case(run, unskipped=3, skipped=1)
289- 289+ 
290- @SupportedDevices(['Ascend910B'])290+ @SupportedDevices(['Ascend910B'])
291- def test_grad_scaling_penalty_2(self):291+ def test_grad_scaling_penalty_2(self):
292- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):292+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
293- for i, (input_data, target) in enumerate(data):293+ for i, (input_data, target) in enumerate(data):
294- optimizer.zero_grad()294+ optimizer.zero_grad()
295- output = model(input_data)295+ output = model(input_data)
296- loss = loss_fn(output, target)296+ loss = loss_fn(output, target)
297- 297+ 
298- if try_scaling_api:298+ if try_scaling_api:
299- grad_params = torch.autograd.grad(scaler.scale(loss),299+ grad_params = torch.autograd.grad(scaler.scale(loss),
300- model.parameters(), create_graph=True)300+ model.parameters(), create_graph=True)
301- inv_scale = 1. / scaler.get_scale()301+ inv_scale = 1. / scaler.get_scale()
302- grad_params = [p * inv_scale for p in grad_params]302+ grad_params = [p * inv_scale for p in grad_params]
303- else:303+ else:
304- grad_params = torch.autograd.grad(loss, model.parameters(), create_graph=True)304+ grad_params = torch.autograd.grad(loss, model.parameters(), create_graph=True)
305- 305+ 
306- grad_norm = 0306+ grad_norm = 0
307- for grad in grad_params:307+ for grad in grad_params:
308- grad_norm += grad.pow(2).sum()308+ grad_norm += grad.pow(2).sum()
309- grad_norm = grad_norm.sqrt()309+ grad_norm = grad_norm.sqrt()
310- loss = loss + grad_norm310+ loss = loss + grad_norm
311- 311+ 
312- if try_scaling_api:312+ if try_scaling_api:
313- scaler.scale(loss).backward()313+ scaler.scale(loss).backward()
314- if i == skip_iter and scaler.is_enabled():314+ if i == skip_iter and scaler.is_enabled():
315- make_device_overflow_2(model)315+ make_device_overflow_2(model)
316- scaler.step(optimizer)316+ scaler.step(optimizer)
317- scaler.update()317+ scaler.update()
318- else:318+ else:
319- loss.backward()319+ loss.backward()
320- if (not scaler.is_enabled()) or (i != skip_iter):320+ if (not scaler.is_enabled()) or (i != skip_iter):
321- optimizer.step()321+ optimizer.step()
322- 322+ 
323- self._run_scaling_case(run, unskipped=3, skipped=1)323+ self._run_scaling_case(run, unskipped=3, skipped=1)
324- 324+ 
325- def test_grad_scaling_accumulation(self):325+ def test_grad_scaling_accumulation(self):
326- def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):326+ def run(data, model, optimizer, scaler, loss_fn, skip_iter, try_scaling_api):
327- iters_to_accumulate = 2327+ iters_to_accumulate = 2
328- for i, (input_data, target) in enumerate(data):328+ for i, (input_data, target) in enumerate(data):
329- output = model(input_data)329+ output = model(input_data)
330- loss = loss_fn(output, target)330+ loss = loss_fn(output, target)
331- loss = loss / iters_to_accumulate331+ loss = loss / iters_to_accumulate
332- if try_scaling_api:332+ if try_scaling_api:
333- scaler.scale(loss).backward()333+ scaler.scale(loss).backward()
334- else:334+ else:
335- loss.backward()335+ loss.backward()
336- if (i + 1) % iters_to_accumulate == 0:336+ if (i + 1) % iters_to_accumulate == 0:
337- if try_scaling_api:337+ if try_scaling_api:
338- scaler.step(optimizer)338+ scaler.step(optimizer)
339- scaler.update()339+ scaler.update()
340- optimizer.zero_grad()340+ optimizer.zero_grad()
341- else:341+ else:
342- optimizer.step()342+ optimizer.step()
343- optimizer.zero_grad()343+ optimizer.zero_grad()
344- 344+ 
345- self._run_scaling_case(run, unskipped=2, skipped=0)345+ self._run_scaling_case(run, unskipped=2, skipped=0)
346- 346+ 
347- @SupportedDevices(['Ascend910A', 'Ascend910P'])347+ @SupportedDevices(['Ascend910A', 'Ascend910P'])
348- def test_grad_scaling_multiple_1(self):348+ def test_grad_scaling_multiple_1(self):
349- # Tests gradient scaling with 2 models and 2 optimizers that both receive gradients from 2 losses.349+ # Tests gradient scaling with 2 models and 2 optimizers that both receive gradients from 2 losses.
350- # Some of the logic here cannot reuse the generic helper functions created for the 1-optimizer cases.350+ # Some of the logic here cannot reuse the generic helper functions created for the 1-optimizer cases.
351- for enabled in True, False:351+ for enabled in True, False:
352- mod_control0, mod_scaling0, opt_control0, opt_scaling0, data, loss_fn, skip_iter = \352+ mod_control0, mod_scaling0, opt_control0, opt_scaling0, data, loss_fn, skip_iter = \
353- self._create_scaling_case()353+ self._create_scaling_case()
354- mod_control1, mod_scaling1, opt_control1, opt_scaling1 = \354+ mod_control1, mod_scaling1, opt_control1, opt_scaling1 = \
355- self._create_scaling_models_optimizers()355+ self._create_scaling_models_optimizers()
356- 356+ 
357- scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)357+ scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)
358- 358+ 
359- def run(model0, model1, optimizer0, optimizer1, try_scaling_api):359+ def run(model0, model1, optimizer0, optimizer1, try_scaling_api):
360- for i, (input_data, target) in enumerate(data):360+ for i, (input_data, target) in enumerate(data):
361- optimizer0.zero_grad()361+ optimizer0.zero_grad()
362- optimizer1.zero_grad()362+ optimizer1.zero_grad()
363- output0 = model0(input_data)363+ output0 = model0(input_data)
364- output1 = model1(input_data)364+ output1 = model1(input_data)
365- loss0 = loss_fn(0.3 * output0 + 0.7 * output1, target)365+ loss0 = loss_fn(0.3 * output0 + 0.7 * output1, target)
366- loss1 = loss_fn(0.6 * output0 - 0.4 * output1, target)366+ loss1 = loss_fn(0.6 * output0 - 0.4 * output1, target)
367- 367+ 
368- if try_scaling_api:368+ if try_scaling_api:
369- scaler.scale(loss0).backward(retain_graph=True)369+ scaler.scale(loss0).backward(retain_graph=True)
370- scaler.scale(loss1).backward()370+ scaler.scale(loss1).backward()
371- if i == skip_iter and scaler.is_enabled():371+ if i == skip_iter and scaler.is_enabled():
372- make_device_overflow_1()372+ make_device_overflow_1()
373- 373+ 
374- # As an additional stress test, separately unscale for one of the optimizers.374+ # As an additional stress test, separately unscale for one of the optimizers.
375- scaler.unscale_(optimizer0)375+ scaler.unscale_(optimizer0)
376- 376+ 
377- scaler.step(optimizer0)377+ scaler.step(optimizer0)
378- scaler.step(optimizer1)378+ scaler.step(optimizer1)
379- scaler.update()379+ scaler.update()
380- else:380+ else:
381- loss0.backward(retain_graph=True)381+ loss0.backward(retain_graph=True)
382- loss1.backward()382+ loss1.backward()
383- if (not scaler.is_enabled()) or (i != skip_iter):383+ if (not scaler.is_enabled()) or (i != skip_iter):
384- optimizer0.step()384+ optimizer0.step()
385- optimizer1.step()385+ optimizer1.step()
386- 386+ 
387- run(mod_control0, mod_control1, opt_control0, opt_control1, False)387+ run(mod_control0, mod_control1, opt_control0, opt_control1, False)
388- run(mod_scaling0, mod_scaling1, opt_scaling0, opt_scaling1, True)388+ run(mod_scaling0, mod_scaling1, opt_scaling0, opt_scaling1, True)
389- 389+ 
390- # The loss scale should have been multiplied by the growth factor 3 times and the backoff factor once.390+ # The loss scale should have been multiplied by the growth factor 3 times and the backoff factor once.
391- self.assertTrue(scaler.get_scale() == (128. * scaler.get_growth_factor()**3 *391+ self.assertTrue(scaler.get_scale() == (128. * scaler.get_growth_factor()**3 *
392- scaler.get_backoff_factor()**1) if enabled else 1.0)392+ scaler.get_backoff_factor()**1) if enabled else 1.0)
393- 393+ 
394- for c, s in zip(chain(mod_control0.parameters(), mod_control1.parameters()),394+ for c, s in zip(chain(mod_control0.parameters(), mod_control1.parameters()),
395- chain(mod_scaling0.parameters(), mod_scaling1.parameters())):395+ chain(mod_scaling0.parameters(), mod_scaling1.parameters())):
396- c = c.cpu().to(torch.float).detach().numpy()396+ c = c.cpu().to(torch.float).detach().numpy()
397- s = s.cpu().to(torch.float).detach().numpy()397+ s = s.cpu().to(torch.float).detach().numpy()
398- self.assertRtolEqual(c, s, 1e-7)398+ self.assertRtolEqual(c, s, 1e-7)
399- 399+ 
400- @SupportedDevices(['Ascend910B'])400+ @SupportedDevices(['Ascend910B'])
401- def test_grad_scaling_multiple_2(self):401+ def test_grad_scaling_multiple_2(self):
402- # Tests gradient scaling with 2 models and 2 optimizers that both receive gradients from 2 losses.402+ # Tests gradient scaling with 2 models and 2 optimizers that both receive gradients from 2 losses.
403- # Some of the logic here cannot reuse the generic helper functions created for the 1-optimizer cases.403+ # Some of the logic here cannot reuse the generic helper functions created for the 1-optimizer cases.
404- for enabled in True, False:404+ for enabled in True, False:
405- mod_control0, mod_scaling0, opt_control0, opt_scaling0, data, loss_fn, skip_iter = \405+ mod_control0, mod_scaling0, opt_control0, opt_scaling0, data, loss_fn, skip_iter = \
406- self._create_scaling_case()406+ self._create_scaling_case()
407- mod_control1, mod_scaling1, opt_control1, opt_scaling1 = \407+ mod_control1, mod_scaling1, opt_control1, opt_scaling1 = \
408- self._create_scaling_models_optimizers()408+ self._create_scaling_models_optimizers()
409- 409+ 
410- scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)410+ scaler = GradScaler(init_scale=128., growth_factor=2.0, enabled=enabled, growth_interval=1)
411- 411+ 
412- def run(model0, model1, optimizer0, optimizer1, try_scaling_api):412+ def run(model0, model1, optimizer0, optimizer1, try_scaling_api):
413- for i, (input_data, target) in enumerate(data):413+ for i, (input_data, target) in enumerate(data):
414- optimizer0.zero_grad()414+ optimizer0.zero_grad()
415- optimizer1.zero_grad()415+ optimizer1.zero_grad()
416- output0 = model0(input_data)416+ output0 = model0(input_data)
417- output1 = model1(input_data)417+ output1 = model1(input_data)
418- loss0 = loss_fn(0.3 * output0 + 0.7 * output1, target)418+ loss0 = loss_fn(0.3 * output0 + 0.7 * output1, target)
419- loss1 = loss_fn(0.6 * output0 - 0.4 * output1, target)419+ loss1 = loss_fn(0.6 * output0 - 0.4 * output1, target)
420- 420+ 
421- if try_scaling_api:421+ if try_scaling_api:
422- scaler.scale(loss0).backward(retain_graph=True)422+ scaler.scale(loss0).backward(retain_graph=True)
423- scaler.scale(loss1).backward()423+ scaler.scale(loss1).backward()
424- if i == skip_iter and scaler.is_enabled():424+ if i == skip_iter and scaler.is_enabled():
425- make_device_overflow_2(model0)425+ make_device_overflow_2(model0)
426- 426+ 
427- # As an additional stress test, separately unscale for one of the optimizers.427+ # As an additional stress test, separately unscale for one of the optimizers.
428- scaler.unscale_(optimizer0)428+ scaler.unscale_(optimizer0)
429- 429+ 
430- scaler.step(optimizer0)430+ scaler.step(optimizer0)
431- scaler.step(optimizer1)431+ scaler.step(optimizer1)
432- scaler.update()432+ scaler.update()
433- else:433+ else:
434- loss0.backward(retain_graph=True)434+ loss0.backward(retain_graph=True)
435- loss1.backward()435+ loss1.backward()
436- if (not scaler.is_enabled()) or (i != skip_iter):436+ if (not scaler.is_enabled()) or (i != skip_iter):
437- optimizer0.step()437+ optimizer0.step()
438- optimizer1.step()438+ optimizer1.step()
439- 439+ 
440- run(mod_control0, mod_control1, opt_control0, opt_control1, False)440+ run(mod_control0, mod_control1, opt_control0, opt_control1, False)
441- run(mod_scaling0, mod_scaling1, opt_scaling0, opt_scaling1, True)441+ run(mod_scaling0, mod_scaling1, opt_scaling0, opt_scaling1, True)
442- 442+ 
443- # The loss scale should have been multiplied by the growth factor 3 times and the backoff factor once.443+ # The loss scale should have been multiplied by the growth factor 3 times and the backoff factor once.
444- self.assertTrue(scaler.get_scale() == (128. * scaler.get_growth_factor() ** 3 *444+ self.assertTrue(scaler.get_scale() == (128. * scaler.get_growth_factor() ** 3 *
445- scaler.get_backoff_factor() ** 1) if enabled else 1.0)445+ scaler.get_backoff_factor() ** 1) if enabled else 1.0)
446- 446+ 
447- for c, s in zip(chain(mod_control0.parameters(), mod_control1.parameters()),447+ for c, s in zip(chain(mod_control0.parameters(), mod_control1.parameters()),
448- chain(mod_scaling0.parameters(), mod_scaling1.parameters())):448+ chain(mod_scaling0.parameters(), mod_scaling1.parameters())):
449- c = c.cpu().to(torch.float).detach().numpy()449+ c = c.cpu().to(torch.float).detach().numpy()
450- s = s.cpu().to(torch.float).detach().numpy()450+ s = s.cpu().to(torch.float).detach().numpy()
451- self.assertRtolEqual(c, s, 1e-7)451+ self.assertRtolEqual(c, s, 1e-7)
452- 452+ 
453- def test_autocast_custom_enabled(self):453+ def test_autocast_custom_enabled(self):
454- class MyMM(torch.autograd.Function):454+ class MyMM(torch.autograd.Function):
455- @staticmethod455+ @staticmethod
456- @torch.npu.amp.custom_fwd456+ @torch.npu.amp.custom_fwd
457- def forward(ctx, a, b):457+ def forward(ctx, a, b):
458- self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))458+ self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))
459- self.assertTrue(a.dtype is torch.float32)459+ self.assertTrue(a.dtype is torch.float32)
460- self.assertTrue(b.dtype is torch.float32)460+ self.assertTrue(b.dtype is torch.float32)
461- self.assertTrue(torch.npu.is_autocast_enabled())461+ self.assertTrue(torch.npu.is_autocast_enabled())
462- ctx.save_for_backward(a, b)462+ ctx.save_for_backward(a, b)
463- return a.mm(b)463+ return a.mm(b)
464- 464+ 
465- @staticmethod465+ @staticmethod
466- @torch.npu.amp.custom_bwd466+ @torch.npu.amp.custom_bwd
467- def backward(ctx, grad):467+ def backward(ctx, grad):
468- self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))468+ self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))
469- self.assertTrue(torch.npu.is_autocast_enabled())469+ self.assertTrue(torch.npu.is_autocast_enabled())
470- a, b = ctx.saved_tensors470+ a, b = ctx.saved_tensors
471- return grad.mm(b.t()), a.t().mm(grad)471+ return grad.mm(b.t()), a.t().mm(grad)
472- 472+ 
473- mymm = MyMM.apply473+ mymm = MyMM.apply
474- 474+ 
475- x = torch.randn((8, 8), device="npu", dtype=torch.float32, requires_grad=True)475+ x = torch.randn((8, 8), device="npu", dtype=torch.float32, requires_grad=True)
476- y = torch.randn((8, 8), device="npu", dtype=torch.float32, requires_grad=True)476+ y = torch.randn((8, 8), device="npu", dtype=torch.float32, requires_grad=True)
477- 477+ 
478- with torch.npu.amp.autocast():478+ with torch.npu.amp.autocast():
479- output = mymm(x, y)479+ output = mymm(x, y)
480- self.assertTrue(output.dtype is torch.float16)480+ self.assertTrue(output.dtype is torch.float16)
481- loss = output.sum()481+ loss = output.sum()
482- loss.backward()482+ loss.backward()
483- 483+ 
484- def test_autocast_custom_cast_inputs(self):484+ def test_autocast_custom_cast_inputs(self):
485- class MyMM(torch.autograd.Function):485+ class MyMM(torch.autograd.Function):
486- @staticmethod486+ @staticmethod
487- @torch.npu.amp.custom_fwd(cast_inputs=torch.float32)487+ @torch.npu.amp.custom_fwd(cast_inputs=torch.float32)
488- def forward(ctx, a, container, expect_type):488+ def forward(ctx, a, container, expect_type):
489- self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))489+ self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))
490- b = container[1][0]490+ b = container[1][0]
491- self.assertTrue(a.dtype is expect_type)491+ self.assertTrue(a.dtype is expect_type)
492- self.assertTrue(b.dtype is expect_type)492+ self.assertTrue(b.dtype is expect_type)
493- self.assertFalse(torch.npu.is_autocast_enabled())493+ self.assertFalse(torch.npu.is_autocast_enabled())
494- ctx.save_for_backward(a, b)494+ ctx.save_for_backward(a, b)
495- return a.mm(b)495+ return a.mm(b)
496- 496+ 
497- @staticmethod497+ @staticmethod
498- @torch.npu.amp.custom_bwd498+ @torch.npu.amp.custom_bwd
499- def backward(ctx, grad):499+ def backward(ctx, grad):
500- self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))500+ self.assertTrue(ctx._dtype is torch.get_autocast_dtype("npu"))
501- a, b = ctx.saved_tensors501+ a, b = ctx.saved_tensors
502- return grad.mm(b.t()), None, None502+ return grad.mm(b.t()), None, None
503- 503+ 
504- mymm = MyMM.apply504+ mymm = MyMM.apply
505- 505+ 
506- x = torch.randn((8, 8), device="npu", dtype=torch.float16, requires_grad=True)506+ x = torch.randn((8, 8), device="npu", dtype=torch.float16, requires_grad=True)
507- # Puts one input tensor in a nested container. y's contained Tensor won't receive a gradient,507+ # Puts one input tensor in a nested container. y's contained Tensor won't receive a gradient,
508- # because torch.autograd.Function can't hand gradients back to non-Tensor forward arguments.508+ # because torch.autograd.Function can't hand gradients back to non-Tensor forward arguments.
509- # Sets requires_grad=False explicitly so we don't lie about expecting a gradient.509+ # Sets requires_grad=False explicitly so we don't lie about expecting a gradient.
510- y = (0, {0: torch.randn((8, 8), device="npu", dtype=torch.float16, requires_grad=False)})510+ y = (0, {0: torch.randn((8, 8), device="npu", dtype=torch.float16, requires_grad=False)})
511- 511+ 
512- with torch.autocast('npu', ):512+ with torch.autocast('npu', ):
513- output = mymm(x, y, torch.float32)513+ output = mymm(x, y, torch.float32)
514- self.assertTrue(output.dtype is torch.float32)514+ self.assertTrue(output.dtype is torch.float32)
515- loss = output.sum()515+ loss = output.sum()
516- loss.backward()516+ loss.backward()
517- 517+ 
518- # Tests if custom_fwd becomes a no-op when mymm runs outside an autocast-enabled region.518+ # Tests if custom_fwd becomes a no-op when mymm runs outside an autocast-enabled region.
519- output = mymm(x, y, torch.float16)519+ output = mymm(x, y, torch.float16)
520- self.assertTrue(output.dtype is torch.float16)520+ self.assertTrue(output.dtype is torch.float16)
521- loss = output.sum()521+ loss = output.sum()
522- loss.backward()522+ loss.backward()
523- 523+ 
524- 524+ 
525-if __name__ == "__main__":525+if __name__ == "__main__":
526- run_tests()526+ run_tests()
Mtest/npu/test_copy.py+137-137
@@ -1,138 +1,138 @@
1-import unittest1+import unittest
2-import torch2+import torch
3-import numpy as np3+import numpy as np
4- 4+ 
5-import torch_npu5+import torch_npu
6-from torch_npu.testing.testcase import TestCase, run_tests6+from torch_npu.testing.testcase import TestCase, run_tests
7-from torch_npu.testing.common_utils import create_common_tensor7+from torch_npu.testing.common_utils import create_common_tensor
8- 8+ 
9- 9+ 
10-class TestCopyKernelMemoryFormat(TestCase):10+class TestCopyKernelMemoryFormat(TestCase):
11- def test_h2d_copy_contiguous_tensor(self):11+ def test_h2d_copy_contiguous_tensor(self):
12- dtype_list = [np.float16, np.float32, np.int32, np.int64]12+ dtype_list = [np.float16, np.float32, np.int32, np.int64]
13- shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]]13+ shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]]
14- shape_format = [14+ shape_format = [
15- [dtype, 2, shape] for dtype in dtype_list for shape in shape_list15+ [dtype, 2, shape] for dtype in dtype_list for shape in shape_list
16- ]16+ ]
17- 17+
18- for item in shape_format:18+ for item in shape_format:
19- cpu_input, npu_input = create_common_tensor(item, -100, 100)19+ cpu_input, npu_input = create_common_tensor(item, -100, 100)
20- npu_input_copy = cpu_input.npu()20+ npu_input_copy = cpu_input.npu()
21- self.assertRtolEqual(npu_input_copy.cpu().numpy(), cpu_input.numpy())21+ self.assertRtolEqual(npu_input_copy.cpu().numpy(), cpu_input.numpy())
22- 22+ 
23- def test_h2d_copy_non_contiguous_tensor(self):23+ def test_h2d_copy_non_contiguous_tensor(self):
24- dtype_list = [np.float16, np.float32]24+ dtype_list = [np.float16, np.float32]
25- shape_list = [[32, 64], [16, 32, 64]]25+ shape_list = [[32, 64], [16, 32, 64]]
26- shape_format = [26+ shape_format = [
27- [dtype, 2, shape] for dtype in dtype_list for shape in shape_list27+ [dtype, 2, shape] for dtype in dtype_list for shape in shape_list
28- ]28+ ]
29- 29+
30- for item in shape_format:30+ for item in shape_format:
31- cpu_input, npu_input = create_common_tensor(item, -100, 100)31+ cpu_input, npu_input = create_common_tensor(item, -100, 100)
32- cpu_transposed = cpu_input.transpose(-1, -2)32+ cpu_transposed = cpu_input.transpose(-1, -2)
33- npu_transposed = cpu_transposed.npu()33+ npu_transposed = cpu_transposed.npu()
34- npu_contiguous = npu_transposed.contiguous()34+ npu_contiguous = npu_transposed.contiguous()
35- self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_transposed.contiguous().numpy())35+ self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_transposed.contiguous().numpy())
36- 36+ 
37- def test_d2h_copy_contiguous_tensor(self):37+ def test_d2h_copy_contiguous_tensor(self):
38- dtype_list = [np.float16, np.float32, np.int32, np.int64]38+ dtype_list = [np.float16, np.float32, np.int32, np.int64]
39- shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]]39+ shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]]
40- shape_format = [40+ shape_format = [
41- [dtype, 2, shape] for dtype in dtype_list for shape in shape_list41+ [dtype, 2, shape] for dtype in dtype_list for shape in shape_list
42- ]42+ ]
43- 43+
44- for item in shape_format:44+ for item in shape_format:
45- cpu_input, npu_input = create_common_tensor(item, -100, 100)45+ cpu_input, npu_input = create_common_tensor(item, -100, 100)
46- cpu_output = npu_input.cpu()46+ cpu_output = npu_input.cpu()
47- self.assertRtolEqual(cpu_output.numpy(), cpu_input.numpy())47+ self.assertRtolEqual(cpu_output.numpy(), cpu_input.numpy())
48- 48+ 
49- def test_d2h_copy_non_contiguous_tensor(self):49+ def test_d2h_copy_non_contiguous_tensor(self):
50- dtype_list = [np.float16, np.float32]50+ dtype_list = [np.float16, np.float32]
51- shape_list = [[32, 64], [16, 32, 64]]51+ shape_list = [[32, 64], [16, 32, 64]]
52- shape_format = [52+ shape_format = [
53- [dtype, 2, shape] for dtype in dtype_list for shape in shape_list53+ [dtype, 2, shape] for dtype in dtype_list for shape in shape_list
54- ]54+ ]
55- 55+
56- for item in shape_format:56+ for item in shape_format:
57- cpu_input, npu_input = create_common_tensor(item, -100, 100)57+ cpu_input, npu_input = create_common_tensor(item, -100, 100)
58- npu_transposed = npu_input.transpose(-1, -2)58+ npu_transposed = npu_input.transpose(-1, -2)
59- cpu_output = npu_transposed.cpu()59+ cpu_output = npu_transposed.cpu()
60- self.assertRtolEqual(cpu_output.numpy(), cpu_input.transpose(-1, -2).contiguous().numpy())60+ self.assertRtolEqual(cpu_output.numpy(), cpu_input.transpose(-1, -2).contiguous().numpy())
61- 61+ 
62- def test_h2d_copy_different_dtype(self):62+ def test_h2d_copy_different_dtype(self):
63- src_dtype_list = [np.float32, np.float16]63+ src_dtype_list = [np.float32, np.float16]
64- dst_dtype_list = [torch.float16, torch.float32]64+ dst_dtype_list = [torch.float16, torch.float32]
65- shape = [32, 64]65+ shape = [32, 64]
66- 66+
67- for src_dtype, dst_dtype in zip(src_dtype_list, dst_dtype_list):67+ for src_dtype, dst_dtype in zip(src_dtype_list, dst_dtype_list):
68- cpu_input = torch.randn(shape, dtype=torch.float32) * 10068+ cpu_input = torch.randn(shape, dtype=torch.float32) * 100
69- cpu_input = cpu_input.to(torch.from_numpy(np.array([])).dtype if src_dtype == np.float32 else torch.float16)69+ cpu_input = cpu_input.to(torch.from_numpy(np.array([])).dtype if src_dtype == np.float32 else torch.float16)
70- 70+
71- npu_input = cpu_input.npu()71+ npu_input = cpu_input.npu()
72- npu_output = npu_input.to(dst_dtype)72+ npu_output = npu_input.to(dst_dtype)
73- 73+
74- cpu_output = cpu_input.to(dst_dtype)74+ cpu_output = cpu_input.to(dst_dtype)
75- self.assertRtolEqual(npu_output.cpu().numpy(), cpu_output.numpy())75+ self.assertRtolEqual(npu_output.cpu().numpy(), cpu_output.numpy())
76- 76+ 
77- def test_d2h_copy_different_dtype(self):77+ def test_d2h_copy_different_dtype(self):
78- dtype_pairs = [78+ dtype_pairs = [
79- (np.float16, torch.float32),79+ (np.float16, torch.float32),
80- (np.float32, torch.float16),80+ (np.float32, torch.float16),
81- ]81+ ]
82- shape = [32, 64]82+ shape = [32, 64]
83- 83+
84- for src_dtype, dst_dtype in dtype_pairs:84+ for src_dtype, dst_dtype in dtype_pairs:
85- cpu_input, npu_input = create_common_tensor([src_dtype, 0, shape], -100, 100)85+ cpu_input, npu_input = create_common_tensor([src_dtype, 0, shape], -100, 100)
86- cpu_output = npu_input.cpu().to(dst_dtype)86+ cpu_output = npu_input.cpu().to(dst_dtype)
87- 87+
88- expected = cpu_input.to(dst_dtype)88+ expected = cpu_input.to(dst_dtype)
89- self.assertRtolEqual(cpu_output.numpy(), expected.numpy())89+ self.assertRtolEqual(cpu_output.numpy(), expected.numpy())
90- 90+ 
91- def test_h2d_copy_slice_tensor(self):91+ def test_h2d_copy_slice_tensor(self):
92- shape = [64, 128]92+ shape = [64, 128]
93- cpu_input = torch.randn(shape)93+ cpu_input = torch.randn(shape)
94- 94+
95- cpu_slice = cpu_input[10:30, 20:60]95+ cpu_slice = cpu_input[10:30, 20:60]
96- npu_slice = cpu_slice.npu()96+ npu_slice = cpu_slice.npu()
97- 97+
98- npu_contiguous = npu_slice.contiguous()98+ npu_contiguous = npu_slice.contiguous()
99- self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_slice.contiguous().numpy())99+ self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_slice.contiguous().numpy())
100- 100+ 
101- def test_d2h_copy_slice_tensor(self):101+ def test_d2h_copy_slice_tensor(self):
102- shape = [64, 128]102+ shape = [64, 128]
103- cpu_input = torch.randn(shape)103+ cpu_input = torch.randn(shape)
104- npu_input = cpu_input.npu()104+ npu_input = cpu_input.npu()
105- 105+
106- npu_slice = npu_input[10:30, 20:60]106+ npu_slice = npu_input[10:30, 20:60]
107- cpu_slice = npu_slice.cpu()107+ cpu_slice = npu_slice.cpu()
108- self.assertRtolEqual(cpu_slice.numpy(), cpu_input[10:30, 20:60].contiguous().numpy())108+ self.assertRtolEqual(cpu_slice.numpy(), cpu_input[10:30, 20:60].contiguous().numpy())
109- 109+ 
110- def test_d2h_copy_broadcast_tensor(self):110+ def test_d2h_copy_broadcast_tensor(self):
111- shape = [1, 64, 1]111+ shape = [1, 64, 1]
112- cpu_input = torch.randn(shape)112+ cpu_input = torch.randn(shape)
113- npu_input = cpu_input.npu()113+ npu_input = cpu_input.npu()
114- 114+
115- npu_broadcast = npu_input.expand(4, 64, 128)115+ npu_broadcast = npu_input.expand(4, 64, 128)
116- cpu_output = npu_broadcast.cpu()116+ cpu_output = npu_broadcast.cpu()
117- self.assertRtolEqual(cpu_output.numpy(), cpu_input.expand(4, 64, 128).contiguous().numpy())117+ self.assertRtolEqual(cpu_output.numpy(), cpu_input.expand(4, 64, 128).contiguous().numpy())
118- 118+ 
119- def test_h2d_copy_permute_tensor(self):119+ def test_h2d_copy_permute_tensor(self):
120- shape = [32, 64, 128]120+ shape = [32, 64, 128]
121- cpu_input = torch.randn(shape)121+ cpu_input = torch.randn(shape)
122- cpu_permuted = cpu_input.permute(2, 0, 1)122+ cpu_permuted = cpu_input.permute(2, 0, 1)
123- npu_permuted = cpu_permuted.npu()123+ npu_permuted = cpu_permuted.npu()
124- npu_contiguous = npu_permuted.contiguous()124+ npu_contiguous = npu_permuted.contiguous()
125- self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_permuted.contiguous().numpy())125+ self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_permuted.contiguous().numpy())
126- 126+ 
127- def test_d2h_copy_permute_tensor(self):127+ def test_d2h_copy_permute_tensor(self):
128- shape = [32, 64, 128]128+ shape = [32, 64, 128]
129- cpu_input = torch.randn(shape)129+ cpu_input = torch.randn(shape)
130- npu_input = cpu_input.npu()130+ npu_input = cpu_input.npu()
131- 131+
132- npu_permuted = npu_input.permute(2, 0, 1) 132+ npu_permuted = npu_input.permute(2, 0, 1)
133- cpu_output = npu_permuted.cpu()133+ cpu_output = npu_permuted.cpu()
134- self.assertRtolEqual(cpu_output.numpy(), cpu_input.permute(2, 0, 1).contiguous().numpy())134+ self.assertRtolEqual(cpu_output.numpy(), cpu_input.permute(2, 0, 1).contiguous().numpy())
135- 135+ 
136- 136+ 
137-if __name__ == "__main__":137+if __name__ == "__main__":
138 run_tests()138 run_tests()
Mtest/npu/test_expandable_segments.py+105-105
@@ -1,105 +1,105 @@
1-import os1+import os
2-import gc2+import gc
3-import unittest3+import unittest
4- 4+ 
5-import torch5+import torch
6-import torch_npu6+import torch_npu
7-from torch_npu.testing.testcase import TestCase, run_tests7+from torch_npu.testing.testcase import TestCase, run_tests
8-from torch.testing._internal.common_utils import TestCase, run_tests, TEST_PRIVATEUSE18+from torch.testing._internal.common_utils import TestCase, run_tests, TEST_PRIVATEUSE1
9- 9+ 
10-os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:True"10+os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:True"
11- 11+ 
12-device_name = torch_npu.npu.get_device_name(0)12+device_name = torch_npu.npu.get_device_name(0)
13- 13+ 
14- 14+ 
15-class Test_expandable_segments(TestCase):15+class Test_expandable_segments(TestCase):
16- @unittest.skipIf(device_name == "Ascend910B4", "Skip when device name is Ascend910B4")16+ @unittest.skipIf(device_name == "Ascend910B4", "Skip when device name is Ascend910B4")
17- def test_empty_virt_addr_cache(self):17+ def test_empty_virt_addr_cache(self):
18- gc.collect()18+ gc.collect()
19- torch_npu.npu.empty_cache()19+ torch_npu.npu.empty_cache()
20- prev = 020+ prev = 0
21- 21+ 
22- x = torch.empty((7500, 1024, 1024), device="npu")22+ x = torch.empty((7500, 1024, 1024), device="npu")
23- del x23+ del x
24- last_r = torch_npu.npu.memory_reserved()24+ last_r = torch_npu.npu.memory_reserved()
25- 25+ 
26- torch_npu.npu.empty_virt_addr_cache()26+ torch_npu.npu.empty_virt_addr_cache()
27- new_r = torch_npu.npu.memory_reserved()27+ new_r = torch_npu.npu.memory_reserved()
28- self.assertEqual(new_r, prev)28+ self.assertEqual(new_r, prev)
29- self.assertEqual(torch_npu.npu.max_memory_reserved(), last_r)29+ self.assertEqual(torch_npu.npu.max_memory_reserved(), last_r)
30- 30+ 
31- # test re-alloc after empty virtual address31+ # test re-alloc after empty virtual address
32- try:32+ try:
33- y = torch.empty((7500, 1024, 1024), device="npu")33+ y = torch.empty((7500, 1024, 1024), device="npu")
34- self.assertGreater(torch_npu.npu.memory_allocated(), prev)34+ self.assertGreater(torch_npu.npu.memory_allocated(), prev)
35- finally:35+ finally:
36- if y is not None:36+ if y is not None:
37- del y37+ del y
38- self.assertEqual(torch_npu.npu.memory_allocated(), prev)38+ self.assertEqual(torch_npu.npu.memory_allocated(), prev)
39- torch_npu.npu.empty_virt_addr_cache()39+ torch_npu.npu.empty_virt_addr_cache()
40- # empty unmapped physical handles with empty_cache()40+ # empty unmapped physical handles with empty_cache()
41- torch_npu.npu.empty_cache()41+ torch_npu.npu.empty_cache()
42- self.assertEqual(torch_npu.npu.memory_reserved(), prev)42+ self.assertEqual(torch_npu.npu.memory_reserved(), prev)
43- 43+ 
44- @unittest.skipIf(TEST_PRIVATEUSE1, "NPU not available for graph capture")44+ @unittest.skipIf(TEST_PRIVATEUSE1, "NPU not available for graph capture")
45- def test_set_segment_state_to_checkpoint_when_expandable_segments(self):45+ def test_set_segment_state_to_checkpoint_when_expandable_segments(self):
46- def tensor_metadata(x):46+ def tensor_metadata(x):
47- return {47+ return {
48- "nbytes": x.untyped_storage().nbytes(),48+ "nbytes": x.untyped_storage().nbytes(),
49- "data_ptr": x.untyped_storage().data_ptr(),49+ "data_ptr": x.untyped_storage().data_ptr(),
50- "size": x.shape,50+ "size": x.shape,
51- "stride": x.stride(),51+ "stride": x.stride(),
52- "dtype": x.dtype,52+ "dtype": x.dtype,
53- "device": x.device,53+ "device": x.device,
54- "storage_offset": x.storage_offset(), }54+ "storage_offset": x.storage_offset(), }
55- 55+ 
56- def reconstruct_from_tensor_metadata(metadata):56+ def reconstruct_from_tensor_metadata(metadata):
57- s = torch._C._construct_storage_from_data_pointer(57+ s = torch._C._construct_storage_from_data_pointer(
58- metadata["data_ptr"], metadata["device"], metadata["nbytes"])58+ metadata["data_ptr"], metadata["device"], metadata["nbytes"])
59- t = torch.empty([0], device=metadata["device"], dtype=metadata["dtype"])59+ t = torch.empty([0], device=metadata["device"], dtype=metadata["dtype"])
60- t.set_(source=s, storage_offset=metadata["storage_offset"],60+ t.set_(source=s, storage_offset=metadata["storage_offset"],
61- size=metadata["size"], stride=metadata["stride"], )61+ size=metadata["size"], stride=metadata["stride"], )
62- return t62+ return t
63- 63+ 
64- def cudagraphify(fn, inputs, pool, stream):64+ def cudagraphify(fn, inputs, pool, stream):
65- torch.npu.synchronize()65+ torch.npu.synchronize()
66- gc.collect()66+ gc.collect()
67- torch.npu.empty_cache()67+ torch.npu.empty_cache()
68- 68+ 
69- graph = torch.npu.NPUGraph()69+ graph = torch.npu.NPUGraph()
70- with torch.npu.graph(graph, stream=stream, pool=pool):70+ with torch.npu.graph(graph, stream=stream, pool=pool):
71- static_outputs = fn(*inputs)71+ static_outputs = fn(*inputs)
72- return graph, static_outputs72+ return graph, static_outputs
73- 73+ 
74- def foo(x, idx):74+ def foo(x, idx):
75- r1 = x.expand([1, 2097152 // 8]).sqrt()75+ r1 = x.expand([1, 2097152 // 8]).sqrt()
76- r2 = x.expand([idx, 2097152]).clone()76+ r2 = x.expand([idx, 2097152]).clone()
77- return r1, r277+ return r1, r2
78- 78+ 
79- # init79+ # init
80- pool_id = torch.npu.graph_pool_handle()80+ pool_id = torch.npu.graph_pool_handle()
81- com_stream = torch.npu.Stream()81+ com_stream = torch.npu.Stream()
82- com_device = torch_npu.npu.current_device()82+ com_device = torch_npu.npu.current_device()
83- inp = torch.tensor([7]).npu()83+ inp = torch.tensor([7]).npu()
84- 84+ 
85- # start capture graph185+ # start capture graph1
86- graph1, outputs1 = cudagraphify(foo, [inp, 1], pool=pool_id, stream=com_stream)86+ graph1, outputs1 = cudagraphify(foo, [inp, 1], pool=pool_id, stream=com_stream)
87- graph1_state = torch_npu._C._npu_getCheckpointState(com_device, pool_id)87+ graph1_state = torch_npu._C._npu_getCheckpointState(com_device, pool_id)
88- output1_metadata = [tensor_metadata(t) for t in outputs1]88+ output1_metadata = [tensor_metadata(t) for t in outputs1]
89- outputs1 = None89+ outputs1 = None
90- 90+ 
91- # start capture graph291+ # start capture graph2
92- graph2, outputs2 = cudagraphify(foo, [inp, 2], pool=pool_id, stream=com_stream)92+ graph2, outputs2 = cudagraphify(foo, [inp, 2], pool=pool_id, stream=com_stream)
93- graph2_state = torch_npu._C._npu_getCheckpointState(com_device, pool_id)93+ graph2_state = torch_npu._C._npu_getCheckpointState(com_device, pool_id)
94- graph2.replay()94+ graph2.replay()
95- outputs2 = None95+ outputs2 = None
96- 96+ 
97- # replay graph197+ # replay graph1
98- graph1.replay()98+ graph1.replay()
99- reconstructed_tensors1 = [reconstruct_from_tensor_metadata(metadata) for metadata in output1_metadata]99+ reconstructed_tensors1 = [reconstruct_from_tensor_metadata(metadata) for metadata in output1_metadata]
100- output1_new_storage = [output.untyped_storage()._cdata for output in reconstructed_tensors1]100+ output1_new_storage = [output.untyped_storage()._cdata for output in reconstructed_tensors1]
101- torch_npu._C._npu_setCheckpointPoolState(com_device, graph1_state, [], output1_new_storage)101+ torch_npu._C._npu_setCheckpointPoolState(com_device, graph1_state, [], output1_new_storage)
102- 102+ 
103- 103+ 
104-if __name__ == '__main__':104+if __name__ == '__main__':
105- run_tests()105+ run_tests()
Mtest/npu/test_graph_tree.py+1210-1210
@@ -1,1210 +1,1210 @@
1-import os1+import os
2- 2+ 
3-os.environ["ASCEND_LAUNCH_BLOCKING"] = "0"3+os.environ["ASCEND_LAUNCH_BLOCKING"] = "0"
4- 4+ 
5-from unittest.mock import patch, MagicMock, call, ANY5+from unittest.mock import patch, MagicMock, call, ANY
6-import weakref6+import weakref
7-import pytest7+import pytest
8-import torch8+import torch
9-import torch_npu9+import torch_npu
10-from torch_npu.npu._graph_tree import (10+from torch_npu.npu._graph_tree import (
11- check_memory_pool,11+ check_memory_pool,
12- clear_cublass_cache,12+ clear_cublass_cache,
13- clear_cublas_manager,13+ clear_cublas_manager,
14- disable_conv_cache_emptying,14+ disable_conv_cache_emptying,
15- enable_history_recording,15+ enable_history_recording,
16- format_tb,16+ format_tb,
17- npugraphify,17+ npugraphify,
18- npugraphify_impl,18+ npugraphify_impl,
19- TreeManagerContainer,19+ TreeManagerContainer,
20- StorageWeakRefWrapper,20+ StorageWeakRefWrapper,
21- NPUWarmupNode,21+ NPUWarmupNode,
22- CompilationMode,22+ CompilationMode,
23- get_container,23+ get_container,
24- get_block_addrs,24+ get_block_addrs,
25- get_manager,25+ get_manager,
26- get_npugraph_segments,26+ get_npugraph_segments,
27- reset_npugraph_trees,27+ reset_npugraph_trees,
28- local,28+ local,
29- OutputAliasInfo,29+ OutputAliasInfo,
30- UnaliasedStorage,30+ UnaliasedStorage,
31- AliasesPriorGraphOutput,31+ AliasesPriorGraphOutput,
32- AliasesNewOutput,32+ AliasesNewOutput,
33- NPUGraphNode,33+ NPUGraphNode,
34- WrappedFunction,34+ WrappedFunction,
35- NPUGraphTreeManager,35+ NPUGraphTreeManager,
36- ExecutionState,36+ ExecutionState,
37- FunctionID,37+ FunctionID,
38- GraphID,38+ GraphID,
39-)39+)
40-from torch_npu.testing.testcase import TestCase, run_tests40+from torch_npu.testing.testcase import TestCase, run_tests
41- 41+ 
42- 42+ 
43-device = "npu:0"43+device = "npu:0"
44-torch.npu.set_device(device)44+torch.npu.set_device(device)
45- 45+ 
46- 46+ 
47-class TestCublasCacheManagement(TestCase):47+class TestCublasCacheManagement(TestCase):
48- @patch("torch_npu.npu._graph_tree.clear_cublass_cache")48+ @patch("torch_npu.npu._graph_tree.clear_cublass_cache")
49- def test_clear_cublas_manager_context(self, mock_clear):49+ def test_clear_cublas_manager_context(self, mock_clear):
50- with clear_cublas_manager():50+ with clear_cublas_manager():
51- mock_clear.assert_called_once()51+ mock_clear.assert_called_once()
52- mock_clear.reset_mock()52+ mock_clear.reset_mock()
53- mock_clear.assert_called_once()53+ mock_clear.assert_called_once()
54- 54+ 
55- 55+ 
56-class TestDisableConvCache(TestCase):56+class TestDisableConvCache(TestCase):
57- def test_disable_conv_cache_emptying(self):57+ def test_disable_conv_cache_emptying(self):
58- with disable_conv_cache_emptying():58+ with disable_conv_cache_emptying():
59- pass # No operation, just ensure no exceptions59+ pass # No operation, just ensure no exceptions
60- 60+ 
61- 61+ 
62-class TestHistoryRecording(TestCase):62+class TestHistoryRecording(TestCase):
63- @patch("torch.npu.memory._record_memory_history")63+ @patch("torch.npu.memory._record_memory_history")
64- def test_enable_history_recording(self, mock_record):64+ def test_enable_history_recording(self, mock_record):
65- original_state = torch_npu._C._npu_isHistoryEnabled()65+ original_state = torch_npu._C._npu_isHistoryEnabled()
66- with enable_history_recording():66+ with enable_history_recording():
67- if not original_state:67+ if not original_state:
68- mock_record.assert_called_once()68+ mock_record.assert_called_once()
69- else:69+ else:
70- mock_record.assert_not_called()70+ mock_record.assert_not_called()
71- mock_record.assert_any_call(None)71+ mock_record.assert_any_call(None)
72- 72+ 
73- 73+ 
74-class TestNpuGraphFunctions(TestCase):74+class TestNpuGraphFunctions(TestCase):
75- def setUp(self):75+ def setUp(self):
76- # Reset global state before each test76+ # Reset global state before each test
77- reset_npugraph_trees()77+ reset_npugraph_trees()
78- 78+ 
79- @patch("torch_npu.npu._graph_tree.TreeManagerContainer")79+ @patch("torch_npu.npu._graph_tree.TreeManagerContainer")
80- def test_get_manager(self, mock_container):80+ def test_get_manager(self, mock_container):
81- # Test manager creation81+ # Test manager creation
82- mock_container.return_value.get_tree_manager.return_value = "mock_manager"82+ mock_container.return_value.get_tree_manager.return_value = "mock_manager"
83- manager = get_manager(0)83+ manager = get_manager(0)
84- self.assertEqual(manager, "mock_manager")84+ self.assertEqual(manager, "mock_manager")
85- 85+ 
86- # Test no-creation path86+ # Test no-creation path
87- manager = get_manager(0, create_if_none_exists=False)87+ manager = get_manager(0, create_if_none_exists=False)
88- mock_container.return_value.get_tree_manager.assert_called_once()88+ mock_container.return_value.get_tree_manager.assert_called_once()
89- 89+ 
90- @patch("torch_npu.npu._graph_tree.npugraphify")90+ @patch("torch_npu.npu._graph_tree.npugraphify")
91- @patch("torch._inductor.compile_fx.align_inputs_from_check_idxs")91+ @patch("torch._inductor.compile_fx.align_inputs_from_check_idxs")
92- def test_npugraphify_impl(self, mock_align, mock_npugraphify):92+ def test_npugraphify_impl(self, mock_align, mock_npugraphify):
93- # Setup mock model and inputs93+ # Setup mock model and inputs
94- mock_model = MagicMock()94+ mock_model = MagicMock()
95- inputs = [1, torch.tensor([2]), 3]95+ inputs = [1, torch.tensor([2]), 3]
96- static_idxs = (1,)96+ static_idxs = (1,)
97- 97+ 
98- # Test caching behavior98+ # Test caching behavior
99- impl = npugraphify_impl(mock_model, inputs, static_idxs)99+ impl = npugraphify_impl(mock_model, inputs, static_idxs)
100- 100+ 
101- # First call101+ # First call
102- mock_npugraphify.return_value = (lambda x: "output1", "output1")102+ mock_npugraphify.return_value = (lambda x: "output1", "output1")
103- result = impl(inputs)103+ result = impl(inputs)
104- self.assertEqual(result, "output1")104+ self.assertEqual(result, "output1")
105- 105+ 
106- # Second call with same int keys106+ # Second call with same int keys
107- result = impl(inputs)107+ result = impl(inputs)
108- self.assertEqual(result, "output1")108+ self.assertEqual(result, "output1")
109- mock_npugraphify.assert_called_once()109+ mock_npugraphify.assert_called_once()
110- 110+ 
111- @patch("torch_npu.npu._graph_tree.get_container")111+ @patch("torch_npu.npu._graph_tree.get_container")
112- def test_npugraphify(self, mock_container):112+ def test_npugraphify(self, mock_container):
113- # Setup mock manager113+ # Setup mock manager
114- mock_manager = MagicMock()114+ mock_manager = MagicMock()
115- mock_container.return_value.get_tree_manager.return_value = mock_manager115+ mock_container.return_value.get_tree_manager.return_value = mock_manager
116- 116+ 
117- # Test valid mode combinations117+ # Test valid mode combinations
118- model = MagicMock()118+ model = MagicMock()
119- inputs = [torch.tensor([1])]119+ inputs = [torch.tensor([1])]
120- 120+ 
121- # Test forward mode121+ # Test forward mode
122- npugraphify(122+ npugraphify(
123- model, inputs, (), device_index=0, is_backward=False, is_inference=False123+ model, inputs, (), device_index=0, is_backward=False, is_inference=False
124- )124+ )
125- mock_manager.add_function.assert_called_with(125+ mock_manager.add_function.assert_called_with(
126- model, inputs, (), None, CompilationMode.FORWARD, (), (), ()126+ model, inputs, (), None, CompilationMode.FORWARD, (), (), ()
127- )127+ )
128- 128+ 
129- # Test backward mode129+ # Test backward mode
130- mock_manager.reset_mock()130+ mock_manager.reset_mock()
131- npugraphify(131+ npugraphify(
132- model, inputs, (), device_index=0, is_backward=True, is_inference=False132+ model, inputs, (), device_index=0, is_backward=True, is_inference=False
133- )133+ )
134- mock_manager.add_function.assert_called_with(134+ mock_manager.add_function.assert_called_with(
135- model, inputs, (), None, CompilationMode.BACKWARD, (), (), ()135+ model, inputs, (), None, CompilationMode.BACKWARD, (), (), ()
136- )136+ )
137- 137+ 
138- # Test invalid mode combination138+ # Test invalid mode combination
139- with self.assertRaises(RuntimeError):139+ with self.assertRaises(RuntimeError):
140- npugraphify(140+ npugraphify(
141- model, inputs, (), device_index=0, is_backward=True, is_inference=True141+ model, inputs, (), device_index=0, is_backward=True, is_inference=True
142- )142+ )
143- 143+ 
144- 144+ 
145-class TestTreeManagerContainer(TestCase):145+class TestTreeManagerContainer(TestCase):
146- def setUp(self):146+ def setUp(self):
147- self.container = TreeManagerContainer(0)147+ self.container = TreeManagerContainer(0)
148- 148+ 
149- def test_initial_state(self):149+ def test_initial_state(self):
150- self.assertIsNone(self.container.tree_manager)150+ self.assertIsNone(self.container.tree_manager)
151- self.assertEqual(self.container.live_npugraphify_fns, 0)151+ self.assertEqual(self.container.live_npugraphify_fns, 0)
152- 152+ 
153- def test_add_strong_reference(self):153+ def test_add_strong_reference(self):
154- self.container.add_strong_reference(lambda: None)154+ self.container.add_strong_reference(lambda: None)
155- # Simulate finalization of fn155+ # Simulate finalization of fn
156- finalizer = weakref.finalize(156+ finalizer = weakref.finalize(
157- lambda: None,157+ lambda: None,
158- self.container.finalize_npugraphify_fn, # Object to monitor # Callback158+ self.container.finalize_npugraphify_fn, # Object to monitor # Callback
159- )159+ )
160- finalizer.atexit = False # Prevent finalizer from running at exit160+ finalizer.atexit = False # Prevent finalizer from running at exit
161- 161+ 
162- # Simulate finalization162+ # Simulate finalization
163- finalizer()163+ finalizer()
164- # If all references are gone, tree_manager should be None164+ # If all references are gone, tree_manager should be None
165- self.container._finalize_tree_manager = MagicMock()165+ self.container._finalize_tree_manager = MagicMock()
166- self.container._finalize_tree_manager()166+ self.container._finalize_tree_manager()
167- self.container._finalize_tree_manager.assert_called_once()167+ self.container._finalize_tree_manager.assert_called_once()
168- 168+ 
169- def test_get_tree_manager(self):169+ def test_get_tree_manager(self):
170- with patch("torch_npu.npu.graphs.NPUGraph.capture_begin"), patch(170+ with patch("torch_npu.npu.graphs.NPUGraph.capture_begin"), patch(
171- "torch_npu.npu.graphs.NPUGraph.capture_end"171+ "torch_npu.npu.graphs.NPUGraph.capture_end"
172- ):172+ ):
173- manager = self.container.get_tree_manager()173+ manager = self.container.get_tree_manager()
174- self.assertIsNotNone(manager)174+ self.assertIsNotNone(manager)
175- self.assertIs(manager, self.container.get_tree_manager()) # Same instance175+ self.assertIs(manager, self.container.get_tree_manager()) # Same instance
176- 176+ 
177- 177+ 
178-class TestStorageWeakRefWrapper(TestCase):178+class TestStorageWeakRefWrapper(TestCase):
179- def test_storage_ref(self):179+ def test_storage_ref(self):
180- tensor = torch.tensor([1], device="npu")180+ tensor = torch.tensor([1], device="npu")
181- wrapper = StorageWeakRefWrapper(tensor)181+ wrapper = StorageWeakRefWrapper(tensor)
182- self.assertEqual(wrapper.data_ptr(), tensor.untyped_storage().data_ptr())182+ self.assertEqual(wrapper.data_ptr(), tensor.untyped_storage().data_ptr())
183- del tensor183+ del tensor
184- # Storage might still be alive due to Python's ref counting; force GC184+ # Storage might still be alive due to Python's ref counting; force GC
185- import gc185+ import gc
186- 186+ 
187- gc.collect()187+ gc.collect()
188- self.assertTrue(wrapper.expired())188+ self.assertTrue(wrapper.expired())
189- 189+ 
190- 190+ 
191-class TestNPUWarmupNode(TestCase):191+class TestNPUWarmupNode(TestCase):
192- @patch("torch_npu.npu._graph_tree.StorageWeakRefWrapper")192+ @patch("torch_npu.npu._graph_tree.StorageWeakRefWrapper")
193- @patch("torch_npu.npu._graph_tree.check_memory_pool")193+ @patch("torch_npu.npu._graph_tree.check_memory_pool")
194- def test_run_captures_outputs(self, mock_check, mock_wrapper):194+ def test_run_captures_outputs(self, mock_check, mock_wrapper):
195- mock_model = MagicMock(return_value=[torch.tensor([2], device="npu")])195+ mock_model = MagicMock(return_value=[torch.tensor([2], device="npu")])
196- wrapped_fn = MagicMock(model=mock_model, constants=[])196+ wrapped_fn = MagicMock(model=mock_model, constants=[])
197- stream = torch.npu.Stream()197+ stream = torch.npu.Stream()
198- node = NPUWarmupNode(198+ node = NPUWarmupNode(
199- wrapped_fn,199+ wrapped_fn,
200- parent=None,200+ parent=None,
201- npu_graphs_pool=(0, 0),201+ npu_graphs_pool=(0, 0),
202- existing_npu_graph=None,202+ existing_npu_graph=None,
203- device_index=0,203+ device_index=0,
204- stack_traces=None,204+ stack_traces=None,
205- stream=stream,205+ stream=stream,
206- already_warm=False,206+ already_warm=False,
207- graph_id=1,207+ graph_id=1,
208- )208+ )
209- outputs = node.run([])209+ outputs = node.run([])
210- self.assertEqual(len(node.outputs_weakrefs), 1)210+ self.assertEqual(len(node.outputs_weakrefs), 1)
211- 211+ 
212- 212+ 
213-class TestTreeManagerIntegration(TestCase):213+class TestTreeManagerIntegration(TestCase):
214- def test_get_container_singleton_per_device(self):214+ def test_get_container_singleton_per_device(self):
215- container1 = get_container(0)215+ container1 = get_container(0)
216- container2 = get_container(0)216+ container2 = get_container(0)
217- self.assertIs(container1, container2)217+ self.assertIs(container1, container2)
218- container3 = get_container(1)218+ container3 = get_container(1)
219- self.assertIsNot(container1, container3)219+ self.assertIsNot(container1, container3)
220- 220+ 
221- def test_reset_npugraph_trees(self):221+ def test_reset_npugraph_trees(self):
222- get_container(0) # Initialize a container222+ get_container(0) # Initialize a container
223- reset_npugraph_trees()223+ reset_npugraph_trees()
224- container_dict = getattr(local, "npu_tree_manager_containers", {})224+ container_dict = getattr(local, "npu_tree_manager_containers", {})
225- self.assertEqual(len(container_dict), 0)225+ self.assertEqual(len(container_dict), 0)
226- 226+ 
227- 227+ 
228-@pytest.fixture228+@pytest.fixture
229-def mock_wrapped_function():229+def mock_wrapped_function():
230- def model_side_effect(inputs):230+ def model_side_effect(inputs):
231- # Clear inputs list while preserving reference231+ # Clear inputs list while preserving reference
232- inputs[:] = []232+ inputs[:] = []
233- return []233+ return []
234- 234+ 
235- return MagicMock(235+ return MagicMock(
236- spec=WrappedFunction,236+ spec=WrappedFunction,
237- static_input_idxs=[0],237+ static_input_idxs=[0],
238- constants=[],238+ constants=[],
239- model=MagicMock(side_effect=model_side_effect),239+ model=MagicMock(side_effect=model_side_effect),
240- )240+ )
241- 241+ 
242- 242+ 
243-@pytest.fixture243+@pytest.fixture
244-def mock_parent_node():244+def mock_parent_node():
245- parent = MagicMock(spec=NPUGraphNode)245+ parent = MagicMock(spec=NPUGraphNode)
246- parent.outputs_weakrefs = []246+ parent.outputs_weakrefs = []
247- parent.path_weakrefs = []247+ parent.path_weakrefs = []
248- parent.parent = None248+ parent.parent = None
249- parent.stack_traces = []249+ parent.stack_traces = []
250- parent.recorded_liveness_after_graph = []250+ parent.recorded_liveness_after_graph = []
251- return parent251+ return parent
252- 252+ 
253- 253+ 
254-@pytest.fixture254+@pytest.fixture
255-def basic_npu_graph_node(mock_wrapped_function, mock_parent_node):255+def basic_npu_graph_node(mock_wrapped_function, mock_parent_node):
256- with patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager"), patch(256+ with patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager"), patch(
257- "torch_npu.npu._graph_tree.check_memory_pool"257+ "torch_npu.npu._graph_tree.check_memory_pool"
258- ), patch("torch_npu._C._npu_getCheckpointState"):258+ ), patch("torch_npu._C._npu_getCheckpointState"):
259- return NPUGraphNode(259+ return NPUGraphNode(
260- wrapped_function=mock_wrapped_function,260+ wrapped_function=mock_wrapped_function,
261- graph_id=1,261+ graph_id=1,
262- parent=mock_parent_node,262+ parent=mock_parent_node,
263- inputs=[torch.tensor([1.0], device="npu")],263+ inputs=[torch.tensor([1.0], device="npu")],
264- npu_graphs_pool=(0, 0),264+ npu_graphs_pool=(0, 0),
265- device_index=0,265+ device_index=0,
266- stack_traces=None,266+ stack_traces=None,
267- stream=torch.npu.Stream(),267+ stream=torch.npu.Stream(),
268- )268+ )
269- 269+ 
270- 270+ 
271-class TestOutputAliasInfo(TestCase):271+class TestOutputAliasInfo(TestCase):
272- def test_aliases_prior_graph_output_validation(self):272+ def test_aliases_prior_graph_output_validation(self):
273- with pytest.raises(RuntimeError):273+ with pytest.raises(RuntimeError):
274- AliasesPriorGraphOutput("invalid_index")274+ AliasesPriorGraphOutput("invalid_index")
275- 275+ 
276- def test_aliases_new_output_validation(self):276+ def test_aliases_new_output_validation(self):
277- with pytest.raises(RuntimeError):277+ with pytest.raises(RuntimeError):
278- AliasesNewOutput("not_an_int")278+ AliasesNewOutput("not_an_int")
279- 279+ 
280- 280+ 
281-class TestNPUGraphNode:281+class TestNPUGraphNode:
282- def tearDown(self):282+ def tearDown(self):
283- torch_npu._C._npu_endAllocateCurrentStreamToPool(0, (0, 0))283+ torch_npu._C._npu_endAllocateCurrentStreamToPool(0, (0, 0))
284- torch_npu._C._npu_releasePool(0, (0, 0))284+ torch_npu._C._npu_releasePool(0, (0, 0))
285- 285+ 
286- def test_initialization(self, mock_wrapped_function, mock_parent_node):286+ def test_initialization(self, mock_wrapped_function, mock_parent_node):
287- inputs = [torch.tensor([1.0], device="npu")]287+ inputs = [torch.tensor([1.0], device="npu")]
288- with patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager"), patch(288+ with patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager"), patch(
289- "torch_npu.npu._graph_tree.check_memory_pool"289+ "torch_npu.npu._graph_tree.check_memory_pool"
290- ), patch("torch_npu._C._npu_getCheckpointState"):290+ ), patch("torch_npu._C._npu_getCheckpointState"):
291- node = NPUGraphNode(291+ node = NPUGraphNode(
292- wrapped_function=mock_wrapped_function,292+ wrapped_function=mock_wrapped_function,
293- graph_id=1,293+ graph_id=1,
294- parent=mock_parent_node,294+ parent=mock_parent_node,
295- inputs=inputs,295+ inputs=inputs,
296- npu_graphs_pool=(0, 0),296+ npu_graphs_pool=(0, 0),
297- device_index=0,297+ device_index=0,
298- stack_traces=None,298+ stack_traces=None,
299- stream=torch.npu.Stream(),299+ stream=torch.npu.Stream(),
300- )300+ )
301- 301+ 
302- assert node.id == 1302+ assert node.id == 1
303- assert node.device == 0303+ assert node.device == 0
304- assert node.parent == mock_parent_node304+ assert node.parent == mock_parent_node
305- assert node.graph is not None305+ assert node.graph is not None
306- 306+ 
307- def test_invalid_input_type(self, mock_wrapped_function):307+ def test_invalid_input_type(self, mock_wrapped_function):
308- with pytest.raises(RuntimeError):308+ with pytest.raises(RuntimeError):
309- NPUGraphNode(309+ NPUGraphNode(
310- wrapped_function=mock_wrapped_function,310+ wrapped_function=mock_wrapped_function,
311- graph_id=1,311+ graph_id=1,
312- parent=None,312+ parent=None,
313- inputs="not_a_list",313+ inputs="not_a_list",
314- npu_graphs_pool=(0, 0),314+ npu_graphs_pool=(0, 0),
315- device_index=0,315+ device_index=0,
316- stack_traces=None,316+ stack_traces=None,
317- stream=torch.npu.Stream(),317+ stream=torch.npu.Stream(),
318- )318+ )
319- 319+ 
320- @patch("torch_npu.npu._graph_tree.check_memory_pool")320+ @patch("torch_npu.npu._graph_tree.check_memory_pool")
321- def test_record_method(self, mock_check, basic_npu_graph_node):321+ def test_record_method(self, mock_check, basic_npu_graph_node):
322- def model_side_effect(inputs):322+ def model_side_effect(inputs):
323- # Clear inputs list while preserving reference323+ # Clear inputs list while preserving reference
324- inputs[:] = []324+ inputs[:] = []
325- return []325+ return []
326- 326+ 
327- mock_model = MagicMock(side_effect=model_side_effect)327+ mock_model = MagicMock(side_effect=model_side_effect)
328- mock_inputs = [torch.tensor([1.0], device="npu")]328+ mock_inputs = [torch.tensor([1.0], device="npu")]
329- 329+ 
330- with patch("torch_npu.npu._graph_tree.clear_cublas_manager"), patch(330+ with patch("torch_npu.npu._graph_tree.clear_cublas_manager"), patch(
331- "torch_npu.npu._graph_tree.get_history_recording"331+ "torch_npu.npu._graph_tree.get_history_recording"
332- ), patch("torch_npu.npu.graphs.NPUGraph.capture_begin"), patch(332+ ), patch("torch_npu.npu.graphs.NPUGraph.capture_begin"), patch(
333- "torch_npu.npu.graphs.NPUGraph.capture_end"333+ "torch_npu.npu.graphs.NPUGraph.capture_end"
334- ), patch(334+ ), patch(
335- "torch_npu._C._npu_getCheckpointState"335+ "torch_npu._C._npu_getCheckpointState"
336- ), patch(336+ ), patch(
337- "torch._dynamo.utils.preserve_rng_state"337+ "torch._dynamo.utils.preserve_rng_state"
338- ):338+ ):
339- 339+ 
340- outputs = basic_npu_graph_node._record(mock_model, mock_inputs)340+ outputs = basic_npu_graph_node._record(mock_model, mock_inputs)
341- 341+ 
342- mock_model.assert_called_once_with(mock_inputs)342+ mock_model.assert_called_once_with(mock_inputs)
343- assert basic_npu_graph_node.recording_outputs == outputs343+ assert basic_npu_graph_node.recording_outputs == outputs
344- 344+ 
345- def test_reconstruct_outputs(self, basic_npu_graph_node):345+ def test_reconstruct_outputs(self, basic_npu_graph_node):
346- # Setup mock metadata and storage info346+ # Setup mock metadata and storage info
347- basic_npu_graph_node.outputs_metadata = [347+ basic_npu_graph_node.outputs_metadata = [
348- {348+ {
349- "nbytes": 4,349+ "nbytes": 4,
350- "data_ptr": 1234,350+ "data_ptr": 1234,
351- "size": (1,),351+ "size": (1,),
352- "stride": (1,),352+ "stride": (1,),
353- "dtype": torch.float32,353+ "dtype": torch.float32,
354- "device": "npu",354+ "device": "npu",
355- "storage_offset": 0,355+ "storage_offset": 0,
356- }356+ }
357- ]357+ ]
358- basic_npu_graph_node.output_weakrefs = [MagicMock()]358+ basic_npu_graph_node.output_weakrefs = [MagicMock()]
359- basic_npu_graph_node.output_storage_alias = [UnaliasedStorage]359+ basic_npu_graph_node.output_storage_alias = [UnaliasedStorage]
360- basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]360+ basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]
361- 361+ 
362- with patch(362+ with patch(
363- "torch_npu._C._construct_NPU_Tensor_From_Storage_And_Metadata"363+ "torch_npu._C._construct_NPU_Tensor_From_Storage_And_Metadata"
364- ) as mock_construct:364+ ) as mock_construct:
365- outputs = basic_npu_graph_node.reconstruct_outputs()365+ outputs = basic_npu_graph_node.reconstruct_outputs()
366- assert len(outputs) == 1366+ assert len(outputs) == 1
367- 367+ 
368- def test_reconstruct_outputs_with_format(self, basic_npu_graph_node):368+ def test_reconstruct_outputs_with_format(self, basic_npu_graph_node):
369- # Setup mock metadata and storage info369+ # Setup mock metadata and storage info
370- basic_npu_graph_node.outputs_metadata = [370+ basic_npu_graph_node.outputs_metadata = [
371- {371+ {
372- "nbytes": 4,372+ "nbytes": 4,
373- "data_ptr": 1234,373+ "data_ptr": 1234,
374- "size": (1,),374+ "size": (1,),
375- "stride": (1,),375+ "stride": (1,),
376- "dtype": torch.float32,376+ "dtype": torch.float32,
377- "device": "npu",377+ "device": "npu",
378- "npu_format": 29,378+ "npu_format": 29,
379- "storage_offset": 0,379+ "storage_offset": 0,
380- }380+ }
381- ]381+ ]
382- basic_npu_graph_node.output_weakrefs = [MagicMock()]382+ basic_npu_graph_node.output_weakrefs = [MagicMock()]
383- basic_npu_graph_node.output_storage_alias = [UnaliasedStorage]383+ basic_npu_graph_node.output_storage_alias = [UnaliasedStorage]
384- basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]384+ basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]
385- 385+ 
386- with patch(386+ with patch(
387- "torch_npu._C._construct_NPU_Tensor_From_Storage_And_Metadata"387+ "torch_npu._C._construct_NPU_Tensor_From_Storage_And_Metadata"
388- ) as mock_construct:388+ ) as mock_construct:
389- outputs = basic_npu_graph_node.reconstruct_outputs()389+ outputs = basic_npu_graph_node.reconstruct_outputs()
390- assert len(outputs) == 1390+ assert len(outputs) == 1
391- 391+ 
392- def test_aliased_output_reconstruction(self, basic_npu_graph_node):392+ def test_aliased_output_reconstruction(self, basic_npu_graph_node):
393- basic_npu_graph_node.outputs_metadata = [393+ basic_npu_graph_node.outputs_metadata = [
394- {394+ {
395- "nbytes": 4,395+ "nbytes": 4,
396- "data_ptr": 1234,396+ "data_ptr": 1234,
397- "size": (1,),397+ "size": (1,),
398- "stride": (1,),398+ "stride": (1,),
399- "dtype": torch.float32,399+ "dtype": torch.float32,
400- "device": "npu",400+ "device": "npu",
401- "storage_offset": 0,401+ "storage_offset": 0,
402- }402+ }
403- ]403+ ]
404- basic_npu_graph_node.output_storage_alias = [AliasesPriorGraphOutput((0, 0))]404+ basic_npu_graph_node.output_storage_alias = [AliasesPriorGraphOutput((0, 0))]
405- basic_npu_graph_node.outputs_weakrefs = [MagicMock()]405+ basic_npu_graph_node.outputs_weakrefs = [MagicMock()]
406- basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]406+ basic_npu_graph_node.cached_tensor_outputs = [MagicMock()]
407- 407+ 
408- with patch("torch_npu.npu._graph_tree.maybe_deref") as mock_maybe_deref:408+ with patch("torch_npu.npu._graph_tree.maybe_deref") as mock_maybe_deref:
409- mock_maybe_deref.return_value = (MagicMock(), 1234)409+ mock_maybe_deref.return_value = (MagicMock(), 1234)
410- outputs = basic_npu_graph_node.reconstruct_outputs()410+ outputs = basic_npu_graph_node.reconstruct_outputs()
411- assert len(outputs) == 1411+ assert len(outputs) == 1
412- 412+ 
413- def test_liveness_tracking(self, basic_npu_graph_node):413+ def test_liveness_tracking(self, basic_npu_graph_node):
414- mock_ref = MagicMock()414+ mock_ref = MagicMock()
415- basic_npu_graph_node.path_weakrefs = [[mock_ref]]415+ basic_npu_graph_node.path_weakrefs = [[mock_ref]]
416- 416+ 
417- with patch("torch_npu.npu._graph_tree.is_live") as mock_is_live:417+ with patch("torch_npu.npu._graph_tree.is_live") as mock_is_live:
418- mock_is_live.return_value = True418+ mock_is_live.return_value = True
419- liveness = basic_npu_graph_node._get_liveness(419+ liveness = basic_npu_graph_node._get_liveness(
420- basic_npu_graph_node.path_weakrefs420+ basic_npu_graph_node.path_weakrefs
421- )421+ )
422- assert liveness == [[True]]422+ assert liveness == [[True]]
423- 423+ 
424- def test_child_management(self, basic_npu_graph_node):424+ def test_child_management(self, basic_npu_graph_node):
425- mock_child = MagicMock()425+ mock_child = MagicMock()
426- basic_npu_graph_node.add_child("test_func", mock_child)426+ basic_npu_graph_node.add_child("test_func", mock_child)
427- assert "test_func" in basic_npu_graph_node.children427+ assert "test_func" in basic_npu_graph_node.children
428- assert mock_child in basic_npu_graph_node.children["test_func"]428+ assert mock_child in basic_npu_graph_node.children["test_func"]
429- 429+ 
430- def test_invalid_run_conditions(self, basic_npu_graph_node):430+ def test_invalid_run_conditions(self, basic_npu_graph_node):
431- basic_npu_graph_node.graph = None431+ basic_npu_graph_node.graph = None
432- with pytest.raises(RuntimeError):432+ with pytest.raises(RuntimeError):
433- basic_npu_graph_node.run_graph()433+ basic_npu_graph_node.run_graph()
434- 434+ 
435- def test_storage_metadata_handling(self, basic_npu_graph_node):435+ def test_storage_metadata_handling(self, basic_npu_graph_node):
436- tensor = torch.tensor([1.0], device="npu")436+ tensor = torch.tensor([1.0], device="npu")
437- metadata = basic_npu_graph_node._tensor_metadata(tensor)437+ metadata = basic_npu_graph_node._tensor_metadata(tensor)
438- 438+ 
439- assert metadata["data_ptr"] == tensor.untyped_storage().data_ptr()439+ assert metadata["data_ptr"] == tensor.untyped_storage().data_ptr()
440- assert metadata["size"] == tensor.shape440+ assert metadata["size"] == tensor.shape
441- 441+ 
442- @patch("torch.npu.synchronize")442+ @patch("torch.npu.synchronize")
443- @patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager")443+ @patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager")
444- def test_input_processing(self, mock_pool_manager, mock_sync, basic_npu_graph_node):444+ def test_input_processing(self, mock_pool_manager, mock_sync, basic_npu_graph_node):
445- inputs = [torch.tensor([1.0], device="npu")]445+ inputs = [torch.tensor([1.0], device="npu")]
446- processed = basic_npu_graph_node._allocate_and_copy_recording_inputs(inputs)446+ processed = basic_npu_graph_node._allocate_and_copy_recording_inputs(inputs)
447- assert len(processed) == 1447+ assert len(processed) == 1
448- assert isinstance(processed[0], torch.Tensor)448+ assert isinstance(processed[0], torch.Tensor)
449- 449+ 
450- def test_check_invariants(self, basic_npu_graph_node):450+ def test_check_invariants(self, basic_npu_graph_node):
451- mock_inputs = [torch.tensor([1.0], device="npu")]451+ mock_inputs = [torch.tensor([1.0], device="npu")]
452- basic_npu_graph_node.static_input_data_ptrs = [mock_inputs[0].data_ptr()]452+ basic_npu_graph_node.static_input_data_ptrs = [mock_inputs[0].data_ptr()]
453- basic_npu_graph_node.npugraph_managed_idxs = [0]453+ basic_npu_graph_node.npugraph_managed_idxs = [0]
454- 454+ 
455- assert basic_npu_graph_node.check_invariants(mock_inputs)455+ assert basic_npu_graph_node.check_invariants(mock_inputs)
456- 456+ 
457- def test_descendant_count(self, basic_npu_graph_node):457+ def test_descendant_count(self, basic_npu_graph_node):
458- mock_child = MagicMock(num_descendants=lambda: 0)458+ mock_child = MagicMock(num_descendants=lambda: 0)
459- basic_npu_graph_node.children["test"] = [mock_child]459+ basic_npu_graph_node.children["test"] = [mock_child]
460- assert basic_npu_graph_node.num_descendants() == 1460+ assert basic_npu_graph_node.num_descendants() == 1
461- 461+ 
462- def test_prepare_alias_info_metadata_int(self, basic_npu_graph_node):462+ def test_prepare_alias_info_metadata_int(self, basic_npu_graph_node):
463- result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(463+ result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
464- MagicMock(), 42464+ MagicMock(), 42
465- )465+ )
466- assert result is None466+ assert result is None
467- 467+ 
468- def test_prepare_alias_info_unaliased_storage(self, basic_npu_graph_node):468+ def test_prepare_alias_info_unaliased_storage(self, basic_npu_graph_node):
469- result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(469+ result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
470- UnaliasedStorage, {"meta": "data"}470+ UnaliasedStorage, {"meta": "data"}
471- )471+ )
472- assert result is None472+ assert result is None
473- 473+ 
474- def test_prepare_alias_info_aliases_prior_graph_valid(self, basic_npu_graph_node):474+ def test_prepare_alias_info_aliases_prior_graph_valid(self, basic_npu_graph_node):
475- mock_ref = MagicMock()475+ mock_ref = MagicMock()
476- basic_npu_graph_node.path_weakrefs = [[mock_ref, mock_ref]]476+ basic_npu_graph_node.path_weakrefs = [[mock_ref, mock_ref]]
477- alias_info = AliasesPriorGraphOutput((0, 1))477+ alias_info = AliasesPriorGraphOutput((0, 1))
478- 478+ 
479- with patch("torch.UntypedStorage._new_with_weak_ptr") as mock_new:479+ with patch("torch.UntypedStorage._new_with_weak_ptr") as mock_new:
480- result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(480+ result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
481- alias_info, {"meta": "data"}481+ alias_info, {"meta": "data"}
482- )482+ )
483- mock_new.assert_called_once_with(mock_ref())483+ mock_new.assert_called_once_with(mock_ref())
484- assert result == mock_new.return_value484+ assert result == mock_new.return_value
485- 485+ 
486- def test_prepare_alias_info_aliases_prior_graph_none_ref(486+ def test_prepare_alias_info_aliases_prior_graph_none_ref(
487- self, basic_npu_graph_node487+ self, basic_npu_graph_node
488- ):488+ ):
489- basic_npu_graph_node.path_weakrefs = [[None, None]]489+ basic_npu_graph_node.path_weakrefs = [[None, None]]
490- alias_info = AliasesPriorGraphOutput((0, 1))490+ alias_info = AliasesPriorGraphOutput((0, 1))
491- 491+ 
492- with pytest.raises(RuntimeError):492+ with pytest.raises(RuntimeError):
493- basic_npu_graph_node.prepare_alias_info_for_tensor_construction(493+ basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
494- alias_info, {"meta": "data"}494+ alias_info, {"meta": "data"}
495- )495+ )
496- 496+ 
497- def test_prepare_alias_info_aliases_new_output(self, basic_npu_graph_node):497+ def test_prepare_alias_info_aliases_new_output(self, basic_npu_graph_node):
498- alias_info = AliasesNewOutput(123)498+ alias_info = AliasesNewOutput(123)
499- result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(499+ result = basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
500- alias_info, {"meta": "data"}500+ alias_info, {"meta": "data"}
501- )501+ )
502- assert result == 123502+ assert result == 123
503- 503+ 
504- def test_prepare_alias_info_invalid_type(self, basic_npu_graph_node):504+ def test_prepare_alias_info_invalid_type(self, basic_npu_graph_node):
505- with pytest.raises(RuntimeError):505+ with pytest.raises(RuntimeError):
506- basic_npu_graph_node.prepare_alias_info_for_tensor_construction(506+ basic_npu_graph_node.prepare_alias_info_for_tensor_construction(
507- "invalid_type", {"meta": "data"}507+ "invalid_type", {"meta": "data"}
508- )508+ )
509- 509+ 
510- # Tests for prepare_storages_for_construction510+ # Tests for prepare_storages_for_construction
511- def test_prepare_storages_mixed_aliases(self, basic_npu_graph_node):511+ def test_prepare_storages_mixed_aliases(self, basic_npu_graph_node):
512- basic_npu_graph_node.output_storage_alias = [512+ basic_npu_graph_node.output_storage_alias = [
513- UnaliasedStorage,513+ UnaliasedStorage,
514- AliasesNewOutput(123),514+ AliasesNewOutput(123),
515- AliasesPriorGraphOutput((0, 1)),515+ AliasesPriorGraphOutput((0, 1)),
516- ]516+ ]
517- basic_npu_graph_node.outputs_metadata = [None, {}, {}]517+ basic_npu_graph_node.outputs_metadata = [None, {}, {}]
518- basic_npu_graph_node.path_weakrefs = [[None, MagicMock(), MagicMock()]]518+ basic_npu_graph_node.path_weakrefs = [[None, MagicMock(), MagicMock()]]
519- 519+ 
520- with patch("torch.UntypedStorage._new_with_weak_ptr"):520+ with patch("torch.UntypedStorage._new_with_weak_ptr"):
521- results = basic_npu_graph_node.prepare_storages_for_construction()521+ results = basic_npu_graph_node.prepare_storages_for_construction()
522- 522+ 
523- assert len(results) == 3523+ assert len(results) == 3
524- assert results[0] is None524+ assert results[0] is None
525- assert results[1] == 123525+ assert results[1] == 123
526- 526+ 
527- # Tests for debug_assert_invariants527+ # Tests for debug_assert_invariants
528- def test_debug_assert_invariants_valid(self, basic_npu_graph_node):528+ def test_debug_assert_invariants_valid(self, basic_npu_graph_node):
529- from torch._inductor import config529+ from torch._inductor import config
530- 530+ 
531- config.triton.fast_path_cudagraph_asserts = True531+ config.triton.fast_path_cudagraph_asserts = True
532- expected_liveness = [[], [True, False]]532+ expected_liveness = [[], [True, False]]
533- newly_dead = [(1, 1)]533+ newly_dead = [(1, 1)]
534- ref = MagicMock(return_value=None)534+ ref = MagicMock(return_value=None)
535- basic_npu_graph_node.outputs_weakrefs = [None, ref]535+ basic_npu_graph_node.outputs_weakrefs = [None, ref]
536- basic_npu_graph_node.parent.outputs_weakrefs = []536+ basic_npu_graph_node.parent.outputs_weakrefs = []
537- basic_npu_graph_node.path_weakrefs = [537+ basic_npu_graph_node.path_weakrefs = [
538- basic_npu_graph_node.parent.outputs_weakrefs,538+ basic_npu_graph_node.parent.outputs_weakrefs,
539- basic_npu_graph_node.outputs_weakrefs,539+ basic_npu_graph_node.outputs_weakrefs,
540- ]540+ ]
541- 541+ 
542- # Should not raise542+ # Should not raise
543- with patch("torch_npu.npu._graph_tree.get_block_addrs"):543+ with patch("torch_npu.npu._graph_tree.get_block_addrs"):
544- basic_npu_graph_node.debug_assert_invariants(expected_liveness, newly_dead)544+ basic_npu_graph_node.debug_assert_invariants(expected_liveness, newly_dead)
545- config.triton.fast_path_cudagraph_asserts = False545+ config.triton.fast_path_cudagraph_asserts = False
546- 546+ 
547- def test_debug_assert_invariants_dead_ref_alive(self, basic_npu_graph_node):547+ def test_debug_assert_invariants_dead_ref_alive(self, basic_npu_graph_node):
548- from torch._inductor import config548+ from torch._inductor import config
549- 549+ 
550- config.triton.fast_path_cudagraph_asserts = True550+ config.triton.fast_path_cudagraph_asserts = True
551- expected_liveness = [[False]]551+ expected_liveness = [[False]]
552- newly_dead = [(0, 0)]552+ newly_dead = [(0, 0)]
553- basic_npu_graph_node.path_weakrefs = [553+ basic_npu_graph_node.path_weakrefs = [
554- [MagicMock(return_value=("ptr", 123))]554+ [MagicMock(return_value=("ptr", 123))]
555- ] # Live ref555+ ] # Live ref
556- 556+ 
557- with pytest.raises(RuntimeError):557+ with pytest.raises(RuntimeError):
558- basic_npu_graph_node.debug_assert_invariants(expected_liveness, newly_dead)558+ basic_npu_graph_node.debug_assert_invariants(expected_liveness, newly_dead)
559- config.triton.fast_path_cudagraph_asserts = False559+ config.triton.fast_path_cudagraph_asserts = False
560- 560+ 
561- # Tests for _initialize_cached_tensors561+ # Tests for _initialize_cached_tensors
562- def test_initialize_cached_tensors_valid(self, basic_npu_graph_node):562+ def test_initialize_cached_tensors_valid(self, basic_npu_graph_node):
563- basic_npu_graph_node.output_storage_alias = [UnaliasedStorage, UnaliasedStorage]563+ basic_npu_graph_node.output_storage_alias = [UnaliasedStorage, UnaliasedStorage]
564- basic_npu_graph_node.outputs_metadata = [564+ basic_npu_graph_node.outputs_metadata = [
565- {"dtype": torch.float},565+ {"dtype": torch.float},
566- {"dtype": torch.int},566+ {"dtype": torch.int},
567- ]567+ ]
568- basic_npu_graph_node.unaliased_in_all_paths = [True, False]568+ basic_npu_graph_node.unaliased_in_all_paths = [True, False]
569- basic_npu_graph_node.outputs_weakrefs = [None, None]569+ basic_npu_graph_node.outputs_weakrefs = [None, None]
570- 570+ 
571- with patch.object(basic_npu_graph_node, "create_storage"), patch(571+ with patch.object(basic_npu_graph_node, "create_storage"), patch(
572- "torch_npu._C._add_cached_tensor"572+ "torch_npu._C._add_cached_tensor"
573- ), patch.object(573+ ), patch.object(
574- basic_npu_graph_node, "_reconstruct_from_tensor_metadata"574+ basic_npu_graph_node, "_reconstruct_from_tensor_metadata"
575- ) as mock_reconstruct:575+ ) as mock_reconstruct:
576- 576+ 
577- mock_reconstruct.return_value = torch.tensor([1.0], device="npu:0")577+ mock_reconstruct.return_value = torch.tensor([1.0], device="npu:0")
578- basic_npu_graph_node._initialize_cached_tensors()578+ basic_npu_graph_node._initialize_cached_tensors()
579- 579+ 
580- assert len(basic_npu_graph_node.cached_tensor_outputs) == 2580+ assert len(basic_npu_graph_node.cached_tensor_outputs) == 2
581- assert basic_npu_graph_node.cached_tensor_outputs[0] is not None581+ assert basic_npu_graph_node.cached_tensor_outputs[0] is not None
582- assert len(basic_npu_graph_node.outputs_weakrefs) == 2582+ assert len(basic_npu_graph_node.outputs_weakrefs) == 2
583- 583+ 
584- def test_initialize_cached_tensors_invalid_storage_info(self, basic_npu_graph_node):584+ def test_initialize_cached_tensors_invalid_storage_info(self, basic_npu_graph_node):
585- basic_npu_graph_node.output_storage_alias = ["invalid"]585+ basic_npu_graph_node.output_storage_alias = ["invalid"]
586- basic_npu_graph_node.unaliased_in_all_paths = [True]586+ basic_npu_graph_node.unaliased_in_all_paths = [True]
587- 587+ 
588- basic_npu_graph_node._initialize_cached_tensors()588+ basic_npu_graph_node._initialize_cached_tensors()
589- 589+ 
590- 590+ 
591-@patch("torch_npu.npu.graphs.NPUGraph.replay")591+@patch("torch_npu.npu.graphs.NPUGraph.replay")
592-@patch("torch_npu.npu._graph_tree.check_memory_pool")592+@patch("torch_npu.npu._graph_tree.check_memory_pool")
593-@patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager")593+@patch("torch_npu.npu._graph_tree._use_npu_memory_pool_manager")
594-class TestNPUGraphNodeRun(TestCase):594+class TestNPUGraphNodeRun(TestCase):
595- def setUp(self):595+ def setUp(self):
596- """Initialize common test components and configurations"""596+ """Initialize common test components and configurations"""
597- self.device = "npu:0"597+ self.device = "npu:0"
598- 598+ 
599- def model_side_effect(inputs):599+ def model_side_effect(inputs):
600- # Clear inputs list while preserving reference600+ # Clear inputs list while preserving reference
601- inputs[:] = []601+ inputs[:] = []
602- return []602+ return []
603- 603+ 
604- self.wrapped_function = MagicMock(604+ self.wrapped_function = MagicMock(
605- spec=WrappedFunction,605+ spec=WrappedFunction,
606- static_input_idxs=[0],606+ static_input_idxs=[0],
607- constants=[],607+ constants=[],
608- model=MagicMock(side_effect=model_side_effect),608+ model=MagicMock(side_effect=model_side_effect),
609- )609+ )
610- self.graph_id = 1610+ self.graph_id = 1
611- self.npu_graphs_pool = (0, 0)611+ self.npu_graphs_pool = (0, 0)
612- self.stream = torch.npu.Stream(device=self.device)612+ self.stream = torch.npu.Stream(device=self.device)
613- 613+ 
614- # Create test tensors614+ # Create test tensors
615- self.static_input = torch.randn(615+ self.static_input = torch.randn(
616- 3, 3, device=self.device616+ 3, 3, device=self.device
617- ) # Static input (parameter-like)617+ ) # Static input (parameter-like)
618- self.dynamic_input = torch.randn(2, 2, device=self.device) # Dynamic input618+ self.dynamic_input = torch.randn(2, 2, device=self.device) # Dynamic input
619- 619+ 
620- def _create_node(self, inputs, parent=None):620+ def _create_node(self, inputs, parent=None):
621- """Helper to create NPUGraphNode instance"""621+ """Helper to create NPUGraphNode instance"""
622- with patch("torch_npu._C._npu_getCheckpointState"), patch(622+ with patch("torch_npu._C._npu_getCheckpointState"), patch(
623- "torch_npu.npu.graphs.NPUGraph.capture_begin"623+ "torch_npu.npu.graphs.NPUGraph.capture_begin"
624- ), patch("torch_npu.npu.graphs.NPUGraph.capture_end"):624+ ), patch("torch_npu.npu.graphs.NPUGraph.capture_end"):
625- return NPUGraphNode(625+ return NPUGraphNode(
626- wrapped_function=self.wrapped_function,626+ wrapped_function=self.wrapped_function,
627- graph_id=self.graph_id,627+ graph_id=self.graph_id,
628- parent=parent,628+ parent=parent,
629- inputs=inputs,629+ inputs=inputs,
630- npu_graphs_pool=self.npu_graphs_pool,630+ npu_graphs_pool=self.npu_graphs_pool,
631- device_index=0,631+ device_index=0,
632- stack_traces=None,632+ stack_traces=None,
633- stream=self.stream,633+ stream=self.stream,
634- )634+ )
635- 635+ 
636- @patch.object(NPUGraphNode, "run_graph")636+ @patch.object(NPUGraphNode, "run_graph")
637- def test_static_input_optimization(637+ def test_static_input_optimization(
638- self, mock_run_graph, mock_pool, mock_check, mock_replay638+ self, mock_run_graph, mock_pool, mock_check, mock_replay
639- ):639+ ):
640- """Verify static inputs bypass copy operations"""640+ """Verify static inputs bypass copy operations"""
641- # Mark all inputs as static641+ # Mark all inputs as static
642- self.wrapped_function.static_input_idxs = [0, 1]642+ self.wrapped_function.static_input_idxs = [0, 1]
643- node = self._create_node([self.static_input, self.static_input.clone()])643+ node = self._create_node([self.static_input, self.static_input.clone()])
644- 644+ 
645- # Execute with cloned inputs645+ # Execute with cloned inputs
646- node.run([self.static_input.clone(), self.static_input.clone()])646+ node.run([self.static_input.clone(), self.static_input.clone()])
647- 647+ 
648- # Validate no copy operations occurred648+ # Validate no copy operations occurred
649- self.assertEqual(mock_run_graph.call_count, 1)649+ self.assertEqual(mock_run_graph.call_count, 1)
650- 650+ 
651- @patch.object(NPUGraphNode, "reconstruct_outputs")651+ @patch.object(NPUGraphNode, "reconstruct_outputs")
652- def test_output_reconstruction_flow(652+ def test_output_reconstruction_flow(
653- self, mock_reconstruct, mock_pool, mock_check, mock_replay653+ self, mock_reconstruct, mock_pool, mock_check, mock_replay
654- ):654+ ):
655- """Test full output reconstruction pipeline"""655+ """Test full output reconstruction pipeline"""
656- # Configure mock reconstruction656+ # Configure mock reconstruction
657- expected_output = torch.tensor([1.0], device=self.device)657+ expected_output = torch.tensor([1.0], device=self.device)
658- mock_reconstruct.return_value = [expected_output]658+ mock_reconstruct.return_value = [expected_output]
659- 659+ 
660- node = self._create_node([self.static_input])660+ node = self._create_node([self.static_input])
661- outputs = node.run([self.static_input.clone()])661+ outputs = node.run([self.static_input.clone()])
662- 662+ 
663- # Validate outputs663+ # Validate outputs
664- self.assertEqual(outputs, [expected_output])664+ self.assertEqual(outputs, [expected_output])
665- mock_reconstruct.assert_called_once()665+ mock_reconstruct.assert_called_once()
666- 666+ 
667- @patch("torch._foreach_copy_")667+ @patch("torch._foreach_copy_")
668- def test_batched_copy_optimization(668+ def test_batched_copy_optimization(
669- self, mock_batched_copy, mock_pool, mock_check, mock_replay669+ self, mock_batched_copy, mock_pool, mock_check, mock_replay
670- ):670+ ):
671- """Verify batched copy operations for efficiency"""671+ """Verify batched copy operations for efficiency"""
672- # Configure multiple dynamic inputs672+ # Configure multiple dynamic inputs
673- self.wrapped_function.static_input_idxs = []673+ self.wrapped_function.static_input_idxs = []
674- inputs = [torch.randn(2, 2, device=self.device) for _ in range(3)]674+ inputs = [torch.randn(2, 2, device=self.device) for _ in range(3)]
675- new_inputs = [t.clone() for t in inputs]675+ new_inputs = [t.clone() for t in inputs]
676- node = self._create_node(inputs)676+ node = self._create_node(inputs)
677- 677+ 
678- # Execute with new inputs678+ # Execute with new inputs
679- node.run(new_inputs)679+ node.run(new_inputs)
680- 680+ 
681- # Validate single batched copy call681+ # Validate single batched copy call
682- args, _ = mock_batched_copy.call_args682+ args, _ = mock_batched_copy.call_args
683- self.assertEqual(len(args[0]), 3)683+ self.assertEqual(len(args[0]), 3)
684- 684+ 
685- def test_memory_cleanup_after_execution(self, mock_pool, mock_check, mock_replay):685+ def test_memory_cleanup_after_execution(self, mock_pool, mock_check, mock_replay):
686- """Validate input list cleanup post-execution"""686+ """Validate input list cleanup post-execution"""
687- initial_inputs = [self.static_input.clone(), self.dynamic_input.clone()]687+ initial_inputs = [self.static_input.clone(), self.dynamic_input.clone()]
688- input_copy = [t.clone() for t in initial_inputs]688+ input_copy = [t.clone() for t in initial_inputs]
689- node = self._create_node(initial_inputs)689+ node = self._create_node(initial_inputs)
690- 690+ 
691- # Execute and verify cleanup691+ # Execute and verify cleanup
692- node.run(input_copy)692+ node.run(input_copy)
693- self.assertEqual(len(input_copy), 0)693+ self.assertEqual(len(input_copy), 0)
694- 694+ 
695- 695+ 
696-class TestGetNpugraphSegments(TestCase):696+class TestGetNpugraphSegments(TestCase):
697- @patch('torch.npu.memory_snapshot') 697+ @patch('torch.npu.memory_snapshot')
698- def test_get_npugraph_segments(self, mock_snapshot): 698+ def test_get_npugraph_segments(self, mock_snapshot):
699- mock_snapshot.return_value = [699+ mock_snapshot.return_value = [
700- {"segment_pool_id": (0, 1), "address": 1000, "blocks": []},700+ {"segment_pool_id": (0, 1), "address": 1000, "blocks": []},
701- {"segment_pool_id": (0, 0), "address": 2000, "blocks": []},701+ {"segment_pool_id": (0, 0), "address": 2000, "blocks": []},
702- {"segment_pool_id": (0, 1), "address": 3000, "blocks": []},702+ {"segment_pool_id": (0, 1), "address": 3000, "blocks": []},
703- ] 703+ ]
704- result = get_npugraph_segments((0, 1)) 704+ result = get_npugraph_segments((0, 1))
705- self.assertEqual(len(result), 2) 705+ self.assertEqual(len(result), 2)
706- mock_snapshot.assert_called_once_with()706+ mock_snapshot.assert_called_once_with()
707- 707+ 
708- 708+ 
709-class TestGetBlockAddrs(TestCase):709+class TestGetBlockAddrs(TestCase):
710- @patch('torch_npu.npu._graph_tree.get_npugraph_segments')710+ @patch('torch_npu.npu._graph_tree.get_npugraph_segments')
711- def test_get_block_addrs_live_only(self, mock_segments):711+ def test_get_block_addrs_live_only(self, mock_segments):
712- mock_segments.return_value = [712+ mock_segments.return_value = [
713- {713+ {
714- "segment_pool_id": (0, 0),714+ "segment_pool_id": (0, 0),
715- "address": 1000,715+ "address": 1000,
716- "blocks": [716+ "blocks": [
717- {"state": "active_allocated", "size": 100},717+ {"state": "active_allocated", "size": 100},
718- {"state": "inactivate", "size": 200},718+ {"state": "inactivate", "size": 200},
719- {"state": "active_allocated", "size": 300},719+ {"state": "active_allocated", "size": 300},
720- ]720+ ]
721- },721+ },
722- {722+ {
723- "segment_pool_id": (0, 0),723+ "segment_pool_id": (0, 0),
724- "address": 2000,724+ "address": 2000,
725- "blocks": [725+ "blocks": [
726- {"state": "active_allocated", "size": 50},726+ {"state": "active_allocated", "size": 50},
727- {"state": "inactivate", "size": 150},727+ {"state": "inactivate", "size": 150},
728- ]728+ ]
729- }729+ }
730- ]730+ ]
731- result = get_block_addrs((0, 0), live_only=True)731+ result = get_block_addrs((0, 0), live_only=True)
732- self.assertEqual(result, [1000, 1300, 2000])732+ self.assertEqual(result, [1000, 1300, 2000])
733- mock_segments.assert_called_once_with((0, 0))733+ mock_segments.assert_called_once_with((0, 0))
734- 734+ 
735- @patch('torch_npu.npu._graph_tree.get_npugraph_segments')735+ @patch('torch_npu.npu._graph_tree.get_npugraph_segments')
736- def test_get_block_addrs_all_blocks(self, mock_segments):736+ def test_get_block_addrs_all_blocks(self, mock_segments):
737- mock_segments.return_value = [737+ mock_segments.return_value = [
738- {738+ {
739- "segment_pool_id": (0, 0),739+ "segment_pool_id": (0, 0),
740- "address": 1000,740+ "address": 1000,
741- "blocks": [741+ "blocks": [
742- {"state": "active_allocated", "size": 100},742+ {"state": "active_allocated", "size": 100},
743- {"state": "inactivate", "size": 200},743+ {"state": "inactivate", "size": 200},
744- ]744+ ]
745- }745+ }
746- ]746+ ]
747- result = get_block_addrs((0, 0), live_only=False)747+ result = get_block_addrs((0, 0), live_only=False)
748- self.assertEqual(result, [1000, 1100])748+ self.assertEqual(result, [1000, 1100])
749- mock_segments.assert_called_once_with((0, 0))749+ mock_segments.assert_called_once_with((0, 0))
750- 750+ 
751- 751+ 
752-class TestFormatTb(TestCase):752+class TestFormatTb(TestCase):
753- def test_format_tb(self):753+ def test_format_tb(self):
754- frames = [754+ frames = [
755- {"filename": "/path/to/file.py", "line": 42, "name": "test_function"},755+ {"filename": "/path/to/file.py", "line": 42, "name": "test_function"},
756- {"filename": "/path/to/module.py", "line": 100, "name": "helper_method"},756+ {"filename": "/path/to/module.py", "line": 100, "name": "helper_method"},
757- ]757+ ]
758- result = format_tb(frames)758+ result = format_tb(frames)
759- self.assertIn("/path/to/file.py", result)759+ self.assertIn("/path/to/file.py", result)
760- self.assertIn("test_function", result)760+ self.assertIn("test_function", result)
761- self.assertIn("/path/to/module.py", result)761+ self.assertIn("/path/to/module.py", result)
762- self.assertIn("helper_method", result)762+ self.assertIn("helper_method", result)
763- self.assertIn("line 100", result)763+ self.assertIn("line 100", result)
764- 764+ 
765- 765+ 
766-class TestCheckMemoryPool(TestCase):766+class TestCheckMemoryPool(TestCase):
767- @patch('torch_npu._C._npu_checkPoolLiveAllocations')767+ @patch('torch_npu._C._npu_checkPoolLiveAllocations')
768- def test_check_memory_pool_fast_path_pass(self, mock_check):768+ def test_check_memory_pool_fast_path_pass(self, mock_check):
769- mock_check.return_value = True769+ mock_check.return_value = True
770- 770+ 
771- mock_storage1 = MagicMock(spec=StorageWeakRefWrapper)771+ mock_storage1 = MagicMock(spec=StorageWeakRefWrapper)
772- mock_storage1.data_ptr.return_value = 1001772+ mock_storage1.data_ptr.return_value = 1001
773- mock_storage1.return_value = True773+ mock_storage1.return_value = True
774- 774+ 
775- mock_storage2 = MagicMock(spec=StorageWeakRefWrapper)775+ mock_storage2 = MagicMock(spec=StorageWeakRefWrapper)
776- mock_storage2.data_ptr.return_value = 1002776+ mock_storage2.data_ptr.return_value = 1002
777- mock_storage2.return_value = True777+ mock_storage2.return_value = True
778- 778+ 
779- check_memory_pool("npu:0", (0, 0), [mock_storage1, mock_storage2])779+ check_memory_pool("npu:0", (0, 0), [mock_storage1, mock_storage2])
780- mock_check.assert_called_once_with(780+ mock_check.assert_called_once_with(
781- "npu:0", (0, 0), {1001, 1002}781+ "npu:0", (0, 0), {1001, 1002}
782- )782+ )
783- 783+ 
784- @patch('torch_npu._C._npu_checkPoolLiveAllocations')784+ @patch('torch_npu._C._npu_checkPoolLiveAllocations')
785- @patch('torch_npu.npu._graph_tree.get_npugraph_segments')785+ @patch('torch_npu.npu._graph_tree.get_npugraph_segments')
786- @patch('torch_npu.npu._graph_tree.format_tb')786+ @patch('torch_npu.npu._graph_tree.format_tb')
787- @patch('gc.collect')787+ @patch('gc.collect')
788- def test_check_memory_pool_slow_path_all_match(788+ def test_check_memory_pool_slow_path_all_match(
789- self, mock_gc, mock_format_tb, mock_segments, mock_check789+ self, mock_gc, mock_format_tb, mock_segments, mock_check
790- ):790+ ):
791- mock_check.return_value = False791+ mock_check.return_value = False
792- mock_segments.return_value = [792+ mock_segments.return_value = [
793- {793+ {
794- "segment_pool_id": (0, 0),794+ "segment_pool_id": (0, 0),
795- "address": 1000,795+ "address": 1000,
796- "blocks": [796+ "blocks": [
797- {"state": "active_allocated", "size": 100, "frames": []},797+ {"state": "active_allocated", "size": 100, "frames": []},
798- {"state": "inactivate", "size": 200},798+ {"state": "inactivate", "size": 200},
799- ]799+ ]
800- }800+ }
801- ]801+ ]
802- mock_storage = MagicMock(spec=StorageWeakRefWrapper)802+ mock_storage = MagicMock(spec=StorageWeakRefWrapper)
803- mock_storage.data_ptr.return_value = 1000803+ mock_storage.data_ptr.return_value = 1000
804- mock_storage.return_value = True804+ mock_storage.return_value = True
805- check_memory_pool("npu:0", (0, 0), [mock_storage])805+ check_memory_pool("npu:0", (0, 0), [mock_storage])
806- mock_gc.assert_called_once_with()806+ mock_gc.assert_called_once_with()
807- mock_segments.assert_called_once_with((0, 0))807+ mock_segments.assert_called_once_with((0, 0))
808- mock_format_tb.assert_not_called()808+ mock_format_tb.assert_not_called()
809- 809+ 
810- @patch('torch_npu._C._npu_checkPoolLiveAllocations')810+ @patch('torch_npu._C._npu_checkPoolLiveAllocations')
811- @patch('torch_npu.npu._graph_tree.get_npugraph_segments')811+ @patch('torch_npu.npu._graph_tree.get_npugraph_segments')
812- @patch('torch_npu.npu._graph_tree.format_tb')812+ @patch('torch_npu.npu._graph_tree.format_tb')
813- @patch('gc.collect')813+ @patch('gc.collect')
814- def test_check_memory_pool_slow_path_unallocated_storage(814+ def test_check_memory_pool_slow_path_unallocated_storage(
815- self, mock_gc, mock_format_tb, mock_segments, mock_check815+ self, mock_gc, mock_format_tb, mock_segments, mock_check
816- ):816+ ):
817- mock_check.return_value = False817+ mock_check.return_value = False
818- mock_segments.return_value = [818+ mock_segments.return_value = [
819- {819+ {
820- "segment_pool_id": (0, 0),820+ "segment_pool_id": (0, 0),
821- "address": 2000,821+ "address": 2000,
822- "blocks": [822+ "blocks": [
823- {"state": "active_allocated", "size": 100, "frames": []},823+ {"state": "active_allocated", "size": 100, "frames": []},
824- ]824+ ]
825- }825+ }
826- ]826+ ]
827- mock_storage = MagicMock(spec=StorageWeakRefWrapper)827+ mock_storage = MagicMock(spec=StorageWeakRefWrapper)
828- mock_storage.data_ptr.return_value = 1000828+ mock_storage.data_ptr.return_value = 1000
829- mock_storage.return_value = True829+ mock_storage.return_value = True
830- with self.assertRaisesRegex(830+ with self.assertRaisesRegex(
831- RuntimeError, r"These storage data ptrs are not allocated in pool \(0, 0\) but should be \{1000\}"831+ RuntimeError, r"These storage data ptrs are not allocated in pool \(0, 0\) but should be \{1000\}"
832- ):832+ ):
833- check_memory_pool("npu:0", (0, 0), [mock_storage])833+ check_memory_pool("npu:0", (0, 0), [mock_storage])
834- 834+ 
835- @patch('torch_npu._C._npu_checkPoolLiveAllocations')835+ @patch('torch_npu._C._npu_checkPoolLiveAllocations')
836- @patch('torch_npu.npu._graph_tree.get_npugraph_segments')836+ @patch('torch_npu.npu._graph_tree.get_npugraph_segments')
837- @patch('torch_npu.npu._graph_tree.format_tb')837+ @patch('torch_npu.npu._graph_tree.format_tb')
838- @patch('gc.collect')838+ @patch('gc.collect')
839- def test_check_memory_pool_slow_path_unaccounted_blocks(839+ def test_check_memory_pool_slow_path_unaccounted_blocks(
840- self, mock_gc, mock_format_tb, mock_segments, mock_check840+ self, mock_gc, mock_format_tb, mock_segments, mock_check
841- ):841+ ):
842- mock_check.return_value = False842+ mock_check.return_value = False
843- mock_segments.return_value = [843+ mock_segments.return_value = [
844- {844+ {
845- "segment_pool_id": (0, 0),845+ "segment_pool_id": (0, 0),
846- "address": 1000,846+ "address": 1000,
847- "blocks": [847+ "blocks": [
848- {"state": "active_allocated", "size": 100, "frames": [848+ {"state": "active_allocated", "size": 100, "frames": [
849- {"filename": "/path/to/file.py", "line": 42, "name": "allocate_func"}849+ {"filename": "/path/to/file.py", "line": 42, "name": "allocate_func"}
850- ]},850+ ]},
851- ]851+ ]
852- }852+ }
853- ]853+ ]
854- live_storages = []854+ live_storages = []
855- mock_format_tb.return_value = "Formatted Traceback"855+ mock_format_tb.return_value = "Formatted Traceback"
856- with self.assertRaisesRegex(856+ with self.assertRaisesRegex(
857- RuntimeError, "These live storage data ptrs are in the npugraph pool but not accounted for"857+ RuntimeError, "These live storage data ptrs are in the npugraph pool but not accounted for"
858- ):858+ ):
859- check_memory_pool("npu:0", (0, 0), live_storages)859+ check_memory_pool("npu:0", (0, 0), live_storages)
860- 860+ 
861- def test_check_memory_pool_invalid_input(self):861+ def test_check_memory_pool_invalid_input(self):
862- invalid_storages = [1, 2, 3]862+ invalid_storages = [1, 2, 3]
863- with self.assertRaisesRegex(863+ with self.assertRaisesRegex(
864- RuntimeError, r"check all\(isinstance\(elem, StorageWeakRefWrapper\) for elem in live_storages_ptrs\) fail"864+ RuntimeError, r"check all\(isinstance\(elem, StorageWeakRefWrapper\) for elem in live_storages_ptrs\) fail"
865- ):865+ ):
866- check_memory_pool("npu:0", (0, 0), invalid_storages)866+ check_memory_pool("npu:0", (0, 0), invalid_storages)
867- 867+ 
868- 868+ 
869-class TestNPUGraphTreeManager:869+class TestNPUGraphTreeManager:
870- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager._run')870+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager._run')
871- def test_run_forward_mode(self, mock_run):871+ def test_run_forward_mode(self, mock_run):
872- manager = NPUGraphTreeManager(0)872+ manager = NPUGraphTreeManager(0)
873- manager.id_to_mode[FunctionID(1)] = CompilationMode.FORWARD873+ manager.id_to_mode[FunctionID(1)] = CompilationMode.FORWARD
874- result = manager.run([torch.tensor([1.0])], FunctionID(1))874+ result = manager.run([torch.tensor([1.0])], FunctionID(1))
875- mock_run.assert_called_once_with([torch.tensor([1.0])], FunctionID(1))875+ mock_run.assert_called_once_with([torch.tensor([1.0])], FunctionID(1))
876- self.assertTrue(manager.running_forwards_with_pending_backwards)876+ self.assertTrue(manager.running_forwards_with_pending_backwards)
877- self.assertTrue(result == mock_run.return_value)877+ self.assertTrue(result == mock_run.return_value)
878- 878+ 
879- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager._run')879+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager._run')
880- def test_run_backward_mode(self, mock_run):880+ def test_run_backward_mode(self, mock_run):
881- manager = NPUGraphTreeManager(0)881+ manager = NPUGraphTreeManager(0)
882- manager.id_to_mode[FunctionID(1)] = CompilationMode.BACKWARD882+ manager.id_to_mode[FunctionID(1)] = CompilationMode.BACKWARD
883- result = manager.run([torch.tensor([1.0])], FunctionID(1))883+ result = manager.run([torch.tensor([1.0])], FunctionID(1))
884- mock_run.assert_called_once_with([torch.tensor([1.0])], FunctionID(1))884+ mock_run.assert_called_once_with([torch.tensor([1.0])], FunctionID(1))
885- self.assertFalse(manager.running_forwards_with_pending_backwards)885+ self.assertFalse(manager.running_forwards_with_pending_backwards)
886- self.assertTrue(result == mock_run.return_value)886+ self.assertTrue(result == mock_run.return_value)
887- 887+ 
888- def test_set_to_running_backward(self):888+ def test_set_to_running_backward(self):
889- manager = NPUGraphTreeManager(0)889+ manager = NPUGraphTreeManager(0)
890- manager.running_forwards_with_pending_backwards = True890+ manager.running_forwards_with_pending_backwards = True
891- manager.set_to_running_backward()891+ manager.set_to_running_backward()
892- self.assertFalse(manager.running_forwards_with_pending_backwards)892+ self.assertFalse(manager.running_forwards_with_pending_backwards)
893- 893+ 
894- def test_shutdown(self):894+ def test_shutdown(self):
895- manager = NPUGraphTreeManager(0)895+ manager = NPUGraphTreeManager(0)
896- mock_node1 = MagicMock()896+ mock_node1 = MagicMock()
897- mock_node2 = MagicMock()897+ mock_node2 = MagicMock()
898- mock_node3 = MagicMock()898+ mock_node3 = MagicMock()
899- manager.roots = {FunctionID(1): [mock_node1]}899+ manager.roots = {FunctionID(1): [mock_node1]}
900- mock_node1.children = {FunctionID(2): [mock_node2]}900+ mock_node1.children = {FunctionID(2): [mock_node2]}
901- mock_node2.children = {FunctionID(3): [mock_node3]}901+ mock_node2.children = {FunctionID(3): [mock_node3]}
902- manager.shutdown()902+ manager.shutdown()
903- mock_node1.remove_node_cached_tensors.assert_called_once_with()903+ mock_node1.remove_node_cached_tensors.assert_called_once_with()
904- mock_node2.remove_node_cached_tensors.assert_called_once_with()904+ mock_node2.remove_node_cached_tensors.assert_called_once_with()
905- mock_node3.remove_node_cached_tensors.assert_called_once_with()905+ mock_node3.remove_node_cached_tensors.assert_called_once_with()
906- assert mock_node1.graph is None906+ assert mock_node1.graph is None
907- assert mock_node2.graph is None907+ assert mock_node2.graph is None
908- assert mock_node3.graph is None908+ assert mock_node3.graph is None
909- assert manager.graph is None909+ assert manager.graph is None
910- assert manager.roots is None910+ assert manager.roots is None
911- assert manager.current_node is None911+ assert manager.current_node is None
912- 912+ 
913- @patch('torch.npu.synchronize')913+ @patch('torch.npu.synchronize')
914- @patch('torch_npu.npu._graph_tree.NPUGraphNode')914+ @patch('torch_npu.npu._graph_tree.NPUGraphNode')
915- def test_record_function(self, mock_node, mock_synchronize):915+ def test_record_function(self, mock_node, mock_synchronize):
916- manager = NPUGraphTreeManager(0)916+ manager = NPUGraphTreeManager(0)
917- manager.ids_to_funcs[FunctionID(1)] = MagicMock()917+ manager.ids_to_funcs[FunctionID(1)] = MagicMock()
918- manager.ids_to_stack_traces[FunctionID(1)] = "stack_trace"918+ manager.ids_to_stack_traces[FunctionID(1)] = "stack_trace"
919- manager.npu_graphs_thread_pool = "pool_handle"919+ manager.npu_graphs_thread_pool = "pool_handle"
920- manager.device_index = 0920+ manager.device_index = 0
921- manager.stream = MagicMock()921+ manager.stream = MagicMock()
922- 922+
923- # 设置模拟返回值923+ # 设置模拟返回值
924- mock_node_instance = MagicMock()924+ mock_node_instance = MagicMock()
925- mock_node.return_value = mock_node_instance925+ mock_node.return_value = mock_node_instance
926- mock_node_instance.run_first_inputs.return_value = [torch.tensor([1.0])]926+ mock_node_instance.run_first_inputs.return_value = [torch.tensor([1.0])]
927- 927+
928- # 执行测试928+ # 执行测试
929- result = manager.record_function([torch.tensor([1.0])], FunctionID(1))929+ result = manager.record_function([torch.tensor([1.0])], FunctionID(1))
930- 930+
931- # 验证调用931+ # 验证调用
932- mock_synchronize.assert_any_call()932+ mock_synchronize.assert_any_call()
933- mock_node.assert_called_once_with(933+ mock_node.assert_called_once_with(
934- manager.ids_to_funcs[FunctionID(1)],934+ manager.ids_to_funcs[FunctionID(1)],
935- ANY, # graph_id935+ ANY, # graph_id
936- None, # parent936+ None, # parent
937- [torch.tensor([1.0])],937+ [torch.tensor([1.0])],
938- "pool_handle",938+ "pool_handle",
939- 0,939+ 0,
940- "stack_trace",940+ "stack_trace",
941- manager.stream941+ manager.stream
942- )942+ )
943- assert isinstance(mock_node.call_args[0][1], GraphID)943+ assert isinstance(mock_node.call_args[0][1], GraphID)
944- assert manager.current_node == mock_node_instance944+ assert manager.current_node == mock_node_instance
945- assert manager.path_state == ExecutionState.RECORDING945+ assert manager.path_state == ExecutionState.RECORDING
946- assert result == [torch.tensor([1.0])]946+ assert result == [torch.tensor([1.0])]
947- 947+ 
948- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.update_generation')948+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.update_generation')
949- def test_execute_node(self, mock_update_gen):949+ def test_execute_node(self, mock_update_gen):
950- manager = NPUGraphTreeManager(0)950+ manager = NPUGraphTreeManager(0)
951- mock_node = MagicMock()951+ mock_node = MagicMock()
952- mock_node.run.return_value = [torch.tensor([1.0])]952+ mock_node.run.return_value = [torch.tensor([1.0])]
953- 953+
954- # 执行测试954+ # 执行测试
955- result = manager.execute_node(mock_node, [torch.tensor([1.0])])955+ result = manager.execute_node(mock_node, [torch.tensor([1.0])])
956- 956+
957- # 验证调用957+ # 验证调用
958- mock_update_gen.assert_called_once_with()958+ mock_update_gen.assert_called_once_with()
959- assert manager.current_node == mock_node959+ assert manager.current_node == mock_node
960- assert manager.path_state == ExecutionState.EXECUTION960+ assert manager.path_state == ExecutionState.EXECUTION
961- assert result == [torch.tensor([1.0])]961+ assert result == [torch.tensor([1.0])]
962- 962+ 
963- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.update_generation')963+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.update_generation')
964- @patch('torch_npu.npu._graph_tree.NPUWarmupNode')964+ @patch('torch_npu.npu._graph_tree.NPUWarmupNode')
965- def test_run_eager(self, mock_warmup_node, mock_update_gen):965+ def test_run_eager(self, mock_warmup_node, mock_update_gen):
966- manager = NPUGraphTreeManager(0)966+ manager = NPUGraphTreeManager(0)
967- manager.ids_to_funcs[FunctionID(1)] = MagicMock()967+ manager.ids_to_funcs[FunctionID(1)] = MagicMock()
968- manager.ids_to_stack_traces[FunctionID(1)] = "stack_trace"968+ manager.ids_to_stack_traces[FunctionID(1)] = "stack_trace"
969- manager.npu_graphs_thread_pool = "pool_handle"969+ manager.npu_graphs_thread_pool = "pool_handle"
970- manager.graph = MagicMock()970+ manager.graph = MagicMock()
971- manager.device_index = 0971+ manager.device_index = 0
972- manager.stream = MagicMock()972+ manager.stream = MagicMock()
973- 973+
974- # 设置模拟返回值974+ # 设置模拟返回值
975- mock_node_instance = MagicMock()975+ mock_node_instance = MagicMock()
976- mock_warmup_node.return_value = mock_node_instance976+ mock_warmup_node.return_value = mock_node_instance
977- mock_node_instance.run.return_value = [torch.tensor([1.0])]977+ mock_node_instance.run.return_value = [torch.tensor([1.0])]
978- 978+
979- # 执行测试979+ # 执行测试
980- result = manager.run_eager([torch.tensor([1.0])], FunctionID(1))980+ result = manager.run_eager([torch.tensor([1.0])], FunctionID(1))
981- 981+
982- # 验证调用982+ # 验证调用
983- mock_update_gen.assert_called_once_with()983+ mock_update_gen.assert_called_once_with()
984- mock_warmup_node.assert_called_once_with(984+ mock_warmup_node.assert_called_once_with(
985- manager.ids_to_funcs[FunctionID(1)],985+ manager.ids_to_funcs[FunctionID(1)],
986- None,986+ None,
987- "pool_handle",987+ "pool_handle",
988- manager.graph,988+ manager.graph,
989- 0,989+ 0,
990- "stack_trace",990+ "stack_trace",
991- manager.stream,991+ manager.stream,
992- False,992+ False,
993- GraphID(-1),993+ GraphID(-1),
994- )994+ )
995- assert manager.current_node == mock_node_instance995+ assert manager.current_node == mock_node_instance
996- assert manager.path_state == ExecutionState.WARMUP996+ assert manager.path_state == ExecutionState.WARMUP
997- assert result == [torch.tensor([1.0])]997+ assert result == [torch.tensor([1.0])]
998- 998+ 
999- def test_new_graph_id(self):999+ def test_new_graph_id(self):
1000- manager = NPUGraphTreeManager(0)1000+ manager = NPUGraphTreeManager(0)
1001- id1 = manager.new_graph_id()1001+ id1 = manager.new_graph_id()
1002- id2 = manager.new_graph_id()1002+ id2 = manager.new_graph_id()
1003- assert isinstance(id1, GraphID)1003+ assert isinstance(id1, GraphID)
1004- assert isinstance(id2, GraphID)1004+ assert isinstance(id2, GraphID)
1005- assert id1 != id21005+ assert id1 != id2
1006- 1006+ 
1007- def test_new_func_id(self):1007+ def test_new_func_id(self):
1008- manager = NPUGraphTreeManager(0)1008+ manager = NPUGraphTreeManager(0)
1009- id1 = manager.new_func_id()1009+ id1 = manager.new_func_id()
1010- id2 = manager.new_func_id()1010+ id2 = manager.new_func_id()
1011- assert isinstance(id1, FunctionID)1011+ assert isinstance(id1, FunctionID)
1012- assert isinstance(id2, FunctionID)1012+ assert isinstance(id2, FunctionID)
1013- assert id1 != id21013+ assert id1 != id2
1014- 1014+ 
1015- def test_in_recording_property(self):1015+ def test_in_recording_property(self):
1016- manager = NPUGraphTreeManager(0)1016+ manager = NPUGraphTreeManager(0)
1017- manager.path_state = ExecutionState.NONE1017+ manager.path_state = ExecutionState.NONE
1018- assert manager.in_recording is False1018+ assert manager.in_recording is False
1019- manager.path_state = ExecutionState.RECORDING1019+ manager.path_state = ExecutionState.RECORDING
1020- assert manager.in_recording is True1020+ assert manager.in_recording is True
1021- 1021+ 
1022- def test_in_warmup_property(self):1022+ def test_in_warmup_property(self):
1023- manager = NPUGraphTreeManager(0)1023+ manager = NPUGraphTreeManager(0)
1024- manager.path_state = ExecutionState.NONE1024+ manager.path_state = ExecutionState.NONE
1025- assert manager.in_warmup is False1025+ assert manager.in_warmup is False
1026- manager.path_state = ExecutionState.WARMUP1026+ manager.path_state = ExecutionState.WARMUP
1027- assert manager.in_warmup is True1027+ assert manager.in_warmup is True
1028- 1028+ 
1029- def test_get_roots(self):1029+ def test_get_roots(self):
1030- manager = NPUGraphTreeManager(0)1030+ manager = NPUGraphTreeManager(0)
1031- mock_node1 = MagicMock()1031+ mock_node1 = MagicMock()
1032- mock_node2 = MagicMock()1032+ mock_node2 = MagicMock()
1033- manager.roots = {1033+ manager.roots = {
1034- FunctionID(1): [mock_node1],1034+ FunctionID(1): [mock_node1],
1035- FunctionID(2): [mock_node2]1035+ FunctionID(2): [mock_node2]
1036- }1036+ }
1037- roots = list(manager.get_roots())1037+ roots = list(manager.get_roots())
1038- assert roots == [mock_node1, mock_node2]1038+ assert roots == [mock_node1, mock_node2]
1039- 1039+ 
1040- def test_current_node_property_and_setter(self):1040+ def test_current_node_property_and_setter(self):
1041- manager = NPUGraphTreeManager(0)1041+ manager = NPUGraphTreeManager(0)
1042- assert manager.current_node is None1042+ assert manager.current_node is None
1043- assert manager.path_state == ExecutionState.NONE1043+ assert manager.path_state == ExecutionState.NONE
1044- mock_node = MagicMock()1044+ mock_node = MagicMock()
1045- manager.current_node = mock_node1045+ manager.current_node = mock_node
1046- assert manager.current_node == mock_node1046+ assert manager.current_node == mock_node
1047- assert manager._current_node == mock_node1047+ assert manager._current_node == mock_node
1048- manager.current_node = None1048+ manager.current_node = None
1049- assert manager.current_node is None1049+ assert manager.current_node is None
1050- assert manager.path_state == ExecutionState.NONE1050+ assert manager.path_state == ExecutionState.NONE
1051- 1051+ 
1052- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')1052+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')
1053- def test_update_generation(self, mock_get_gen):1053+ def test_update_generation(self, mock_get_gen):
1054- manager = NPUGraphTreeManager(0)1054+ manager = NPUGraphTreeManager(0)
1055- mock_get_gen.return_value = 51055+ mock_get_gen.return_value = 5
1056- manager.update_generation()1056+ manager.update_generation()
1057- assert manager.current_gen == 51057+ assert manager.current_gen == 5
1058- mock_get_gen.assert_called_once_with()1058+ mock_get_gen.assert_called_once_with()
1059- 1059+ 
1060- @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 3)1060+ @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 3)
1061- def test_get_curr_generation_mark_step(self):1061+ def test_get_curr_generation_mark_step(self):
1062- result = NPUGraphTreeManager.get_curr_generation()1062+ result = NPUGraphTreeManager.get_curr_generation()
1063- assert result == 31063+ assert result == 3
1064- 1064+ 
1065- @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 0)1065+ @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 0)
1066- @patch('torch_npu.npu._graph_tree.GenerationTracker.generation', 5)1066+ @patch('torch_npu.npu._graph_tree.GenerationTracker.generation', 5)
1067- def test_get_curr_generation_generation_tracker(self):1067+ def test_get_curr_generation_generation_tracker(self):
1068- result = NPUGraphTreeManager.get_curr_generation()1068+ result = NPUGraphTreeManager.get_curr_generation()
1069- assert result == 51069+ assert result == 5
1070- 1070+ 
1071- @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 3)1071+ @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 3)
1072- def test_user_invoked_mark_step_true(self):1072+ def test_user_invoked_mark_step_true(self):
1073- result = NPUGraphTreeManager.user_invoked_mark_step()1073+ result = NPUGraphTreeManager.user_invoked_mark_step()
1074- assert result is True1074+ assert result is True
1075- 1075+ 
1076- @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 0)1076+ @patch('torch_npu.npu._graph_tree.MarkStepBox.mark_step_counter', 0)
1077- def test_user_invoked_mark_step_false(self):1077+ def test_user_invoked_mark_step_false(self):
1078- result = NPUGraphTreeManager.user_invoked_mark_step()1078+ result = NPUGraphTreeManager.user_invoked_mark_step()
1079- assert result is False1079+ assert result is False
1080- 1080+ 
1081- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1081+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1082- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.user_invoked_mark_step')1082+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.user_invoked_mark_step')
1083- def test_can_start_new_generation_true_user_mark_step(1083+ def test_can_start_new_generation_true_user_mark_step(
1084- self, mock_user_mark_step, mock_in_new_invocation1084+ self, mock_user_mark_step, mock_in_new_invocation
1085- ):1085+ ):
1086- manager = NPUGraphTreeManager(0)1086+ manager = NPUGraphTreeManager(0)
1087- mock_in_new_invocation.return_value = True1087+ mock_in_new_invocation.return_value = True
1088- mock_user_mark_step.return_value = True1088+ mock_user_mark_step.return_value = True
1089- result = manager.can_start_new_generation1089+ result = manager.can_start_new_generation
1090- 1090+ 
1091- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1091+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1092- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.user_invoked_mark_step')1092+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.user_invoked_mark_step')
1093- def test_can_start_new_generation_true_no_pending_backwards(1093+ def test_can_start_new_generation_true_no_pending_backwards(
1094- self, mock_user_mark_step, mock_in_new_invocation1094+ self, mock_user_mark_step, mock_in_new_invocation
1095- ):1095+ ):
1096- manager = NPUGraphTreeManager(0)1096+ manager = NPUGraphTreeManager(0)
1097- manager.running_forwards_with_pending_backwards = False1097+ manager.running_forwards_with_pending_backwards = False
1098- mock_in_new_invocation.return_value = True1098+ mock_in_new_invocation.return_value = True
1099- mock_user_mark_step.return_value = False1099+ mock_user_mark_step.return_value = False
1100- result = manager.can_start_new_generation()1100+ result = manager.can_start_new_generation()
1101- assert result is True1101+ assert result is True
1102- 1102+ 
1103- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1103+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1104- def test_can_start_new_generation_false_pending_backwards(1104+ def test_can_start_new_generation_false_pending_backwards(
1105- self, mock_in_new_invocation1105+ self, mock_in_new_invocation
1106- ):1106+ ):
1107- manager = NPUGraphTreeManager(0)1107+ manager = NPUGraphTreeManager(0)
1108- manager.running_forwards_with_pending_backwards = True1108+ manager.running_forwards_with_pending_backwards = True
1109- mock_in_new_invocation.return_value = True1109+ mock_in_new_invocation.return_value = True
1110- result = manager.can_start_new_generation()1110+ result = manager.can_start_new_generation()
1111- assert result is False1111+ assert result is False
1112- 1112+ 
1113- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1113+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1114- def test_can_start_new_generation_false_not_new_invocation(1114+ def test_can_start_new_generation_false_not_new_invocation(
1115- self, mock_in_new_invocation1115+ self, mock_in_new_invocation
1116- ):1116+ ):
1117- manager = NPUGraphTreeManager(0)1117+ manager = NPUGraphTreeManager(0)
1118- mock_in_new_invocation.return_value = False1118+ mock_in_new_invocation.return_value = False
1119- result = manager.can_start_new_generation()1119+ result = manager.can_start_new_generation()
1120- assert result is False1120+ assert result is False
1121- 1121+ 
1122- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')1122+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')
1123- def test_in_new_torch_compile_invocation_true(self, mock_get_gen):1123+ def test_in_new_torch_compile_invocation_true(self, mock_get_gen):
1124- manager = NPUGraphTreeManager(0)1124+ manager = NPUGraphTreeManager(0)
1125- manager.current_gen = 11125+ manager.current_gen = 1
1126- mock_get_gen.return_value = 21126+ mock_get_gen.return_value = 2
1127- result = manager.in_new_torch_compile_invocation()1127+ result = manager.in_new_torch_compile_invocation()
1128- assert result is True1128+ assert result is True
1129- 1129+ 
1130- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')1130+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.get_curr_generation')
1131- def test_in_new_torch_compile_invocation_false(self, mock_get_gen):1131+ def test_in_new_torch_compile_invocation_false(self, mock_get_gen):
1132- manager = NPUGraphTreeManager(0)1132+ manager = NPUGraphTreeManager(0)
1133- manager.current_gen = 11133+ manager.current_gen = 1
1134- mock_get_gen.return_value = 11134+ mock_get_gen.return_value = 1
1135- result = manager.in_new_torch_compile_invocation()1135+ result = manager.in_new_torch_compile_invocation()
1136- assert result is False1136+ assert result is False
1137- 1137+ 
1138- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1138+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1139- @patch('warnings.warn')1139+ @patch('warnings.warn')
1140- def test_check_warn_on_unable_to_start_executing_no_warn(1140+ def test_check_warn_on_unable_to_start_executing_no_warn(
1141- self, mock_warn, mock_in_new_invocation1141+ self, mock_warn, mock_in_new_invocation
1142- ):1142+ ):
1143- manager = NPUGraphTreeManager(0)1143+ manager = NPUGraphTreeManager(0)
1144- mock_in_new_invocation.return_value = False1144+ mock_in_new_invocation.return_value = False
1145- manager.check_warn_on_unable_to_start_executing(FunctionID(1))1145+ manager.check_warn_on_unable_to_start_executing(FunctionID(1))
1146- mock_warn.assert_not_called()1146+ mock_warn.assert_not_called()
1147- 1147+ 
1148- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1148+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1149- @patch('warnings.warn')1149+ @patch('warnings.warn')
1150- def test_check_warn_on_unable_to_start_executing_already_warned(1150+ def test_check_warn_on_unable_to_start_executing_already_warned(
1151- self, mock_warn, mock_in_new_invocation1151+ self, mock_warn, mock_in_new_invocation
1152- ):1152+ ):
1153- manager = NPUGraphTreeManager(0)1153+ manager = NPUGraphTreeManager(0)
1154- manager.warned_functions.add(FunctionID(1))1154+ manager.warned_functions.add(FunctionID(1))
1155- mock_in_new_invocation.return_value = True1155+ mock_in_new_invocation.return_value = True
1156- manager.check_warn_on_unable_to_start_executing(FunctionID(1))1156+ manager.check_warn_on_unable_to_start_executing(FunctionID(1))
1157- mock_warn.assert_not_called()1157+ mock_warn.assert_not_called()
1158- 1158+ 
1159- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1159+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1160- @patch('warnings.warn')1160+ @patch('warnings.warn')
1161- def test_check_warn_on_unable_to_start_executing_no_repeated_pattern(1161+ def test_check_warn_on_unable_to_start_executing_no_repeated_pattern(
1162- self, mock_warn, mock_in_new_invocation1162+ self, mock_warn, mock_in_new_invocation
1163- ):1163+ ):
1164- manager = NPUGraphTreeManager(0)1164+ manager = NPUGraphTreeManager(0)
1165- mock_in_new_invocation.return_value = True1165+ mock_in_new_invocation.return_value = True
1166- 1166+
1167- mock_node = MagicMock()1167+ mock_node = MagicMock()
1168- mock_node._path_from_root = [MagicMock()]1168+ mock_node._path_from_root = [MagicMock()]
1169- mock_node._path_from_root[0].wrapped_function.id = FunctionID(2)1169+ mock_node._path_from_root[0].wrapped_function.id = FunctionID(2)
1170- mock_node.wrapped_function.id = FunctionID(1)1170+ mock_node.wrapped_function.id = FunctionID(1)
1171- manager.current_node = mock_node1171+ manager.current_node = mock_node
1172- manager.check_warn_on_unable_to_start_executing(FunctionID(1))1172+ manager.check_warn_on_unable_to_start_executing(FunctionID(1))
1173- mock_warn.assert_not_called()1173+ mock_warn.assert_not_called()
1174- 1174+ 
1175- @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')1175+ @patch('torch_npu.npu._graph_tree.NPUGraphTreeManager.in_new_torch_compile_invocation')
1176- @patch('warnings.warn')1176+ @patch('warnings.warn')
1177- def test_check_warn_on_unable_to_start_executing_warn(1177+ def test_check_warn_on_unable_to_start_executing_warn(
1178- self, mock_warn, mock_in_new_invocation1178+ self, mock_warn, mock_in_new_invocation
1179- ):1179+ ):
1180- manager = NPUGraphTreeManager(0)1180+ manager = NPUGraphTreeManager(0)
1181- mock_in_new_invocation.return_value = True1181+ mock_in_new_invocation.return_value = True
1182- 1182+
1183- mock_node1 = MagicMock()1183+ mock_node1 = MagicMock()
1184- mock_node1.wrapped_function.id = FunctionID(1)1184+ mock_node1.wrapped_function.id = FunctionID(1)
1185- mock_node1.parent = MagicMock()1185+ mock_node1.parent = MagicMock()
1186- mock_node1.parent.wrapped_function.id = FunctionID(0)1186+ mock_node1.parent.wrapped_function.id = FunctionID(0)
1187- 1187+
1188- mock_node2 = MagicMock()1188+ mock_node2 = MagicMock()
1189- mock_node2.wrapped_function.id = FunctionID(1)1189+ mock_node2.wrapped_function.id = FunctionID(1)
1190- mock_node2.parent = MagicMock()1190+ mock_node2.parent = MagicMock()
1191- mock_node2.parent.wrapped_function.id = FunctionID(0)1191+ mock_node2.parent.wrapped_function.id = FunctionID(0)
1192- 1192+
1193- mock_current_node = MagicMock()1193+ mock_current_node = MagicMock()
1194- mock_current_node.wrapped_function.id = FunctionID(1)1194+ mock_current_node.wrapped_function.id = FunctionID(1)
1195- mock_current_node.parent = MagicMock()1195+ mock_current_node.parent = MagicMock()
1196- mock_current_node.parent.wrapped_function.id = FunctionID(0)1196+ mock_current_node.parent.wrapped_function.id = FunctionID(0)
1197- 1197+
1198- mock_current_node._path_from_root = [mock_node1, mock_node2]1198+ mock_current_node._path_from_root = [mock_node1, mock_node2]
1199- manager.current_node = mock_current_node1199+ manager.current_node = mock_current_node
1200- manager.check_warn_on_unable_to_start_executing(FunctionID(1))1200+ manager.check_warn_on_unable_to_start_executing(FunctionID(1))
1201- mock_warn.assert_called_once_with(1201+ mock_warn.assert_called_once_with(
1202- "Unable to hit fast path of NPUGraphs because of pending, uninvoked backwards. "1202+ "Unable to hit fast path of NPUGraphs because of pending, uninvoked backwards. "
1203- "Consider running with torch.no_grad() or using torch.compiler.npugraph_mark_step_begin() "1203+ "Consider running with torch.no_grad() or using torch.compiler.npugraph_mark_step_begin() "
1204- "before each model invocation"1204+ "before each model invocation"
1205- )1205+ )
1206- assert FunctionID(1) in manager.warned_functions1206+ assert FunctionID(1) in manager.warned_functions
1207- 1207+ 
1208- 1208+ 
1209-if __name__ == "__main__":1209+if __name__ == "__main__":
1210- run_tests()1210+ run_tests()
Mtest/npu/test_kernel_check.py+0-1
@@ -118,4 +118,3 @@ class TestKernelCheck(TestCase):
118 118 
119if __name__ == "__main__":119if __name__ == "__main__":
120 run_tests()120 run_tests()
121- 
Mtest/npu/test_pin_memory_host_register.py+19-19
@@ -1,19 +1,19 @@
1-import unittest1+import unittest
2- 2+ 
3-import os3+import os
4-import torch4+import torch
5-from torch.testing._internal.common_utils import TestCase, run_tests5+from torch.testing._internal.common_utils import TestCase, run_tests
6-from torch_npu.npu.utils import get_cann_version6+from torch_npu.npu.utils import get_cann_version
7- 7+ 
8- 8+ 
9-class TestPinMemoryHostRegister(TestCase):9+class TestPinMemoryHostRegister(TestCase):
10- 10+ 
11- @unittest.skipUnless(get_cann_version("RUNTIME") >= "8.5.0" and get_cann_version(module="DRIVER") >= "25.5.0", "This feature is not supported in older versions.")11+ @unittest.skipUnless(get_cann_version("RUNTIME") >= "8.5.0" and get_cann_version(module="DRIVER") >= "25.5.0", "This feature is not supported in older versions.")
12- def test_pin_memory_host_register(self):12+ def test_pin_memory_host_register(self):
13- os.environ["PYTORCH_NPU_ALLOC_CONF"] = "pinned_mem_register:True"13+ os.environ["PYTORCH_NPU_ALLOC_CONF"] = "pinned_mem_register:True"
14- cpu_tensor = torch.ones([2, 3])14+ cpu_tensor = torch.ones([2, 3])
15- pin_tensor = cpu_tensor.pin_memory()15+ pin_tensor = cpu_tensor.pin_memory()
16- self.assertTrue(pin_tensor.is_pinned())16+ self.assertTrue(pin_tensor.is_pinned())
17- 17+ 
18-if __name__ == '__main__':18+if __name__ == '__main__':
19- run_tests()19+ run_tests()
Mtest/npu/test_save_async.py+119-119
@@ -1,119 +1,119 @@
1-import os1+import os
2-import time2+import time
3-import copy3+import copy
4- 4+ 
5-import torch5+import torch
6-import torch.nn as nn6+import torch.nn as nn
7-import torch.optim as optim7+import torch.optim as optim
8- 8+ 
9-import torch_npu9+import torch_npu
10-from torch_npu.testing.testcase import TestCase, run_tests10+from torch_npu.testing.testcase import TestCase, run_tests
11-from torch_npu.utils._path_manager import PathManager11+from torch_npu.utils._path_manager import PathManager
12- 12+ 
13- 13+ 
14-class TestAsyncSave(TestCase):14+class TestAsyncSave(TestCase):
15- test_save_path = os.path.join(15+ test_save_path = os.path.join(
16- os.path.realpath(os.path.dirname(__file__)), "test_save_async")16+ os.path.realpath(os.path.dirname(__file__)), "test_save_async")
17- 17+ 
18- @classmethod18+ @classmethod
19- def setUpClass(cls):19+ def setUpClass(cls):
20- PathManager.make_dir_safety(TestAsyncSave.test_save_path)20+ PathManager.make_dir_safety(TestAsyncSave.test_save_path)
21- 21+ 
22- @classmethod22+ @classmethod
23- def tearDownClass(cls):23+ def tearDownClass(cls):
24- PathManager.remove_path_safety(TestAsyncSave.test_save_path)24+ PathManager.remove_path_safety(TestAsyncSave.test_save_path)
25- 25+
26- def wait_for_save_completion(self, file_path, timeout_sec=60, poll_interval_sec=0.5):26+ def wait_for_save_completion(self, file_path, timeout_sec=60, poll_interval_sec=0.5):
27- start_time = time.time()27+ start_time = time.time()
28- 28+ 
29- while time.time() - start_time < timeout_sec:29+ while time.time() - start_time < timeout_sec:
30- if os.path.exists(file_path):30+ if os.path.exists(file_path):
31- current_size = os.path.getsize(file_path)31+ current_size = os.path.getsize(file_path)
32- time.sleep(poll_interval_sec)32+ time.sleep(poll_interval_sec)
33- new_size = os.path.getsize(file_path)33+ new_size = os.path.getsize(file_path)
34- 34+ 
35- if current_size == new_size:35+ if current_size == new_size:
36- return True36+ return True
37- else:37+ else:
38- time.sleep(poll_interval_sec)38+ time.sleep(poll_interval_sec)
39- 39+ 
40- return False40+ return False
41- 41+ 
42- def test_save_async_tensor(self):42+ def test_save_async_tensor(self):
43- save_tensor = torch.rand(1024, dtype=torch.float32).npu()43+ save_tensor = torch.rand(1024, dtype=torch.float32).npu()
44- async_save_path = os.path.join(TestAsyncSave.test_save_path, "async_save_tensor.pt")44+ async_save_path = os.path.join(TestAsyncSave.test_save_path, "async_save_tensor.pt")
45- torch_npu.utils.save_async(save_tensor, async_save_path)45+ torch_npu.utils.save_async(save_tensor, async_save_path)
46- 46+
47- if self.wait_for_save_completion(async_save_path):47+ if self.wait_for_save_completion(async_save_path):
48- tensor_async = torch.load(async_save_path, weights_only=False)48+ tensor_async = torch.load(async_save_path, weights_only=False)
49- self.assertEqual(tensor_async, save_tensor)49+ self.assertEqual(tensor_async, save_tensor)
50- else:50+ else:
51- self.assertTrue(False, f"{async_save_path} is not exist!")51+ self.assertTrue(False, f"{async_save_path} is not exist!")
52- 52+
53- def test_save_async(self):53+ def test_save_async(self):
54- loss1 = [1.6099495, 1.6099086, 1.6098710]54+ loss1 = [1.6099495, 1.6099086, 1.6098710]
55- loss2 = []55+ loss2 = []
56- model_list = []56+ model_list = []
57- checkpoint_list = []57+ checkpoint_list = []
58- model_origin = nn.Sequential(58+ model_origin = nn.Sequential(
59- nn.Linear(100, 50),59+ nn.Linear(100, 50),
60- nn.ReLU(),60+ nn.ReLU(),
61- nn.Linear(50, 20),61+ nn.Linear(50, 20),
62- nn.ReLU(),62+ nn.ReLU(),
63- nn.Linear(20, 5),63+ nn.Linear(20, 5),
64- nn.ReLU()64+ nn.ReLU()
65- )65+ )
66- 66+ 
67- input_data = torch.ones(6400, 100).npu()67+ input_data = torch.ones(6400, 100).npu()
68- labels = torch.arange(5).repeat(1280).npu()68+ labels = torch.arange(5).repeat(1280).npu()
69- 69+ 
70- criterion = nn.CrossEntropyLoss()70+ criterion = nn.CrossEntropyLoss()
71- model = model_origin.npu()71+ model = model_origin.npu()
72- optimerizer = optim.SGD(model.parameters(), lr=0.1)72+ optimerizer = optim.SGD(model.parameters(), lr=0.1)
73- for step in range(3):73+ for step in range(3):
74- outputs = model(input_data)74+ outputs = model(input_data)
75- loss = criterion(outputs, labels)75+ loss = criterion(outputs, labels)
76- 76+ 
77- optimerizer.zero_grad()77+ optimerizer.zero_grad()
78- loss.backward()78+ loss.backward()
79- 79+ 
80- optimerizer.step()80+ optimerizer.step()
81- 81+
82- loss2.append(loss)82+ loss2.append(loss)
83- checkpoint = {83+ checkpoint = {
84- "model": model.state_dict(),84+ "model": model.state_dict(),
85- "optimizer": optimerizer.state_dict()85+ "optimizer": optimerizer.state_dict()
86- }86+ }
87- checkpoint_list.append(copy.deepcopy(checkpoint))87+ checkpoint_list.append(copy.deepcopy(checkpoint))
88- model_list.append(copy.deepcopy(model))88+ model_list.append(copy.deepcopy(model))
89- checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{step}.path")89+ checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{step}.path")
90- model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{step}.path")90+ model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{step}.path")
91- torch_npu.utils.save_async(checkpoint, checkpoint_async_path, model=model)91+ torch_npu.utils.save_async(checkpoint, checkpoint_async_path, model=model)
92- torch_npu.utils.save_async(model, model_async_path, model=model)92+ torch_npu.utils.save_async(model, model_async_path, model=model)
93- 93+ 
94- for i in range(3):94+ for i in range(3):
95- self.assertEqual(loss1[i], loss2[i].item())95+ self.assertEqual(loss1[i], loss2[i].item())
96- checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{i}.path")96+ checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{i}.path")
97- if self.wait_for_save_completion(checkpoint_async_path):97+ if self.wait_for_save_completion(checkpoint_async_path):
98- checkpoint_async = torch.load(checkpoint_async_path, weights_only=False)98+ checkpoint_async = torch.load(checkpoint_async_path, weights_only=False)
99- self.assertEqual(checkpoint_list[i], checkpoint_async, prec=2e-3)99+ self.assertEqual(checkpoint_list[i], checkpoint_async, prec=2e-3)
100- else:100+ else:
101- self.assertTrue(False, f"{checkpoint_async_path} is not exist!")101+ self.assertTrue(False, f"{checkpoint_async_path} is not exist!")
102- model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{i}.path")102+ model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{i}.path")
103- if self.wait_for_save_completion(model_async_path):103+ if self.wait_for_save_completion(model_async_path):
104- model_async = torch.load(model_async_path, weights_only=False)104+ model_async = torch.load(model_async_path, weights_only=False)
105- else:105+ else:
106- self.assertTrue(False, f"{model_async_path} is not exist!")106+ self.assertTrue(False, f"{model_async_path} is not exist!")
107- state_dict_sync = model_list[i].state_dict()107+ state_dict_sync = model_list[i].state_dict()
108- state_dict_async = model_async.state_dict()108+ state_dict_async = model_async.state_dict()
109- 109+ 
110- key_sync = sorted(state_dict_sync.keys())110+ key_sync = sorted(state_dict_sync.keys())
111- key_async = sorted(state_dict_async.keys())111+ key_async = sorted(state_dict_async.keys())
112- 112+ 
113- self.assertEqual(key_sync, key_async)113+ self.assertEqual(key_sync, key_async)
114- for key in key_async:114+ for key in key_async:
115- self.assertEqual(state_dict_async[key], state_dict_sync[key], prec=2e-3)115+ self.assertEqual(state_dict_async[key], state_dict_sync[key], prec=2e-3)
116- 116+ 
117-if __name__ == '__main__':117+if __name__ == '__main__':
118- torch.npu.set_device(0)118+ torch.npu.set_device(0)
119- run_tests()119+ run_tests()
Mtest/profiler/analysis/test_profiler_config.py+0-1
@@ -81,4 +81,3 @@ class TestProfilerConfig(TestCase):
81 81 
82if __name__ == "__main__":82if __name__ == "__main__":
83 run_tests()83 run_tests()
84- 
Mtest/profiler/test_non_intrusive_profile.py+0-1
@@ -39,4 +39,3 @@ class TestNoneInstrusiveProfile(TestCase):
39 39 
40if __name__ == "__main__":40if __name__ == "__main__":
41 run_tests()41 run_tests()
42- 
Mtest/test_fake_tensor.py+2016-2016
Mtest/unsupported_test_cases/.pytorch-disabled-tests.json+3-1
@@ -32074,5 +32074,7 @@
32074 "test_softmax_forward_64bit_indexing_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],32074 "test_softmax_forward_64bit_indexing_npu (__main__.TestNNDeviceTypePRIVATEUSE1)": ["", [""]],
32075 "test_grad_with_v_schedule (__main__.TestScheduleLowering)": ["", [""]],32075 "test_grad_with_v_schedule (__main__.TestScheduleLowering)": ["", [""]],
32076 "test_npu_roi_align_1 (__main__.TestPsRoiPooling)": ["", [""]],32076 "test_npu_roi_align_1 (__main__.TestPsRoiPooling)": ["", [""]],
32077- "test_silu (__main__.TestActivations)": ["", [""]]32077+ "test_silu (__main__.TestActivations)": ["", [""]],
32078+ "test_data_parallel_rnn (__main__.TestDataParallel)": ["", ["Disabled during A1 to A2 chip transition"]],
32079+ "test_alltoall_single_2p_size_dist (__main__.HcclAlltoAllSingleTest)": ["", ["Disabled during A1 to A2 chip transition"]]
32078}32080}
Mtools/flight_recorder/components/builder.py+307-307
@@ -1,307 +1,307 @@
1-import argparse1+import argparse
2-import ast2+import ast
3-import os3+import os
4-from typing import Any4+from typing import Any
5- 5+ 
6-from tools.flight_recorder.components.fr_logger import FlightRecorderLogger6+from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
7-from tools.flight_recorder.components.types import (7+from tools.flight_recorder.components.types import (
8- Collective,8+ Collective,
9- Database,9+ Database,
10- EntryState,10+ EntryState,
11- Group,11+ Group,
12- MatchStateRecord,12+ MatchStateRecord,
13- Membership,13+ Membership,
14- HCCLCall,14+ HCCLCall,
15- Op,15+ Op,
16- Traceback,16+ Traceback,
17-)17+)
18-from tools.flight_recorder.components.utils import (18+from tools.flight_recorder.components.utils import (
19- ProcessGroupData,19+ ProcessGroupData,
20- align_trace_from_beginning,20+ align_trace_from_beginning,
21- check_current_entry_match,21+ check_current_entry_match,
22- check_no_missing_dump_files,22+ check_no_missing_dump_files,
23- check_version,23+ check_version,
24- EntryContext,24+ EntryContext,
25- error_analysis,25+ error_analysis,
26- get_version_detail,26+ get_version_detail,
27- just_print_entries,27+ just_print_entries,
28-)28+)
29- 29+ 
30- 30+ 
31-# Set up logging31+# Set up logging
32-logger: FlightRecorderLogger = FlightRecorderLogger()32+logger: FlightRecorderLogger = FlightRecorderLogger()
33- 33+ 
34- 34+ 
35-try:35+try:
36- from tabulate import tabulate36+ from tabulate import tabulate
37-except ModuleNotFoundError:37+except ModuleNotFoundError:
38- logger.warning("tabulate is not installed. Proceeding without it.")38+ logger.warning("tabulate is not installed. Proceeding without it.")
39- 39+ 
40- # Define a no-op tabulate function40+ # Define a no-op tabulate function
41- def tabulate(data: Any, headers: Any = None) -> Any: # type: ignore[misc]41+ def tabulate(data: Any, headers: Any = None) -> Any: # type: ignore[misc]
42- return data42+ return data
43- 43+ 
44- 44+ 
45-"""45+"""
O
OopenLiBingCI5月17日

此条代码评论区间+38至+45

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
46-Flat DB builder46+Flat DB builder
47-"""47+"""
48- 48+ 
49- 49+ 
50-def build_groups_memberships(50+def build_groups_memberships(
51- pg_config: Any,51+ pg_config: Any,
52-) -> tuple[52+) -> tuple[
53- list[Group],53+ list[Group],
54- dict[Any, Group],54+ dict[Any, Group],
55- list[Membership],55+ list[Membership],
56- dict[str, set[Any]],56+ dict[str, set[Any]],
57- dict[tuple[str, int], str],57+ dict[tuple[str, int], str],
58-]:58+]:
59- """59+ """
60- pg_config: {60+ pg_config: {
61- global_rank: {61+ global_rank: {
62- (pg_guid, desc, ranks)62+ (pg_guid, desc, ranks)
63- }63+ }
64- }64+ }
65- 65+ 
66- `pg_guid` is a system generated id, but depending on the mode of PG creation it could be a globally incrementing int66+ `pg_guid` is a system generated id, but depending on the mode of PG creation it could be a globally incrementing int
67- or a hash of the ranks. See `_process_group_name` in distributed_c10d.py.67+ or a hash of the ranks. See `_process_group_name` in distributed_c10d.py.
68- `desc` is provided by the user (optionally) and should be 'meaningful' (e.g. TP/PP/DP group)68+ `desc` is provided by the user (optionally) and should be 'meaningful' (e.g. TP/PP/DP group)
69- `ranks` is a list of the 'global ranks' that are members of the PG.69+ `ranks` is a list of the 'global ranks' that are members of the PG.
70- 70+ 
71- (pg_guid, desc, ranks) tuples are appended lazily to the flight buffer when `getHCCLComm` is called on a PG and71+ (pg_guid, desc, ranks) tuples are appended lazily to the flight buffer when `getHCCLComm` is called on a PG and
72- the `enabled_` flag is true for that PG.72+ the `enabled_` flag is true for that PG.
73- - the order of calling (init_process_group, new_group, etc) does not affect the order of the tuples in the list73+ - the order of calling (init_process_group, new_group, etc) does not affect the order of the tuples in the list
74- 74+ 
75- Returns:75+ Returns:
76- `groups`: a groups table where each row is a Group namedtuple.76+ `groups`: a groups table where each row is a Group namedtuple.
77- `_groups`: a dict that is indexed by pg_guid with Group namedtuple as value.77+ `_groups`: a dict that is indexed by pg_guid with Group namedtuple as value.
78- `memberships`: a membership table where each row is a Membership namedtuple.78+ `memberships`: a membership table where each row is a Membership namedtuple.
79- `_memberships`: a dict that is indexed by pg_guid with set of ranks (int) as value.79+ `_memberships`: a dict that is indexed by pg_guid with set of ranks (int) as value.
80- `_pg_guids`: a dict that is indexed by (pg_uid, global_rank) with pg_guid as value.80+ `_pg_guids`: a dict that is indexed by (pg_uid, global_rank) with pg_guid as value.
81- """81+ """
82- # flat lists for return82+ # flat lists for return
83- groups = []83+ groups = []
84- memberships = []84+ memberships = []
85- 85+ 
86- # dicts for faster cross-rank validation86+ # dicts for faster cross-rank validation
87- _groups = {}87+ _groups = {}
88- _memberships = {}88+ _memberships = {}
89- _pg_guids = {}89+ _pg_guids = {}
90- for global_rank in pg_config:90+ for global_rank in pg_config:
91- for pg_uid in pg_config[global_rank]:91+ for pg_uid in pg_config[global_rank]:
92- desc = pg_config[global_rank][pg_uid]["desc"]92+ desc = pg_config[global_rank][pg_uid]["desc"]
93- ranks = ast.literal_eval(pg_config[global_rank][pg_uid]["ranks"])93+ ranks = ast.literal_eval(pg_config[global_rank][pg_uid]["ranks"])
94- # With the adoption of the split_group API, we can have multiple PGs with the same pg_guid (PG Name)94+ # With the adoption of the split_group API, we can have multiple PGs with the same pg_guid (PG Name)
95- # So we need to add the hash of all its ranks within the PG as well.95+ # So we need to add the hash of all its ranks within the PG as well.
96- # Also guid must be a string because `_process_group_name` returns a string.96+ # Also guid must be a string because `_process_group_name` returns a string.
97- pg_guid = pg_uid + str(hash(frozenset(ranks)))97+ pg_guid = pg_uid + str(hash(frozenset(ranks)))
98- _pg_guids[(pg_uid, global_rank)] = pg_guid98+ _pg_guids[(pg_uid, global_rank)] = pg_guid
99- if isinstance(ranks, str):99+ if isinstance(ranks, str):
100- ranks = ast.literal_eval(ranks)100+ ranks = ast.literal_eval(ranks)
101- if pg_guid not in _groups:101+ if pg_guid not in _groups:
102- groups.append(Group(id=pg_guid, desc=desc, size=len(ranks)))102+ groups.append(Group(id=pg_guid, desc=desc, size=len(ranks)))
103- for rank in ranks:103+ for rank in ranks:
104- memberships.append(Membership(group_id=pg_guid, global_rank=rank))104+ memberships.append(Membership(group_id=pg_guid, global_rank=rank))
105- _groups[pg_guid] = groups[-1]105+ _groups[pg_guid] = groups[-1]
106- _memberships[pg_guid] = set(ranks)106+ _memberships[pg_guid] = set(ranks)
107- else:107+ else:
108- # validation across ranks108+ # validation across ranks
109- if _groups[pg_guid].desc != desc:109+ if _groups[pg_guid].desc != desc:
110- raise ValueError(110+ raise ValueError(
111- f"Description mismatch for group {pg_guid}: "111+ f"Description mismatch for group {pg_guid}: "
112- f"expected '{desc}', got '{_groups[pg_guid].desc}'"112+ f"expected '{desc}', got '{_groups[pg_guid].desc}'"
113- )113+ )
114- 114+ 
115- if _memberships[pg_guid] != set(ranks):115+ if _memberships[pg_guid] != set(ranks):
116- raise ValueError(116+ raise ValueError(
117- f"Membership mismatch for group {pg_guid}: "117+ f"Membership mismatch for group {pg_guid}: "
118- f"expected {set(ranks)}, got {_memberships[pg_guid]}"118+ f"expected {set(ranks)}, got {_memberships[pg_guid]}"
119- )119+ )
120- 120+ 
121- return groups, _groups, memberships, _memberships, _pg_guids121+ return groups, _groups, memberships, _memberships, _pg_guids
122- 122+ 
123- 123+ 
124-def build_collectives(124+def build_collectives(
125- all_entries: dict[int, list[dict[str, Any]]],125+ all_entries: dict[int, list[dict[str, Any]]],
126- _groups: dict[str, Group],126+ _groups: dict[str, Group],
127- _memberships: dict[str, set[Any]],127+ _memberships: dict[str, set[Any]],
128- _pg_guids: dict[tuple[str, int], str],128+ _pg_guids: dict[tuple[str, int], str],
129- version: str,129+ version: str,
130-) -> tuple[list[Traceback], list[Collective], list[HCCLCall]]:130+) -> tuple[list[Traceback], list[Collective], list[HCCLCall]]:
131- """131+ """
132- groups, memberships are the non-flat dicts that are indexable132+ groups, memberships are the non-flat dicts that are indexable
133- all_entries is a raw dict from the original dumps:133+ all_entries is a raw dict from the original dumps:
134- 134+ 
135- all_entries: {135+ all_entries: {
136- global_rank: [136+ global_rank: [
137- {137+ {
138- record_id: ordered id of the event in the trace buffer138+ record_id: ordered id of the event in the trace buffer
139- pg_id: ProcessGroupHCCL::uid_139+ pg_id: ProcessGroupHCCL::uid_
140- *note: `pg_id` corresponds to nothing in groups table140+ *note: `pg_id` corresponds to nothing in groups table
141- process_group: (pg_name, desc)141+ process_group: (pg_name, desc)
142- *note: `pg_name`, `desc` corresponds to `pg_id`, `desc` in groups table142+ *note: `pg_name`, `desc` corresponds to `pg_id`, `desc` in groups table
143- collective_seq_id: ordered id for collective operations and coalesced group operations143+ collective_seq_id: ordered id for collective operations and coalesced group operations
144- p2p_seq_id: ordered id for point-to-point operations144+ p2p_seq_id: ordered id for point-to-point operations
145- op_id: ordered id including individual ops inside coalescing group145+ op_id: ordered id including individual ops inside coalescing group
146- profiling_name: descriptive name of the operation146+ profiling_name: descriptive name of the operation
147- 'time_created_ns',147+ 'time_created_ns',
148- 'input_sizes',148+ 'input_sizes',
149- 'output_sizes',149+ 'output_sizes',
150- 'state',150+ 'state',
151- 'time_discovered_started_ns',151+ 'time_discovered_started_ns',
152- 'time_discovered_completed_ns',152+ 'time_discovered_completed_ns',
153- 'retired',153+ 'retired',
154- 'frames',154+ 'frames',
155- }155+ }
156- ]156+ ]
157- }157+ }
158- """158+ """
159- tracebacks: list[Traceback] = []159+ tracebacks: list[Traceback] = []
160- 160+ 
161- collectives: list[Collective] = []161+ collectives: list[Collective] = []
162- hccl_calls: list[HCCLCall] = []162+ hccl_calls: list[HCCLCall] = []
163- 163+ 
164- # once we find one mismatch, we stop pairing up collectives since the pairing is possibly incorrect164+ # once we find one mismatch, we stop pairing up collectives since the pairing is possibly incorrect
165- # instead, just record the remaining ops as HCCLCalls165+ # instead, just record the remaining ops as HCCLCalls
166- mismatch = {_groups[g].id: 0 for g in _groups}166+ mismatch = {_groups[g].id: 0 for g in _groups}
167- MISMATCH_TAIL = 10167+ MISMATCH_TAIL = 10
168- 168+ 
169- # For best effort partial analysis.169+ # For best effort partial analysis.
170- dumps_ranks = set()170+ dumps_ranks = set()
171- for key in all_entries.keys():171+ for key in all_entries.keys():
172- try:172+ try:
173- dumps_ranks.add(int(key))173+ dumps_ranks.add(int(key))
174- except ValueError as e:174+ except ValueError as e:
175- raise ValueError(f"Cannot extract rank from '{key}") from e175+ raise ValueError(f"Cannot extract rank from '{key}") from e
176- """176+ """
177- - it doesn't matter what order I put collectives/hcclops into their table. we can later on re-sort it by start time177+ - it doesn't matter what order I put collectives/hcclops into their table. we can later on re-sort it by start time
178- - there could be multiple options for the "first" collective to pair up (rank 0,1 might do a bcast while rank 2,3 do a bcast)178+ - there could be multiple options for the "first" collective to pair up (rank 0,1 might do a bcast while rank 2,3 do a bcast)
179- - within a group, the first collective must be the same on all ranks in the group, then it can be marked as a179+ - within a group, the first collective must be the same on all ranks in the group, then it can be marked as a
180- collective and removed180+ collective and removed
181- """181+ """
182- while all_entries:182+ while all_entries:
183- # we greedily match collectives, starting arbitrarily with the trace from the first rank183+ # we greedily match collectives, starting arbitrarily with the trace from the first rank
184- # later, if we exhaust the first rank, we continue with the next 'first rank'184+ # later, if we exhaust the first rank, we continue with the next 'first rank'
185- rank_iter = iter(all_entries)185+ rank_iter = iter(all_entries)
186- first_rank = next(rank_iter)186+ first_rank = next(rank_iter)
187- other_ranks = list(rank_iter)187+ other_ranks = list(rank_iter)
188- 188+ 
189- if len(all_entries[first_rank]) == 0:189+ if len(all_entries[first_rank]) == 0:
190- all_entries.pop(first_rank)190+ all_entries.pop(first_rank)
191- continue191+ continue
192- 192+ 
193- # lets match the first collective! we need to know which ranks are involved, and ensure that this same193+ # lets match the first collective! we need to know which ranks are involved, and ensure that this same
194- # collective is also the first one on those ranks within that group194+ # collective is also the first one on those ranks within that group
195- entries = all_entries[first_rank]195+ entries = all_entries[first_rank]
196- current_entry = entries[0]196+ current_entry = entries[0]
197- 197+ 
198- desc = current_entry["process_group"][1] if current_entry["process_group"][1] else "default_pg"198+ desc = current_entry["process_group"][1] if current_entry["process_group"][1] else "default_pg"
199- # For db build and logs printing, we want to use the original pg_name, not the hash one.199+ # For db build and logs printing, we want to use the original pg_name, not the hash one.
200- original_pg_name = current_entry["process_group"][0]200+ original_pg_name = current_entry["process_group"][0]
201- pg_name = _pg_guids[(original_pg_name, first_rank)]201+ pg_name = _pg_guids[(original_pg_name, first_rank)]
202- expected_ranks = set(_memberships[pg_name])202+ expected_ranks = set(_memberships[pg_name])
203- entry_state = EntryState(current_entry, expected_ranks)203+ entry_state = EntryState(current_entry, expected_ranks)
204- match_record = MatchStateRecord(204+ match_record = MatchStateRecord(
205- expected_ranks=expected_ranks,205+ expected_ranks=expected_ranks,
206- other_ranks=other_ranks,206+ other_ranks=other_ranks,
207- entry_state=entry_state,207+ entry_state=entry_state,
208- candidate_ranks={first_rank},208+ candidate_ranks={first_rank},
209- candidate_idx={},209+ candidate_idx={},
210- found_ranks=set(),210+ found_ranks=set(),
211- found_idx={},211+ found_idx={},
212- errors=set(),212+ errors=set(),
213- )213+ )
214- 214+ 
215- check_current_entry_match(215+ check_current_entry_match(
216- all_entries=all_entries,216+ all_entries=all_entries,
217- current_entry=current_entry,217+ current_entry=current_entry,
218- _memberships=_memberships,218+ _memberships=_memberships,
219- pg_data=ProcessGroupData(pg_guids=_pg_guids, pg_name=pg_name, desc=desc, mismatch=mismatch),219+ pg_data=ProcessGroupData(pg_guids=_pg_guids, pg_name=pg_name, desc=desc, mismatch=mismatch),
220- match_record=match_record,220+ match_record=match_record,
221- )221+ )
222- 222+ 
223- # Use heuristics to decide what type of errors and error messages we should print.223+ # Use heuristics to decide what type of errors and error messages we should print.
224- error_analysis(224+ error_analysis(
225- entry_context=EntryContext(all_entries, current_entry, dumps_ranks, first_rank),225+ entry_context=EntryContext(all_entries, current_entry, dumps_ranks, first_rank),
226- match_record=match_record,226+ match_record=match_record,
227- mismatch=mismatch,227+ mismatch=mismatch,
228- version=get_version_detail(version),228+ version=get_version_detail(version),
229- pg_name=pg_name,229+ pg_name=pg_name,
230- )230+ )
231- # at this point there are 3 possibilities231+ # at this point there are 3 possibilities
232- # 1. we found a match on all the ranks that are members of the group232+ # 1. we found a match on all the ranks that are members of the group
233- # -> we create a Collective and remove the individual entries from their original lists233+ # -> we create a Collective and remove the individual entries from their original lists
234- if match_record.found_ranks == expected_ranks and mismatch[pg_name] == 0:234+ if match_record.found_ranks == expected_ranks and mismatch[pg_name] == 0:
235- collectives.append(match_record.entry_state.to_collective(len(collectives)))235+ collectives.append(match_record.entry_state.to_collective(len(collectives)))
236- idx_map = {r: match_record.found_idx[r] if r != first_rank else 0 for r in match_record.found_ranks}236+ idx_map = {r: match_record.found_idx[r] if r != first_rank else 0 for r in match_record.found_ranks}
237- hccl_calls.extend(237+ hccl_calls.extend(
238- match_record.entry_state.to_hccl_call(all_entries, idx_map, len(hccl_calls), collectives[-1].id)238+ match_record.entry_state.to_hccl_call(all_entries, idx_map, len(hccl_calls), collectives[-1].id)
239- )239+ )
240- 240+ 
241- # 2. we found a partial match but some ranks are missing241+ # 2. we found a partial match but some ranks are missing
242- # 3. we found no match242+ # 3. we found no match
243- else:243+ else:
244- logger.debug("appending a non-matching collective")244+ logger.debug("appending a non-matching collective")
245- idx_map = {r: match_record.candidate_idx[r] if r != first_rank else 0 for r in match_record.candidate_ranks}245+ idx_map = {r: match_record.candidate_idx[r] if r != first_rank else 0 for r in match_record.candidate_ranks}
246- collectives.append(246+ collectives.append(
247- match_record.entry_state.to_collective(247+ match_record.entry_state.to_collective(
248- len(collectives),248+ len(collectives),
249- errors=match_record.errors,249+ errors=match_record.errors,
250- idx_map=idx_map,250+ idx_map=idx_map,
251- all_entries=all_entries,251+ all_entries=all_entries,
252- )252+ )
253- )253+ )
254- hccl_calls.extend(match_record.entry_state.to_hccl_call(all_entries, idx_map, len(hccl_calls), None))254+ hccl_calls.extend(match_record.entry_state.to_hccl_call(all_entries, idx_map, len(hccl_calls), None))
255- 255+ 
256- if mismatch[pg_name] > MISMATCH_TAIL:256+ if mismatch[pg_name] > MISMATCH_TAIL:
257- logger.error("Too many mismatches for process_group %s: %s aborting", pg_name, desc)257+ logger.error("Too many mismatches for process_group %s: %s aborting", pg_name, desc)
258- break258+ break
259- return tracebacks, collectives, hccl_calls259+ return tracebacks, collectives, hccl_calls
260- 260+ 
261- 261+ 
262-def build_db(details: dict[str, dict[str, Any]], args: argparse.Namespace, version: str) -> Database:262+def build_db(details: dict[str, dict[str, Any]], args: argparse.Namespace, version: str) -> Database:
263- if args.verbose:263+ if args.verbose:
264- os.environ["FR_TRACE_VERBOSE_OUTPUT"] = "1"264+ os.environ["FR_TRACE_VERBOSE_OUTPUT"] = "1"
265- # temporary state used for building database265+ # temporary state used for building database
266- entries = {}266+ entries = {}
267- pg_config = {}267+ pg_config = {}
268- version_by_ranks = {}268+ version_by_ranks = {}
269- for rank, dump in details.items():269+ for rank, dump in details.items():
270- entries[rank] = dump["entries"]270+ entries[rank] = dump["entries"]
271- version_by_ranks[rank] = dump["version"]271+ version_by_ranks[rank] = dump["version"]
272- pg_config[rank] = dump["pg_config"]272+ pg_config[rank] = dump["pg_config"]
273- 273+ 
274- # Ensure version is consistent across all ranks.274+ # Ensure version is consistent across all ranks.
275- check_version(version_by_ranks, version)275+ check_version(version_by_ranks, version)
276- entries = align_trace_from_beginning(entries)276+ entries = align_trace_from_beginning(entries)
277- 277+ 
278- # flattened database278+ # flattened database
279- groups, _groups, memberships, _memberships, _pg_guids = build_groups_memberships(pg_config)279+ groups, _groups, memberships, _memberships, _pg_guids = build_groups_memberships(pg_config)
280- logger.debug("built groups, memberships")280+ logger.debug("built groups, memberships")
281- 281+ 
282- if not args.allow_incomplete_ranks:282+ if not args.allow_incomplete_ranks:
283- check_no_missing_dump_files(entries, memberships)283+ check_no_missing_dump_files(entries, memberships)
284- 284+ 
285- if args.just_print_entries:285+ if args.just_print_entries:
286- just_print_entries(entries, _groups, _memberships, _pg_guids, args)286+ just_print_entries(entries, _groups, _memberships, _pg_guids, args)
287- return None287+ return None
288- 288+ 
289- tracebacks, collectives, hccl_calls = build_collectives(entries, _groups, _memberships, _pg_guids, version)289+ tracebacks, collectives, hccl_calls = build_collectives(entries, _groups, _memberships, _pg_guids, version)
290- logger.debug("built collectives, hccl_calls")290+ logger.debug("built collectives, hccl_calls")
291- if args.verbose:291+ if args.verbose:
292- logger.debug("Groups")292+ logger.debug("Groups")
293- logger.debug(tabulate(groups, headers=Group._fields))293+ logger.debug(tabulate(groups, headers=Group._fields))
294- logger.debug("Memberships")294+ logger.debug("Memberships")
295- logger.debug(tabulate(memberships, headers=Membership._fields))295+ logger.debug(tabulate(memberships, headers=Membership._fields))
296- logger.debug("Collectives")296+ logger.debug("Collectives")
297- logger.debug(tabulate(collectives, headers=Collective._fields))297+ logger.debug(tabulate(collectives, headers=Collective._fields))
298- logger.debug("HCCLCalls")298+ logger.debug("HCCLCalls")
299- logger.debug(tabulate(hccl_calls, headers=HCCLCall._fields))299+ logger.debug(tabulate(hccl_calls, headers=HCCLCall._fields))
300- db = Database(300+ db = Database(
301- tracebacks=tracebacks,301+ tracebacks=tracebacks,
302- collectives=collectives,302+ collectives=collectives,
303- hcclcalls=hccl_calls,303+ hcclcalls=hccl_calls,
304- groups=groups,304+ groups=groups,
305- memberships=memberships,305+ memberships=memberships,
306- )306+ )
307- return db307+ return db
Mtools/flight_recorder/components/config_manager.py+71-71
@@ -1,71 +1,71 @@
1-import argparse1+import argparse
2-import logging2+import logging
3-from collections.abc import Sequence3+from collections.abc import Sequence
4-from typing import Optional4+from typing import Optional
5- 5+ 
6-from tools.flight_recorder.components.fr_logger import FlightRecorderLogger6+from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
7- 7+ 
8- 8+ 
9-logger: FlightRecorderLogger = FlightRecorderLogger()9+logger: FlightRecorderLogger = FlightRecorderLogger()
10- 10+ 
11- 11+ 
12-class JobConfig:12+class JobConfig:
13- """13+ """
14- A helper class to manage the script configuration.14+ A helper class to manage the script configuration.
15- """15+ """
16- 16+ 
17- def __init__(self: "JobConfig"):17+ def __init__(self: "JobConfig"):
18- self.parser = argparse.ArgumentParser(description="PyTorch Flight recorder analyzing script.")18+ self.parser = argparse.ArgumentParser(description="PyTorch Flight recorder analyzing script.")
19- self.parser.add_argument(19+ self.parser.add_argument(
20- "trace_dir",20+ "trace_dir",
21- nargs="?",21+ nargs="?",
22- help="Directory containing one trace file per rank, named with <prefix>_<rank>.",22+ help="Directory containing one trace file per rank, named with <prefix>_<rank>.",
23- )23+ )
24- self.parser.add_argument(24+ self.parser.add_argument(
25- "--selected-ranks",25+ "--selected-ranks",
26- default=None,26+ default=None,
27- nargs="+",27+ nargs="+",
28- type=int,28+ type=int,
29- help="List of ranks we want to show traces for.",29+ help="List of ranks we want to show traces for.",
30- )30+ )
31- self.parser.add_argument(31+ self.parser.add_argument(
32- "--allow-incomplete-ranks",32+ "--allow-incomplete-ranks",
33- action="store_true",33+ action="store_true",
34- help=(34+ help=(
35- "FR trace require all ranks to have dumps for analysis. "35+ "FR trace require all ranks to have dumps for analysis. "
36- "This flag allows best-effort partial analysis of results "36+ "This flag allows best-effort partial analysis of results "
37- "and printing of collected data."37+ "and printing of collected data."
38- ),38+ ),
39- )39+ )
40- self.parser.add_argument(40+ self.parser.add_argument(
41- "--pg-filters",41+ "--pg-filters",
42- default=None,42+ default=None,
43- nargs="+",43+ nargs="+",
44- type=str,44+ type=str,
45- help=(45+ help=(
46- "List of filter strings, it could be pg name or pg desc. "46+ "List of filter strings, it could be pg name or pg desc. "
47- "If specified, only show traces for the given pg."47+ "If specified, only show traces for the given pg."
48- ),48+ ),
49- )49+ )
50- self.parser.add_argument("-o", "--output", default=None)50+ self.parser.add_argument("-o", "--output", default=None)
51- self.parser.add_argument(51+ self.parser.add_argument(
52- "-p",52+ "-p",
53- "--prefix",53+ "--prefix",
54- help=(54+ help=(
55- "Common filename prefix to strip such that rank can be extracted. "55+ "Common filename prefix to strip such that rank can be extracted. "
56- "If not specified, will attempt to infer a common prefix."56+ "If not specified, will attempt to infer a common prefix."
57- ),57+ ),
58- default=None,58+ default=None,
59- )59+ )
60- self.parser.add_argument("-j", "--just_print_entries", action="store_true")60+ self.parser.add_argument("-j", "--just_print_entries", action="store_true")
61- self.parser.add_argument("-v", "--verbose", action="store_true")61+ self.parser.add_argument("-v", "--verbose", action="store_true")
62- 62+ 
63- def parse_args(self: "JobConfig", args: Optional[Sequence[str]]) -> argparse.Namespace:63+ def parse_args(self: "JobConfig", args: Optional[Sequence[str]]) -> argparse.Namespace:
64- args = self.parser.parse_args(args)64+ args = self.parser.parse_args(args)
65- if args.selected_ranks is not None and not args.just_print_entries:65+ if args.selected_ranks is not None and not args.just_print_entries:
66- raise ValueError("Cannot use --selected-ranks without --just-print-entries")66+ raise ValueError("Cannot use --selected-ranks without --just-print-entries")
67- if args.pg_filters is not None and not args.just_print_entries:67+ if args.pg_filters is not None and not args.just_print_entries:
68- raise ValueError("Cannot use --pg-filters without --just-print-entries")68+ raise ValueError("Cannot use --pg-filters without --just-print-entries")
69- if args.verbose:69+ if args.verbose:
70- logger.set_log_level(logging.DEBUG)70+ logger.set_log_level(logging.DEBUG)
71- return args71+ return args
Mtools/flight_recorder/components/fr_logger.py+42-42
@@ -1,42 +1,42 @@
1-import logging1+import logging
2-from typing import Any, Callable, Optional2+from typing import Any, Callable, Optional
3- 3+ 
4- 4+ 
5-class FlightRecorderLogger:5+class FlightRecorderLogger:
6- _instance: Optional[Any] = None6+ _instance: Optional[Any] = None
7- logger: logging.Logger7+ logger: logging.Logger
8- 8+ 
9- def __init__(self) -> None:9+ def __init__(self) -> None:
10- self.logger: logging.Logger = logging.getLogger("Flight Recorder")10+ self.logger: logging.Logger = logging.getLogger("Flight Recorder")
11- 11+ 
12- def __new__(cls) -> Any:12+ def __new__(cls) -> Any:
13- if cls._instance is None:13+ if cls._instance is None:
14- cls._instance = super().__new__(cls)14+ cls._instance = super().__new__(cls)
15- cls._instance.logger = logging.getLogger("Flight Recorder")15+ cls._instance.logger = logging.getLogger("Flight Recorder")
16- cls._instance.logger.setLevel(logging.INFO)16+ cls._instance.logger.setLevel(logging.INFO)
17- ch = logging.StreamHandler()17+ ch = logging.StreamHandler()
18- cls._instance.logger.addHandler(ch)18+ cls._instance.logger.addHandler(ch)
19- return cls._instance19+ return cls._instance
20- 20+ 
21- def set_log_level(self, level: int) -> None:21+ def set_log_level(self, level: int) -> None:
22- self.logger.setLevel(level)22+ self.logger.setLevel(level)
23- 23+ 
24- @property24+ @property
25- def debug(self) -> Callable[..., None]:25+ def debug(self) -> Callable[..., None]:
26- return self.logger.debug26+ return self.logger.debug
27- 27+ 
28- @property28+ @property
29- def info(self) -> Callable[..., None]:29+ def info(self) -> Callable[..., None]:
30- return self.logger.info30+ return self.logger.info
31- 31+ 
32- @property32+ @property
33- def warning(self) -> Callable[..., None]:33+ def warning(self) -> Callable[..., None]:
34- return self.logger.warning34+ return self.logger.warning
35- 35+ 
36- @property36+ @property
37- def error(self) -> Callable[..., None]:37+ def error(self) -> Callable[..., None]:
38- return self.logger.error38+ return self.logger.error
39- 39+ 
40- @property40+ @property
41- def critical(self) -> Callable[..., None]:41+ def critical(self) -> Callable[..., None]:
42- return self.logger.critical42+ return self.logger.critical
Mtools/flight_recorder/components/loader.py+83-83
@@ -1,83 +1,83 @@
1-import os1+import os
2-import pickle2+import pickle
3-import re3+import re
4-from collections import defaultdict4+from collections import defaultdict
5- 5+ 
6-from tools.flight_recorder.components.fr_logger import FlightRecorderLogger6+from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
7-from tools.flight_recorder.components.utils import get_valid_read_path7+from tools.flight_recorder.components.utils import get_valid_read_path
8- 8+ 
9-MAX_DEPTH = 39+MAX_DEPTH = 3
10- 10+ 
11-logger: FlightRecorderLogger = FlightRecorderLogger()11+logger: FlightRecorderLogger = FlightRecorderLogger()
12- 12+ 
13-SAFE_CLASSES = {13+SAFE_CLASSES = {
14- # Built-in security type14+ # Built-in security type
15- "builtins": {"str", "int", "float", "list", "dict", "tuple"},15+ "builtins": {"str", "int", "float", "list", "dict", "tuple"},
16-}16+}
17- 17+ 
18-exp = re.compile(r"^([a-zA-Z0-9_]{0,100}?)(\d+)$")18+exp = re.compile(r"^([a-zA-Z0-9_]{0,100}?)(\d+)$")
19- 19+ 
20- 20+ 
21-class SafeUnpickler(pickle.Unpickler):21+class SafeUnpickler(pickle.Unpickler):
22- def find_class(self, module, name):22+ def find_class(self, module, name):
23- # Check if the module and class are in the whitelist23+ # Check if the module and class are in the whitelist
24- if module in SAFE_CLASSES and name in SAFE_CLASSES[module]:24+ if module in SAFE_CLASSES and name in SAFE_CLASSES[module]:
25- return super().find_class(module, name)25+ return super().find_class(module, name)
26- raise pickle.UnpicklingError(f"Forbidden class: {module}.{name}")26+ raise pickle.UnpicklingError(f"Forbidden class: {module}.{name}")
27- 27+ 
28- 28+ 
29-def read_dump(prefix, filename):29+def read_dump(prefix, filename):
30- basename = os.path.basename(filename)30+ basename = os.path.basename(filename)
31- try:31+ try:
32- rank = int(basename[len(prefix):])32+ rank = int(basename[len(prefix):])
33- except ValueError as e:33+ except ValueError as e:
34- raise ValueError(f"Cannot extract rank from '{basename}' with prefix '{prefix}'.") from e34+ raise ValueError(f"Cannot extract rank from '{basename}' with prefix '{prefix}'.") from e
35- filename = get_valid_read_path(filename)35+ filename = get_valid_read_path(filename)
36- try:36+ try:
37- with open(filename, "rb") as infile:37+ with open(filename, "rb") as infile:
38- dump = SafeUnpickler(infile).load()38+ dump = SafeUnpickler(infile).load()
39- except Exception as e:39+ except Exception as e:
40- logger.error(f"Failed to load data from {filename}: {e}")40+ logger.error(f"Failed to load data from {filename}: {e}")
41- return rank, dump41+ return rank, dump
42- 42+ 
43- 43+ 
44-def determine_prefix(files):44+def determine_prefix(files):
45- possible_prefixes: defaultdict[str, set[int]] = defaultdict(set)45+ possible_prefixes: defaultdict[str, set[int]] = defaultdict(set)
46- for f in files:46+ for f in files:
47- m = exp.search(f)47+ m = exp.search(f)
48- if m:48+ if m:
49- p, r = m.groups()49+ p, r = m.groups()
50- possible_prefixes[p].add(int(r))50+ possible_prefixes[p].add(int(r))
51- if len(possible_prefixes) == 1:51+ if len(possible_prefixes) == 1:
52- prefix = next(iter(possible_prefixes))52+ prefix = next(iter(possible_prefixes))
53- return prefix53+ return prefix
54- else:54+ else:
55- raise ValueError(55+ raise ValueError(
56- "Unable to automatically determine the common prefix for the trace file names. "56+ "Unable to automatically determine the common prefix for the trace file names. "
57- "Please specify --prefix argument manually"57+ "Please specify --prefix argument manually"
58- )58+ )
59- 59+ 
60- 60+ 
61-def read_dir(args):61+def read_dir(args):
62- """Load recorder data for all ranks"""62+ """Load recorder data for all ranks"""
63- prefix = args.prefix63+ prefix = args.prefix
64- path = args.trace_dir64+ path = args.trace_dir
65- details = {}65+ details = {}
66- version = ""66+ version = ""
67- for root, _, files in os.walk(path):67+ for root, _, files in os.walk(path):
68- current_depth = root.count(os.sep) - path.count(os.sep)68+ current_depth = root.count(os.sep) - path.count(os.sep)
69- if current_depth > MAX_DEPTH:69+ if current_depth > MAX_DEPTH:
70- logger.error("The current file depth has exceeded the maximum depth limit, which is set to {MAX_DEPTH}.")70+ logger.error("The current file depth has exceeded the maximum depth limit, which is set to {MAX_DEPTH}.")
71- break71+ break
72- if prefix is None:72+ if prefix is None:
73- prefix = determine_prefix(files)73+ prefix = determine_prefix(files)
74- for f in files:74+ for f in files:
75- if "py_traceback" in f:75+ if "py_traceback" in f:
76- continue76+ continue
77- if f.find(prefix) != 0:77+ if f.find(prefix) != 0:
78- continue78+ continue
79- rank, dump = read_dump(prefix, os.path.join(root, f))79+ rank, dump = read_dump(prefix, os.path.join(root, f))
80- details[rank] = dump80+ details[rank] = dump
81- if not version:81+ if not version:
82- version = str(details[rank]["version"])82+ version = str(details[rank]["version"])
83- return details, version83+ return details, version
Mtools/flight_recorder/components/types.py+550-550
@@ -1,550 +1,550 @@
1-import math1+import math
2-import os2+import os
3-from enum import auto, Enum3+from enum import auto, Enum
4-from typing import (4+from typing import (
5- _eval_type,5+ _eval_type,
6- Any,6+ Any,
7- Generic,7+ Generic,
8- NamedTuple,8+ NamedTuple,
9- Optional,9+ Optional,
10- TypeVar,10+ TypeVar,
11-)11+)
12- 12+ 
13-from tools.flight_recorder.components.fr_logger import FlightRecorderLogger13+from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
14- 14+ 
15- 15+ 
16-T = TypeVar("T", bound=NamedTuple)16+T = TypeVar("T", bound=NamedTuple)
17- 17+ 
18- 18+ 
19-class Ref(Generic[T]):19+class Ref(Generic[T]):
20- pass20+ pass
21- 21+ 
22- 22+ 
23-class TypeInfo(NamedTuple):23+class TypeInfo(NamedTuple):
24- name: str24+ name: str
25- fields: list[tuple[str, type]] # type: ignore[type-arg]25+ fields: list[tuple[str, type]] # type: ignore[type-arg]
26- 26+ 
27- @classmethod27+ @classmethod
28- def from_type(cls, c: T) -> "TypeInfo":28+ def from_type(cls, c: T) -> "TypeInfo":
O
OopenLiBingCI5月17日

此条代码评论区间+23至+28

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
29- if hasattr(c, "__name__"):29+ if hasattr(c, "__name__"):
30- name = c.__name__30+ name = c.__name__
31- else:31+ else:
32- name = str(c)32+ name = str(c)
33- return cls(33+ return cls(
34- name,34+ name,
35- [(f, _eval_type(c.__annotations__[f], globals(), {})) for f in c._fields],35+ [(f, _eval_type(c.__annotations__[f], globals(), {})) for f in c._fields],
36- )36+ )
37- 37+ 
38- 38+ 
39-class MatchState(Enum):39+class MatchState(Enum):
40- """40+ """
41- Enum representing the possible states of matching for collective operations.41+ Enum representing the possible states of matching for collective operations.
42- 42+ 
43- - FULLY_MATCHED: Indicates that all aspects of the collective operations match.43+ - FULLY_MATCHED: Indicates that all aspects of the collective operations match.
44- - COLLECTIVE_TYPE_MISMATCH: The types of the collective operations differ.44+ - COLLECTIVE_TYPE_MISMATCH: The types of the collective operations differ.
45- - SIZE_OR_SYNTAX_MISMATCH: There is a mismatch in input/output sizes or violation of collective syntax.45+ - SIZE_OR_SYNTAX_MISMATCH: There is a mismatch in input/output sizes or violation of collective syntax.
46- - COLLECTIVE_STATE_MISMATCH:46+ - COLLECTIVE_STATE_MISMATCH:
47- The states of the collective not same, such as one finished while another just started or scheduled.47+ The states of the collective not same, such as one finished while another just started or scheduled.
48- - COLLECTIVE_DTYPE_MISMATCH: The data types of the collective input/output differ.48+ - COLLECTIVE_DTYPE_MISMATCH: The data types of the collective input/output differ.
49- - UNDECIDED:49+ - UNDECIDED:
50- The match status is ambiguous or cannot be determined, e.g., we might need to check all ranks for alltoall_base.50+ The match status is ambiguous or cannot be determined, e.g., we might need to check all ranks for alltoall_base.
51- """51+ """
52- 52+ 
53- FULLY_MATCHED = auto()53+ FULLY_MATCHED = auto()
54- COLLECTIVE_TYPE_MISMATCH = auto()54+ COLLECTIVE_TYPE_MISMATCH = auto()
55- SIZE_OR_SYNTAX_MISMATCH = auto()55+ SIZE_OR_SYNTAX_MISMATCH = auto()
56- COLLECTIVE_STATE_MISMATCH = auto()56+ COLLECTIVE_STATE_MISMATCH = auto()
57- COLLECTIVE_DTYPE_MISMATCH = auto()57+ COLLECTIVE_DTYPE_MISMATCH = auto()
58- UNDECIDED = auto()58+ UNDECIDED = auto()
59- 59+ 
60- 60+ 
61-class MatchInfo:61+class MatchInfo:
62- """62+ """
63- Aside from the match state, we also store some dynamic info for the match such as the culprit rank63+ Aside from the match state, we also store some dynamic info for the match such as the culprit rank
64- or collective state that caused the mismatch.64+ or collective state that caused the mismatch.
65- """65+ """
66- 66+ 
67- def __init__(self, state: MatchState, culprit: Optional[str] = None) -> None:67+ def __init__(self, state: MatchState, culprit: Optional[str] = None) -> None:
68- self._state = state68+ self._state = state
69- self.culprit = culprit69+ self.culprit = culprit
70- 70+ 
71- def __str__(self) -> str:71+ def __str__(self) -> str:
72- details = f", {self.culprit}" if getattr(self, "culprit", None) else ""72+ details = f", {self.culprit}" if getattr(self, "culprit", None) else ""
73- return f"Error type: {self._state.name}{details}"73+ return f"Error type: {self._state.name}{details}"
74- 74+ 
75- @property75+ @property
76- def state(self) -> MatchState:76+ def state(self) -> MatchState:
77- return self._state77+ return self._state
78- 78+ 
79- 79+ 
80-class Group(NamedTuple):80+class Group(NamedTuple):
81- id: str81+ id: str
82- desc: str82+ desc: str
83- size: int83+ size: int
84- 84+ 
85- 85+ 
86-class Membership(NamedTuple):86+class Membership(NamedTuple):
87- group_id: str87+ group_id: str
88- global_rank: int88+ global_rank: int
89- 89+ 
90- 90+ 
91-class Traceback(NamedTuple):91+class Traceback(NamedTuple):
92- id: int92+ id: int
93- frames: str93+ frames: str
94- 94+ 
95- 95+ 
96-class Collective(NamedTuple):96+class Collective(NamedTuple):
97- id: int97+ id: int
98- group_id: str98+ group_id: str
99- pass_check: bool99+ pass_check: bool
100- collective_seq_id: int100+ collective_seq_id: int
101- p2p_seq_id: int101+ p2p_seq_id: int
102- record_id: int102+ record_id: int
103- pg_desc: str103+ pg_desc: str
104- collective_name: str104+ collective_name: str
105- input_sizes: list[list[int]]105+ input_sizes: list[list[int]]
106- output_sizes: list[list[int]]106+ output_sizes: list[list[int]]
107- expected_ranks: set[int]107+ expected_ranks: set[int]
108- collective_state: str108+ collective_state: str
109- collective_frames: list[dict[str, str]]109+ collective_frames: list[dict[str, str]]
110- input_numel: Optional[int] = None110+ input_numel: Optional[int] = None
111- output_numel: Optional[int] = None111+ output_numel: Optional[int] = None
112- missing_ranks: Optional[set[int]] = None112+ missing_ranks: Optional[set[int]] = None
113- mismatch_collectives: Optional[dict[int, "Collective"]] = None113+ mismatch_collectives: Optional[dict[int, "Collective"]] = None
114- type_of_mismatch: Optional[MatchInfo] = None114+ type_of_mismatch: Optional[MatchInfo] = None
115- 115+ 
116- 116+ 
117-class HCCLCall(NamedTuple):117+class HCCLCall(NamedTuple):
118- id: int118+ id: int
119- collective_id: Ref[Collective]119+ collective_id: Ref[Collective]
120- group_id: str120+ group_id: str
121- global_rank: int # technically Ref[Process] once we have it121+ global_rank: int # technically Ref[Process] once we have it
122- traceback_id: Ref[Traceback]122+ traceback_id: Ref[Traceback]
123- collective_type: str123+ collective_type: str
124- sizes: list[list[int]]124+ sizes: list[list[int]]
125- 125+ 
126- 126+ 
127-class Database(NamedTuple):127+class Database(NamedTuple):
128- groups: list[Group]128+ groups: list[Group]
129- memberships: list[Membership]129+ memberships: list[Membership]
130- tracebacks: list[Traceback]130+ tracebacks: list[Traceback]
131- collectives: list[Collective]131+ collectives: list[Collective]
132- hcclcalls: list[HCCLCall]132+ hcclcalls: list[HCCLCall]
133- 133+ 
134- 134+ 
135-types = [135+types = [
136- TypeInfo.from_type(t) # type: ignore[type-var]136+ TypeInfo.from_type(t) # type: ignore[type-var]
137- for t in globals().values()137+ for t in globals().values()
138- if (isinstance(t, type) and issubclass(t, tuple) and hasattr(t, "_fields") and t is not TypeInfo)138+ if (isinstance(t, type) and issubclass(t, tuple) and hasattr(t, "_fields") and t is not TypeInfo)
O
OopenLiBingCI5月17日

此条代码评论区间+132至+138

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
139-]139+]
140- 140+ 
141- 141+ 
142-COLLECTIVES = {142+COLLECTIVES = {
143- "broadcast",143+ "broadcast",
144- "_broadcast_oop",144+ "_broadcast_oop",
145- "reduce",145+ "reduce",
146- "_reduce_oop",146+ "_reduce_oop",
147- "all_gather",147+ "all_gather",
148- "all_reduce",148+ "all_reduce",
149- "_all_gather_base",149+ "_all_gather_base",
150- "all_gather_into_tensor_coalesced",150+ "all_gather_into_tensor_coalesced",
151- "reduce_scatter",151+ "reduce_scatter",
152- "reduce_scatter_tensor_coalesced",152+ "reduce_scatter_tensor_coalesced",
153- "_reduce_scatter_base",153+ "_reduce_scatter_base",
154- "gather",154+ "gather",
155- "scatter",155+ "scatter",
156- "all_to_all",156+ "all_to_all",
157- "all_reduce_barrier",157+ "all_reduce_barrier",
158- "allreduce_coalesced",158+ "allreduce_coalesced",
159- "ALLGATHER_coalesced",159+ "ALLGATHER_coalesced",
160- "REDUCE_SCATTER_coalesced",160+ "REDUCE_SCATTER_coalesced",
161-}161+}
162- 162+ 
163-P2P = {163+P2P = {
164- "send",164+ "send",
165- "recv",165+ "recv",
166-}166+}
167- 167+ 
168- 168+ 
169-class EntryState:169+class EntryState:
170- """170+ """
171- Util class to keep track of the state of an entry and standardize the way we171+ Util class to keep track of the state of an entry and standardize the way we
172- log the error info during analysis.172+ log the error info during analysis.
173- """173+ """
174- 174+ 
175- def __init__(self, entry: dict[str, Any], expected_ranks: set[int]) -> None:175+ def __init__(self, entry: dict[str, Any], expected_ranks: set[int]) -> None:
176- self.pg_name = entry["process_group"][0]176+ self.pg_name = entry["process_group"][0]
177- self.desc = entry["process_group"][1]177+ self.desc = entry["process_group"][1]
178- self.pg_desc = f"{self.pg_name}:{self.desc}" if self.desc != "undefined" else self.pg_name178+ self.pg_desc = f"{self.pg_name}:{self.desc}" if self.desc != "undefined" else self.pg_name
179- self.profiling_name = entry["profiling_name"]179+ self.profiling_name = entry["profiling_name"]
180- self.collective_seq_id = entry["collective_seq_id"]180+ self.collective_seq_id = entry["collective_seq_id"]
181- self.p2p_seq_id = entry["p2p_seq_id"]181+ self.p2p_seq_id = entry["p2p_seq_id"]
182- self.record_id = entry["record_id"]182+ self.record_id = entry["record_id"]
183- self.input_sizes = entry["input_sizes"]183+ self.input_sizes = entry["input_sizes"]
184- self.output_sizes = entry["output_sizes"]184+ self.output_sizes = entry["output_sizes"]
185- self.collective_state = entry["state"]185+ self.collective_state = entry["state"]
186- self.collective_frames = entry.get("frames", [])186+ self.collective_frames = entry.get("frames", [])
187- self.expected_ranks = expected_ranks187+ self.expected_ranks = expected_ranks
188- self.missing_ranks: set[int]188+ self.missing_ranks: set[int]
189- self.input_numel: int189+ self.input_numel: int
190- self.output_numel: int190+ self.output_numel: int
191- self.errors: set[tuple[int, MatchInfo]]191+ self.errors: set[tuple[int, MatchInfo]]
192- 192+ 
193- 193+ 
194- def log(194+ def log(
195- self,195+ self,
196- logger: FlightRecorderLogger,196+ logger: FlightRecorderLogger,
197- logger_msg: str,197+ logger_msg: str,
198- frame_formatter: Any,198+ frame_formatter: Any,
199- additional_info: dict = None,199+ additional_info: dict = None,
200- ) -> None:200+ ) -> None:
201- logger.info(201+ logger.info(
202- logger_msg,202+ logger_msg,
203- self.collective_seq_id,203+ self.collective_seq_id,
204- )204+ )
205- logger.info("internal record id: %s", self.record_id)205+ logger.info("internal record id: %s", self.record_id)
206- logger.info("group info: %s", self.pg_desc)206+ logger.info("group info: %s", self.pg_desc)
207- logger.info("collective: %s", self.profiling_name)207+ logger.info("collective: %s", self.profiling_name)
208- if additional_info and "missing_ranks" in additional_info:208+ if additional_info and "missing_ranks" in additional_info:
209- missing_ranks = additional_info["missing_ranks"]209+ missing_ranks = additional_info["missing_ranks"]
210- self.missing_ranks = missing_ranks210+ self.missing_ranks = missing_ranks
211- logger.info("missing ranks: %s", missing_ranks)211+ logger.info("missing ranks: %s", missing_ranks)
212- if additional_info and "total_numel" in additional_info:212+ if additional_info and "total_numel" in additional_info:
213- total_numel = additional_info["total_numel"]213+ total_numel = additional_info["total_numel"]
214- self.input_numel = total_numel[0]214+ self.input_numel = total_numel[0]
215- self.output_numel = total_numel[1]215+ self.output_numel = total_numel[1]
216- logger.info("total input numel: %d", total_numel[0])216+ logger.info("total input numel: %d", total_numel[0])
217- logger.info("total output numel: %d", total_numel[1])217+ logger.info("total output numel: %d", total_numel[1])
218- logger.info("input sizes: %s", self.input_sizes)218+ logger.info("input sizes: %s", self.input_sizes)
219- logger.info("output sizes: %s", self.output_sizes)219+ logger.info("output sizes: %s", self.output_sizes)
220- logger.info("world size: %d", len(self.expected_ranks))220+ logger.info("world size: %d", len(self.expected_ranks))
221- logger.info("expected ranks: %s", str(self.expected_ranks))221+ logger.info("expected ranks: %s", str(self.expected_ranks))
222- logger.info("collective state: %s", self.collective_state)222+ logger.info("collective state: %s", self.collective_state)
223- if additional_info and "errors" in additional_info:223+ if additional_info and "errors" in additional_info:
224- errors = additional_info["errors"]224+ errors = additional_info["errors"]
225- self.errors = errors225+ self.errors = errors
226- error_msg = ", ".join(f"Culprit rank {error[0]}; {str(error[1])}" for error in errors)226+ error_msg = ", ".join(f"Culprit rank {error[0]}; {str(error[1])}" for error in errors)
227- logger.info("error msg: %s", error_msg)227+ logger.info("error msg: %s", error_msg)
228- logger.info("collective stack trace: \n %s", frame_formatter(self.collective_frames))228+ logger.info("collective stack trace: \n %s", frame_formatter(self.collective_frames))
229- 229+ 
230- def to_collective(230+ def to_collective(
231- self,231+ self,
232- collective_id: int,232+ collective_id: int,
233- errors: Optional[set[tuple[int, MatchInfo]]] = None,233+ errors: Optional[set[tuple[int, MatchInfo]]] = None,
234- idx_map: Optional[dict[int, int]] = None,234+ idx_map: Optional[dict[int, int]] = None,
235- all_entries: Optional[dict[int, list[dict[str, Any]]]] = None,235+ all_entries: Optional[dict[int, list[dict[str, Any]]]] = None,
236- ) -> Collective:236+ ) -> Collective:
237- if not errors:237+ if not errors:
238- return Collective(238+ return Collective(
239- id=collective_id,239+ id=collective_id,
240- group_id=self.pg_name,240+ group_id=self.pg_name,
241- record_id=self.record_id,241+ record_id=self.record_id,
242- pg_desc=self.pg_desc,242+ pg_desc=self.pg_desc,
243- pass_check=True,243+ pass_check=True,
244- collective_seq_id=self.collective_seq_id,244+ collective_seq_id=self.collective_seq_id,
245- p2p_seq_id=self.p2p_seq_id,245+ p2p_seq_id=self.p2p_seq_id,
246- collective_name=self.profiling_name,246+ collective_name=self.profiling_name,
247- input_sizes=self.input_sizes,247+ input_sizes=self.input_sizes,
248- output_sizes=self.output_sizes,248+ output_sizes=self.output_sizes,
249- expected_ranks=self.expected_ranks,249+ expected_ranks=self.expected_ranks,
250- collective_state=self.collective_state,250+ collective_state=self.collective_state,
251- collective_frames=self.collective_frames,251+ collective_frames=self.collective_frames,
252- missing_ranks=getattr(self, "missing_ranks", None),252+ missing_ranks=getattr(self, "missing_ranks", None),
253- )253+ )
254- else:254+ else:
255- if idx_map is None:255+ if idx_map is None:
256- raise ValueError("idx_map cannot be None")256+ raise ValueError("idx_map cannot be None")
257- if all_entries is None:257+ if all_entries is None:
258- raise ValueError("all_entries cannot be None")258+ raise ValueError("all_entries cannot be None")
259- mismatch_collectives = {}259+ mismatch_collectives = {}
260- for rank, error in errors:260+ for rank, error in errors:
261- idx = idx_map[rank]261+ idx = idx_map[rank]
262- entry = all_entries[rank][idx]262+ entry = all_entries[rank][idx]
263- desc = entry["process_group"][1]263+ desc = entry["process_group"][1]
264- pg_name = entry["process_group"][0]264+ pg_name = entry["process_group"][0]
265- mismatch_collectives[rank] = Collective(265+ mismatch_collectives[rank] = Collective(
266- id=collective_id,266+ id=collective_id,
267- group_id=entry["process_group"][0],267+ group_id=entry["process_group"][0],
268- record_id=entry["record_id"],268+ record_id=entry["record_id"],
269- pg_desc=f"{pg_name}:{desc}" if desc != "undefined" else pg_name,269+ pg_desc=f"{pg_name}:{desc}" if desc != "undefined" else pg_name,
270- pass_check=False,270+ pass_check=False,
271- collective_seq_id=entry["collective_seq_id"],271+ collective_seq_id=entry["collective_seq_id"],
272- p2p_seq_id=entry["p2p_seq_id"],272+ p2p_seq_id=entry["p2p_seq_id"],
273- collective_name=entry["profiling_name"],273+ collective_name=entry["profiling_name"],
274- input_sizes=entry["input_sizes"],274+ input_sizes=entry["input_sizes"],
275- output_sizes=entry["output_sizes"],275+ output_sizes=entry["output_sizes"],
276- expected_ranks=self.expected_ranks,276+ expected_ranks=self.expected_ranks,
277- collective_state=entry["state"],277+ collective_state=entry["state"],
278- collective_frames=entry.get("frames", []),278+ collective_frames=entry.get("frames", []),
279- type_of_mismatch=error,279+ type_of_mismatch=error,
280- )280+ )
281- return Collective(281+ return Collective(
282- id=collective_id,282+ id=collective_id,
283- group_id=self.pg_name,283+ group_id=self.pg_name,
284- record_id=self.record_id,284+ record_id=self.record_id,
285- pg_desc=self.pg_desc,285+ pg_desc=self.pg_desc,
286- pass_check=False,286+ pass_check=False,
287- collective_seq_id=self.collective_seq_id,287+ collective_seq_id=self.collective_seq_id,
288- p2p_seq_id=self.p2p_seq_id,288+ p2p_seq_id=self.p2p_seq_id,
289- collective_name=self.profiling_name,289+ collective_name=self.profiling_name,
290- input_sizes=self.input_sizes,290+ input_sizes=self.input_sizes,
291- output_sizes=self.output_sizes,291+ output_sizes=self.output_sizes,
292- expected_ranks=self.expected_ranks,292+ expected_ranks=self.expected_ranks,
293- collective_state=self.collective_state,293+ collective_state=self.collective_state,
294- collective_frames=self.collective_frames,294+ collective_frames=self.collective_frames,
295- input_numel=self.input_numel if hasattr(self, "input_numel") else None,295+ input_numel=self.input_numel if hasattr(self, "input_numel") else None,
296- output_numel=self.output_numel if hasattr(self, "output_numel") else None,296+ output_numel=self.output_numel if hasattr(self, "output_numel") else None,
297- missing_ranks=self.missing_ranks if hasattr(self, "missing_ranks") else None,297+ missing_ranks=self.missing_ranks if hasattr(self, "missing_ranks") else None,
298- mismatch_collectives=mismatch_collectives,298+ mismatch_collectives=mismatch_collectives,
299- )299+ )
300- 300+ 
301- def to_hccl_call(301+ def to_hccl_call(
302- self,302+ self,
303- all_entries: dict[int, list[dict[str, Any]]],303+ all_entries: dict[int, list[dict[str, Any]]],
304- idx_map: dict[int, int],304+ idx_map: dict[int, int],
305- hccl_call_id: int,305+ hccl_call_id: int,
306- collective_id: Any,306+ collective_id: Any,
307- ) -> list[HCCLCall]:307+ ) -> list[HCCLCall]:
308- result = []308+ result = []
309- for i, k in idx_map.items():309+ for i, k in idx_map.items():
310- all_entries[i].pop(k)310+ all_entries[i].pop(k)
311- result.append(311+ result.append(
312- HCCLCall(312+ HCCLCall(
313- id=hccl_call_id,313+ id=hccl_call_id,
314- collective_id=collective_id,314+ collective_id=collective_id,
315- group_id=self.pg_name, # type: ignore[arg-type]315+ group_id=self.pg_name, # type: ignore[arg-type]
316- global_rank=i,316+ global_rank=i,
317- traceback_id=0, # type: ignore[arg-type]317+ traceback_id=0, # type: ignore[arg-type]
O
OopenLiBingCI5月17日

此条代码评论区间+313至+317

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
318- collective_type=self.profiling_name,318+ collective_type=self.profiling_name,
319- sizes=self.input_sizes,319+ sizes=self.input_sizes,
O
OopenLiBingCI5月17日

此条代码评论区间+315至+319

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
320- )320+ )
321- )321+ )
322- hccl_call_id += 1322+ hccl_call_id += 1
323- return result323+ return result
324- 324+ 
325- 325+ 
326-class Op:326+class Op:
327- """Parses relevant info about operation out of 'event' dict327+ """Parses relevant info about operation out of 'event' dict
328- 328+ 
329- examples of supported `profiling_name`s:329+ examples of supported `profiling_name`s:
330- hccl:broadcast330+ hccl:broadcast
331- hccl:send 1->2331+ hccl:send 1->2
332- hccl:recv 3<-0332+ hccl:recv 3<-0
333- """333+ """
334- MISSING_FRAMES_ERR = "Event missing 'frames' field or empty frames array"334+ MISSING_FRAMES_ERR = "Event missing 'frames' field or empty frames array"
335- INVALID_FRAME_ERR = "Frame[0] missing 'name' field"335+ INVALID_FRAME_ERR = "Frame[0] missing 'name' field"
336- 336+ 
337- 337+ 
338- def __init__(self, event: dict[Any, Any], memberships: dict[str, set[Any]], pg_name: str):338+ def __init__(self, event: dict[Any, Any], memberships: dict[str, set[Any]], pg_name: str):
339- 339+ 
340- frames = event.get("frames")340+ frames = event.get("frames")
341- if not frames: 341+ if not frames:
342- raise ValueError(self.MISSING_FRAMES_ERR)342+ raise ValueError(self.MISSING_FRAMES_ERR)
343- first_frame = frames[0] if len(frames) > 0 else None343+ first_frame = frames[0] if len(frames) > 0 else None
344- if not first_frame:344+ if not first_frame:
345- raise ValueError(self.MISSING_FRAMES_ERR)345+ raise ValueError(self.MISSING_FRAMES_ERR)
346- self.profiling_name = first_frame.get("name")346+ self.profiling_name = first_frame.get("name")
347- if self.profiling_name is None:347+ if self.profiling_name is None:
348- raise ValueError(self.INVALID_FRAME_ERR)348+ raise ValueError(self.INVALID_FRAME_ERR)
349- parts = self.profiling_name.split(":")349+ parts = self.profiling_name.split(":")
350- self.type = parts[0]350+ self.type = parts[0]
351- meta = parts[1] if len(parts) == 2 else None351+ meta = parts[1] if len(parts) == 2 else None
352- self.state = event.get("state")352+ self.state = event.get("state")
353- self.pg_name, self.pg_desc = event.get("process_group")353+ self.pg_name, self.pg_desc = event.get("process_group")
354- if type == "send":354+ if type == "send":
355- s, d = meta.split("->")355+ s, d = meta.split("->")
356- self._src, self._dst = int(s), int(d)356+ self._src, self._dst = int(s), int(d)
357- elif type == "recv":357+ elif type == "recv":
358- d, s = meta.split("<-")358+ d, s = meta.split("<-")
359- self._dst, self._src = int(d), int(s)359+ self._dst, self._src = int(d), int(s)
360- else:360+ else:
361- self._src, self._dst = -1, -1361+ self._src, self._dst = -1, -1
362- self._init_global_src_dst(memberships[pg_name])362+ self._init_global_src_dst(memberships[pg_name])
363- self.pg_size = len(memberships[pg_name])363+ self.pg_size = len(memberships[pg_name])
364- if type in P2P | COLLECTIVES:364+ if type in P2P | COLLECTIVES:
365- self.input_sizes = event.get("input_sizes")365+ self.input_sizes = event.get("input_sizes")
366- self.output_sizes = event.get("output_sizes")366+ self.output_sizes = event.get("output_sizes")
367- else:367+ else:
368- self.input_sizes, self.output_sizes = None, None368+ self.input_sizes, self.output_sizes = None, None
369- self.collective_seq_id = event.get("collective_seq_id")369+ self.collective_seq_id = event.get("collective_seq_id")
370- self.p2p_seq_id = event.get("p2p_seq_id")370+ self.p2p_seq_id = event.get("p2p_seq_id")
371- self.input_dtypes = event.get("input_dtypes")371+ self.input_dtypes = event.get("input_dtypes")
372- self.output_dtypes = event.get("output_dtypes")372+ self.output_dtypes = event.get("output_dtypes")
373- self.time_created_ns = event.get("time_created_ns")373+ self.time_created_ns = event.get("time_created_ns")
374- self.collective_frames = event.get("frames", [])374+ self.collective_frames = event.get("frames", [])
375- self.is_verbose = os.getenv("FR_TRACE_VERBOSE_OUTPUT", "0") == "1"375+ self.is_verbose = os.getenv("FR_TRACE_VERBOSE_OUTPUT", "0") == "1"
376- 376+ 
377- def _init_global_src_dst(self, pg_ranks: set[Any]) -> None:377+ def _init_global_src_dst(self, pg_ranks: set[Any]) -> None:
378- pg_ranks = sorted(pg_ranks)378+ pg_ranks = sorted(pg_ranks)
379- self._src_g = pg_ranks[self._src] if self._src is not None else None379+ self._src_g = pg_ranks[self._src] if self._src is not None else None
380- self._dst_g = pg_ranks[self._dst] if self._dst is not None else None380+ self._dst_g = pg_ranks[self._dst] if self._dst is not None else None
381- 381+ 
382- @property382+ @property
383- def src(self) -> int:383+ def src(self) -> int:
384- if self.type not in P2P:384+ if self.type not in P2P:
385- raise ValueError(f"Can't get src of non-p2p op (type: {self.type})")385+ raise ValueError(f"Can't get src of non-p2p op (type: {self.type})")
386- return self._src386+ return self._src
387- 387+ 
388- @property388+ @property
389- def dst(self) -> int:389+ def dst(self) -> int:
390- if self.type not in P2P:390+ if self.type not in P2P:
391- raise ValueError(f"Can't get dst of non-p2p op (type: {self.type})")391+ raise ValueError(f"Can't get dst of non-p2p op (type: {self.type})")
392- return self._dst392+ return self._dst
393- 393+ 
394- def __repr__(self) -> str:394+ def __repr__(self) -> str:
395- p2p_info = ""395+ p2p_info = ""
396- if self.type in P2P:396+ if self.type in P2P:
397- p2p_info = f"s={self._src_g} d={self._dst_g}"397+ p2p_info = f"s={self._src_g} d={self._dst_g}"
398- if self.is_verbose:398+ if self.is_verbose:
399- verbose_info = (399+ verbose_info = (
400- f"timestamp_created={self.time_created_ns}",400+ f"timestamp_created={self.time_created_ns}",
401- p2p_info,401+ p2p_info,
402- f"input_sizes={self.input_sizes}",402+ f"input_sizes={self.input_sizes}",
403- f"output_sizes={self.output_sizes}",403+ f"output_sizes={self.output_sizes}",
404- f"input_dtypes={self.input_dtypes}",404+ f"input_dtypes={self.input_dtypes}",
405- f"output_dtypes={self.output_dtypes}",405+ f"output_dtypes={self.output_dtypes}",
406- "collective_seq_id | p2p_seq_id=" f"{self.p2p_seq_id if self.type in P2P else self.collective_seq_id}",406+ "collective_seq_id | p2p_seq_id=" f"{self.p2p_seq_id if self.type in P2P else self.collective_seq_id}",
407- f"pg_name={self.pg_name}",407+ f"pg_name={self.pg_name}",
408- f"pg_description={self.pg_desc}",408+ f"pg_description={self.pg_desc}",
409- f"pg_size={self.pg_size}",409+ f"pg_size={self.pg_size}",
410- f"state={self.state}",410+ f"state={self.state}",
411- )411+ )
412- return f"{self.type}({', '.join(s for s in verbose_info if s)})"412+ return f"{self.type}({', '.join(s for s in verbose_info if s)})"
413- return f"{self.type}(%sinput_sizes={self.input_sizes}, state={self.state})" % (413+ return f"{self.type}(%sinput_sizes={self.input_sizes}, state={self.state})" % (
414- f"{p2p_info}, " if p2p_info else ""414+ f"{p2p_info}, " if p2p_info else ""
415- )415+ )
416- 416+ 
417- def has_different_dtypes_and_non_empty_sizes(self, other):417+ def has_different_dtypes_and_non_empty_sizes(self, other):
418- """418+ """
419- Check if the input/output dtypes are different and the sizes are non-empty.419+ Check if the input/output dtypes are different and the sizes are non-empty.
420- """420+ """
421- # Check if input/output dtypes are different and sizes are non-empty421+ # Check if input/output dtypes are different and sizes are non-empty
422- condition1 = set(self.input_dtypes) != set(self.output_dtypes) and self.input_sizes[0] and self.output_sizes[0]422+ condition1 = set(self.input_dtypes) != set(self.output_dtypes) and self.input_sizes[0] and self.output_sizes[0]
423- condition2 = set(self.input_dtypes) != set(other.input_dtypes) and self.input_sizes[0] and other.input_sizes[0]423+ condition2 = set(self.input_dtypes) != set(other.input_dtypes) and self.input_sizes[0] and other.input_sizes[0]
424- condition3 = (424+ condition3 = (
425- set(self.input_dtypes) != set(other.output_dtypes) and self.input_sizes[0] and other.output_sizes[0]425+ set(self.input_dtypes) != set(other.output_dtypes) and self.input_sizes[0] and other.output_sizes[0]
426- )426+ )
427- return condition1 or condition2 or condition3427+ return condition1 or condition2 or condition3
428- 428+ 
429- def match(self, other: "Op") -> MatchInfo:429+ def match(self, other: "Op") -> MatchInfo:
430- if self.type == "send":430+ if self.type == "send":
431- return (431+ return (
432- MatchInfo(MatchState.FULLY_MATCHED)432+ MatchInfo(MatchState.FULLY_MATCHED)
433- if (433+ if (
434- other.type == "recv"434+ other.type == "recv"
435- and self.src == other.src435+ and self.src == other.src
436- and self.dst == other.dst436+ and self.dst == other.dst
437- and self.input_sizes == other.output_sizes437+ and self.input_sizes == other.output_sizes
438- )438+ )
439- else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)439+ else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)
440- )440+ )
441- elif self.type == "recv":441+ elif self.type == "recv":
442- return (442+ return (
443- MatchInfo(MatchState.FULLY_MATCHED)443+ MatchInfo(MatchState.FULLY_MATCHED)
444- if (444+ if (
445- other.type == "send"445+ other.type == "send"
446- and self.src == other.src446+ and self.src == other.src
447- and self.dst == other.dst447+ and self.dst == other.dst
448- and self.output_sizes == other.input_sizes448+ and self.output_sizes == other.input_sizes
449- )449+ )
450- else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)450+ else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)
451- )451+ )
452- elif self.type in COLLECTIVES:452+ elif self.type in COLLECTIVES:
453- if self.type != other.type:453+ if self.type != other.type:
454- return MatchInfo(454+ return MatchInfo(
455- MatchState.COLLECTIVE_TYPE_MISMATCH,455+ MatchState.COLLECTIVE_TYPE_MISMATCH,
456- f"Expected collective type: '{self.type}' does not match found collective type: '{other.type}'",456+ f"Expected collective type: '{self.type}' does not match found collective type: '{other.type}'",
457- )457+ )
458- if self.state != other.state:458+ if self.state != other.state:
459- return MatchInfo(459+ return MatchInfo(
460- MatchState.COLLECTIVE_STATE_MISMATCH,460+ MatchState.COLLECTIVE_STATE_MISMATCH,
461- f"Expected state: '{self.state}' does not match found state: '{other.state}'",461+ f"Expected state: '{self.state}' does not match found state: '{other.state}'",
462- )462+ )
463- if self.has_different_dtypes_and_non_empty_sizes(self):463+ if self.has_different_dtypes_and_non_empty_sizes(self):
464- return MatchInfo(464+ return MatchInfo(
465- MatchState.COLLECTIVE_DTYPE_MISMATCH,465+ MatchState.COLLECTIVE_DTYPE_MISMATCH,
466- f"Expected dtypes: '{set(self.input_dtypes)}' does not "466+ f"Expected dtypes: '{set(self.input_dtypes)}' does not "
467- f"match found dtype: '{set(self.output_dtypes)}/"467+ f"match found dtype: '{set(self.output_dtypes)}/"
468- f"{set(other.input_dtypes)}/{set(other.output_dtypes)}'",468+ f"{set(other.input_dtypes)}/{set(other.output_dtypes)}'",
469- )469+ )
470- if self.type == "all_to_all":470+ if self.type == "all_to_all":
471- return MatchInfo(MatchState.UNDECIDED)471+ return MatchInfo(MatchState.UNDECIDED)
472- if self.type != "scatter" and self.input_sizes != other.input_sizes:472+ if self.type != "scatter" and self.input_sizes != other.input_sizes:
473- return MatchInfo(473+ return MatchInfo(
474- MatchState.SIZE_OR_SYNTAX_MISMATCH,474+ MatchState.SIZE_OR_SYNTAX_MISMATCH,
475- f"Expected input sizes: '{self.input_sizes}' does not match found input sizes: "475+ f"Expected input sizes: '{self.input_sizes}' does not match found input sizes: "
476- f"'{other.input_sizes}'",476+ f"'{other.input_sizes}'",
477- )477+ )
478- if self.type != "gather" and self.output_sizes != other.output_sizes:478+ if self.type != "gather" and self.output_sizes != other.output_sizes:
479- return MatchInfo(479+ return MatchInfo(
480- MatchState.SIZE_OR_SYNTAX_MISMATCH,480+ MatchState.SIZE_OR_SYNTAX_MISMATCH,
481- f"Expected output sizes: '{self.output_sizes}' does not match found output sizes: "481+ f"Expected output sizes: '{self.output_sizes}' does not match found output sizes: "
482- f"'{other.output_sizes}'",482+ f"'{other.output_sizes}'",
483- )483+ )
484- if self.type in ["all_reduce", "allreduce_coalesced"] and self.input_sizes != other.output_sizes:484+ if self.type in ["all_reduce", "allreduce_coalesced"] and self.input_sizes != other.output_sizes:
485- return MatchInfo(485+ return MatchInfo(
486- MatchState.SIZE_OR_SYNTAX_MISMATCH,486+ MatchState.SIZE_OR_SYNTAX_MISMATCH,
487- f"Expected input sizes: '{self.input_sizes}' does not match found output sizes: '{other.output_sizes}'",487+ f"Expected input sizes: '{self.input_sizes}' does not match found output sizes: '{other.output_sizes}'",
488- )488+ )
489- if self.type in [489+ if self.type in [
490- "all_gather",490+ "all_gather",
491- "all_gather_base",491+ "all_gather_base",
492- "all_gather_into_tensor_coalesced",492+ "all_gather_into_tensor_coalesced",
493- ] and not (math.prod(other.output_sizes[0]) == math.prod(self.input_sizes[0]) * self.pg_size):493+ ] and not (math.prod(other.output_sizes[0]) == math.prod(self.input_sizes[0]) * self.pg_size):
494- return MatchInfo(494+ return MatchInfo(
495- MatchState.SIZE_OR_SYNTAX_MISMATCH,495+ MatchState.SIZE_OR_SYNTAX_MISMATCH,
496- f"Found input numel '{math.prod(other.input_sizes[0])} * pg size {self.pg_size}' "496+ f"Found input numel '{math.prod(other.input_sizes[0])} * pg size {self.pg_size}' "
497- f"does not match output numel '{math.prod(other.output_sizes[0])}'",497+ f"does not match output numel '{math.prod(other.output_sizes[0])}'",
498- )498+ )
499- if self.type in [499+ if self.type in [
500- "reduce_scatter",500+ "reduce_scatter",
501- "_reduce_scatter_base",501+ "_reduce_scatter_base",
502- "reduce_scatter_tensor_coalesced",502+ "reduce_scatter_tensor_coalesced",
503- ] and not (math.prod(other.input_sizes[0]) == math.prod(self.output_sizes[0]) * self.pg_size):503+ ] and not (math.prod(other.input_sizes[0]) == math.prod(self.output_sizes[0]) * self.pg_size):
504- return MatchInfo(504+ return MatchInfo(
505- MatchState.SIZE_OR_SYNTAX_MISMATCH,505+ MatchState.SIZE_OR_SYNTAX_MISMATCH,
506- f"Found input numel '{math.prod(other.input_sizes[0])}' does not match output numel "506+ f"Found input numel '{math.prod(other.input_sizes[0])}' does not match output numel "
507- f"'{math.prod(other.output_sizes[0])} * pg size {self.pg_size}'",507+ f"'{math.prod(other.output_sizes[0])} * pg size {self.pg_size}'",
508- )508+ )
509- elif self.type in [509+ elif self.type in [
510- "coalesced",510+ "coalesced",
511- "ALLGATHER_coalesced",511+ "ALLGATHER_coalesced",
512- "REDUCE_SCATTER_coalesced",512+ "REDUCE_SCATTER_coalesced",
513- ]:513+ ]:
514- return (514+ return (
515- MatchInfo(MatchState.FULLY_MATCHED)515+ MatchInfo(MatchState.FULLY_MATCHED)
516- if (other.type == self.type)516+ if (other.type == self.type)
517- else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)517+ else MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)
518- )518+ )
519- return MatchInfo(MatchState.FULLY_MATCHED)519+ return MatchInfo(MatchState.FULLY_MATCHED)
520- 520+ 
521- 521+ 
522-class MatchStateRecord:522+class MatchStateRecord:
523- def __init__(523+ def __init__(
524- self,524+ self,
525- expected_ranks: set[int],525+ expected_ranks: set[int],
526- other_ranks: list[int],526+ other_ranks: list[int],
527- entry_state: EntryState,527+ entry_state: EntryState,
528- candidate_ranks: set[int],528+ candidate_ranks: set[int],
529- candidate_idx: dict[int, int],529+ candidate_idx: dict[int, int],
530- found_ranks: set[int],530+ found_ranks: set[int],
531- found_idx: dict[int, int],531+ found_idx: dict[int, int],
532- errors: set[tuple[int, MatchInfo]],532+ errors: set[tuple[int, MatchInfo]],
533- ) -> None:533+ ) -> None:
534- self.expected_ranks = expected_ranks534+ self.expected_ranks = expected_ranks
535- self.other_ranks = other_ranks535+ self.other_ranks = other_ranks
536- self.entry_state = entry_state536+ self.entry_state = entry_state
537- self.candidate_ranks = candidate_ranks537+ self.candidate_ranks = candidate_ranks
538- self.candidate_idx = candidate_idx538+ self.candidate_idx = candidate_idx
539- self.found_ranks = found_ranks539+ self.found_ranks = found_ranks
540- self.found_idx = found_idx540+ self.found_idx = found_idx
541- self.errors = errors541+ self.errors = errors
542- self.has_undecided_case = False542+ self.has_undecided_case = False
543- 543+ 
544- def reset_for_coalesced(self, entry_state: EntryState, candidate_ranks: set[int]) -> None:544+ def reset_for_coalesced(self, entry_state: EntryState, candidate_ranks: set[int]) -> None:
545- self.entry_state = entry_state545+ self.entry_state = entry_state
546- self.candidate_ranks = candidate_ranks546+ self.candidate_ranks = candidate_ranks
547- self.candidate_idx = {}547+ self.candidate_idx = {}
548- self.found_ranks = set()548+ self.found_ranks = set()
549- self.found_idx = {}549+ self.found_idx = {}
550- self.errors = set()550+ self.errors = set()
Mtools/flight_recorder/components/utils.py+402-402
@@ -1,402 +1,402 @@
1-__all__ = []1+__all__ = []
2- 2+ 
3-import argparse3+import argparse
4-import math4+import math
5-from typing import Any5+from typing import Any
6-import os6+import os
7-import re7+import re
8-import sys8+import sys
9-import stat9+import stat
10- 10+ 
11-from tools.flight_recorder.components.fr_logger import FlightRecorderLogger11+from tools.flight_recorder.components.fr_logger import FlightRecorderLogger
12-from tools.flight_recorder.components.types import (12+from tools.flight_recorder.components.types import (
13- Group,13+ Group,
14- MatchInfo,14+ MatchInfo,
15- MatchState,15+ MatchState,
16- MatchStateRecord,16+ MatchStateRecord,
17- Membership,17+ Membership,
18- Op,18+ Op,
19-)19+)
20- 20+ 
21-logger: FlightRecorderLogger = FlightRecorderLogger()21+logger: FlightRecorderLogger = FlightRecorderLogger()
22- 22+ 
23-try:23+try:
24- from tabulate import tabulate24+ from tabulate import tabulate
25-except ModuleNotFoundError:25+except ModuleNotFoundError:
26- logger.debug("tabulate is not installed. Proceeding without it.")26+ logger.debug("tabulate is not installed. Proceeding without it.")
27- 27+ 
28-PATH_WHITE_LIST_REGEX = re.compile(r"[^_A-Za-z0-9/.-]")28+PATH_WHITE_LIST_REGEX = re.compile(r"[^_A-Za-z0-9/.-]")
29-MAX_READ_FILE_SIZE_4G = 4294967296 # 4G, 4 * 1024 * 1024 * 102429+MAX_READ_FILE_SIZE_4G = 4294967296 # 4G, 4 * 1024 * 1024 * 1024
30-MAX_READ_FILE_SIZE_32G = 34359738368 # 32G, 32 * 1024 * 1024 * 102430+MAX_READ_FILE_SIZE_32G = 34359738368 # 32G, 32 * 1024 * 1024 * 1024
31-MAX_READ_FILE_SIZE_512G = 549755813888 # 512G, 512 * 1024 * 1024 * 102431+MAX_READ_FILE_SIZE_512G = 549755813888 # 512G, 512 * 1024 * 1024 * 1024
32- 32+ 
33-# group not writable, others no permission, max stat is 75033+# group not writable, others no permission, max stat is 750
34-WRITE_FILE_NOT_PERMITTED_STAT = stat.S_IWGRP | stat.S_IWOTH | stat.S_IROTH | stat.S_IXOTH34+WRITE_FILE_NOT_PERMITTED_STAT = stat.S_IWGRP | stat.S_IWOTH | stat.S_IROTH | stat.S_IXOTH
35-# group not writable, others not writable, max stat is 75535+# group not writable, others not writable, max stat is 755
36-READ_FILE_NOT_PERMITTED_STAT = stat.S_IWGRP | stat.S_IWOTH36+READ_FILE_NOT_PERMITTED_STAT = stat.S_IWGRP | stat.S_IWOTH
37- 37+ 
38- 38+ 
39-def type_to_str(value_type):39+def type_to_str(value_type):
40- return " or ".join(ii.__name__ for ii in value_type) if isinstance(value_type, tuple) else value_type.__name__40+ return " or ".join(ii.__name__ for ii in value_type) if isinstance(value_type, tuple) else value_type.__name__
41- 41+ 
42- 42+ 
43-def check_type(value, value_type, param_name="value"):43+def check_type(value, value_type, param_name="value"):
44- if not isinstance(value, value_type):44+ if not isinstance(value, value_type):
45- raise TypeError("{} must be {}, not {}.".format(param_name, type_to_str(value_type), type(value).__name__))45+ raise TypeError("{} must be {}, not {}.".format(param_name, type_to_str(value_type), type(value).__name__))
46- 46+ 
47- 47+ 
48-def get_valid_path(path):48+def get_valid_path(path):
49- check_type(path, str, "path")49+ check_type(path, str, "path")
50- if not path or len(path) == 0:50+ if not path or len(path) == 0:
51- raise ValueError("The value of the path cannot be empty.")51+ raise ValueError("The value of the path cannot be empty.")
52- if PATH_WHITE_LIST_REGEX.search(path): # Check special char52+ if PATH_WHITE_LIST_REGEX.search(path): # Check special char
53- raise ValueError("Input path contains invalid characters.") # Not printing out the path value for invalid char53+ raise ValueError("Input path contains invalid characters.") # Not printing out the path value for invalid char
54- path = os.path.expanduser(path)54+ path = os.path.expanduser(path)
55- if os.path.islink(os.path.abspath(path)): # when checking link, get rid of the "/" at the path tail if any55+ if os.path.islink(os.path.abspath(path)): # when checking link, get rid of the "/" at the path tail if any
56- raise ValueError("The value of the path cannot be a symbolic link: {}.".format(path))56+ raise ValueError("The value of the path cannot be a symbolic link: {}.".format(path))
57- 57+ 
58- real_path = os.path.realpath(path)58+ real_path = os.path.realpath(path)
59- 59+ 
60- if len(real_path) > 4096:60+ if len(real_path) > 4096:
61- raise ValueError("The length of file path should be less than 4096.")61+ raise ValueError("The length of file path should be less than 4096.")
62- 62+ 
63- if real_path != path and PATH_WHITE_LIST_REGEX.search(real_path): # Check special char again63+ if real_path != path and PATH_WHITE_LIST_REGEX.search(real_path): # Check special char again
64- raise ValueError("Input path contains invalid characters.") # Not printing out the path value for invalid char64+ raise ValueError("Input path contains invalid characters.") # Not printing out the path value for invalid char
65- 65+ 
66- return real_path66+ return real_path
67- 67+ 
68- 68+ 
69-def is_belong_to_user_or_group(file_stat):69+def is_belong_to_user_or_group(file_stat):
70- return file_stat.st_uid == os.getuid() or file_stat.st_gid in os.getgroups()70+ return file_stat.st_uid == os.getuid() or file_stat.st_gid in os.getgroups()
71- 71+ 
72- 72+ 
73-def get_valid_read_path(path, size_max=MAX_READ_FILE_SIZE_4G, check_user_stat=True, is_dir=False):73+def get_valid_read_path(path, size_max=MAX_READ_FILE_SIZE_4G, check_user_stat=True, is_dir=False):
74- real_path = get_valid_path(path)74+ real_path = get_valid_path(path)
75- if not is_dir and not os.path.isfile(real_path):75+ if not is_dir and not os.path.isfile(real_path):
76- raise ValueError("The path {} doesn't exists or not a file.".format(path))76+ raise ValueError("The path {} doesn't exists or not a file.".format(path))
77- if is_dir and not os.path.isdir(real_path):77+ if is_dir and not os.path.isdir(real_path):
78- raise ValueError("The path {} doesn't exists or not a directory.".format(path))78+ raise ValueError("The path {} doesn't exists or not a directory.".format(path))
79- 79+ 
80- file_stat = os.stat(real_path)80+ file_stat = os.stat(real_path)
81- if check_user_stat and not sys.platform.startswith("win") and not is_belong_to_user_or_group(file_stat):81+ if check_user_stat and not sys.platform.startswith("win") and not is_belong_to_user_or_group(file_stat):
82- raise ValueError("The file {} doesn't belong to the current user or group.".format(path))82+ raise ValueError("The file {} doesn't belong to the current user or group.".format(path))
83- if check_user_stat and os.stat(path).st_mode & READ_FILE_NOT_PERMITTED_STAT > 0:83+ if check_user_stat and os.stat(path).st_mode & READ_FILE_NOT_PERMITTED_STAT > 0:
84- raise ValueError("The file {} is group writable, or is others writable.".format(path))84+ raise ValueError("The file {} is group writable, or is others writable.".format(path))
85- if not os.access(real_path, os.R_OK) or file_stat.st_mode & stat.S_IRUSR == 0: # At least been 40085+ if not os.access(real_path, os.R_OK) or file_stat.st_mode & stat.S_IRUSR == 0: # At least been 400
86- raise ValueError("Current user doesn't have read permission to the file {}.".format(path))86+ raise ValueError("Current user doesn't have read permission to the file {}.".format(path))
87- if not is_dir and size_max > 0 and file_stat.st_size > size_max:87+ if not is_dir and size_max > 0 and file_stat.st_size > size_max:
88- raise ValueError("The file {} exceeds size limitation of {}.".format(path, size_max))88+ raise ValueError("The file {} exceeds size limitation of {}.".format(path, size_max))
89- return real_path89+ return real_path
90- 90+ 
91- 91+ 
92-def check_write_directory(dir_name, check_user_stat=True):92+def check_write_directory(dir_name, check_user_stat=True):
93- real_dir_name = get_valid_path(dir_name)93+ real_dir_name = get_valid_path(dir_name)
94- if not os.path.isdir(real_dir_name):94+ if not os.path.isdir(real_dir_name):
95- raise ValueError("The file writen directory {} doesn't exists.".format(dir_name))95+ raise ValueError("The file writen directory {} doesn't exists.".format(dir_name))
96- 96+ 
97- file_stat = os.stat(real_dir_name)97+ file_stat = os.stat(real_dir_name)
98- if check_user_stat and not sys.platform.startswith("win") and not is_belong_to_user_or_group(file_stat):98+ if check_user_stat and not sys.platform.startswith("win") and not is_belong_to_user_or_group(file_stat):
99- raise ValueError("The file writen directory {} doesn't belong to the current user or group.".format(dir_name))99+ raise ValueError("The file writen directory {} doesn't belong to the current user or group.".format(dir_name))
100- if not os.access(real_dir_name, os.W_OK):100+ if not os.access(real_dir_name, os.W_OK):
101- raise ValueError("Current user doesn't have writen permission to file writen directory {}.".format(dir_name))101+ raise ValueError("Current user doesn't have writen permission to file writen directory {}.".format(dir_name))
102- 102+ 
103- 103+ 
104-def get_valid_write_path(path, check_user_stat=True, is_dir=False, warn_exists=True):104+def get_valid_write_path(path, check_user_stat=True, is_dir=False, warn_exists=True):
105- real_path = get_valid_path(path)105+ real_path = get_valid_path(path)
106- real_path_dir = real_path if is_dir else os.path.dirname(real_path)106+ real_path_dir = real_path if is_dir else os.path.dirname(real_path)
107- check_write_directory(real_path_dir, check_user_stat=check_user_stat)107+ check_write_directory(real_path_dir, check_user_stat=check_user_stat)
108- 108+ 
109- if not is_dir and os.path.exists(real_path):109+ if not is_dir and os.path.exists(real_path):
110- if os.path.isdir(real_path):110+ if os.path.isdir(real_path):
111- raise ValueError("The file {} exist and is a directory.".format(path))111+ raise ValueError("The file {} exist and is a directory.".format(path))
112- if check_user_stat and os.stat(real_path).st_uid != os.getuid(): # Has to be exactly belonging to current user112+ if check_user_stat and os.stat(real_path).st_uid != os.getuid(): # Has to be exactly belonging to current user
113- raise ValueError("The file {} doesn't belong to the current user.".format(path))113+ raise ValueError("The file {} doesn't belong to the current user.".format(path))
114- if check_user_stat and os.stat(real_path).st_mode & WRITE_FILE_NOT_PERMITTED_STAT > 0:114+ if check_user_stat and os.stat(real_path).st_mode & WRITE_FILE_NOT_PERMITTED_STAT > 0:
115- raise ValueError("The file {} permission for others is not 0, or is group writable.".format(path))115+ raise ValueError("The file {} permission for others is not 0, or is group writable.".format(path))
116- if not os.access(real_path, os.W_OK):116+ if not os.access(real_path, os.W_OK):
117- raise ValueError("The file {} exist and not writable.".format(path))117+ raise ValueError("The file {} exist and not writable.".format(path))
118- if warn_exists:118+ if warn_exists:
119- logger.warning("%s already exist. The original file will be overwritten.", path)119+ logger.warning("%s already exist. The original file will be overwritten.", path)
120- return real_path120+ return real_path
121- 121+ 
122- 122+ 
123-def format_frame(frame: dict[str, str]) -> str:123+def format_frame(frame: dict[str, str]) -> str:
124- name = frame["name"]124+ name = frame["name"]
125- filename = frame["filename"]125+ filename = frame["filename"]
126- line = frame["line"]126+ line = frame["line"]
127- return f"{name} at {filename}:{line}"127+ return f"{name} at {filename}:{line}"
128- 128+ 
129- 129+ 
130-def format_frames(frames: list[dict[str, str]]) -> str:130+def format_frames(frames: list[dict[str, str]]) -> str:
131- formatted_frames = []131+ formatted_frames = []
132- for frame in frames:132+ for frame in frames:
133- formatted_frames.append(format_frame(frame))133+ formatted_frames.append(format_frame(frame))
134- return "\n".join(formatted_frames)134+ return "\n".join(formatted_frames)
135- 135+ 
136- 136+ 
137-def match_one_event(137+def match_one_event(
138- event_a: dict[Any, Any],138+ event_a: dict[Any, Any],
139- event_b: dict[Any, Any],139+ event_b: dict[Any, Any],
140- memberships: dict[str, set[Any]],140+ memberships: dict[str, set[Any]],
141- pg_name: str,141+ pg_name: str,
142-) -> MatchInfo:142+) -> MatchInfo:
143- op_a = Op(event_a, memberships, pg_name)143+ op_a = Op(event_a, memberships, pg_name)
144- op_b = Op(event_b, memberships, pg_name)144+ op_b = Op(event_b, memberships, pg_name)
145- return op_a.match(op_b)145+ return op_a.match(op_b)
146- 146+ 
147- 147+ 
148-def check_size_alltoall(alltoall_cases: list[dict[str, Any]]) -> tuple[bool, int, int]:148+def check_size_alltoall(alltoall_cases: list[dict[str, Any]]) -> tuple[bool, int, int]:
149- input_numel = 0149+ input_numel = 0
150- output_numel = 0150+ output_numel = 0
151- for e in alltoall_cases:151+ for e in alltoall_cases:
152- input_numel += math.prod(e["input_sizes"][0])152+ input_numel += math.prod(e["input_sizes"][0])
153- output_numel += math.prod(e["output_sizes"][0])153+ output_numel += math.prod(e["output_sizes"][0])
154- return input_numel != output_numel, input_numel, output_numel154+ return input_numel != output_numel, input_numel, output_numel
155- 155+ 
156- 156+ 
157-class ProcessGroupData:157+class ProcessGroupData:
158- def __init__(self, pg_guids: dict[tuple[str, int], str], pg_name: str, desc: str, mismatch: dict[str, int]):158+ def __init__(self, pg_guids: dict[tuple[str, int], str], pg_name: str, desc: str, mismatch: dict[str, int]):
159- self.pg_guids, self.pg_name, self.desc, self.mismatch = pg_guids, pg_name, desc, mismatch159+ self.pg_guids, self.pg_name, self.desc, self.mismatch = pg_guids, pg_name, desc, mismatch
160- 160+ 
161- 161+ 
162-def check_current_entry_match(162+def check_current_entry_match(
163- all_entries: dict[int, list[dict[str, Any]]],163+ all_entries: dict[int, list[dict[str, Any]]],
164- current_entry: dict[str, Any],164+ current_entry: dict[str, Any],
165- _memberships: dict[str, set[Any]],165+ _memberships: dict[str, set[Any]],
166- pg_data: ProcessGroupData,166+ pg_data: ProcessGroupData,
167- match_record: MatchStateRecord,167+ match_record: MatchStateRecord,
168-) -> None:168+) -> None:
169- pg_guids, pg_name, mismatch, desc = pg_data.pg_guids, pg_data.pg_name, pg_data.mismatch, pg_data.desc169+ pg_guids, pg_name, mismatch, desc = pg_data.pg_guids, pg_data.pg_name, pg_data.mismatch, pg_data.desc
170- for rank in match_record.expected_ranks.intersection(set(match_record.other_ranks)):170+ for rank in match_record.expected_ranks.intersection(set(match_record.other_ranks)):
171- for entry_idx, entry in enumerate(all_entries[rank]):171+ for entry_idx, entry in enumerate(all_entries[rank]):
172- # step over ops from other PGs172+ # step over ops from other PGs
173- # only check match state when seq_id matches173+ # only check match state when seq_id matches
174- if (174+ if (
175- pg_guids[(entry["process_group"][0], rank)] == pg_name175+ pg_guids[(entry["process_group"][0], rank)] == pg_name
176- and entry["collective_seq_id"] == match_record.entry_state.collective_seq_id176+ and entry["collective_seq_id"] == match_record.entry_state.collective_seq_id
177- ):177+ ):
178- match_info = match_one_event(current_entry, entry, _memberships, pg_name)178+ match_info = match_one_event(current_entry, entry, _memberships, pg_name)
179- if match_info.state in [MatchState.FULLY_MATCHED, MatchState.UNDECIDED] and mismatch[pg_name] == 0:179+ if match_info.state in [MatchState.FULLY_MATCHED, MatchState.UNDECIDED] and mismatch[pg_name] == 0:
180- match_record.found_ranks.add(rank)180+ match_record.found_ranks.add(rank)
181- match_record.found_idx[rank] = entry_idx181+ match_record.found_idx[rank] = entry_idx
182- match_record.has_undecided_case = match_info.state == MatchState.UNDECIDED182+ match_record.has_undecided_case = match_info.state == MatchState.UNDECIDED
183- else:183+ else:
184- match_record.candidate_ranks.add(rank)184+ match_record.candidate_ranks.add(rank)
185- match_record.candidate_idx[rank] = entry_idx185+ match_record.candidate_idx[rank] = entry_idx
186- if match_info.state not in [186+ if match_info.state not in [
187- MatchState.FULLY_MATCHED,187+ MatchState.FULLY_MATCHED,
188- MatchState.UNDECIDED,188+ MatchState.UNDECIDED,
189- ]:189+ ]:
190- match_record.errors.add((rank, match_info))190+ match_record.errors.add((rank, match_info))
191- break191+ break
192- 192+ 
193- 193+ 
194-class EntryContext:194+class EntryContext:
195- def __init__(self, all_entries, current_entry, dumps_ranks, first_rank):195+ def __init__(self, all_entries, current_entry, dumps_ranks, first_rank):
196- self.all_entries = all_entries196+ self.all_entries = all_entries
197- self.current_entry = current_entry197+ self.current_entry = current_entry
198- self.dumps_ranks = dumps_ranks198+ self.dumps_ranks = dumps_ranks
199- self.first_rank = first_rank199+ self.first_rank = first_rank
200- 200+ 
201- 201+ 
202-def error_analysis(202+def error_analysis(
203- entry_context: EntryContext,203+ entry_context: EntryContext,
204- match_record: MatchStateRecord, # all204+ match_record: MatchStateRecord, # all
205- mismatch: dict[str, int], # all205+ mismatch: dict[str, int], # all
206- version: tuple[int, int], # 2206+ version: tuple[int, int], # 2
207- pg_name: str, # all, mismatch207+ pg_name: str, # all, mismatch
208-) -> None:208+) -> None:
209- all_entries = entry_context.all_entries209+ all_entries = entry_context.all_entries
210- current_entry = entry_context.current_entry210+ current_entry = entry_context.current_entry
211- dumps_ranks = entry_context.dumps_ranks211+ dumps_ranks = entry_context.dumps_ranks
212- first_rank = entry_context.first_rank212+ first_rank = entry_context.first_rank
213- major_v, minor_v = version[0], version[1]213+ major_v, minor_v = version[0], version[1]
214- # case one: not every rank join the collective or in the flight recorder.214+ # case one: not every rank join the collective or in the flight recorder.
215- if (215+ if (
216- match_record.candidate_ranks | match_record.found_ranks216+ match_record.candidate_ranks | match_record.found_ranks
217- ) != match_record.expected_ranks and match_record.expected_ranks - (217+ ) != match_record.expected_ranks and match_record.expected_ranks - (
218- match_record.candidate_ranks | match_record.found_ranks218+ match_record.candidate_ranks | match_record.found_ranks
219- ) <= dumps_ranks:219+ ) <= dumps_ranks:
220- mismatch[pg_name] += 1220+ mismatch[pg_name] += 1
221- logger_msg = "Not all ranks joining collective, sequence number: %s"221+ logger_msg = "Not all ranks joining collective, sequence number: %s"
222- missing_ranks = match_record.expected_ranks - (match_record.candidate_ranks | match_record.found_ranks)222+ missing_ranks = match_record.expected_ranks - (match_record.candidate_ranks | match_record.found_ranks)
223- match_record.entry_state.log(223+ match_record.entry_state.log(
224- logger, logger_msg, format_frames, additional_info={"missing_ranks": missing_ranks}224+ logger, logger_msg, format_frames, additional_info={"missing_ranks": missing_ranks}
225- )225+ )
226- match_record.candidate_ranks.update(match_record.found_ranks)226+ match_record.candidate_ranks.update(match_record.found_ranks)
227- match_record.candidate_idx.update(match_record.found_idx)227+ match_record.candidate_idx.update(match_record.found_idx)
228- match_record.found_idx.clear()228+ match_record.found_idx.clear()
229- match_record.found_ranks.clear()229+ match_record.found_ranks.clear()
230- elif len(match_record.candidate_ranks) == 1 and dumps_ranks == match_record.expected_ranks:230+ elif len(match_record.candidate_ranks) == 1 and dumps_ranks == match_record.expected_ranks:
231- # case two: alltoall or alltoall_base case.231+ # case two: alltoall or alltoall_base case.
232- if match_record.has_undecided_case:232+ if match_record.has_undecided_case:
233- alltoall_cases = [current_entry] + [233+ alltoall_cases = [current_entry] + [
234- all_entries[rank][match_record.found_idx[rank]] for rank in match_record.found_ranks234+ all_entries[rank][match_record.found_idx[rank]] for rank in match_record.found_ranks
235- ]235+ ]
236- fail_check, total_input_numel, total_output_numel = check_size_alltoall(alltoall_cases)236+ fail_check, total_input_numel, total_output_numel = check_size_alltoall(alltoall_cases)
237- if major_v <= 2 and minor_v <= 3:237+ if major_v <= 2 and minor_v <= 3:
238- # We don't log the input/output sizes for alltoall before v2.4,238+ # We don't log the input/output sizes for alltoall before v2.4,
239- # so we don't consider the size mismatch as an error for now.239+ # so we don't consider the size mismatch as an error for now.
240- fail_check = False240+ fail_check = False
241- if fail_check:241+ if fail_check:
242- # When we see errors in all_to_all, it's hard to tell which rank is the source of the error.242+ # When we see errors in all_to_all, it's hard to tell which rank is the source of the error.
243- mismatch[pg_name] += 1243+ mismatch[pg_name] += 1
244- logger_msg = "Input/output mismatch in the collective sequence number: %s"244+ logger_msg = "Input/output mismatch in the collective sequence number: %s"
245- match_record.entry_state.log(245+ match_record.entry_state.log(
246- logger,246+ logger,
247- logger_msg,247+ logger_msg,
248- format_frames,248+ format_frames,
249- additional_info={"total_numel": (total_input_numel, total_output_numel)},249+ additional_info={"total_numel": (total_input_numel, total_output_numel)},
250- )250+ )
251- match_record.candidate_ranks.update(match_record.found_ranks)251+ match_record.candidate_ranks.update(match_record.found_ranks)
252- match_record.candidate_idx.update(match_record.found_idx)252+ match_record.candidate_idx.update(match_record.found_idx)
253- match_record.found_idx.clear()253+ match_record.found_idx.clear()
254- match_record.found_ranks.clear()254+ match_record.found_ranks.clear()
255- match_record.errors.add((first_rank, MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)))255+ match_record.errors.add((first_rank, MatchInfo(MatchState.SIZE_OR_SYNTAX_MISMATCH)))
256- else:256+ else:
257- match_record.found_ranks.update(match_record.candidate_ranks)257+ match_record.found_ranks.update(match_record.candidate_ranks)
258- match_record.found_idx.update(match_record.candidate_idx)258+ match_record.found_idx.update(match_record.candidate_idx)
259- match_record.candidate_idx.clear()259+ match_record.candidate_idx.clear()
260- match_record.candidate_ranks.clear()260+ match_record.candidate_ranks.clear()
261- # case three: all joined and everything matches on all ranks.261+ # case three: all joined and everything matches on all ranks.
262- else:262+ else:
263- match_record.found_ranks.update(match_record.candidate_ranks)263+ match_record.found_ranks.update(match_record.candidate_ranks)
264- match_record.found_idx.update(match_record.candidate_idx)264+ match_record.found_idx.update(match_record.candidate_idx)
265- match_record.candidate_idx.clear()265+ match_record.candidate_idx.clear()
266- match_record.candidate_ranks.clear()266+ match_record.candidate_ranks.clear()
267- # case four: mismatch cases due to not same type, size mismatch or state mismatch.267+ # case four: mismatch cases due to not same type, size mismatch or state mismatch.
268- elif len(match_record.errors) > 0:268+ elif len(match_record.errors) > 0:
269- mismatch[pg_name] += 1269+ mismatch[pg_name] += 1
270- logger_msg = "Collective sequence number: %s has errors"270+ logger_msg = "Collective sequence number: %s has errors"
271- match_record.entry_state.log(logger, logger_msg, format_frames, errors=match_record.errors)271+ match_record.entry_state.log(logger, logger_msg, format_frames, errors=match_record.errors)
272- match_record.candidate_ranks.update(match_record.found_ranks)272+ match_record.candidate_ranks.update(match_record.found_ranks)
273- match_record.candidate_idx.update(match_record.found_idx)273+ match_record.candidate_idx.update(match_record.found_idx)
274- match_record.found_idx.clear()274+ match_record.found_idx.clear()
275- match_record.found_ranks.clear()275+ match_record.found_ranks.clear()
276- # partial analysis case when we cannot decide what's wrong with this collective entry.276+ # partial analysis case when we cannot decide what's wrong with this collective entry.
277- else:277+ else:
278- match_record.candidate_ranks.update(match_record.found_ranks)278+ match_record.candidate_ranks.update(match_record.found_ranks)
279- match_record.candidate_idx.update(match_record.found_idx)279+ match_record.candidate_idx.update(match_record.found_idx)
280- match_record.found_idx.clear()280+ match_record.found_idx.clear()
281- match_record.found_ranks.clear()281+ match_record.found_ranks.clear()
282- if match_record.expected_ranks - dumps_ranks:282+ if match_record.expected_ranks - dumps_ranks:
283- mismatch[pg_name] += 1283+ mismatch[pg_name] += 1
284- logger.info(284+ logger.info(
285- "We cannot decide what's wrong with this collective entry "285+ "We cannot decide what's wrong with this collective entry "
286- "because we missed FR dumps from ranks (%s) so we don't have enough "286+ "because we missed FR dumps from ranks (%s) so we don't have enough "
287- "information. If you want to debug further use -j to dump all raw trace",287+ "information. If you want to debug further use -j to dump all raw trace",
288- str(match_record.expected_ranks - dumps_ranks),288+ str(match_record.expected_ranks - dumps_ranks),
289- )289+ )
290- else:290+ else:
291- logger.info(291+ logger.info(
292- "No errors found for this collective entry, There could be some "292+ "No errors found for this collective entry, There could be some "
293- "other reasons why we see collective timeout."293+ "other reasons why we see collective timeout."
294- )294+ )
295- 295+ 
296- 296+ 
297-def just_print_entries(297+def just_print_entries(
298- all_entries: dict[int, list[dict[str, Any]]],298+ all_entries: dict[int, list[dict[str, Any]]],
299- _groups: dict[str, Group],299+ _groups: dict[str, Group],
300- _memberships: dict[str, set[Any]],300+ _memberships: dict[str, set[Any]],
301- _pg_guids: dict[tuple[str, int], str],301+ _pg_guids: dict[tuple[str, int], str],
302- args: argparse.Namespace,302+ args: argparse.Namespace,
303-) -> None:303+) -> None:
304- rows = []304+ rows = []
305- ranks = sorted(all_entries.keys())305+ ranks = sorted(all_entries.keys())
306- headers = [f"Rank {rank}" for rank in ranks if args.selected_ranks is None or rank in args.selected_ranks]306+ headers = [f"Rank {rank}" for rank in ranks if args.selected_ranks is None or rank in args.selected_ranks]
307- progress = True307+ progress = True
308- while progress:308+ while progress:
309- progress = False309+ progress = False
310- row = []310+ row = []
311- for rank in ranks:311+ for rank in ranks:
312- if args.selected_ranks is not None and rank not in args.selected_ranks:312+ if args.selected_ranks is not None and rank not in args.selected_ranks:
313- continue313+ continue
314- if len(all_entries[rank]) == 0:314+ if len(all_entries[rank]) == 0:
315- row.append("")315+ row.append("")
316- else:316+ else:
317- entry = all_entries[rank].pop(0)317+ entry = all_entries[rank].pop(0)
318- pg_name = _pg_guids[(entry["process_group"][0], rank)]318+ pg_name = _pg_guids[(entry["process_group"][0], rank)]
319- if (319+ if (
320- args.pg_filters is None320+ args.pg_filters is None
321- or entry["process_group"][1] in args.pg_filters321+ or entry["process_group"][1] in args.pg_filters
322- or entry["process_group"][0] in args.pg_filters322+ or entry["process_group"][0] in args.pg_filters
323- ):323+ ):
324- row.append(str(Op(entry, _memberships, pg_name)))324+ row.append(str(Op(entry, _memberships, pg_name)))
325- else:325+ else:
326- row.append("")326+ row.append("")
327- progress = True327+ progress = True
328- if progress:328+ if progress:
329- rows.append(row)329+ rows.append(row)
330- 330+ 
331- logger.info(tabulate(rows, headers=headers))331+ logger.info(tabulate(rows, headers=headers))
332- 332+ 
333- 333+ 
334-def check_no_missing_dump_files(entries: dict[int, Any], memberships: list[Membership]) -> None:334+def check_no_missing_dump_files(entries: dict[int, Any], memberships: list[Membership]) -> None:
335- all_ranks = {int(m.global_rank) for m in memberships}335+ all_ranks = {int(m.global_rank) for m in memberships}
336- 336+ 
337- dumps_ranks = {int(key) for key in entries.keys()}337+ dumps_ranks = {int(key) for key in entries.keys()}
338- missing_ranks = all_ranks - dumps_ranks338+ missing_ranks = all_ranks - dumps_ranks
339- if missing_ranks:339+ if missing_ranks:
340- raise ValueError(340+ raise ValueError(
341- f"Missing dump files for {len(missing_ranks)} ranks: {sorted(missing_ranks)}\n"341+ f"Missing dump files for {len(missing_ranks)} ranks: {sorted(missing_ranks)}\n"
342- f"Expected ranks: {sorted(all_ranks)}\n"342+ f"Expected ranks: {sorted(all_ranks)}\n"
343- f"Found dumps for: {sorted(dumps_ranks)}"343+ f"Found dumps for: {sorted(dumps_ranks)}"
344- )344+ )
345- 345+ 
346- 346+ 
347-def check_version(version_by_ranks: dict[str, str], expected_version: str) -> None:347+def check_version(version_by_ranks: dict[str, str], expected_version: str) -> None:
348- for rank, actual_version in version_by_ranks.items():348+ for rank, actual_version in version_by_ranks.items():
349- if actual_version != expected_version:349+ if actual_version != expected_version:
350- raise ValueError(f"Version mismatch at rank {rank}: " f"expected {expected_version}, got {actual_version}")350+ raise ValueError(f"Version mismatch at rank {rank}: " f"expected {expected_version}, got {actual_version}")
351- 351+ 
352- 352+ 
353-def get_version_detail(version_str: str) -> tuple[int, int]:353+def get_version_detail(version_str: str) -> tuple[int, int]:
354- parts = version_str.split(".")354+ parts = version_str.split(".")
355- if len(parts) != 2:355+ if len(parts) != 2:
356- raise ValueError(f"Invalid version format: expected 'X.Y', got '{version_str}'")356+ raise ValueError(f"Invalid version format: expected 'X.Y', got '{version_str}'")
357- 357+ 
358- try:358+ try:
359- major, minor = int(parts[0]), int(parts[1])359+ major, minor = int(parts[0]), int(parts[1])
360- except ValueError as e:360+ except ValueError as e:
361- raise ValueError(f"Version components must be integers: '{version_str}'") from e361+ raise ValueError(f"Version components must be integers: '{version_str}'") from e
362- 362+ 
363- return major, minor363+ return major, minor
364- 364+ 
365- 365+ 
366-def align_trace_from_beginning(366+def align_trace_from_beginning(
367- entries: dict[int, list[dict[str, Any]]],367+ entries: dict[int, list[dict[str, Any]]],
368-) -> dict[int, list[dict[str, Any]]]:368+) -> dict[int, list[dict[str, Any]]]:
369- """369+ """
370- Align the trace entries by record ID for entries.370+ Align the trace entries by record ID for entries.
371- This function takes a dictionary of rank names to lists of trace entries as input.371+ This function takes a dictionary of rank names to lists of trace entries as input.
372- Each trace entry is a dictionary containing information about a collective operation,372+ Each trace entry is a dictionary containing information about a collective operation,
373- including its unique identifier (`record_id` is monotonically increasing as we write into the ring buffer).373+ including its unique identifier (`record_id` is monotonically increasing as we write into the ring buffer).
374- The function finds the largest starting point across all ranks by taking the maximum374+ The function finds the largest starting point across all ranks by taking the maximum
375- `record_id` value of the first entry in each rank. Finally, it filters out any375+ `record_id` value of the first entry in each rank. Finally, it filters out any
376- entries with `record_id` values less than the maximum starting point.376+ entries with `record_id` values less than the maximum starting point.
377- The function returns the updated dictionary of sorted and filtered trace entries.377+ The function returns the updated dictionary of sorted and filtered trace entries.
378- 378+ 
379- Args:379+ Args:
380- entries (Dict[str, List[Dict[str, Any]]]): A dictionary of rank names to lists of trace entries.380+ entries (Dict[str, List[Dict[str, Any]]]): A dictionary of rank names to lists of trace entries.
381- 381+ 
382- Returns:382+ Returns:
383- entries (Dict[str, List[Dict[str, Any]]]): Entries sorted by record ID and filtered by the maximum starting point.383+ entries (Dict[str, List[Dict[str, Any]]]): Entries sorted by record ID and filtered by the maximum starting point.
384- """384+ """
385- 385+ 
386- maximum_starting_record_id = 0386+ maximum_starting_record_id = 0
387- for rank in entries:387+ for rank in entries:
388- # Although this is a ring buffer, we already sort the entries by `record_id` when dumping, we just388+ # Although this is a ring buffer, we already sort the entries by `record_id` when dumping, we just
389- # need to find the largest starting point. For example, if the buffer has the following entries:389+ # need to find the largest starting point. For example, if the buffer has the following entries:
390- # Rank 0: [0, 1, 2, 3, 4, 5, 6]390+ # Rank 0: [0, 1, 2, 3, 4, 5, 6]
391- # Rank 1: [1, 2, 3, 4, 5, 6, 7]391+ # Rank 1: [1, 2, 3, 4, 5, 6, 7]
392- # Rank 2: [2, 3, 4, 5, 6, 7, 8]392+ # Rank 2: [2, 3, 4, 5, 6, 7, 8]
393- # Rank 3: [0, 1, 2, 3, 4, 5, None]393+ # Rank 3: [0, 1, 2, 3, 4, 5, None]
394- # Then we should start from collective 2 not 0 because any collective before,394+ # Then we should start from collective 2 not 0 because any collective before,
395- # we don't have complete records from all ranks so we need to ignore them.395+ # we don't have complete records from all ranks so we need to ignore them.
396- first_record_id = entries[rank][0]["record_id"]396+ first_record_id = entries[rank][0]["record_id"]
397- maximum_starting_record_id = max(maximum_starting_record_id, first_record_id)397+ maximum_starting_record_id = max(maximum_starting_record_id, first_record_id)
398- 398+ 
399- for rank in entries:399+ for rank in entries:
400- entries[rank] = [entry for entry in entries[rank] if entry["record_id"] >= maximum_starting_record_id]400+ entries[rank] = [entry for entry in entries[rank] if entry["record_id"] >= maximum_starting_record_id]
401- 401+ 
402- return entries402+ return entries
Mtools/flight_recorder/fr_trace.py+26-26
@@ -1,26 +1,26 @@
1-from collections.abc import Sequence1+from collections.abc import Sequence
2-from typing import Optional2+from typing import Optional
3-import pickle3+import pickle
4- 4+ 
5-from tools.flight_recorder.components.builder import build_db5+from tools.flight_recorder.components.builder import build_db
6-from tools.flight_recorder.components.config_manager import JobConfig6+from tools.flight_recorder.components.config_manager import JobConfig
7-from tools.flight_recorder.components.loader import read_dir7+from tools.flight_recorder.components.loader import read_dir
8-from tools.flight_recorder.components.types import types8+from tools.flight_recorder.components.types import types
9-from tools.flight_recorder.components.utils import get_valid_read_path, get_valid_write_path9+from tools.flight_recorder.components.utils import get_valid_read_path, get_valid_write_path
10- 10+ 
11- 11+ 
12-def main(args: Optional[Sequence[str]] = None) -> None:12+def main(args: Optional[Sequence[str]] = None) -> None:
13- config = JobConfig()13+ config = JobConfig()
14- args = config.parse_args(args)14+ args = config.parse_args(args)
15- get_valid_read_path(args.trace_dir, is_dir=True)15+ get_valid_read_path(args.trace_dir, is_dir=True)
16- 16+ 
17- details, version = read_dir(args)17+ details, version = read_dir(args)
18- db = build_db(details, args, version)18+ db = build_db(details, args, version)
19- if args.output:19+ if args.output:
20- args.output = get_valid_write_path(args.output)20+ args.output = get_valid_write_path(args.output)
21- with open(args.output, "wb") as f:21+ with open(args.output, "wb") as f:
22- pickle.dump((types, db), f)22+ pickle.dump((types, db), f)
23- 23+ 
24- 24+ 
25-if __name__ == "__main__":25+if __name__ == "__main__":
26- main()26+ main()
Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/meta_kernel.py+522-522
@@ -1,523 +1,523 @@
1-from itertools import count1+from itertools import count
2-from typing import List, Union, Optional, Tuple, Any, Dict2+from typing import List, Union, Optional, Tuple, Any, Dict
3-import os3+import os
4-import sympy4+import sympy
5-import textwrap5+import textwrap
6- 6+ 
7-import torch7+import torch
8-import torch.fx8+import torch.fx
9- 9+ 
10-from torch._functorch.aot_autograd import set_model_name, get_aot_compilation_context10+from torch._functorch.aot_autograd import set_model_name, get_aot_compilation_context
11-from torch._inductor.codegen.simd import (11+from torch._inductor.codegen.simd import (
12- log,12+ log,
13- SIMDKernel,13+ SIMDKernel,
14- SIMDKernelFeatures,14+ SIMDKernelFeatures,
15- MultiKernel,15+ MultiKernel,
16- code_hash16+ code_hash
17-)17+)
18-from torch._inductor.codegen.triton import (18+from torch._inductor.codegen.triton import (
19- SIMDScheduling,19+ SIMDScheduling,
20- FixedTritonConfig,20+ FixedTritonConfig,
21-)21+)
22-from torch._inductor import config, ir, scheduler, metrics22+from torch._inductor import config, ir, scheduler, metrics
23-from torch._inductor.codecache import get_path23+from torch._inductor.codecache import get_path
24-from torch._dynamo.utils import counters24+from torch._dynamo.utils import counters
25-from torch._inductor.codegen.common import (25+from torch._inductor.codegen.common import (
26- IndentedBuffer,26+ IndentedBuffer,
27- Kernel,27+ Kernel,
28-)28+)
29-from torch._inductor.virtualized import V29+from torch._inductor.virtualized import V
30-from torch._inductor.codegen.triton import (30+from torch._inductor.codegen.triton import (
31- TritonKernel31+ TritonKernel
32-)32+)
33-from torch._inductor.utils import (33+from torch._inductor.utils import (
34- get_fused_kernel_name,34+ get_fused_kernel_name,
35- get_kernel_metadata,35+ get_kernel_metadata,
36-)36+)
37-from torch._inductor.scheduler import Scheduler37+from torch._inductor.scheduler import Scheduler
38-from torch_mlir.compiler_utils import OutputType38+from torch_mlir.compiler_utils import OutputType
39- 39+ 
40-from torch.fx.experimental.proxy_tensor import make_fx40+from torch.fx.experimental.proxy_tensor import make_fx
41-from torch._dynamo.device_interface import get_interface_for_device41+from torch._dynamo.device_interface import get_interface_for_device
42-from ..torch_mlir_patch import stateless_fx_import42+from ..torch_mlir_patch import stateless_fx_import
43-from ...npu.inductor_patch.lowering import map_strings_to_operators43+from ...npu.inductor_patch.lowering import map_strings_to_operators
44-from ...npu.utils import (44+from ...npu.utils import (
45- parse_fx_example_inputs,45+ parse_fx_example_inputs,
46- npu_cast_to_prim_cast,46+ npu_cast_to_prim_cast,
47- get_fx_graph_code,47+ get_fx_graph_code,
48- scalarize_tensor_ops_on_scalars,48+ scalarize_tensor_ops_on_scalars,
49- to_folder,49+ to_folder,
50- modify_gm_for_acc_comp,50+ modify_gm_for_acc_comp,
51- get_num_call_functions,51+ get_num_call_functions,
52- is_fx_dynamic,52+ is_fx_dynamic,
53- view_to_reshape53+ view_to_reshape
54-)54+)
55-from ... import config as anir_config55+from ... import config as anir_config
56-from ...npu.inductor_patch.lowering import merge_fx_graphs56+from ...npu.inductor_patch.lowering import merge_fx_graphs
57- 57+ 
58- 58+ 
59-id_iter = count()59+id_iter = count()
60- 60+ 
61- 61+ 
62-class NpuTritonKernel(TritonKernel):62+class NpuTritonKernel(TritonKernel):
63- def __init__(self, 63+ def __init__(self,
64- tiling: Dict[str, sympy.Expr],64+ tiling: Dict[str, sympy.Expr],
65- min_elem_per_thread=0,65+ min_elem_per_thread=0,
66- optimize_mask=True,66+ optimize_mask=True,
67- fixed_config: Optional[FixedTritonConfig] = None,67+ fixed_config: Optional[FixedTritonConfig] = None,
68- **kwargs,68+ **kwargs,
69- ):69+ ):
70- super().__init__(70+ super().__init__(
71- tiling,71+ tiling,
72- min_elem_per_thread=min_elem_per_thread,72+ min_elem_per_thread=min_elem_per_thread,
73- optimize_mask=optimize_mask,73+ optimize_mask=optimize_mask,
74- fixed_config = fixed_config,74+ fixed_config = fixed_config,
75- **kwargs,75+ **kwargs,
76- )76+ )
77- 77+ 
78- @staticmethod78+ @staticmethod
79- def inductor_meta_common():79+ def inductor_meta_common():
80- return {}80+ return {}
81- 81+
82- def call_kernel(self, call_args, name: str):82+ def call_kernel(self, call_args, name: str):
83- wrapper = V.graph.wrapper_code83+ wrapper = V.graph.wrapper_code
84- for call_arg in call_args:84+ for call_arg in call_args:
85- if call_arg.startswith('_uwu_'):85+ if call_arg.startswith('_uwu_'):
86- expression = map_strings_to_operators(call_arg)86+ expression = map_strings_to_operators(call_arg)
87- wrapper.writeline(f'{call_arg} = {expression}')87+ wrapper.writeline(f'{call_arg} = {expression}')
88- if len(call_args) > 0:88+ if len(call_args) > 0:
89- wrapper.generate_kernel_call(89+ wrapper.generate_kernel_call(
90- name,90+ name,
91- call_args,91+ call_args,
92- )92+ )
93- 93+ 
94- 94+ 
95-def _nc_key(nc):95+def _nc_key(nc):
96- if isinstance(nc, dict):96+ if isinstance(nc, dict):
97- return (tuple(nc.get("inputs", [])), tuple(nc.get("outputs", [])))97+ return (tuple(nc.get("inputs", [])), tuple(nc.get("outputs", [])))
98- return tuple(nc)98+ return tuple(nc)
99- 99+ 
100- 100+ 
101-def find_common_positions(list1, list2):101+def find_common_positions(list1, list2):
102- common_elements = set(list1) & set(list2)102+ common_elements = set(list1) & set(list2)
103- merged_list = list1 + list2103+ merged_list = list1 + list2
104- positions = [index for index, element in enumerate(merged_list) if element in common_elements]104+ positions = [index for index, element in enumerate(merged_list) if element in common_elements]
105- 105+ 
106- return merged_list, positions106+ return merged_list, positions
107- 107+ 
108- 108+ 
109-def refresh_input_meta_with_buffer_layout(node):109+def refresh_input_meta_with_buffer_layout(node):
110- val = node.meta.get('val')110+ val = node.meta.get('val')
111- try:111+ try:
112- layout = getattr(V.graph.try_get_buffer(node.target), "layout", None)112+ layout = getattr(V.graph.try_get_buffer(node.target), "layout", None)
113- except Exception:113+ except Exception:
114- layout = None114+ layout = None
115- if not torch.is_tensor(val) or layout is None:115+ if not torch.is_tensor(val) or layout is None:
116- return val116+ return val
117- 117+ 
118- size, stride = tuple(layout.size), tuple(layout.stride)118+ size, stride = tuple(layout.size), tuple(layout.stride)
119- if any(not isinstance(x, int) for x in (*size, *stride)):119+ if any(not isinstance(x, int) for x in (*size, *stride)):
120- return val120+ return val
121- 121+ 
122- if val.size() != size or val.stride() != stride:122+ if val.size() != size or val.stride() != stride:
123- with V.graph.fake_mode:123+ with V.graph.fake_mode:
124- val = torch.empty_strided(124+ val = torch.empty_strided(
125- size,125+ size,
126- stride,126+ stride,
127- dtype=val.dtype,127+ dtype=val.dtype,
128- device=val.device,128+ device=val.device,
129- requires_grad=val.requires_grad,129+ requires_grad=val.requires_grad,
130- )130+ )
131- node.meta['val'] = val131+ node.meta['val'] = val
132- return val132+ return val
133- 133+ 
134- 134+ 
135-def create_fx_from_snodes_by_traced_graph(snodes: List[scheduler.SchedulerNode], triton_kernel: TritonKernel):135+def create_fx_from_snodes_by_traced_graph(snodes: List[scheduler.SchedulerNode], triton_kernel: TritonKernel):
136- call_inputs = []136+ call_inputs = []
137- for snode in snodes:137+ for snode in snodes:
138- snode.node.data.traced_graph.last_node.name = snode.node.get_name()138+ snode.node.data.traced_graph.last_node.name = snode.node.get_name()
139- if len(snodes) == 1:139+ if len(snodes) == 1:
140- traced_graph = snodes[0].node.data.traced_graph140+ traced_graph = snodes[0].node.data.traced_graph
141- else:141+ else:
142- traced_graph = merge_fx_graphs([snode.node.data.traced_graph for snode in snodes])142+ traced_graph = merge_fx_graphs([snode.node.data.traced_graph for snode in snodes])
143- inputs = []143+ inputs = []
144- for node in traced_graph.graph.nodes:144+ for node in traced_graph.graph.nodes:
145- if node.op == 'placeholder':145+ if node.op == 'placeholder':
146- call_inputs.append(node.target)146+ call_inputs.append(node.target)
147- inputs.append(refresh_input_meta_with_buffer_layout(node))147+ inputs.append(refresh_input_meta_with_buffer_layout(node))
148- non_contiguous_indices = {}148+ non_contiguous_indices = {}
149- non_contiguous_indices["inputs"] = [i for i, inp in enumerate(inputs) if torch.is_tensor(inp) and not inp.is_contiguous()]149+ non_contiguous_indices["inputs"] = [i for i, inp in enumerate(inputs) if torch.is_tensor(inp) and not inp.is_contiguous()]
150- num_inputs = len(call_inputs)150+ num_inputs = len(call_inputs)
151- call_outputs = []151+ call_outputs = []
152- for snode in snodes:152+ for snode in snodes:
153- if snode.has_aliasing_or_mutation():153+ if snode.has_aliasing_or_mutation():
154- for buf in snode.get_outputs():154+ for buf in snode.get_outputs():
155- if len(buf.get_mutations()):155+ if len(buf.get_mutations()):
156- call_outputs.extend(buf.get_mutations())156+ call_outputs.extend(buf.get_mutations())
157- elif len(buf.get_aliases()):157+ elif len(buf.get_aliases()):
158- call_outputs.append(buf.get_name())158+ call_outputs.append(buf.get_name())
159- elif snode.node.get_name() not in (V.graph.removed_buffers | V.graph.inplaced_to_remove):159+ elif snode.node.get_name() not in (V.graph.removed_buffers | V.graph.inplaced_to_remove):
160- call_outputs.append(snode.node.get_name())160+ call_outputs.append(snode.node.get_name())
161- num_outputs = len(call_outputs)161+ num_outputs = len(call_outputs)
162- call_args, mutated_indices = find_common_positions(call_inputs, call_outputs)162+ call_args, mutated_indices = find_common_positions(call_inputs, call_outputs)
163- outputs = traced_graph.last_node if isinstance(traced_graph.last_node, List) \163+ outputs = traced_graph.last_node if isinstance(traced_graph.last_node, List) \
164- else [traced_graph.last_node]164+ else [traced_graph.last_node]
165- outputs = [output for output in outputs if output.name not in (V.graph.removed_buffers | V.graph.inplaced_to_remove)]165+ outputs = [output for output in outputs if output.name not in (V.graph.removed_buffers | V.graph.inplaced_to_remove)]
166- traced_graph.graph.output(tuple(outputs))166+ traced_graph.graph.output(tuple(outputs))
167- traced_graph.graph.lint()167+ traced_graph.graph.lint()
168- orig_module = torch.nn.Module()168+ orig_module = torch.nn.Module()
169- gm = torch.fx.GraphModule(orig_module, traced_graph.graph)169+ gm = torch.fx.GraphModule(orig_module, traced_graph.graph)
170- gm.recompile()170+ gm.recompile()
171- 171+ 
172- def runnable_gm(*args):172+ def runnable_gm(*args):
173- return torch.fx.Interpreter(gm).run(*args)173+ return torch.fx.Interpreter(gm).run(*args)
174- with V.graph.fake_mode: 174+ with V.graph.fake_mode:
175- gm = make_fx(runnable_gm)(*inputs)175+ gm = make_fx(runnable_gm)(*inputs)
176- view_to_reshape(gm)176+ view_to_reshape(gm)
177- non_contiguous_indices["outputs"] = [i + num_inputs 177+ non_contiguous_indices["outputs"] = [i + num_inputs
178- for i, call_output in enumerate(call_outputs)178+ for i, call_output in enumerate(call_outputs)
179- if not V.graph.try_get_buffer(call_output).layout.is_contiguous()]179+ if not V.graph.try_get_buffer(call_output).layout.is_contiguous()]
180- return (gm, call_args, {"num_outputs": num_outputs, 180+ return (gm, call_args, {"num_outputs": num_outputs,
181- "non_contiguous_indices": non_contiguous_indices, 181+ "non_contiguous_indices": non_contiguous_indices,
182- "mutated_indices": mutated_indices, })182+ "mutated_indices": mutated_indices, })
183- 183+ 
184- 184+ 
185-class NpuMetaKernel(Kernel):185+class NpuMetaKernel(Kernel):
186- def __init__(self, gm: torch.fx.GraphModule, snodes: list[scheduler.SchedulerNode], call_args: list[str], non_contiguous_indices: list[int], num_outputs: list[int] = None, mutated_indices: list[int] = None):186+ def __init__(self, gm: torch.fx.GraphModule, snodes: list[scheduler.SchedulerNode], call_args: list[str], non_contiguous_indices: list[int], num_outputs: list[int] = None, mutated_indices: list[int] = None):
187- super().__init__()187+ super().__init__()
188- if gm is None:188+ if gm is None:
189- self._gm = None189+ self._gm = None
190- self._gm_with_prim_cast = None190+ self._gm_with_prim_cast = None
191- self._is_dynamic = False191+ self._is_dynamic = False
192- else:192+ else:
193- self._gm = gm193+ self._gm = gm
194- self._gm_with_prim_cast = self.build_gm_with_prim_cast(gm)194+ self._gm_with_prim_cast = self.build_gm_with_prim_cast(gm)
195- self._is_dynamic = is_fx_dynamic(self._gm)195+ self._is_dynamic = is_fx_dynamic(self._gm)
196- 196+
197- if anir_config.online_acc_comp:197+ if anir_config.online_acc_comp:
198- modify_gm_for_acc_comp(self._gm)198+ modify_gm_for_acc_comp(self._gm)
199- 199+
200- self._snodes = snodes200+ self._snodes = snodes
201- self._call_args = call_args201+ self._call_args = call_args
202- self.non_contiguous_indices = non_contiguous_indices202+ self.non_contiguous_indices = non_contiguous_indices
203- self.num_outputs = num_outputs203+ self.num_outputs = num_outputs
204- self.mutated_indices = mutated_indices204+ self.mutated_indices = mutated_indices
205- 205+ 
206- def get_mlir_output_type(self):206+ def get_mlir_output_type(self):
207- return OutputType.RAW207+ return OutputType.RAW
208- 208+ 
209- def imports_for_benchmark_kernel(self):209+ def imports_for_benchmark_kernel(self):
210- return textwrap.dedent(210+ return textwrap.dedent(
211- """211+ """
212- from torch._dynamo.testing import rand_strided212+ from torch._dynamo.testing import rand_strided
213- {}213+ {}
214- import torch214+ import torch
215- """.format(215+ """.format(
216- V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")216+ V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")
217- )217+ )
218- )218+ )
219- 219+
220- def build_gm_with_prim_cast(self, gm):220+ def build_gm_with_prim_cast(self, gm):
221- return npu_cast_to_prim_cast(gm)221+ return npu_cast_to_prim_cast(gm)
222- 222+ 
223- def codegen_kernel(self, name=None):223+ def codegen_kernel(self, name=None):
224- code = IndentedBuffer()224+ code = IndentedBuffer()
225- 225+ 
226- scalarize_tensor_ops_on_scalars(self._gm_with_prim_cast)226+ scalarize_tensor_ops_on_scalars(self._gm_with_prim_cast)
227- 227+ 
228- set_model_name(f'MODEL_NAME')228+ set_model_name(f'MODEL_NAME')
229- *_, model_name, nth_graph = get_aot_compilation_context()229+ *_, model_name, nth_graph = get_aot_compilation_context()
230- 230+ 
231- import torch_mlir231+ import torch_mlir
232- 232+ 
233- mlir_module = stateless_fx_import(233+ mlir_module = stateless_fx_import(
234- self._gm_with_prim_cast,234+ self._gm_with_prim_cast,
235- model_name=model_name,235+ model_name=model_name,
236- output_type=self.get_mlir_output_type(),236+ output_type=self.get_mlir_output_type(),
237- import_symbolic_shape_expressions=False237+ import_symbolic_shape_expressions=False
238- )238+ )
239- from torch_mlir.compiler_utils import run_pipeline_with_repro_report239+ from torch_mlir.compiler_utils import run_pipeline_with_repro_report
240- 240+ 
241- run_pipeline_with_repro_report(241+ run_pipeline_with_repro_report(
242- mlir_module,242+ mlir_module,
243- f"builtin.module(torch-lower-to-backend-contract)",243+ f"builtin.module(torch-lower-to-backend-contract)",
244- "Lowering TorchFX IR -> Torch Backend IR",244+ "Lowering TorchFX IR -> Torch Backend IR",
245- )245+ )
246- code.splice(f'{str(mlir_module)}')246+ code.splice(f'{str(mlir_module)}')
247- 247+ 
248- return code.getvalue()248+ return code.getvalue()
249- 249+ 
250- def get_call_args(self):250+ def get_call_args(self):
251- return self._call_args251+ return self._call_args
252- 252+ 
253- def call_kernel(self, name: str, node=None):253+ def call_kernel(self, name: str, node=None):
254- wrapper = V.graph.wrapper_code254+ wrapper = V.graph.wrapper_code
255- call_args = self.get_call_args()255+ call_args = self.get_call_args()
256- 256+
257- for call_arg in call_args:257+ for call_arg in call_args:
258- if call_arg.startswith('_uwu'):258+ if call_arg.startswith('_uwu'):
259- expression = map_strings_to_operators(call_arg)259+ expression = map_strings_to_operators(call_arg)
260- wrapper.writeline(f'{call_arg} = {expression}')260+ wrapper.writeline(f'{call_arg} = {expression}')
261- 261+
262- if len(call_args) > 0:262+ if len(call_args) > 0:
263- wrapper.generate_kernel_call(name, call_args)263+ wrapper.generate_kernel_call(name, call_args)
264- 264+ 
265- def codegen_debug_performance(self, fd):265+ def codegen_debug_performance(self, fd):
266- from ...npu.utils import generate_compiler_repro_string, generate_fake_inputs266+ from ...npu.utils import generate_compiler_repro_string, generate_fake_inputs
267- name_to_example_inputs = parse_fx_example_inputs(self._gm)267+ name_to_example_inputs = parse_fx_example_inputs(self._gm)
268- call_args_str = ", ".join(list(name_to_example_inputs.keys()))268+ call_args_str = ", ".join(list(name_to_example_inputs.keys()))
269- fd.write(generate_compiler_repro_string(self._gm))269+ fd.write(generate_compiler_repro_string(self._gm))
270- fd.write("\n")270+ fd.write("\n")
271- fd.write("if __name__ == '__main__':\n")271+ fd.write("if __name__ == '__main__':\n")
272- fd.write(" from torch._inductor.utils import print_performance\n")272+ fd.write(" from torch._inductor.utils import print_performance\n")
273- fd.write(f" with torch.no_grad():\n")273+ fd.write(f" with torch.no_grad():\n")
274- fd.write(generate_fake_inputs(name_to_example_inputs))274+ fd.write(generate_fake_inputs(name_to_example_inputs))
275- fd.write('\n')275+ fd.write('\n')
276- fd.write(f" fn = lambda: mod({call_args_str})\n")276+ fd.write(f" fn = lambda: mod({call_args_str})\n")
277- fd.write(f" print_performance(fn, times=10, repeat=10)\n")277+ fd.write(f" print_performance(fn, times=10, repeat=10)\n")
278- 278+ 
279- 279+ 
280-class NpuMetaScheduling(SIMDScheduling):280+class NpuMetaScheduling(SIMDScheduling):
281- kernel_type = NpuTritonKernel281+ kernel_type = NpuTritonKernel
282- meta_kernel_type = NpuMetaKernel282+ meta_kernel_type = NpuMetaKernel
283- 283+ 
284- def __init__(self, sched: Scheduler):284+ def __init__(self, sched: Scheduler):
285- super().__init__(sched)285+ super().__init__(sched)
286- self.orig_fnode_name_to_fnode = {}286+ self.orig_fnode_name_to_fnode = {}
287- 287+ 
288- def _postprocess_src_code(self, src_code, mlir_kernel, kernel_name):288+ def _postprocess_src_code(self, src_code, mlir_kernel, kernel_name):
289- return src_code, {}289+ return src_code, {}
290- 290+ 
291- def define_kernel(self, src_code, mlir_kernel, traced_graph, mode=None):291+ def define_kernel(self, src_code, mlir_kernel, traced_graph, mode=None):
292- if mode is None:292+ if mode is None:
293- mode = anir_config._get_compile_mode()293+ mode = anir_config._get_compile_mode()
294- 294+
295- wrapper = V.graph.wrapper_code295+ wrapper = V.graph.wrapper_code
296- 296+ 
297- kernel_key = (src_code, tuple(mlir_kernel.non_contiguous_indices))297+ kernel_key = (src_code, tuple(mlir_kernel.non_contiguous_indices))
298- 298+ 
299- if kernel_key in wrapper.src_to_kernel:299+ if kernel_key in wrapper.src_to_kernel:
300- cached_val = wrapper.src_to_kernel[kernel_key]300+ cached_val = wrapper.src_to_kernel[kernel_key]
301- if isinstance(cached_val, str):301+ if isinstance(cached_val, str):
302- return cached_val302+ return cached_val
303- else:303+ else:
304- log.warning(f"Found invalid cache entry for kernel {mlir_kernel}. Recompiling.")304+ log.warning(f"Found invalid cache entry for kernel {mlir_kernel}. Recompiling.")
305- del wrapper.src_to_kernel[kernel_key]305+ del wrapper.src_to_kernel[kernel_key]
306- 306+ 
307- fused_name = (307+ fused_name = (
308- get_fused_kernel_name(mlir_kernel._snodes, config.triton.descriptive_names)308+ get_fused_kernel_name(mlir_kernel._snodes, config.triton.descriptive_names)
309- if config.triton.descriptive_names309+ if config.triton.descriptive_names
310- else ""310+ else ""
311- )311+ )
312- 312+ 
313- if mode in ["complete_fallback", "auto_fallback"]:313+ if mode in ["complete_fallback", "auto_fallback"]:
314- fx_graph_suffix = f"{next(id_iter)}"314+ fx_graph_suffix = f"{next(id_iter)}"
315- prefix = self._get_kernel_prefix()315+ prefix = self._get_kernel_prefix()
316- kernel_name = "_".join([prefix, fused_name, fx_graph_suffix])316+ kernel_name = "_".join([prefix, fused_name, fx_graph_suffix])
317- else:317+ else:
318- kernel_suffix = wrapper.next_kernel_suffix()318+ kernel_suffix = wrapper.next_kernel_suffix()
319- prefix = self._get_kernel_prefix()319+ prefix = self._get_kernel_prefix()
320- kernel_name = "_".join([prefix, fused_name, kernel_suffix])320+ kernel_name = "_".join([prefix, fused_name, kernel_suffix])
321- compile_src_code, extra_kernel_meta = self._postprocess_src_code(src_code, mlir_kernel, kernel_name)321+ compile_src_code, extra_kernel_meta = self._postprocess_src_code(src_code, mlir_kernel, kernel_name)
322- 322+ 
323- current_device = V.graph.get_current_device_or_throw()323+ current_device = V.graph.get_current_device_or_throw()
324- traced_graph_hash = code_hash(traced_graph.print_readable(print_output=False) + kernel_name)324+ traced_graph_hash = code_hash(traced_graph.print_readable(print_output=False) + kernel_name)
325- num_call_functions = get_num_call_functions(mlir_kernel._gm)325+ num_call_functions = get_num_call_functions(mlir_kernel._gm)
326- 326+ 
327- if num_call_functions <= 1 or kernel_name in anir_config.force_fallback_kernel_names:327+ if num_call_functions <= 1 or kernel_name in anir_config.force_fallback_kernel_names:
328- mode = "complete_fallback"328+ mode = "complete_fallback"
329- 329+ 
330- kernel_meta = self._prepare_kernel_meta(330+ kernel_meta = self._prepare_kernel_meta(
331- mlir_kernel, current_device, kernel_name, traced_graph_hash, num_call_functions, compile_src_code331+ mlir_kernel, current_device, kernel_name, traced_graph_hash, num_call_functions, compile_src_code
332- )332+ )
333- kernel_meta.update(extra_kernel_meta)333+ kernel_meta.update(extra_kernel_meta)
334- 334+ 
335- wrapper.src_to_kernel[kernel_key] = kernel_name335+ wrapper.src_to_kernel[kernel_key] = kernel_name
336- 336+
337- subs_name = kernel_name if config.triton.unique_kernel_names else f"{self._get_kernel_prefix()}_"337+ subs_name = kernel_name if config.triton.unique_kernel_names else f"{self._get_kernel_prefix()}_"
338- 338+
339- compile_wrapper = IndentedBuffer()339+ compile_wrapper = IndentedBuffer()
340- metadata_comment = ""340+ metadata_comment = ""
341- 341+ 
342- if mode == "auto_fallback":342+ if mode == "auto_fallback":
343- src_code = src_code.replace("MODEL_NAME", kernel_name)343+ src_code = src_code.replace("MODEL_NAME", kernel_name)
344- self._handle_auto_fallback_mode(344+ self._handle_auto_fallback_mode(
345- compile_wrapper, src_code, kernel_name, subs_name, kernel_meta, wrapper, metadata_comment, mlir_kernel345+ compile_wrapper, src_code, kernel_name, subs_name, kernel_meta, wrapper, metadata_comment, mlir_kernel
346- )346+ )
347- elif mode == "complete_fallback":347+ elif mode == "complete_fallback":
348- self._handle_complete_fallback_mode(348+ self._handle_complete_fallback_mode(
349- compile_wrapper, kernel_name, kernel_meta, wrapper, metadata_comment, mlir_kernel349+ compile_wrapper, kernel_name, kernel_meta, wrapper, metadata_comment, mlir_kernel
350- )350+ )
351- else:351+ else:
352- src_code = src_code.replace("MODEL_NAME", kernel_name)352+ src_code = src_code.replace("MODEL_NAME", kernel_name)
353- self._handle_default_mode(353+ self._handle_default_mode(
354- compile_wrapper, src_code, kernel_name, subs_name, kernel_meta, wrapper, metadata_comment, mlir_kernel354+ compile_wrapper, src_code, kernel_name, subs_name, kernel_meta, wrapper, metadata_comment, mlir_kernel
355- )355+ )
356- 356+ 
357- if mode in ["complete_fallback", "auto_fallback"]:357+ if mode in ["complete_fallback", "auto_fallback"]:
358- self._dump_fx_graph_for_fallback(358+ self._dump_fx_graph_for_fallback(
359- mlir_kernel, current_device, traced_graph_hash, kernel_name, compile_wrapper.getvalue()359+ mlir_kernel, current_device, traced_graph_hash, kernel_name, compile_wrapper.getvalue()
360- )360+ )
361- 361+ 
362- return kernel_name362+ return kernel_name
363- 363+ 
364- def _get_kernel_prefix(self) -> str:364+ def _get_kernel_prefix(self) -> str:
365- return "mlir"365+ return "mlir"
366- 366+ 
367- def _prepare_kernel_meta(self, mlir_kernel, device, name, hash_val, num_calls, src_code) -> Dict[str, Any]:367+ def _prepare_kernel_meta(self, mlir_kernel, device, name, hash_val, num_calls, src_code) -> Dict[str, Any]:
368- device_interface = get_interface_for_device(device.type)368+ device_interface = get_interface_for_device(device.type)
369- device = torch.device(device.type, device.index)369+ device = torch.device(device.type, device.index)
370- device_name = device_interface.get_compute_capability(device)370+ device_name = device_interface.get_compute_capability(device)
371- return {371+ return {
372- 'device_str': device.type,372+ 'device_str': device.type,
373- 'device_index': device.index,373+ 'device_index': device.index,
374- 'device_name': device_name,374+ 'device_name': device_name,
375- 'num_outputs': mlir_kernel.num_outputs,375+ 'num_outputs': mlir_kernel.num_outputs,
376- 'non_contiguous_indices': mlir_kernel.non_contiguous_indices,376+ 'non_contiguous_indices': mlir_kernel.non_contiguous_indices,
377- 'dynamic': mlir_kernel._is_dynamic,377+ 'dynamic': mlir_kernel._is_dynamic,
378- 'mutated_indices': mlir_kernel.mutated_indices,378+ 'mutated_indices': mlir_kernel.mutated_indices,
379- 'traced_graph_cache': anir_config.traced_graph_cache or "traced_graph_cache",379+ 'traced_graph_cache': anir_config.traced_graph_cache or "traced_graph_cache",
380- 'traced_graph_hash': hash_val,380+ 'traced_graph_hash': hash_val,
381- 'num_call_functions': num_calls,381+ 'num_call_functions': num_calls,
382- 'is_reduction': 'linalg.reduce' in src_code,382+ 'is_reduction': 'linalg.reduce' in src_code,
383- 'are_deterministic_algorithms_enabled': torch.are_deterministic_algorithms_enabled(),383+ 'are_deterministic_algorithms_enabled': torch.are_deterministic_algorithms_enabled(),
384- }384+ }
385- 385+ 
386- def _handle_auto_fallback_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment, mlir_kernel=None):386+ def _handle_auto_fallback_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment, mlir_kernel=None):
387- _basename, _, kernel_path = get_path(code_hash(src_code.strip()), "py")387+ _basename, _, kernel_path = get_path(code_hash(src_code.strip()), "py")
388- 388+
389- compile_wrapper.writeline(f"async_compile.{self._get_compile_api()}({subs_name!r}, '''")389+ compile_wrapper.writeline(f"async_compile.{self._get_compile_api()}({subs_name!r}, '''")
390- compile_wrapper.splice(src_code, strip=True)390+ compile_wrapper.splice(src_code, strip=True)
391- compile_wrapper.writeline(f"''', kernel_meta={meta})")391+ compile_wrapper.writeline(f"''', kernel_meta={meta})")
392- 392+ 
393- metadata_comment = f"# kernel path: {kernel_path}"393+ metadata_comment = f"# kernel path: {kernel_path}"
394- 394+ 
395- origins, detailed_origins = get_kernel_metadata(mlir_kernel._snodes, wrapper)395+ origins, detailed_origins = get_kernel_metadata(mlir_kernel._snodes, wrapper)
396- metadata_comment += "\n" + origins + "\n" + detailed_origins396+ metadata_comment += "\n" + origins + "\n" + detailed_origins
397- 397+
398- wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)398+ wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)
399- 399+ 
400- if metrics.is_metric_table_enabled("kernel_metadata"):400+ if metrics.is_metric_table_enabled("kernel_metadata"):
401- metrics.log_kernel_metadata(name, kernel_path, src_code)401+ metrics.log_kernel_metadata(name, kernel_path, src_code)
402- 402+ 
403- def _handle_complete_fallback_mode(self, compile_wrapper, name, meta, wrapper, metadata_comment, mlir_kernel):403+ def _handle_complete_fallback_mode(self, compile_wrapper, name, meta, wrapper, metadata_comment, mlir_kernel):
404- compile_wrapper.writeline(f"async_compile.import_fx({name!r}, kernel_meta={meta})")404+ compile_wrapper.writeline(f"async_compile.import_fx({name!r}, kernel_meta={meta})")
405- metadata_comment = f'"""\n{mlir_kernel._gm.print_readable(print_output=False)}\n"""'405+ metadata_comment = f'"""\n{mlir_kernel._gm.print_readable(print_output=False)}\n"""'
406- wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)406+ wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)
407- 407+ 
408- def _handle_default_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment):408+ def _handle_default_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment):
409- self._handle_auto_fallback_mode(compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment)409+ self._handle_auto_fallback_mode(compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment)
410- 410+ 
411- def _get_compile_api(self) -> str:411+ def _get_compile_api(self) -> str:
412- return "auto_fallback"412+ return "auto_fallback"
413- 413+ 
414- def _dump_fx_graph_for_fallback(self, mlir_kernel, device, graph_hash, kernel_name, compile_code):414+ def _dump_fx_graph_for_fallback(self, mlir_kernel, device, graph_hash, kernel_name, compile_code):
415- 415+
416- cache_root = os.getenv("TORCHINDUCTOR_CACHE_DIR")416+ cache_root = os.getenv("TORCHINDUCTOR_CACHE_DIR")
417- dump_path = os.path.join(417+ dump_path = os.path.join(
418- cache_root, 418+ cache_root,
419- anir_config.traced_graph_cache or "traced_graph_cache", 419+ anir_config.traced_graph_cache or "traced_graph_cache",
420- str(device.index), 420+ str(device.index),
421- graph_hash421+ graph_hash
422- )422+ )
423- 423+
424- if not os.path.exists(dump_path):424+ if not os.path.exists(dump_path):
425- os.makedirs(dump_path, exist_ok=True)425+ os.makedirs(dump_path, exist_ok=True)
426- to_folder(mlir_kernel._gm, dump_path, graph_hash=graph_hash, module_name=graph_hash)426+ to_folder(mlir_kernel._gm, dump_path, graph_hash=graph_hash, module_name=graph_hash)
427- 427+ 
428- if anir_config.fx_subgraph_dump_path is not None:428+ if anir_config.fx_subgraph_dump_path is not None:
429- subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(device.index), kernel_name)429+ subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(device.index), kernel_name)
430- os.makedirs(subgraph_dump_path, exist_ok=True)430+ os.makedirs(subgraph_dump_path, exist_ok=True)
431- 431+
432- num_args = len(mlir_kernel._gm.code.split('forward(', )[1].split(')')[0].split(', ')) - 1432+ num_args = len(mlir_kernel._gm.code.split('forward(', )[1].split(')')[0].split(', ')) - 1
433- fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=False, kernel_code=compile_code, kernel_name=kernel_name)433+ fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=False, kernel_code=compile_code, kernel_name=kernel_name)
434- runnable_fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=True, kernel_code=compile_code, kernel_name=kernel_name)434+ runnable_fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=True, kernel_code=compile_code, kernel_name=kernel_name)
435- 435+
436- with open(os.path.join(subgraph_dump_path, f'{kernel_name}.py'), 'w') as f:436+ with open(os.path.join(subgraph_dump_path, f'{kernel_name}.py'), 'w') as f:
437- f.write(fx_graph_code)437+ f.write(fx_graph_code)
438- with open(os.path.join(subgraph_dump_path, f'runnable_{kernel_name}.py'), 'w') as f:438+ with open(os.path.join(subgraph_dump_path, f'runnable_{kernel_name}.py'), 'w') as f:
439- f.write(runnable_fx_graph_code)439+ f.write(runnable_fx_graph_code)
440- 440+ 
441- 441+ 
442- def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures, nodes):442+ def codegen_node_schedule(self, kernel_features: SIMDKernelFeatures, nodes):
443- node_schedule = kernel_features.node_schedule443+ node_schedule = kernel_features.node_schedule
444- 444+ 
445- tiling = self.select_tiling(445+ tiling = self.select_tiling(
446- node_schedule, kernel_features.numel, kernel_features.reduction_numel446+ node_schedule, kernel_features.numel, kernel_features.reduction_numel
447- )447+ )
448- 448+ 
449- kernels = self.create_kernel_choices(449+ kernels = self.create_kernel_choices(
450- kernel_features, [tiling], {"features": kernel_features}450+ kernel_features, [tiling], {"features": kernel_features}
451- )451+ )
452- for kernel in kernels:452+ for kernel in kernels:
453- super().codegen_node_schedule_with_kernel(node_schedule, kernel)453+ super().codegen_node_schedule_with_kernel(node_schedule, kernel)
454- MultiKernel.merge_workspaces_inplace(kernels)454+ MultiKernel.merge_workspaces_inplace(kernels)
455- for kernel in kernels:455+ for kernel in kernels:
456- V.graph.removed_buffers |= kernel.removed_buffers456+ V.graph.removed_buffers |= kernel.removed_buffers
457- V.graph.inplaced_to_remove |= kernel.inplaced_to_remove457+ V.graph.inplaced_to_remove |= kernel.inplaced_to_remove
458- if not anir_config.traced_graph_cache:458+ if not anir_config.traced_graph_cache:
459- anir_config.traced_graph_cache = "traced_graph_cache"459+ anir_config.traced_graph_cache = "traced_graph_cache"
460- os.makedirs(460+ os.makedirs(
461- os.path.join(os.getenv("TORCHINDUCTOR_CACHE_DIR"), anir_config.traced_graph_cache),461+ os.path.join(os.getenv("TORCHINDUCTOR_CACHE_DIR"), anir_config.traced_graph_cache),
462- exist_ok=True,462+ exist_ok=True,
463- )463+ )
464- traced_graph, call_args, compile_kwargs = create_fx_from_snodes_by_traced_graph(464+ traced_graph, call_args, compile_kwargs = create_fx_from_snodes_by_traced_graph(
465- nodes, kernel465+ nodes, kernel
466- )466+ )
467- mlir_kernel = self.meta_kernel_type(traced_graph, nodes, call_args, **compile_kwargs)467+ mlir_kernel = self.meta_kernel_type(traced_graph, nodes, call_args, **compile_kwargs)
468- with V.set_kernel_handler(mlir_kernel):468+ with V.set_kernel_handler(mlir_kernel):
469- src_code = mlir_kernel.codegen_kernel()469+ src_code = mlir_kernel.codegen_kernel()
470- kernel_name = self.define_kernel(src_code, mlir_kernel, traced_graph)470+ kernel_name = self.define_kernel(src_code, mlir_kernel, traced_graph)
471- log.debug("Generating kernel code with kernel_name: %s", kernel_name)471+ log.debug("Generating kernel code with kernel_name: %s", kernel_name)
472- kernel.kernel_name = kernel_name472+ kernel.kernel_name = kernel_name
473- kernel.code_hash = code_hash(src_code)473+ kernel.code_hash = code_hash(src_code)
474- del kernel474+ del kernel
475- 475+ 
476- final_kernel: Union[SIMDKernel, MultiKernel]476+ final_kernel: Union[SIMDKernel, MultiKernel]
477- if len(kernels) > 1:477+ if len(kernels) > 1:
478- raise RuntimeError("MultiKernel not Implemented for this backend!")478+ raise RuntimeError("MultiKernel not Implemented for this backend!")
479- else:479+ else:
480- (final_kernel,) = kernels480+ (final_kernel,) = kernels
481- 481+ 
482- with V.set_kernel_handler(final_kernel):482+ with V.set_kernel_handler(final_kernel):
483- for node in kernel_features.scheduler_nodes():483+ for node in kernel_features.scheduler_nodes():
484- node.mark_run()484+ node.mark_run()
485- 485+ 
486- self.codegen_comment(node_schedule)486+ self.codegen_comment(node_schedule)
487- final_kernel.call_kernel(call_args, final_kernel.kernel_name)487+ final_kernel.call_kernel(call_args, final_kernel.kernel_name)
488- 488+ 
489- if config.nan_asserts:489+ if config.nan_asserts:
490- final_kernel.codegen_nan_check()490+ final_kernel.codegen_nan_check()
491- if config.warn_mix_layout:491+ if config.warn_mix_layout:
492- final_kernel.warn_mix_layout(kernels[0].kernel_name)492+ final_kernel.warn_mix_layout(kernels[0].kernel_name)
493- 493+ 
494- if (494+ if (
495- V.graph.wrapper_code.supports_intermediate_hooks495+ V.graph.wrapper_code.supports_intermediate_hooks
496- and config.generate_intermediate_hooks496+ and config.generate_intermediate_hooks
497- ):497+ ):
498- live_outs = kernels[0].args.live_output_buffers()498+ live_outs = kernels[0].args.live_output_buffers()
499- for node in kernel_features.scheduler_nodes():499+ for node in kernel_features.scheduler_nodes():
500- name = node.get_name()500+ name = node.get_name()
501- if name not in live_outs:501+ if name not in live_outs:
502- continue502+ continue
503- if node.node is None:503+ if node.node is None:
504- raise RuntimeError("assert node.node is not None")504+ raise RuntimeError("assert node.node is not None")
505- 505+ 
506- origin_node = node.node.get_origin_node()506+ origin_node = node.node.get_origin_node()
507- if origin_node is not None:507+ if origin_node is not None:
508- counters["inductor"]["intermediate_hooks"] += 1508+ counters["inductor"]["intermediate_hooks"] += 1
509- V.graph.wrapper_code.writeline(509+ V.graph.wrapper_code.writeline(
510- f"run_intermediate_hooks({origin_node.name!r}, {name})"510+ f"run_intermediate_hooks({origin_node.name!r}, {name})"
511- )511+ )
512- 512+ 
513- self.scheduler.free_buffers()513+ self.scheduler.free_buffers()
514- 514+ 
515- def codegen_node(self, node: Union[scheduler.SchedulerNode, object]):515+ def codegen_node(self, node: Union[scheduler.SchedulerNode, object]):
516- nodes: List[scheduler.SchedulerNode] = node.get_nodes()516+ nodes: List[scheduler.SchedulerNode] = node.get_nodes()
517- 517+
518- _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group518+ _, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group
519- 519+ 
520- node_schedule = self.generate_node_schedule(nodes, numel, rnumel)520+ node_schedule = self.generate_node_schedule(nodes, numel, rnumel)
521- kernel_features = SIMDKernelFeatures(node_schedule, numel, rnumel)521+ kernel_features = SIMDKernelFeatures(node_schedule, numel, rnumel)
522- 522+
523 return self.codegen_node_schedule(kernel_features, nodes)523 return self.codegen_node_schedule(kernel_features, nodes)
Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/meta_compiler.py+329-329
@@ -1,330 +1,330 @@
1-import os1+import os
2-import sys2+import sys
3-import importlib3+import importlib
4-import shutil4+import shutil
5-from itertools import count5+from itertools import count
6-from typing import Any, Callable, Dict, List, Optional, Tuple, Iterator6+from typing import Any, Callable, Dict, List, Optional, Tuple, Iterator
7- 7+ 
8-import torch8+import torch
9-from torch._inductor.compile_fx import clone_preserve_strides9+from torch._inductor.compile_fx import clone_preserve_strides
10- 10+ 
11-from .. import config as anir_config11+from .. import config as anir_config
12-from .utils import replace_placeholders12+from .utils import replace_placeholders
13- 13+ 
14- 14+ 
15-_dump_id_iter: Iterator[int] = count()15+_dump_id_iter: Iterator[int] = count()
16- 16+ 
17-class MetaCompiler:17+class MetaCompiler:
18- def __init__(18+ def __init__(
19- self,19+ self,
20- kernel_name: str = "",20+ kernel_name: str = "",
21- multiprocess_compile: bool = False,21+ multiprocess_compile: bool = False,
22- no_more_compile: bool = False,22+ no_more_compile: bool = False,
23- kernel_meta: Optional[Dict[str, Any]] = None,23+ kernel_meta: Optional[Dict[str, Any]] = None,
24- autotune: bool = True,24+ autotune: bool = True,
25- ):25+ ):
26- kernel_meta = kernel_meta or {}26+ kernel_meta = kernel_meta or {}
27- 27+ 
28- self.kernel_name = kernel_name28+ self.kernel_name = kernel_name
29- self.kernel_meta = kernel_meta29+ self.kernel_meta = kernel_meta
30- 30+ 
31- self.dynamic = kernel_meta.get("dynamic")31+ self.dynamic = kernel_meta.get("dynamic")
32- self.mutated_indices = kernel_meta.get("mutated_indices") or []32+ self.mutated_indices = kernel_meta.get("mutated_indices") or []
33- self.kernel_hash = kernel_meta.get("kernel_hash")33+ self.kernel_hash = kernel_meta.get("kernel_hash")
34- self.signature = kernel_meta.get("signature")34+ self.signature = kernel_meta.get("signature")
35- self.ranks = kernel_meta.get("ranks")35+ self.ranks = kernel_meta.get("ranks")
36- self.num_outputs = kernel_meta.get("num_outputs")36+ self.num_outputs = kernel_meta.get("num_outputs")
37- self.num_call_functions = kernel_meta.get("num_call_functions")37+ self.num_call_functions = kernel_meta.get("num_call_functions")
38- self.device_index = kernel_meta.get("device_index", 0)38+ self.device_index = kernel_meta.get("device_index", 0)
39- self.traced_graph_hash = kernel_meta.get("traced_graph_hash")39+ self.traced_graph_hash = kernel_meta.get("traced_graph_hash")
40- 40+ 
41- self.multiprocess_compile = multiprocess_compile41+ self.multiprocess_compile = multiprocess_compile
42- self.no_more_compile = no_more_compile42+ self.no_more_compile = no_more_compile
43- self.autotune = autotune43+ self.autotune = autotune
44- 44+ 
45- self._fallback_call: Optional[Callable[..., Any]] = None45+ self._fallback_call: Optional[Callable[..., Any]] = None
46- self._nc_input_indices: Optional[List[int]] = None46+ self._nc_input_indices: Optional[List[int]] = None
47- self._nc_output_indices: Optional[List[int]] = None47+ self._nc_output_indices: Optional[List[int]] = None
48- 48+ 
49- self.autotuned: bool = False49+ self.autotuned: bool = False
50- self.launchers: list[Callable[..., Any]] = []50+ self.launchers: list[Callable[..., Any]] = []
51- self.kernel_paths: list[Any] = []51+ self.kernel_paths: list[Any] = []
52- self.is_fallback_kernels: list[bool] = []52+ self.is_fallback_kernels: list[bool] = []
53- 53+ 
54- def register_launcher(54+ def register_launcher(
55- self,55+ self,
56- launcher: Callable[..., Any],56+ launcher: Callable[..., Any],
57- kernel_path: Any = None,57+ kernel_path: Any = None,
58- is_fallback_kernel: bool = False,58+ is_fallback_kernel: bool = False,
59- ) -> None:59+ ) -> None:
60- """Register a compiled or fallback launcher for runtime execution."""60+ """Register a compiled or fallback launcher for runtime execution."""
61- self.launchers.append(launcher)61+ self.launchers.append(launcher)
62- self.kernel_paths.append(kernel_path)62+ self.kernel_paths.append(kernel_path)
63- self.is_fallback_kernels.append(is_fallback_kernel)63+ self.is_fallback_kernels.append(is_fallback_kernel)
64- 64+ 
65- def get_primary_launcher_index(self) -> int:65+ def get_primary_launcher_index(self) -> int:
66- if not self.launchers:66+ if not self.launchers:
67- raise RuntimeError("No valid launcher")67+ raise RuntimeError("No valid launcher")
68- for idx, is_fallback_kernel in enumerate(self.is_fallback_kernels):68+ for idx, is_fallback_kernel in enumerate(self.is_fallback_kernels):
69- if not is_fallback_kernel:69+ if not is_fallback_kernel:
70- return idx70+ return idx
71- return 071+ return 0
72- 72+ 
73- def get_primary_launcher(self) -> Callable[..., Any]:73+ def get_primary_launcher(self) -> Callable[..., Any]:
74- return self.launchers[self.get_primary_launcher_index()]74+ return self.launchers[self.get_primary_launcher_index()]
75- 75+ 
76- def data_dump(self, *args, dump_path=None):76+ def data_dump(self, *args, dump_path=None):
77- if not dump_path:77+ if not dump_path:
78- dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)78+ dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)
79- data_dump_path = os.path.join(dump_path, 'data.pth')79+ data_dump_path = os.path.join(dump_path, 'data.pth')
80- args_cpu = [arg.cpu() if isinstance(arg, torch.Tensor) else arg for arg in args]80+ args_cpu = [arg.cpu() if isinstance(arg, torch.Tensor) else arg for arg in args]
81- torch.save(args_cpu, data_dump_path)81+ torch.save(args_cpu, data_dump_path)
82- 82+ 
83- def data_dump_fake(self, *args, dump_path=None):83+ def data_dump_fake(self, *args, dump_path=None):
84- if not dump_path:84+ if not dump_path:
85- dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)85+ dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)
86- runable_py_path = os.path.join(dump_path, f'runnable_{self.kernel_name}.py')86+ runable_py_path = os.path.join(dump_path, f'runnable_{self.kernel_name}.py')
87- fake_inputs = [f'rand_strided({arg.shape}, {arg.stride()}, device="{arg.device.type}", dtype={arg.dtype})' \87+ fake_inputs = [f'rand_strided({arg.shape}, {arg.stride()}, device="{arg.device.type}", dtype={arg.dtype})' \
88- if isinstance(arg, torch.Tensor) else str(arg) for arg in args[:-self.num_outputs]]88+ if isinstance(arg, torch.Tensor) else str(arg) for arg in args[:-self.num_outputs]]
89- fake_outputs = [f'empty_strided({arg.shape}, {arg.stride()}, device="{arg.device.type}", dtype={arg.dtype})' \89+ fake_outputs = [f'empty_strided({arg.shape}, {arg.stride()}, device="{arg.device.type}", dtype={arg.dtype})' \
90- if isinstance(arg, torch.Tensor) else str(arg) for arg in args[-self.num_outputs:]]90+ if isinstance(arg, torch.Tensor) else str(arg) for arg in args[-self.num_outputs:]]
91- replacements = {"FAKE_ARGS_PLACEHOLDER": f"args = [{', '.join(fake_inputs + fake_outputs)}]"}91+ replacements = {"FAKE_ARGS_PLACEHOLDER": f"args = [{', '.join(fake_inputs + fake_outputs)}]"}
92- replace_placeholders(runable_py_path, replacements)92+ replace_placeholders(runable_py_path, replacements)
93- 93+ 
94- def fx_subgraph_dump(self, suffix):94+ def fx_subgraph_dump(self, suffix):
95- subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)95+ subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(self.device_index), self.kernel_name)
96- failed_fx_subgraph_dump_path = anir_config.fx_subgraph_dump_path + f'_{suffix}'96+ failed_fx_subgraph_dump_path = anir_config.fx_subgraph_dump_path + f'_{suffix}'
97- failed_subgraph_dump_path = os.path.join(failed_fx_subgraph_dump_path, str(self.device_index), f'{next(_dump_id_iter)}_' + self.kernel_name)97+ failed_subgraph_dump_path = os.path.join(failed_fx_subgraph_dump_path, str(self.device_index), f'{next(_dump_id_iter)}_' + self.kernel_name)
98- if os.path.exists(failed_subgraph_dump_path):98+ if os.path.exists(failed_subgraph_dump_path):
99- shutil.rmtree(failed_subgraph_dump_path)99+ shutil.rmtree(failed_subgraph_dump_path)
100- shutil.copytree(subgraph_dump_path, failed_subgraph_dump_path)100+ shutil.copytree(subgraph_dump_path, failed_subgraph_dump_path)
101- return failed_subgraph_dump_path101+ return failed_subgraph_dump_path
102- 102+
103- def acc_compare_and_dump(self, *args, **kwargs):103+ def acc_compare_and_dump(self, *args, **kwargs):
104- from torch.testing._comparison import _make_mismatch_msg104+ from torch.testing._comparison import _make_mismatch_msg
105- self.register_fx_fallback(self.kernel_meta)105+ self.register_fx_fallback(self.kernel_meta)
106- launcher_fx = self.launchers[1]106+ launcher_fx = self.launchers[1]
107- launcher = self.launchers[0]107+ launcher = self.launchers[0]
108- 108+ 
109- fx_outputs = [clone_preserve_strides(arg).to(torch.float32) if arg.dtype == torch.bfloat16 \109+ fx_outputs = [clone_preserve_strides(arg).to(torch.float32) if arg.dtype == torch.bfloat16 \
110- else clone_preserve_strides(arg) for arg in args[-self.num_outputs:]]110+ else clone_preserve_strides(arg) for arg in args[-self.num_outputs:]]
111- fx_inputs = [clone_preserve_strides(arg) if isinstance(arg, torch.Tensor) else arg for arg in args[:-self.num_outputs]]111+ fx_inputs = [clone_preserve_strides(arg) if isinstance(arg, torch.Tensor) else arg for arg in args[:-self.num_outputs]]
112- fx_inputs = [inp.float() if isinstance(inp, torch.Tensor) and inp.dtype == torch.bfloat16 else inp for inp in fx_inputs]112+ fx_inputs = [inp.float() if isinstance(inp, torch.Tensor) and inp.dtype == torch.bfloat16 else inp for inp in fx_inputs]
113- 113+
114- fx_args = fx_inputs + fx_outputs114+ fx_args = fx_inputs + fx_outputs
115- launcher_fx(*fx_args, **kwargs)115+ launcher_fx(*fx_args, **kwargs)
116- 116+ 
117- if self.dynamic:117+ if self.dynamic:
118- args_new = self.prepare_runtime_args(118+ args_new = self.prepare_runtime_args(
119- list(args),119+ list(args),
120- )120+ )
121- else:121+ else:
122- args_new = args122+ args_new = args
123- 123+
124- output = launcher(*args_new, **kwargs)124+ output = launcher(*args_new, **kwargs)
125- 125+ 
126- has_acc_error = False126+ has_acc_error = False
127- num_inputs = len(args) - self.num_outputs127+ num_inputs = len(args) - self.num_outputs
128- for idx, (actual, expected) in enumerate(zip(args[num_inputs:], fx_outputs)):128+ for idx, (actual, expected) in enumerate(zip(args[num_inputs:], fx_outputs)):
129- if actual.dtype != expected.dtype:129+ if actual.dtype != expected.dtype:
130- expected = expected.to(actual.dtype)130+ expected = expected.to(actual.dtype)
131- acc_comp_tol = anir_config.acc_comp_tol.get(actual.dtype, anir_config.acc_comp_tol['default'])131+ acc_comp_tol = anir_config.acc_comp_tol.get(actual.dtype, anir_config.acc_comp_tol['default'])
132- rtol = acc_comp_tol['rtol']132+ rtol = acc_comp_tol['rtol']
133- atol = acc_comp_tol['atol']133+ atol = acc_comp_tol['atol']
134- matches = torch.isclose(134+ matches = torch.isclose(
135- actual, expected, rtol=rtol, atol=atol, equal_nan=True135+ actual, expected, rtol=rtol, atol=atol, equal_nan=True
136- )136+ )
137- if not matches.all():137+ if not matches.all():
138- abs_diff = abs(actual - expected)138+ abs_diff = abs(actual - expected)
139- rel_diff = abs_diff / abs(expected)139+ rel_diff = abs_diff / abs(expected)
140- rel_diff.masked_fill_(matches, 0)140+ rel_diff.masked_fill_(matches, 0)
141- number_of_elements = matches.numel()141+ number_of_elements = matches.numel()
142- total_mismatches = number_of_elements - int(torch.sum(matches))142+ total_mismatches = number_of_elements - int(torch.sum(matches))
143- extra = (143+ extra = (
144- f"Mismatched elements: {total_mismatches} / {number_of_elements} "144+ f"Mismatched elements: {total_mismatches} / {number_of_elements} "
145- f"({total_mismatches / number_of_elements:.1%})"145+ f"({total_mismatches / number_of_elements:.1%})"
146- )146+ )
147- msg = _make_mismatch_msg(147+ msg = _make_mismatch_msg(
148- default_identifier="Tensor-likes",148+ default_identifier="Tensor-likes",
149- identifier=None,149+ identifier=None,
150- extra=extra,150+ extra=extra,
151- abs_diff=abs_diff.max().item(),151+ abs_diff=abs_diff.max().item(),
152- abs_diff_idx=None,152+ abs_diff_idx=None,
153- atol=atol,153+ atol=atol,
154- rel_diff=rel_diff.max().item(),154+ rel_diff=rel_diff.max().item(),
155- rel_diff_idx=None,155+ rel_diff_idx=None,
156- rtol=rtol,156+ rtol=rtol,
157- )157+ )
158- print(f"Kernel Name: {self.kernel_name}\n{msg}", flush=True)158+ print(f"Kernel Name: {self.kernel_name}\n{msg}", flush=True)
159- has_acc_error = True159+ has_acc_error = True
160- 160+ 
161- del abs_diff161+ del abs_diff
162- del rel_diff162+ del rel_diff
163- del matches163+ del matches
164- del expected164+ del expected
165- 165+
166- if anir_config.fx_subgraph_dump_path:166+ if anir_config.fx_subgraph_dump_path:
167- data = args167+ data = args
168- if has_acc_error:168+ if has_acc_error:
169- data_dump_path = self.fx_subgraph_dump('acc_failed')169+ data_dump_path = self.fx_subgraph_dump('acc_failed')
170- self.data_dump_fake(*data, dump_path=data_dump_path)170+ self.data_dump_fake(*data, dump_path=data_dump_path)
171- del fx_inputs171+ del fx_inputs
172- torch.npu.synchronize()172+ torch.npu.synchronize()
173- self.launchers = [self.launchers[0]]173+ self.launchers = [self.launchers[0]]
174- self.is_fallback_kernels = [self.is_fallback_kernels[0]]174+ self.is_fallback_kernels = [self.is_fallback_kernels[0]]
175- 175+
176- return output176+ return output
177- 177+ 
178- def compile(self, *args, **kwargs):178+ def compile(self, *args, **kwargs):
179- raise NotImplementedError179+ raise NotImplementedError
180- 180+ 
181- def ensure_runtime_ready(self, *args, **kwargs) -> None:181+ def ensure_runtime_ready(self, *args, **kwargs) -> None:
182- if not self.autotuned:182+ if not self.autotuned:
183- if self.autotune:183+ if self.autotune:
184- self.register_fx_fallback(self.kernel_meta)184+ self.register_fx_fallback(self.kernel_meta)
185- self.autotune_to_one_config(*args, **kwargs)185+ self.autotune_to_one_config(*args, **kwargs)
186- self.autotuned = True186+ self.autotuned = True
187- 187+ 
188- def prepare_runtime_args(188+ def prepare_runtime_args(
189- self,189+ self,
190- args_list: List[Any],190+ args_list: List[Any],
191- ) -> List[Any]:191+ ) -> List[Any]:
192- args_new = ()192+ args_new = ()
193- for arg in args_list:193+ for arg in args_list:
194- if not torch.is_tensor(arg):194+ if not torch.is_tensor(arg):
195- args_new = args_new + (arg,)195+ args_new = args_new + (arg,)
196- continue196+ continue
197- args_new = args_new + (arg, arg, 0) + arg.size() + arg.stride()197+ args_new = args_new + (arg, arg, 0) + arg.size() + arg.stride()
198- args_list = list(args_new)198+ args_list = list(args_new)
199- return args_list199+ return args_list
200- 200+ 
201- def autotune_to_one_config(self, *args, **kwargs):201+ def autotune_to_one_config(self, *args, **kwargs):
202- return None202+ return None
203- 203+ 
204- def _normalize_contiguous_args(204+ def _normalize_contiguous_args(
205- self, args: List[Any]205+ self, args: List[Any]
206- ) -> Tuple[List[Any], Optional[List[torch.Tensor]]]:206+ ) -> Tuple[List[Any], Optional[List[torch.Tensor]]]:
207- meta = self.kernel_meta or {}207+ meta = self.kernel_meta or {}
208- num_outputs = int(meta.get("num_outputs") or 0)208+ num_outputs = int(meta.get("num_outputs") or 0)
209- num_call_functions = int(meta.get("num_call_functions") or 0)209+ num_call_functions = int(meta.get("num_call_functions") or 0)
210- 210+ 
211- if num_outputs < 0 or len(args) < num_outputs:211+ if num_outputs < 0 or len(args) < num_outputs:
212- return args, None212+ return args, None
213- 213+ 
214- num_inputs = len(args) - num_outputs214+ num_inputs = len(args) - num_outputs
215- 215+ 
216- if self._nc_input_indices is None:216+ if self._nc_input_indices is None:
217- self._nc_input_indices = []217+ self._nc_input_indices = []
218- if num_call_functions > 0:218+ if num_call_functions > 0:
219- input_slice = args[:num_inputs] if num_outputs > 0 else args219+ input_slice = args[:num_inputs] if num_outputs > 0 else args
220- for idx, arg in enumerate(input_slice):220+ for idx, arg in enumerate(input_slice):
221- if not isinstance(arg, torch.Tensor) or arg.is_contiguous():221+ if not isinstance(arg, torch.Tensor) or arg.is_contiguous():
222- continue222+ continue
223- args[idx] = arg.contiguous()223+ args[idx] = arg.contiguous()
224- self._nc_input_indices.append(idx)224+ self._nc_input_indices.append(idx)
225- else:225+ else:
226- for idx in self._nc_input_indices:226+ for idx in self._nc_input_indices:
227- if idx < len(args) and isinstance(args[idx], torch.Tensor):227+ if idx < len(args) and isinstance(args[idx], torch.Tensor):
228- args[idx] = args[idx].contiguous()228+ args[idx] = args[idx].contiguous()
229- 229+ 
230- if num_outputs == 0:230+ if num_outputs == 0:
231- return args, None231+ return args, None
232- 232+ 
233- original_outputs: Optional[List[torch.Tensor]] = None233+ original_outputs: Optional[List[torch.Tensor]] = None
234- if self._nc_output_indices is None:234+ if self._nc_output_indices is None:
235- self._nc_output_indices = []235+ self._nc_output_indices = []
236- original_outputs = []236+ original_outputs = []
237- for j, out in enumerate(args[num_inputs:]):237+ for j, out in enumerate(args[num_inputs:]):
238- if not isinstance(out, torch.Tensor) or out.is_contiguous():238+ if not isinstance(out, torch.Tensor) or out.is_contiguous():
239- continue239+ continue
240- tmp = torch.empty(out.shape, dtype=out.dtype, device=out.device)240+ tmp = torch.empty(out.shape, dtype=out.dtype, device=out.device)
241- out_idx = num_inputs + j241+ out_idx = num_inputs + j
242- original_outputs.append(out)242+ original_outputs.append(out)
243- args[out_idx] = tmp243+ args[out_idx] = tmp
244- self._nc_output_indices.append(out_idx)244+ self._nc_output_indices.append(out_idx)
245- else:245+ else:
246- original_outputs = []246+ original_outputs = []
247- for out_idx in self._nc_output_indices:247+ for out_idx in self._nc_output_indices:
248- tmp = torch.empty(248+ tmp = torch.empty(
249- args[out_idx].shape, dtype=args[out_idx].dtype, device=args[out_idx].device249+ args[out_idx].shape, dtype=args[out_idx].dtype, device=args[out_idx].device
250- )250+ )
251- original_outputs.append(args[out_idx])251+ original_outputs.append(args[out_idx])
252- args[out_idx] = tmp252+ args[out_idx] = tmp
253- 253+ 
254- return args, original_outputs254+ return args, original_outputs
255- 255+ 
256- def _copy_back_non_contiguous_outputs(256+ def _copy_back_non_contiguous_outputs(
257- self,257+ self,
258- args_list: List[Any],258+ args_list: List[Any],
259- original_outputs: Optional[List[torch.Tensor]],259+ original_outputs: Optional[List[torch.Tensor]],
260- ) -> None:260+ ) -> None:
261- if original_outputs and self._nc_output_indices:261+ if original_outputs and self._nc_output_indices:
262- for orig, idx in zip(original_outputs, self._nc_output_indices):262+ for orig, idx in zip(original_outputs, self._nc_output_indices):
263- orig.copy_(args_list[idx])263+ orig.copy_(args_list[idx])
264- 264+ 
265- 265+ 
266- def _fx_graph_call_factory(self, module: torch.nn.Module, num_outputs: int) -> Callable[..., Any]:266+ def _fx_graph_call_factory(self, module: torch.nn.Module, num_outputs: int) -> Callable[..., Any]:
267- def module_call(*args, **kwargs):267+ def module_call(*args, **kwargs):
268- actual_args = args[:-num_outputs] if num_outputs > 0 else args268+ actual_args = args[:-num_outputs] if num_outputs > 0 else args
269- actual_outputs = module.forward(*actual_args)269+ actual_outputs = module.forward(*actual_args)
270- for out1, out2 in zip(actual_outputs, args[-num_outputs:]):270+ for out1, out2 in zip(actual_outputs, args[-num_outputs:]):
271- if isinstance(out1, torch.Tensor) and not out1.is_contiguous():271+ if isinstance(out1, torch.Tensor) and not out1.is_contiguous():
272- out1 = out1.contiguous()272+ out1 = out1.contiguous()
273- out2.data = out1.data273+ out2.data = out1.data
274- return module_call274+ return module_call
275- 275+ 
276- def _load_traced_graph_model(self) -> torch.nn.Module:276+ def _load_traced_graph_model(self) -> torch.nn.Module:
277- meta = self.kernel_meta or {}277+ meta = self.kernel_meta or {}
278- traced_graph_hash = meta.get("traced_graph_hash")278+ traced_graph_hash = meta.get("traced_graph_hash")
279- if not traced_graph_hash:279+ if not traced_graph_hash:
280- raise RuntimeError("traced_graph_hash missing, cannot build FX fallback")280+ raise RuntimeError("traced_graph_hash missing, cannot build FX fallback")
281- 281+ 
282- traced_graph_cache = meta.get("traced_graph_cache")282+ traced_graph_cache = meta.get("traced_graph_cache")
283- if traced_graph_cache is None:283+ if traced_graph_cache is None:
284- raise RuntimeError("traced_graph_cache missing, cannot locate traced graph dump")284+ raise RuntimeError("traced_graph_cache missing, cannot locate traced graph dump")
285- 285+ 
286- base_cache = os.getenv("TORCHINDUCTOR_CACHE_DIR", "")286+ base_cache = os.getenv("TORCHINDUCTOR_CACHE_DIR", "")
287- dump_root = os.path.join(base_cache, traced_graph_cache)287+ dump_root = os.path.join(base_cache, traced_graph_cache)
288- dump_path = os.path.join(dump_root, str(self.device_index), traced_graph_hash)288+ dump_path = os.path.join(dump_root, str(self.device_index), traced_graph_hash)
289- 289+ 
290- sys.path.append(dump_path)290+ sys.path.append(dump_path)
291- try:291+ try:
292- module = importlib.import_module(traced_graph_hash)292+ module = importlib.import_module(traced_graph_hash)
293- finally:293+ finally:
294- sys.path.remove(dump_path)294+ sys.path.remove(dump_path)
295- 295+ 
296- Model = getattr(module, traced_graph_hash, None)296+ Model = getattr(module, traced_graph_hash, None)
297- if Model is None:297+ if Model is None:
298- raise RuntimeError("Cannot find valid graph module class in traced graph dump")298+ raise RuntimeError("Cannot find valid graph module class in traced graph dump")
299- 299+ 
300- return Model()300+ return Model()
301- 301+ 
302- def register_fx_fallback(self, kernel_meta) -> None:302+ def register_fx_fallback(self, kernel_meta) -> None:
303- model = self._load_traced_graph_model()303+ model = self._load_traced_graph_model()
304- num_outputs = kernel_meta.get("num_outputs", 0)304+ num_outputs = kernel_meta.get("num_outputs", 0)
305- module_call = self._fx_graph_call_factory(model, num_outputs)305+ module_call = self._fx_graph_call_factory(model, num_outputs)
306- self.register_launcher(306+ self.register_launcher(
307- module_call,307+ module_call,
308- kernel_path=self.kernel_name + "_fx_fallback",308+ kernel_path=self.kernel_name + "_fx_fallback",
309- is_fallback_kernel=True,309+ is_fallback_kernel=True,
310- )310+ )
311- 311+ 
312- def run(self, *args, **kwargs):312+ def run(self, *args, **kwargs):
313- args_list = list(args)313+ args_list = list(args)
314- args_list, original_outputs = self._normalize_contiguous_args(args_list)314+ args_list, original_outputs = self._normalize_contiguous_args(args_list)
315- self.ensure_runtime_ready(*args_list, **kwargs)315+ self.ensure_runtime_ready(*args_list, **kwargs)
316- 316+ 
317- launcher_idx = self.get_primary_launcher_index()317+ launcher_idx = self.get_primary_launcher_index()
318- launcher = self.launchers[launcher_idx]318+ launcher = self.launchers[launcher_idx]
319- is_fallback_kernel = self.is_fallback_kernels[launcher_idx]319+ is_fallback_kernel = self.is_fallback_kernels[launcher_idx]
320- if anir_config.online_acc_comp and not is_fallback_kernel:320+ if anir_config.online_acc_comp and not is_fallback_kernel:
321- output = self.acc_compare_and_dump(*args_list, **kwargs)321+ output = self.acc_compare_and_dump(*args_list, **kwargs)
322- self._copy_back_non_contiguous_outputs(args_list, original_outputs)322+ self._copy_back_non_contiguous_outputs(args_list, original_outputs)
323- return output323+ return output
324- if self.dynamic and not is_fallback_kernel:324+ if self.dynamic and not is_fallback_kernel:
325- args_list = self.prepare_runtime_args(325+ args_list = self.prepare_runtime_args(
326- args_list,326+ args_list,
327- )327+ )
328- ret = launcher(*tuple(args_list), **kwargs)328+ ret = launcher(*tuple(args_list), **kwargs)
329- self._copy_back_non_contiguous_outputs(args_list, original_outputs)329+ self._copy_back_non_contiguous_outputs(args_list, original_outputs)
330 return ret330 return ret
Mtorch_npu/_inductor/codegen/common.py+11-11
@@ -1,11 +1,11 @@
1-import torch._inductor.codegen.common1+import torch._inductor.codegen.common
2-from torch._inductor.codegen.common import device_op_overrides_dict2+from torch._inductor.codegen.common import device_op_overrides_dict
3- 3+ 
4- 4+ 
5-def register_device_op_overrides_npu():5+def register_device_op_overrides_npu():
6- if not device_op_overrides_dict:6+ if not device_op_overrides_dict:
7- # remove cuda, xpu device_op_override, add npu device_op_override7+ # remove cuda, xpu device_op_override, add npu device_op_override
8- from torch._inductor.codegen import cpu_device_op_overrides, mps_device_op_overrides # noqa: F4018+ from torch._inductor.codegen import cpu_device_op_overrides, mps_device_op_overrides # noqa: F401
9- from .npu import device_op_overrides # noqa: F4019+ from .npu import device_op_overrides # noqa: F401
10- elif "npu" not in device_op_overrides_dict:10+ elif "npu" not in device_op_overrides_dict:
O
OopenLiBingCI5月17日

此条代码评论区间+6至+10

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

likedislike
11- from .npu import device_op_overrides11+ from .npu import device_op_overrides
O
OopenLiBingCI5月17日

此条代码评论区间+7至+11

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

likedislike
Mtorch_npu/_inductor/fx_passes/pattern_match/npu_fusion_attention_graph.py+0-1
@@ -254,4 +254,3 @@ def register_fa_pass():
254 register_replacement(254 register_replacement(
255 **register_replacement_kwargs,255 **register_replacement_kwargs,
256 )256 )
257- 
Mtorch_npu/csrc/afd/CMakeLists.txt+5-5
@@ -1,6 +1,6 @@
1-FILE(GLOB _AFD_SRCS *.cpp)1+FILE(GLOB _AFD_SRCS *.cpp)
2- 2+ 
3-LIST(APPEND AFD_SRCS ${_AFD_SRCS})3+LIST(APPEND AFD_SRCS ${_AFD_SRCS})
4- 4+ 
5-# Pass to parent5+# Pass to parent
6set(AFD_SRCS ${AFD_SRCS} PARENT_SCOPE)6set(AFD_SRCS ${AFD_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/afd/Init.cpp+44-44
@@ -1,44 +1,44 @@
1-#include "torch_npu/csrc/afd/Init.h"1+#include "torch_npu/csrc/afd/Init.h"
2-#include <torch/csrc/python_headers.h>2+#include <torch/csrc/python_headers.h>
3-#include <torch/csrc/Exceptions.h>3+#include <torch/csrc/Exceptions.h>
4-#include <torch/csrc/utils/object_ptr.h>4+#include <torch/csrc/utils/object_ptr.h>
5-#include <torch/csrc/utils/pybind.h>5+#include <torch/csrc/utils/pybind.h>
6-#include "torch_npu/csrc/afd/ScheduleContext.h"6+#include "torch_npu/csrc/afd/ScheduleContext.h"
7- 7+ 
8-namespace torch_npu {8+namespace torch_npu {
9-namespace afd {9+namespace afd {
10- 10+ 
11-PyObject *afd_init(PyObject * _unused, PyObject * noargs)11+PyObject *afd_init(PyObject * _unused, PyObject * noargs)
12-{12+{
13- auto torch_npu_C_module = THPObjectPtr(PyImport_ImportModule("torch_npu._C"));13+ auto torch_npu_C_module = THPObjectPtr(PyImport_ImportModule("torch_npu._C"));
14- if (!torch_npu_C_module) {14+ if (!torch_npu_C_module) {
15- throw python_error();15+ throw python_error();
16- }16+ }
17- auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();17+ auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();
18- 18+ 
19- auto m = torch_npu_C_m.def_submodule("_afd", "Attention-FFN Disaggregation");19+ auto m = torch_npu_C_m.def_submodule("_afd", "Attention-FFN Disaggregation");
20- auto module = py::handle(m).cast<py::module>();20+ auto module = py::handle(m).cast<py::module>();
21- 21+ 
22- py::class_<ScheduleContextHolder>(module, "ScheduleContextHolder")22+ py::class_<ScheduleContextHolder>(module, "ScheduleContextHolder")
23- .def(py::init<int32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint64_t, uint64_t,23+ .def(py::init<int32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint64_t, uint64_t,
24- uint64_t, uint64_t>())24+ uint64_t, uint64_t>())
25- .def("init", &ScheduleContextHolder::Init)25+ .def("init", &ScheduleContextHolder::Init)
26- .def("get_context_tensor", &ScheduleContextHolder::GetContextTensor)26+ .def("get_context_tensor", &ScheduleContextHolder::GetContextTensor)
27- .def("stop_schedule", &ScheduleContextHolder::StopSchedule)27+ .def("stop_schedule", &ScheduleContextHolder::StopSchedule)
28- .def("get_schedule_context_info", &ScheduleContextHolder::GetScheduleContextInfo);28+ .def("get_schedule_context_info", &ScheduleContextHolder::GetScheduleContextInfo);
29- Py_RETURN_TRUE;29+ Py_RETURN_TRUE;
30-}30+}
31- 31+ 
32-// methods on torch._C32+// methods on torch._C
33-PyMethodDef methods[] = {33+PyMethodDef methods[] = {
34- {"_afd_init", afd_init, METH_NOARGS, nullptr},34+ {"_afd_init", afd_init, METH_NOARGS, nullptr},
35- {nullptr, nullptr, 0, nullptr}35+ {nullptr, nullptr, 0, nullptr}
36-};36+};
37- 37+ 
38-PyMethodDef *python_functions()38+PyMethodDef *python_functions()
39-{39+{
40- return methods;40+ return methods;
41-}41+}
42- 42+ 
43-} // namespace afd43+} // namespace afd
44-} // namespace torch_npu44+} // namespace torch_npu
Mtorch_npu/csrc/afd/Init.h+11-11
@@ -1,12 +1,12 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <torch/csrc/python_headers.h>3+#include <torch/csrc/python_headers.h>
4-#include "torch_npu/csrc/core/npu/NPUMacros.h"4+#include "torch_npu/csrc/core/npu/NPUMacros.h"
5- 5+ 
6-namespace torch_npu {6+namespace torch_npu {
7-namespace afd {7+namespace afd {
8- 8+ 
9-TORCH_NPU_API PyMethodDef *python_functions();9+TORCH_NPU_API PyMethodDef *python_functions();
10- 10+ 
11-} // namespace afd11+} // namespace afd
12} // namespace torch_npu12} // namespace torch_npu
Mtorch_npu/csrc/aten/CMakeLists.txt+21-21
@@ -1,21 +1,21 @@
1-FILE(GLOB _ATEN_SRCS1+FILE(GLOB _ATEN_SRCS
2- *.cpp2+ *.cpp
3- common/*.cpp3+ common/*.cpp
4- mirror/*.cpp4+ mirror/*.cpp
5- ops/*.cpp5+ ops/*.cpp
6- ops/op_api/*.cpp)6+ ops/op_api/*.cpp)
7- 7+ 
8-FILE(GLOB _EXCLUDE8+FILE(GLOB _EXCLUDE
9- VariableTypeEverything.cpp9+ VariableTypeEverything.cpp
10- ADInplaceOrViewTypeEverything.cpp10+ ADInplaceOrViewTypeEverything.cpp
11- python_functionsEverything.cpp11+ python_functionsEverything.cpp
12- RegisterFunctionalizationEverything.cpp)12+ RegisterFunctionalizationEverything.cpp)
13- 13+ 
14-FOREACH(ITEM ${_EXCLUDE})14+FOREACH(ITEM ${_EXCLUDE})
15- LIST(REMOVE_ITEM _ATEN_SRCS ${ITEM})15+ LIST(REMOVE_ITEM _ATEN_SRCS ${ITEM})
16-ENDFOREACH()16+ENDFOREACH()
17- 17+ 
18-LIST(APPEND ATEN_SRCS ${_ATEN_SRCS})18+LIST(APPEND ATEN_SRCS ${_ATEN_SRCS})
19- 19+ 
20-# Pass to parent20+# Pass to parent
21-set(ATEN_SRCS ${ATEN_SRCS} PARENT_SCOPE)21+set(ATEN_SRCS ${ATEN_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/aten/PinMemory.cpp+50-50
@@ -1,50 +1,50 @@
1-#include <ATen/ATen.h>1+#include <ATen/ATen.h>
2-#include <ATen/NativeFunctions.h>2+#include <ATen/NativeFunctions.h>
3-#include <ATen/TensorUtils.h>3+#include <ATen/TensorUtils.h>
4-#include <c10/core/Storage.h>4+#include <c10/core/Storage.h>
5-#include "torch_npu/csrc/core/npu/NPUFunctions.h"5+#include "torch_npu/csrc/core/npu/NPUFunctions.h"
6-#include "torch_npu/csrc/core/npu/NPUException.h"6+#include "torch_npu/csrc/core/npu/NPUException.h"
7-#include "torch_npu/csrc/core/npu/CachingHostAllocator.h"7+#include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
8- 8+ 
9-#define TORCH_ASSERT_ONLY_METHOD_OPERATORS9+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
10-#include <ATen/core/Tensor.h>10+#include <ATen/core/Tensor.h>
11-#include <ATen/core/dispatch/DispatchKeyExtractor.h>11+#include <ATen/core/dispatch/DispatchKeyExtractor.h>
12-#include <torch/library.h>12+#include <torch/library.h>
13- 13+ 
14-#ifndef AT_PER_OPERATOR_HEADERS14+#ifndef AT_PER_OPERATOR_HEADERS
15-#include <ATen/Operators.h>15+#include <ATen/Operators.h>
16-#else16+#else
17-#include <ATen/ops/is_pinned_ops.h>17+#include <ATen/ops/is_pinned_ops.h>
18-#include <ATen/ops/_pin_memory_ops.h>18+#include <ATen/ops/_pin_memory_ops.h>
19- 19+ 
20-${ops_headers}20+${ops_headers}
21-#endif21+#endif
22- 22+ 
23-namespace at_npu {23+namespace at_npu {
24-namespace native {24+namespace native {
25- 25+ 
26-bool is_pinned(const at::Tensor& self, c10::optional<at::Device> device)26+bool is_pinned(const at::Tensor& self, c10::optional<at::Device> device)
27-{27+{
28- // Only CPU tensors can be pinned28+ // Only CPU tensors can be pinned
29- if (!self.is_cpu()) {29+ if (!self.is_cpu()) {
30- return false;30+ return false;
31- }31+ }
32- c10::DispatchKeySet _dk = c10::DispatchKeySet(c10::computeDispatchKey(c10::nullopt, self.layout(), device.value_or(c10::DeviceType::PrivateUse1)));32+ c10::DispatchKeySet _dk = c10::DispatchKeySet(c10::computeDispatchKey(c10::nullopt, self.layout(), device.value_or(c10::DeviceType::PrivateUse1)));
33- return at::_ops::is_pinned::redispatch(_dk, self, device);33+ return at::_ops::is_pinned::redispatch(_dk, self, device);
34-}34+}
35- 35+ 
36-at::Tensor _pin_memory(const at::Tensor& self, c10::optional<at::Device> device)36+at::Tensor _pin_memory(const at::Tensor& self, c10::optional<at::Device> device)
37-{37+{
38- TORCH_CHECK(self.device().is_cpu(), "cannot pin '", self.toString(), "' only dense CPU tensors can be pinned", PTA_ERROR(ErrCode::TYPE));38+ TORCH_CHECK(self.device().is_cpu(), "cannot pin '", self.toString(), "' only dense CPU tensors can be pinned", PTA_ERROR(ErrCode::TYPE));
39- c10::DispatchKeySet _dk = c10::DispatchKeySet(c10::computeDispatchKey(c10::nullopt, self.layout(), device.value_or(c10::DeviceType::PrivateUse1)));39+ c10::DispatchKeySet _dk = c10::DispatchKeySet(c10::computeDispatchKey(c10::nullopt, self.layout(), device.value_or(c10::DeviceType::PrivateUse1)));
40- return at::_ops::_pin_memory::redispatch(_dk, self, device);40+ return at::_ops::_pin_memory::redispatch(_dk, self, device);
41-}41+}
42- 42+ 
43-TORCH_LIBRARY_IMPL(aten, BackendSelect, m)43+TORCH_LIBRARY_IMPL(aten, BackendSelect, m)
44-{44+{
45- m.impl(TORCH_SELECTIVE_NAME("aten::is_pinned"), TORCH_FN(is_pinned));45+ m.impl(TORCH_SELECTIVE_NAME("aten::is_pinned"), TORCH_FN(is_pinned));
46- m.impl(TORCH_SELECTIVE_NAME("aten::_pin_memory"), TORCH_FN(_pin_memory));46+ m.impl(TORCH_SELECTIVE_NAME("aten::_pin_memory"), TORCH_FN(_pin_memory));
47-}47+}
48- 48+ 
49-}49+}
50-}50+}
Mtorch_npu/csrc/aten/ops/op_api/CloneKernelOpApi.cpp+0-1
@@ -45,4 +45,3 @@ at::Tensor NPUNativeOpApiFunctions::clone(const at::Tensor &src, c10::optional<c
45 45 
46} // namespace native46} // namespace native
47} // namespace at_npu47} // namespace at_npu
48- 
Mtorch_npu/csrc/core/CMakeLists.txt+5-5
@@ -1,6 +1,6 @@
1-FILE(GLOB _CORE_SRCS *.cpp npu/*.cpp npu/*/*.cpp)1+FILE(GLOB _CORE_SRCS *.cpp npu/*.cpp npu/*/*.cpp)
2- 2+ 
3-LIST(APPEND CORE_SRCS ${_CORE_SRCS})3+LIST(APPEND CORE_SRCS ${_CORE_SRCS})
4- 4+ 
5-# Pass to parent5+# Pass to parent
6set(CORE_SRCS ${CORE_SRCS} PARENT_SCOPE)6set(CORE_SRCS ${CORE_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/core/NPUStorageImpl.h+0-1
@@ -67,4 +67,3 @@ c10::intrusive_ptr<c10::StorageImpl> make_npu_storage_impl(
67 bool resizable);67 bool resizable);
68 68 
69} // namespace torch_npu69} // namespace torch_npu
70- 
Mtorch_npu/csrc/core/npu/NPUErrorCodes.h+400-400
@@ -1,400 +1,400 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <unordered_map>3+#include <unordered_map>
4- 4+ 
5-namespace c10_npu::acl {5+namespace c10_npu::acl {
6- 6+
7-class AclErrorCode {7+class AclErrorCode {
8-public:8+public:
9- std::unordered_map<int, std::string> error_code_map = {9+ std::unordered_map<int, std::string> error_code_map = {
10- {100000, "Parameter verification failed.\n\10+ {100000, "Parameter verification failed.\n\
11- Check whether the input parameter value of the interface is correct."},11+ Check whether the input parameter value of the interface is correct."},
12- {100001, "ACL uninitialized.\n\12+ {100001, "ACL uninitialized.\n\
13- (1)Check whether the acl.init interface has been invoked for initialization. Ensure that the acl.init \13+ (1)Check whether the acl.init interface has been invoked for initialization. Ensure that the acl.init \
14-interface has been invoked before other pyACL interfaces.\n\14+interface has been invoked before other pyACL interfaces.\n\
15- (2)Check whether the initialization interface of the corresponding function has been invoked, for example, \15+ (2)Check whether the initialization interface of the corresponding function has been invoked, for example, \
16-the acl.mdl.init_dump interface for initializing Dump and the acl.prof.init interface for initializing Profiling."},16+the acl.mdl.init_dump interface for initializing Dump and the acl.prof.init interface for initializing Profiling."},
17- {100002, "Repeated initialization or repeated loading.\n\17+ {100002, "Repeated initialization or repeated loading.\n\
18- Check whether the corresponding interface is repeatedly invoked for initialization or loading."},18+ Check whether the corresponding interface is repeatedly invoked for initialization or loading."},
19- {100003, "Invalid file.\n\19+ {100003, "Invalid file.\n\
20- Check whether the file exists and can be accessed."},20+ Check whether the file exists and can be accessed."},
21- {100004, "Failed to write the file.\n\21+ {100004, "Failed to write the file.\n\
22- Check whether the file path exists and whether the file has the write permission."},22+ Check whether the file path exists and whether the file has the write permission."},
23- {100005, "Invalid file size.\n\23+ {100005, "Invalid file size.\n\
24- Check whether the file size meets the interface requirements."},24+ Check whether the file size meets the interface requirements."},
25- {100006, "Failed to parse the file.\n\25+ {100006, "Failed to parse the file.\n\
26- Check whether the file content is valid."},26+ Check whether the file content is valid."},
27- {100007, "The file is missing parameters.\n\27+ {100007, "The file is missing parameters.\n\
28- Check whether the file content is complete."},28+ Check whether the file content is complete."},
29- {100008, "Invalid file parameter.\n\29+ {100008, "Invalid file parameter.\n\
30- Check whether the parameter values in the file are correct."},30+ Check whether the parameter values in the file are correct."},
31- {100009, "Invalid dump configuration.\n\31+ {100009, "Invalid dump configuration.\n\
32- Check whether the dump configuration in the acl.init interface configuration file is correct. For details, \32+ Check whether the dump configuration in the acl.init interface configuration file is correct. For details, \
33-see 'Preparing Comparison Data > Preparing Offline Model Dump Data Files' in the <Precision Comparison Tool \33+see 'Preparing Comparison Data > Preparing Offline Model Dump Data Files' in the <Precision Comparison Tool \
34-User Guide>."},34+User Guide>."},
35- {100010, "Invalid Profiling configuration.\n\35+ {100010, "Invalid Profiling configuration.\n\
36- Check whether the profiling configuration is correct."},36+ Check whether the profiling configuration is correct."},
37- {100011, "Invalid model ID.\n\37+ {100011, "Invalid model ID.\n\
38- Check whether the model ID is correct and whether the model is correctly loaded."},38+ Check whether the model ID is correct and whether the model is correctly loaded."},
39- {100012, "Failed to deserialize the model.\n\39+ {100012, "Failed to deserialize the model.\n\
40- The model may not match the current version. Convert the model again by referring to the <ATC Tool User \40+ The model may not match the current version. Convert the model again by referring to the <ATC Tool User \
41-Guide>."},41+Guide>."},
42- {100013, "Failed to parse the model.\n\42+ {100013, "Failed to parse the model.\n\
43- The model may not match the current version. Convert the model again by referring to the <ATC Tool User \43+ The model may not match the current version. Convert the model again by referring to the <ATC Tool User \
44-Guide>."},44+Guide>."},
45- {100014, "Failed to read the model.\n\45+ {100014, "Failed to read the model.\n\
46- Check whether the model file exists and can be accessed."},46+ Check whether the model file exists and can be accessed."},
47- {100015, "Invalid model size.\n\47+ {100015, "Invalid model size.\n\
48- The model file is invalid, Convert the model again by referring to the <ATC Tool User Guide>."},48+ The model file is invalid, Convert the model again by referring to the <ATC Tool User Guide>."},
49- {100016, "The model is missing parameters.\n\49+ {100016, "The model is missing parameters.\n\
50- The model may not match the current version. Convert the model again by referring to the <ATC Tool User \50+ The model may not match the current version. Convert the model again by referring to the <ATC Tool User \
51-Guide>."},51+Guide>."},
52- {100017, "The input for the model does not match.\n\52+ {100017, "The input for the model does not match.\n\
53- Check whether the model input is correct."},53+ Check whether the model input is correct."},
54- {100018, "The output of the model does not match.\n\54+ {100018, "The output of the model does not match.\n\
55- Check whether the output of the model is correct."},55+ Check whether the output of the model is correct."},
56- {100019, "Model is not-dynamic.\n\56+ {100019, "Model is not-dynamic.\n\
57- Check whether the current model supports dynamic scenarios. If not, convert the model again by referring \57+ Check whether the current model supports dynamic scenarios. If not, convert the model again by referring \
58-to the <ATC Tool User Guide>."},58+to the <ATC Tool User Guide>."},
59- {100020, "The type of a single operator does not match.\n\59+ {100020, "The type of a single operator does not match.\n\
60- Check whether the operator type is correct."},60+ Check whether the operator type is correct."},
61- {100021, "The input of a single operator does not match.\n\61+ {100021, "The input of a single operator does not match.\n\
62- Check whether the operator input is correct."},62+ Check whether the operator input is correct."},
63- {100022, "The output of a single operator does not match.\n\63+ {100022, "The output of a single operator does not match.\n\
64- Check whether the operator output is correct."},64+ Check whether the operator output is correct."},
65- {100023, "The attributes of a single operator do not match.\n\65+ {100023, "The attributes of a single operator do not match.\n\
66- Check whether the operator attributes are correct."},66+ Check whether the operator attributes are correct."},
67- {100024, "Single operator not found.\n\67+ {100024, "Single operator not found.\n\
68- Check whether the operator type is supported."},68+ Check whether the operator type is supported."},
69- {100025, "Failed to load a single operator.\n\69+ {100025, "Failed to load a single operator.\n\
70- The model may not match the current version. Convert the model again by referring to the <ATC Tool User \70+ The model may not match the current version. Convert the model again by referring to the <ATC Tool User \
71-Guide>."},71+Guide>."},
72- {100026, "Unsupported data type.\n\72+ {100026, "Unsupported data type.\n\
73- Check whether the data type exists or is currently supported."},73+ Check whether the data type exists or is currently supported."},
74- {100027, "The format does not match.\n\74+ {100027, "The format does not match.\n\
75- Check whether the format is correct."},75+ Check whether the format is correct."},
76- {100028, "When the operator interface is compiled in binary selection mode, the operator haven't registered a \76+ {100028, "When the operator interface is compiled in binary selection mode, the operator haven't registered a \
77-selector.\n\77+selector.\n\
78- Check whether the acl.op.register_compile_func interface is invoked to register the operator selector."},78+ Check whether the acl.op.register_compile_func interface is invoked to register the operator selector."},
79- {100029, "During operator compilation, the operator kernel haven't registered.\n\79+ {100029, "During operator compilation, the operator kernel haven't registered.\n\
80- Check whether the acl.op.create_kernel interface is invoked to register the operator kernel."},80+ Check whether the acl.op.create_kernel interface is invoked to register the operator kernel."},
81- {100030, "When the operator interface is compiled in binary selection mode, the operator is registered \81+ {100030, "When the operator interface is compiled in binary selection mode, the operator is registered \
82-repeatedly.\n\82+repeatedly.\n\
83- Check whether the acl.op.register_compile_func interface is repeatedly invoked to register the operator \83+ Check whether the acl.op.register_compile_func interface is repeatedly invoked to register the operator \
84-selector."},84+selector."},
85- {100031, "During operator compilation, the operator kernel is registered repeatedly.\n\85+ {100031, "During operator compilation, the operator kernel is registered repeatedly.\n\
86- Check whether the acl.op.create_kernel interface is repeatedly invoked to register the operator kernel."},86+ Check whether the acl.op.create_kernel interface is repeatedly invoked to register the operator kernel."},
87- {100032, "Invalid queue ID.\n\87+ {100032, "Invalid queue ID.\n\
88- Check whether the queue ID is correct."},88+ Check whether the queue ID is correct."},
89- {100033, "Duplicate subscription.\n\89+ {100033, "Duplicate subscription.\n\
90- Check whether the acl.rt.subscribe_report interface is invoked repeatedly for the same stream."},90+ Check whether the acl.rt.subscribe_report interface is invoked repeatedly for the same stream."},
91- {100034, "The stream is not subscribed.\n\91+ {100034, "The stream is not subscribed.\n\
92- Check whether the acl.rt.subscribe_report interface has been invoked.\n\92+ Check whether the acl.rt.subscribe_report interface has been invoked.\n\
93- [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_STREAM_NO_CB_REG' \93+ [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_STREAM_NO_CB_REG' \
94-return code."},94+return code."},
95- {100035, "The thread is not subscribed.\n\95+ {100035, "The thread is not subscribed.\n\
96- Check whether the acl.rt.subscribe_report interface has been invoked.\n\96+ Check whether the acl.rt.subscribe_report interface has been invoked.\n\
97- [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_THREAD_SUBSCRIBE' \97+ [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_THREAD_SUBSCRIBE' \
98-return code."},98+return code."},
99- {100036, "Waiting for callback time out.\n\99+ {100036, "Waiting for callback time out.\n\
100- (1)Check whether the acl.rt.launch_callback interface has been invoked to deliver the callback task.\n\100+ (1)Check whether the acl.rt.launch_callback interface has been invoked to deliver the callback task.\n\
101- (2)Check whether the timeout interval in the acl.rt.process_report interface is proper.\n\101+ (2)Check whether the timeout interval in the acl.rt.process_report interface is proper.\n\
102- (3)Check whether the callback task has been processed. If the callback task has been processed but the \102+ (3)Check whether the callback task has been processed. If the callback task has been processed but the \
103-acl.rt.process_report interface is still invoked, optimize the code logic.\n\103+acl.rt.process_report interface is still invoked, optimize the code logic.\n\
104- [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_REPORT_TIMEOUT' \104+ [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_REPORT_TIMEOUT' \
105-return code."},105+return code."},
106- {100037, "Repeated deinitialization.\n\106+ {100037, "Repeated deinitialization.\n\
107- Check whether the acl.finalize interface is repeatedly invoked for deinitialization."},107+ Check whether the acl.finalize interface is repeatedly invoked for deinitialization."},
108- {100038, "The static AIPP configuration information does not exist.\n\108+ {100038, "The static AIPP configuration information does not exist.\n\
109- When invoking the 'acl.mdl.get_first_aipp_info' interface, ensure that the index value is correct.\n\109+ When invoking the 'acl.mdl.get_first_aipp_info' interface, ensure that the index value is correct.\n\
110- [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_GE_AIPP_NOT_EXIST' \110+ [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_GE_AIPP_NOT_EXIST' \
111-return code."},111+return code."},
112- {100039, "The dynamic library path configured before running the application is the path of the \112+ {100039, "The dynamic library path configured before running the application is the path of the \
113-compilation stub, not the correct dynamic library path.\n\113+compilation stub, not the correct dynamic library path.\n\
114- Check the configuration of the dynamic library path and ensure that the dynamic library in running mode \114+ Check the configuration of the dynamic library path and ensure that the dynamic library in running mode \
115-is used."},115+is used."},
116- {100040, "Group is not set.\n\116+ {100040, "Group is not set.\n\
117- Check if the aclrtSetGroup interface has been called."},117+ Check if the aclrtSetGroup interface has been called."},
118- {100041, "No corresponding Group is created.\n\118+ {100041, "No corresponding Group is created.\n\
119- Check whether the Group ID set during the interface invocation is within the supported range. \119+ Check whether the Group ID set during the interface invocation is within the supported range. \
120-The value range of the Group ID is [0, (number of groups -1)]. You can invoke the aclrtGetGroupCount interface to \120+The value range of the Group ID is [0, (number of groups -1)]. You can invoke the aclrtGetGroupCount interface to \
121-obtain the number of groups."},121+obtain the number of groups."},
122- {100042, "A profiling data collection task exists.\n\122+ {100042, "A profiling data collection task exists.\n\
123- Check whether multiple profiling configurations are delivered to the same device.\n\123+ Check whether multiple profiling configurations are delivered to the same device.\n\
124- For details, see the Profiling pyACL API. (Performance data is collected and flushed to disks through \124+ For details, see the Profiling pyACL API. (Performance data is collected and flushed to disks through \
125-the Profiling pyACL API.) The code logic is adjusted based on the interface invoking requirements and interface \125+the Profiling pyACL API.) The code logic is adjusted based on the interface invoking requirements and interface \
126-invoking sequence in."},126+invoking sequence in."},
127- {100043, "The acl.prof.init interface is not used for analysis initialization.\n\127+ {100043, "The acl.prof.init interface is not used for analysis initialization.\n\
128- Check the interface invoking sequence for collecting and analyzing data and analyze the pyACL API by \128+ Check the interface invoking sequence for collecting and analyzing data and analyze the pyACL API by \
129-referring to.\n\129+referring to.\n\
130- See (Performance data collected and flushed to disks by analyzing the pyACL API) Description in."},130+ See (Performance data collected and flushed to disks by analyzing the pyACL API) Description in."},
131- {100044, "A task for obtaining dump data exists.\n\131+ {100044, "A task for obtaining dump data exists.\n\
132- Before invoking the 'acl.mdl.init_dump', 'acl.mdl.set_dump', 'acl.mdl.finalize_dump' interfaces to configure \132+ Before invoking the 'acl.mdl.init_dump', 'acl.mdl.set_dump', 'acl.mdl.finalize_dump' interfaces to configure \
133-dump information: Check whether the acl.init interface has been invoked to configure the dump information.\n\133+dump information: Check whether the acl.init interface has been invoked to configure the dump information.\n\
134- If yes, adjust the code logic and retain one method to configure the dump information."},134+ If yes, adjust the code logic and retain one method to configure the dump information."},
135- {100045, "The acl.mdl.init_dump interface is not used to initialize the dump.\n\135+ {100045, "The acl.mdl.init_dump interface is not used to initialize the dump.\n\
136- Check the invoking sequence of the interfaces for obtaining dump data. For details, see the description \136+ Check the invoking sequence of the interfaces for obtaining dump data. For details, see the description \
137-of the acl.mdl.init_dump interface."},137+of the acl.mdl.init_dump interface."},
138- {148046, "Subscribe to the same model repeatedly.\n\138+ {148046, "Subscribe to the same model repeatedly.\n\
139- Check the API invoking sequence. See (Performance data collected and flushed to disks by analyzing \139+ Check the API invoking sequence. See (Performance data collected and flushed to disks by analyzing \
140-the pyACL API) Description in."},140+the pyACL API) Description in."},
141- {148047, "A conflict occurs in invoking the interface for collecting performance data.\n\141+ {148047, "A conflict occurs in invoking the interface for collecting performance data.\n\
142- The interfaces for collecting profiling performance data in the two modes cannot be invoked crossly.\n\142+ The interfaces for collecting profiling performance data in the two modes cannot be invoked crossly.\n\
143- The 'acl.prof.model_subscribe', 'acl.prof.get_op_*', 'acl.prof.model_un_subscribe' interface cannot be \143+ The 'acl.prof.model_subscribe', 'acl.prof.get_op_*', 'acl.prof.model_un_subscribe' interface cannot be \
144-invoked between the 'acl.prof.init' and 'acl.prof.finalize' interfaces.\n\144+invoked between the 'acl.prof.init' and 'acl.prof.finalize' interfaces.\n\
145- The 'acl.prof.init', 'acl.prof.start', 'acl.prof.stop', and 'acl.prof.finalize' interfaces cannot be \145+ The 'acl.prof.init', 'acl.prof.start', 'acl.prof.stop', and 'acl.prof.finalize' interfaces cannot be \
146-invoked between the 'acl.prof.model_subscribe' and 'acl.prof.model_un_subscribe' interfaces."},146+invoked between the 'acl.prof.model_subscribe' and 'acl.prof.model_un_subscribe' interfaces."},
147- {148048, "Invalid operator cache information aging configuration.\n\147+ {148048, "Invalid operator cache information aging configuration.\n\
148- Check the aging configuration of the operator cache information. For details, see the configuration \148+ Check the aging configuration of the operator cache information. For details, see the configuration \
149-description and example in acl.init."},149+description and example in acl.init."},
150- {148049, "The 'ASCEND_OPP_PATH' environment variable is not set, or the value of the environment variable \150+ {148049, "The 'ASCEND_OPP_PATH' environment variable is not set, or the value of the environment variable \
151-is incorrect.\n\151+is incorrect.\n\
152- Check whether the 'ASCEND_OPP_PATH' environment variable is set and whether the value of the environment \152+ Check whether the 'ASCEND_OPP_PATH' environment variable is set and whether the value of the environment \
153-variable is the installation path of the opp software package."},153+variable is the installation path of the opp software package."},
154- {148050, "The operator does not support dynamic Shape.\n\154+ {148050, "The operator does not support dynamic Shape.\n\
155- Check whether the shape of the operator in the single-operator model file is dynamic. If yes, change the \155+ Check whether the shape of the operator in the single-operator model file is dynamic. If yes, change the \
156-shape to a fixed one.\n\156+shape to a fixed one.\n\
157- Check whether the shape of aclTensorDesc is dynamic during operator compilation. If yes, create \157+ Check whether the shape of aclTensorDesc is dynamic during operator compilation. If yes, create \
158-aclTensorDesc based on the fixed shape."},158+aclTensorDesc based on the fixed shape."},
159- {148051, "Related resources have not been released.\n\159+ {148051, "Related resources have not been released.\n\
160- This error code is returned if the related channel is not destroyed when the channel description \160+ This error code is returned if the related channel is not destroyed when the channel description \
161-information is being destroyed. Check whether the channel associated with the channel description is destroyed."},161+information is being destroyed. Check whether the channel associated with the channel description is destroyed."},
162- {148052, "Input image encoding format (such as arithmetic encoding and progressive encoding) that is not \162+ {148052, "Input image encoding format (such as arithmetic encoding and progressive encoding) that is not \
163-supported by the JPEGD function.\n\163+supported by the JPEGD function.\n\
164- When JPEGD image decoding is implemented, only Huffman encoding is supported.\n\164+ When JPEGD image decoding is implemented, only Huffman encoding is supported.\n\
165- The color space of the original image before compression is YUV. The ratio of pixel components is \n\165+ The color space of the original image before compression is YUV. The ratio of pixel components is \n\
166- 4:4:4, 4:2:2, 4:2:0, 4:0:0, or 4:4:0.\n\166+ 4:4:4, 4:2:2, 4:2:0, 4:0:0, or 4:4:0.\n\
167- Arithmetic coding, progressive JPEG, and JPEG2000 are not supported."},167+ Arithmetic coding, progressive JPEG, and JPEG2000 are not supported."},
168- {200000, "Failed to apply for memory.\n\168+ {200000, "Failed to apply for memory.\n\
169- Check the available memory in the hardware environment."},169+ Check the available memory in the hardware environment."},
170- {200001, "The interface does not support this function.\n\170+ {200001, "The interface does not support this function.\n\
171- Check whether the invoked interface is supported."},171+ Check whether the invoked interface is supported."},
172- {200002, "Invalid device.\n\172+ {200002, "Invalid device.\n\
173- Check whether the device exists.\n\173+ Check whether the device exists.\n\
174- [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_INVALID_DEVICEID' \174+ [notice] This return code will be discarded in later versions. Use the 'ACL_ERROR_RT_INVALID_DEVICEID' \
175-return code."},175+return code."},
176- {200003, "The memory address is not aligned.\n\176+ {200003, "The memory address is not aligned.\n\
177- Check whether the memory address meets the interface requirements."},177+ Check whether the memory address meets the interface requirements."},
178- {200004, "The resources do not match.\n\178+ {200004, "The resources do not match.\n\
179- Check whether the correct resources such as Stream and Context are transferred when the interface is \179+ Check whether the correct resources such as Stream and Context are transferred when the interface is \
180-invoked."},180+invoked."},
181- {200005, "Invalid resource handle.\n\181+ {200005, "Invalid resource handle.\n\
182- Check whether the transferred resources such as Stream and Context are destroyed or occupied when the \182+ Check whether the transferred resources such as Stream and Context are destroyed or occupied when the \
183-interface is invoked."},183+interface is invoked."},
184- {200006, "This feature is not supported.\n\184+ {200006, "This feature is not supported.\n\
185- Rectify the fault based on the error information in the ascend log or contact Huawei technical support. For \185+ Rectify the fault based on the error information in the ascend log or contact Huawei technical support. For \
186-details about logs, see the Log Reference."},186+details about logs, see the Log Reference."},
187- {200007, "Unsupported profiling configurations are delivered.\n\187+ {200007, "Unsupported profiling configurations are delivered.\n\
188- Check whether the profiling configuration is correct by referring to the description in \188+ Check whether the profiling configuration is correct by referring to the description in \
189-'acl.prof.create_config'."},189+'acl.prof.create_config'."},
190- {300000, "The storage limit is exceeded.\n\190+ {300000, "The storage limit is exceeded.\n\
191- Check the remaining storage space in the hardware environment."},191+ Check the remaining storage space in the hardware environment."},
192- {500000, "Unknown internal error.\n\192+ {500000, "Unknown internal error.\n\
193- Rectify the fault based on the error information in the ascend log."},193+ Rectify the fault based on the error information in the ascend log."},
194- {500001, "The internal ACL of the system is incorrect.\n\194+ {500001, "The internal ACL of the system is incorrect.\n\
195- Rectify the fault based on the error information in the ascend log."},195+ Rectify the fault based on the error information in the ascend log."},
196- {500002, "A GE error occurs in the system.\n\196+ {500002, "A GE error occurs in the system.\n\
197- Rectify the fault based on the error information in the ascend log."},197+ Rectify the fault based on the error information in the ascend log."},
198- {500003, "The internal RUNTIME of the system is incorrect.\n\198+ {500003, "The internal RUNTIME of the system is incorrect.\n\
199- Rectify the fault based on the error information in the ascend log."},199+ Rectify the fault based on the error information in the ascend log."},
200- {500004, "An internal DRV (Driver) error occurs.\n\200+ {500004, "An internal DRV (Driver) error occurs.\n\
201- Rectify the fault based on the error information in the ascend log."},201+ Rectify the fault based on the error information in the ascend log."},
202- {500005, "Profiling error.\n\202+ {500005, "Profiling error.\n\
203- Rectify the fault based on the error information in the ascend log."},203+ Rectify the fault based on the error information in the ascend log."},
204- /* following return codes is for the internal RUNTIME */204+ /* following return codes is for the internal RUNTIME */
205- {107000, "Parameter verification failed.\n\205+ {107000, "Parameter verification failed.\n\
206- Check whether the input parameters of the interface are correct."},206+ Check whether the input parameters of the interface are correct."},
207- {107001, "Invalid device ID.\n\207+ {107001, "Invalid device ID.\n\
208- Check whether the device ID is valid."},208+ Check whether the device ID is valid."},
209- /* Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! */209+ /* Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! */
210- {107002, "The context is empty.\n\210+ {107002, "The context is empty.\n\
211- Check whether acl.rt.set_context or acl.rt.set_device is called."},211+ Check whether acl.rt.set_context or acl.rt.set_device is called."},
212- {107003, "The stream is not in the current context.\n\212+ {107003, "The stream is not in the current context.\n\
213- Check whether the context where the stream is located is the same as the current context."},213+ Check whether the context where the stream is located is the same as the current context."},
214- {107004, "The model is not in the current context.\n\214+ {107004, "The model is not in the current context.\n\
215- Check whether the loaded model is consistent with the current context."},215+ Check whether the loaded model is consistent with the current context."},
216- {107005, "The stream is not in the current model.\n\216+ {107005, "The stream is not in the current model.\n\
217- Check whether the stream has been bound to the model."},217+ Check whether the stream has been bound to the model."},
218- {107006, "The event timestamp is invalid.\n\218+ {107006, "The event timestamp is invalid.\n\
219- Check whether the event is created."},219+ Check whether the event is created."},
220- {107007, "Invert the event timestamp.\n\220+ {107007, "Invert the event timestamp.\n\
221- Check whether the event is created."},221+ Check whether the event is created."},
222- {107008, "The memory address is not aligned.\n\222+ {107008, "The memory address is not aligned.\n\
223- Check whether the applied memory addresses are aligned. For details about the restrictions on the memory \223+ Check whether the applied memory addresses are aligned. For details about the restrictions on the memory \
224-application interface, see Memory Management."},224+application interface, see Memory Management."},
225- {107009, "Failed to open the file.\n\225+ {107009, "Failed to open the file.\n\
226- Check whether the file exists."},226+ Check whether the file exists."},
227- {107010, "Failed to write the file.\n\227+ {107010, "Failed to write the file.\n\
228- Check whether the file exists or has the write permission."},228+ Check whether the file exists or has the write permission."},
229- {107011, "The stream is not subscribed to or subscribed to repeatedly.\n\229+ {107011, "The stream is not subscribed to or subscribed to repeatedly.\n\
230- Check whether the current stream is subscribed to or repeatedly subscribed to."},230+ Check whether the current stream is subscribed to or repeatedly subscribed to."},
231- {107012, "The thread is not subscribed or subscribed repeatedly.\n\231+ {107012, "The thread is not subscribed or subscribed repeatedly.\n\
232- Check whether the current thread subscribes to or subscribes to the thread repeatedly."},232+ Check whether the current thread subscribes to or subscribes to the thread repeatedly."},
233- {107013, "The group is not set."},233+ {107013, "The group is not set."},
234- {107014, "The corresponding group is not created.\n\234+ {107014, "The corresponding group is not created.\n\
235- Check whether the group ID set when the interface is invoked is within the supported range. The value \235+ Check whether the group ID set when the interface is invoked is within the supported range. The value \
236-range of the group ID is [0, (Number of groups - 1)]."},236+range of the group ID is [0, (Number of groups - 1)]."},
237- {107015, "The stream corresponding to the callback is not registered with the thread.\n\237+ {107015, "The stream corresponding to the callback is not registered with the thread.\n\
238- Check whether the stream has been registered with the thread and whether the acl.rt.subscribe_report \238+ Check whether the stream has been registered with the thread and whether the acl.rt.subscribe_report \
239-interface is invoked."},239+interface is invoked."},
240- {107016, "Invalid memory type.\n\240+ {107016, "Invalid memory type.\n\
241- Check whether the memory type is valid."},241+ Check whether the memory type is valid."},
242- {107017, "Invalid resource handle.\n\242+ {107017, "Invalid resource handle.\n\
243- Check whether the input and used parameters are correct."},243+ Check whether the input and used parameters are correct."},
244- {107018, "The memory type applied for is incorrect.\n\244+ {107018, "The memory type applied for is incorrect.\n\
245- Check whether the input and used memory types are correct."},245+ Check whether the input and used memory types are correct."},
246- {107019, "Task execution timed out.\n\246+ {107019, "Task execution timed out.\n\
247- Re-execute the interface for delivering the task."},247+ Re-execute the interface for delivering the task."},
248- {207000, "This feature is not supported.\n\248+ {207000, "This feature is not supported.\n\
249- Rectify the fault based on the error information in the ascend log."},249+ Rectify the fault based on the error information in the ascend log."},
250- {207001, "Failed to apply for memory.\n\250+ {207001, "Failed to apply for memory.\n\
251- Check the remaining storage space in the hardware environment."},251+ Check the remaining storage space in the hardware environment."},
252- {207002, "Failed to release the memory.\n\252+ {207002, "Failed to release the memory.\n\
253- Rectify the fault based on the error information in the ascend log."},253+ Rectify the fault based on the error information in the ascend log."},
254- {207003, "The operation of the aicore operator overflows.\n\254+ {207003, "The operation of the aicore operator overflows.\n\
255- Check whether the corresponding aicore operator operation overflows."},255+ Check whether the corresponding aicore operator operation overflows."},
256- {207004, "The device is unavailable.\n\256+ {207004, "The device is unavailable.\n\
257- Check whether the device is running properly."},257+ Check whether the device is running properly."},
258- {207005, "Failed to apply for memory.\n\258+ {207005, "Failed to apply for memory.\n\
259- Check the remaining storage space in the hardware environment."},259+ Check the remaining storage space in the hardware environment."},
260- {207006, "You do not have the operation permission.\n\260+ {207006, "You do not have the operation permission.\n\
261- Check whether the permission of the user who runs the application is correct."},261+ Check whether the permission of the user who runs the application is correct."},
262- {207007, "Event resources are insufficient.\n\262+ {207007, "Event resources are insufficient.\n\
263- Check whether the number of events meets the requirements by referring to the description of the \263+ Check whether the number of events meets the requirements by referring to the description of the \
264-acl.rt.create_event interface."},264+acl.rt.create_event interface."},
265- {207008, "Stream resources are insufficient.\n\265+ {207008, "Stream resources are insufficient.\n\
266- Check whether the number of streams meets the requirements by referring to the description of the \266+ Check whether the number of streams meets the requirements by referring to the description of the \
267-acl.rt.create_stream interface."},267+acl.rt.create_stream interface."},
268- {207009, "Notify resources in the system are insufficient.\n\268+ {207009, "Notify resources in the system are insufficient.\n\
269- There are too many concurrent data preprocessing tasks or model inference consumes too many resources. \269+ There are too many concurrent data preprocessing tasks or model inference consumes too many resources. \
270-You are advised to reduce the number of concurrent tasks or uninstall some models."},270+You are advised to reduce the number of concurrent tasks or uninstall some models."},
271- {207010, "Insufficient model resources.\n\271+ {207010, "Insufficient model resources.\n\
272- You are advised to uninstall some models."},272+ You are advised to uninstall some models."},
273- {207011, "Runtime internal resources are insufficient.\n\273+ {207011, "Runtime internal resources are insufficient.\n\
274- Rectify the fault based on the error information in the ascend log."},274+ Rectify the fault based on the error information in the ascend log."},
275- {207012, "The number of queues exceeds the upper limit.\n\275+ {207012, "The number of queues exceeds the upper limit.\n\
276- Destroy unnecessary queues before creating new queues."},276+ Destroy unnecessary queues before creating new queues."},
277- {207013, "The queue is empty.\n\277+ {207013, "The queue is empty.\n\
278- Cannot obtain data from an empty queue. Add data to the queue and then obtain data."},278+ Cannot obtain data from an empty queue. Add data to the queue and then obtain data."},
279- {207014, "The queue is full. \n\279+ {207014, "The queue is full. \n\
280- You cannot add data to a queue that is full. Obtain data from the queue and then add data."},280+ You cannot add data to a queue that is full. Obtain data from the queue and then add data."},
281- {207015, "The queue is initialized repeatedly. \n\281+ {207015, "The queue is initialized repeatedly. \n\
282- You are advised to initialize the queue only once."},282+ You are advised to initialize the queue only once."},
283- {207018, "The memory on the device is exhausted. \n\283+ {207018, "The memory on the device is exhausted. \n\
284- Check the memory usage on the device and properly plan the memory usage based on the memory specifications \284+ Check the memory usage on the device and properly plan the memory usage based on the memory specifications \
285-on the device."},285+on the device."},
286- {507000, "An internal error occurs in the runtime module on the host. \n\286+ {507000, "An internal error occurs in the runtime module on the host. \n\
287- Rectify the fault based on the error information in the ascend log."},287+ Rectify the fault based on the error information in the ascend log."},
288- {507001, "An internal error occurs in the task scheduler module on the device. \n\288+ {507001, "An internal error occurs in the task scheduler module on the device. \n\
289- Rectify the fault based on the error information in the ascend log."},289+ Rectify the fault based on the error information in the ascend log."},
290- {507002, "The number of tasks on the stream reaches the maximum. \n\290+ {507002, "The number of tasks on the stream reaches the maximum. \n\
291- Rectify the fault based on the error information in the ascend log."},291+ Rectify the fault based on the error information in the ascend log."},
292- {507003, "The number of tasks on the stream is empty. \n\292+ {507003, "The number of tasks on the stream is empty. \n\
293- Rectify the fault based on the error information in the ascend log."},293+ Rectify the fault based on the error information in the ascend log."},
294- {507004, "Not all tasks on the stream are executed. \n\294+ {507004, "Not all tasks on the stream are executed. \n\
295- Rectify the fault based on the error information in the ascend log."},295+ Rectify the fault based on the error information in the ascend log."},
296- {507005, "Task execution on the AI CPU is complete. \n\296+ {507005, "Task execution on the AI CPU is complete. \n\
297- Rectify the fault based on the error information in the ascend log."},297+ Rectify the fault based on the error information in the ascend log."},
298- {507006, "The event is not complete. \n\298+ {507006, "The event is not complete. \n\
299- Rectify the fault based on the error information in the ascend log."},299+ Rectify the fault based on the error information in the ascend log."},
300- {507007, "Failed to release the context. \n\300+ {507007, "Failed to release the context. \n\
301- Rectify the fault based on the error information in the ascend log."},301+ Rectify the fault based on the error information in the ascend log."},
302- {507008, "Failed to obtain the SOC version. \n\302+ {507008, "Failed to obtain the SOC version. \n\
303- Rectify the fault based on the error information in the ascend log."},303+ Rectify the fault based on the error information in the ascend log."},
304- {507009, "The task type is not supported. \n\304+ {507009, "The task type is not supported. \n\
305- Rectify the fault based on the error information in the ascend log."},305+ Rectify the fault based on the error information in the ascend log."},
306- {507010, "The task scheduler loses the heartbeat. \n\306+ {507010, "The task scheduler loses the heartbeat. \n\
307- Rectify the fault based on the error information in the ascend log."},307+ Rectify the fault based on the error information in the ascend log."},
308- {507011, "Model execution failed. \n\308+ {507011, "Model execution failed. \n\
309- Rectify the fault based on the error information in the ascend log."},309+ Rectify the fault based on the error information in the ascend log."},
310- {507012, "Failed to obtain the task scheduler message. \n\310+ {507012, "Failed to obtain the task scheduler message. \n\
311- Rectify the fault based on the error information in the ascend log."},311+ Rectify the fault based on the error information in the ascend log."},
312- {507013, "System Direct Memory Access (DMA) hardware execution error. \n\312+ {507013, "System Direct Memory Access (DMA) hardware execution error. \n\
313- Rectify the fault based on the error information in the ascend log."},313+ Rectify the fault based on the error information in the ascend log."},
314- {507014, "The aicore execution times out. \n\314+ {507014, "The aicore execution times out. \n\
315- Rectify the fault based on the error information in the ascend log."},315+ Rectify the fault based on the error information in the ascend log."},
316- {507015, "The aicore execution is abnormal. \n\316+ {507015, "The aicore execution is abnormal. \n\
317- Rectify the fault based on the error information in the ascend log."},317+ Rectify the fault based on the error information in the ascend log."},
318- {507016, "An exception occurs when the aicore trap is executed. \n\318+ {507016, "An exception occurs when the aicore trap is executed. \n\
319- Rectify the fault based on the error information in the ascend log."},319+ Rectify the fault based on the error information in the ascend log."},
320- {507017, "The aicpu execution times out. \n\320+ {507017, "The aicpu execution times out. \n\
321- Rectify the fault based on the error information in the ascend log."},321+ Rectify the fault based on the error information in the ascend log."},
322- {507018, "The aicpu execution is abnormal. \n\322+ {507018, "The aicpu execution is abnormal. \n\
323- Rectify the fault based on the error information in the ascend log."},323+ Rectify the fault based on the error information in the ascend log."},
324- {507019, "The AICPU does not send a response to the task scheduler after data dump. \n\324+ {507019, "The AICPU does not send a response to the task scheduler after data dump. \n\
325- Rectify the fault based on the error information in the ascend log."},325+ Rectify the fault based on the error information in the ascend log."},
326- {507020, "The AIPPU does not send a response to the task scheduler after executing the model. \n\326+ {507020, "The AIPPU does not send a response to the task scheduler after executing the model. \n\
327- Rectify the fault based on the error information in the ascend log."},327+ Rectify the fault based on the error information in the ascend log."},
328- {507021, "The profiling function is abnormal. \n\328+ {507021, "The profiling function is abnormal. \n\
329- Rectify the fault based on the error information in the ascend log."},329+ Rectify the fault based on the error information in the ascend log."},
330- {507022, "The communication between processes is abnormal. \n\330+ {507022, "The communication between processes is abnormal. \n\
331- Rectify the fault based on the error information in the ascend log."},331+ Rectify the fault based on the error information in the ascend log."},
332- {507023, "The model exits. \n\332+ {507023, "The model exits. \n\
333- Rectify the fault based on the error information in the ascend log."},333+ Rectify the fault based on the error information in the ascend log."},
334- {507024, "The operator is being deregistered. \n\334+ {507024, "The operator is being deregistered. \n\
335- Rectify the fault based on the error information in the ascend log."},335+ Rectify the fault based on the error information in the ascend log."},
336- {507025, "The ring buffer function is not initialized. \n\336+ {507025, "The ring buffer function is not initialized. \n\
337- Rectify the fault based on the error information in the ascend log."},337+ Rectify the fault based on the error information in the ascend log."},
338- {507026, "The ring buffer has no data. \n\338+ {507026, "The ring buffer has no data. \n\
339- Rectify the fault based on the error information in the ascend log."},339+ Rectify the fault based on the error information in the ascend log."},
340- {507027, "The kernel in RUNTIME is not registered. \n\340+ {507027, "The kernel in RUNTIME is not registered. \n\
341- Rectify the fault based on the error information in the ascend log."},341+ Rectify the fault based on the error information in the ascend log."},
342- {507028, "Repeatedly register the kernel inside the RUNTIME. \n\342+ {507028, "Repeatedly register the kernel inside the RUNTIME. \n\
343- Rectify the fault based on the error information in the ascend log."},343+ Rectify the fault based on the error information in the ascend log."},
344- {507029, "The debug function failed to be registered. \n\344+ {507029, "The debug function failed to be registered. \n\
345- Rectify the fault based on the error information in the ascend log."},345+ Rectify the fault based on the error information in the ascend log."},
346- {507030, "Deregistration of the debugging function fails. \n\346+ {507030, "Deregistration of the debugging function fails. \n\
347- Rectify the fault based on the error information in the ascend log."},347+ Rectify the fault based on the error information in the ascend log."},
348- {507031, "The tag is not in the current context. \n\348+ {507031, "The tag is not in the current context. \n\
349- Rectify the fault based on the error information in the ascend log."},349+ Rectify the fault based on the error information in the ascend log."},
350- {507032, "The number of registered programs exceeds the upper limit. \n\350+ {507032, "The number of registered programs exceeds the upper limit. \n\
351- Rectify the fault based on the error information in the ascend log."},351+ Rectify the fault based on the error information in the ascend log."},
352- {507033, "Failed to start the device. \n\352+ {507033, "Failed to start the device. \n\
353- Rectify the fault based on the error information in the ascend log."},353+ Rectify the fault based on the error information in the ascend log."},
354- {507034, "Vector core execution timed out. \n\354+ {507034, "Vector core execution timed out. \n\
355- Rectify the fault based on the error information in the ascend log."},355+ Rectify the fault based on the error information in the ascend log."},
356- {507035, "The vector core execution is abnormal. \n\356+ {507035, "The vector core execution is abnormal. \n\
357- Rectify the fault based on the error information in the ascend log."},357+ Rectify the fault based on the error information in the ascend log."},
358- {507036, "An exception occurs when vector core traps are executed. \n\358+ {507036, "An exception occurs when vector core traps are executed. \n\
359- Rectify the fault based on the error information in the ascend log."},359+ Rectify the fault based on the error information in the ascend log."},
360- {507037, "An exception occurred when applying for internal resources of the Runtime. \n\360+ {507037, "An exception occurred when applying for internal resources of the Runtime. \n\
361- Rectify the fault based on the error information in the ascend log."},361+ Rectify the fault based on the error information in the ascend log."},
362- {507038, "An error occurred when modifying the die mode, can not change the die mode. \n\362+ {507038, "An error occurred when modifying the die mode, can not change the die mode. \n\
363- Rectify the fault based on the error information in the ascend log."},363+ Rectify the fault based on the error information in the ascend log."},
364- {507039, "The die cannot be specified in single-die mode. \n\364+ {507039, "The die cannot be specified in single-die mode. \n\
365- Rectify the fault based on the error information in the ascend log."},365+ Rectify the fault based on the error information in the ascend log."},
366- {507040, "The specified die ID is incorrect. \n\366+ {507040, "The specified die ID is incorrect. \n\
367- Rectify the fault based on the error information in the ascend log."},367+ Rectify the fault based on the error information in the ascend log."},
368- {507041, "The die mode is not set. \n\368+ {507041, "The die mode is not set. \n\
369- Rectify the fault based on the error information in the ascend log."},369+ Rectify the fault based on the error information in the ascend log."},
370- {507042, "The aicore trap read out-of-bounds exception. \n\370+ {507042, "The aicore trap read out-of-bounds exception. \n\
371- Rectify the fault based on the error information in the ascend log."},371+ Rectify the fault based on the error information in the ascend log."},
372- {507043, "The aicore trap write out-of-bounds exception. \n\372+ {507043, "The aicore trap write out-of-bounds exception. \n\
373- Rectify the fault based on the error information in the ascend log."},373+ Rectify the fault based on the error information in the ascend log."},
374- {507044, "Vector core trap read out-of-bounds exception. \n\374+ {507044, "Vector core trap read out-of-bounds exception. \n\
375- Rectify the fault based on the error information in the ascend log."},375+ Rectify the fault based on the error information in the ascend log."},
376- {507045, "Vector core trap write out-of-bounds exception. \n\376+ {507045, "Vector core trap write out-of-bounds exception. \n\
377- Rectify the fault based on the error information in the ascend log."},377+ Rectify the fault based on the error information in the ascend log."},
378- {507046, "In the specified timeout waiting event, all tasks in the specified stream are not completed. \n\378+ {507046, "In the specified timeout waiting event, all tasks in the specified stream are not completed. \n\
379- Rectify the fault based on the error information in the ascend log."},379+ Rectify the fault based on the error information in the ascend log."},
380- {507047, "During the specified event synchronization waiting time, the event is not executed completely. \n\380+ {507047, "During the specified event synchronization waiting time, the event is not executed completely. \n\
381- Rectify the fault based on the error information in the ascend log."},381+ Rectify the fault based on the error information in the ascend log."},
382- {507048, "The execution of the internal task times out. \n\382+ {507048, "The execution of the internal task times out. \n\
383- Rectify the fault based on the error information in the ascend log."},383+ Rectify the fault based on the error information in the ascend log."},
384- {507049, "An exception occurs during the execution of an internal task. \n\384+ {507049, "An exception occurs during the execution of an internal task. \n\
385- Rectify the fault based on the error information in the ascend log."},385+ Rectify the fault based on the error information in the ascend log."},
386- {507050, "The trap of the internal task is abnormal. \n\386+ {507050, "The trap of the internal task is abnormal. \n\
387- Rectify the fault based on the error information in the ascend log."},387+ Rectify the fault based on the error information in the ascend log."},
388- {507051, "Messages fail to be sent during data enqueuing. \n\388+ {507051, "Messages fail to be sent during data enqueuing. \n\
389- Rectify the fault based on the error information in the ascend log."},389+ Rectify the fault based on the error information in the ascend log."},
390- {507052, "Memory copy fails during data enqueuing. \n\390+ {507052, "Memory copy fails during data enqueuing. \n\
391- Rectify the fault based on the error information in the ascend log."},391+ Rectify the fault based on the error information in the ascend log."},
392- {507899, "An internal error occurs in the Driver module. \n\392+ {507899, "An internal error occurs in the Driver module. \n\
393- Rectify the fault based on the error information in the ascend log."},393+ Rectify the fault based on the error information in the ascend log."},
394- {507900, "An internal error occurs on the AI CPU module. \n\394+ {507900, "An internal error occurs on the AI CPU module. \n\
395- Rectify the fault based on the error information in the ascend log."},395+ Rectify the fault based on the error information in the ascend log."},
396- {507901, "The internal host device communication (HDC) session is disconnected. \n\396+ {507901, "The internal host device communication (HDC) session is disconnected. \n\
397- Rectify the fault based on the error information in the ascend log."},397+ Rectify the fault based on the error information in the ascend log."},
398- }; /* aclError code */398+ }; /* aclError code */
399-};399+};
400-} /* c10_npu::acl */400+} /* c10_npu::acl */
Mtorch_npu/csrc/core/npu/NpuVariables.cpp+0-1
@@ -115,4 +115,3 @@ bool IsAclnnOnly()
115 return GetSocVersion() >= SocVersion::Ascend950;115 return GetSocVersion() >= SocVersion::Ascend950;
116}116}
117} // namespace c10_npu117} // namespace c10_npu
118- 
Mtorch_npu/csrc/core/npu/NpuVariables.h+0-1
@@ -45,4 +45,3 @@ bool IsBF16Supported();
45bool IsAclnnOnly();45bool IsAclnnOnly();
46} // namespace c10_npu46} // namespace c10_npu
47#endif47#endif
48- 
Mtorch_npu/csrc/distributed/CMakeLists.txt+11-11
@@ -1,12 +1,12 @@
1-if (DEFINED BUILD_LIBTORCH)1+if (DEFINED BUILD_LIBTORCH)
2- FILE(GLOB _DIST_SRCS *.cpp)2+ FILE(GLOB _DIST_SRCS *.cpp)
3- # Exclude Python binding files when building libtorch3+ # Exclude Python binding files when building libtorch
4- list(REMOVE_ITEM _DIST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/Init.cpp")4+ list(REMOVE_ITEM _DIST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/Init.cpp")
5-else()5+else()
6- FILE(GLOB _DIST_SRCS *.cpp rpc/*.cpp symm_mem/*.cpp)6+ FILE(GLOB _DIST_SRCS *.cpp rpc/*.cpp symm_mem/*.cpp)
7-endif()7+endif()
8- 8+ 
9-LIST(APPEND DIST_SRCS ${_DIST_SRCS})9+LIST(APPEND DIST_SRCS ${_DIST_SRCS})
10- 10+ 
11-# Pass to parent11+# Pass to parent
12set(DIST_SRCS ${DIST_SRCS} PARENT_SCOPE)12set(DIST_SRCS ${DIST_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/distributed/HCCLUtils.cpp+311-311
@@ -1,311 +1,311 @@
1-#include <filesystem>1+#include <filesystem>
2-#include <fstream>2+#include <fstream>
3-#include <string>3+#include <string>
4- 4+ 
5-#include <torch/csrc/distributed/c10d/Utils.hpp>5+#include <torch/csrc/distributed/c10d/Utils.hpp>
6- 6+ 
7-#include "torch_npu/csrc/core/npu/interface/HcclInterface.h"7+#include "torch_npu/csrc/core/npu/interface/HcclInterface.h"
8-#include "torch_npu/csrc/distributed/HCCLUtils.hpp"8+#include "torch_npu/csrc/distributed/HCCLUtils.hpp"
9- 9+ 
10- 10+ 
11-namespace c10d_npu {11+namespace c10d_npu {
12-bool isFileExists(const std::string& path)12+bool isFileExists(const std::string& path)
13-{13+{
14- std::filesystem::path filePath(path);14+ std::filesystem::path filePath(path);
15- 15+ 
16- if (!filePath.is_absolute()) {16+ if (!filePath.is_absolute()) {
17- TORCH_CHECK(false, "Path is not absolute.", DIST_ERROR(ErrCode::UNAVAIL))17+ TORCH_CHECK(false, "Path is not absolute.", DIST_ERROR(ErrCode::UNAVAIL))
18- return false;18+ return false;
19- }19+ }
20- 20+ 
21- if (std::filesystem::exists(filePath) && std::filesystem::is_regular_file(filePath)) {21+ if (std::filesystem::exists(filePath) && std::filesystem::is_regular_file(filePath)) {
22- return true;22+ return true;
23- } else {23+ } else {
24- return false;24+ return false;
25- }25+ }
26-}26+}
27- 27+ 
28-bool checkFilePathReadable(const std::string& file)28+bool checkFilePathReadable(const std::string& file)
29-{29+{
30- std::filesystem::path filePath(file);30+ std::filesystem::path filePath(file);
31- 31+ 
32- if (!std::filesystem::exists(filePath)) {32+ if (!std::filesystem::exists(filePath)) {
33- return false;33+ return false;
34- }34+ }
35- 35+ 
36- if (std::filesystem::is_symlink(filePath)) {36+ if (std::filesystem::is_symlink(filePath)) {
37- return false;37+ return false;
38- }38+ }
39- 39+ 
40- if (!std::filesystem::is_regular_file(filePath)) {40+ if (!std::filesystem::is_regular_file(filePath)) {
41- return false;41+ return false;
42- }42+ }
43- 43+ 
44- std::filesystem::perms perms = std::filesystem::status(filePath).permissions();44+ std::filesystem::perms perms = std::filesystem::status(filePath).permissions();
45- if ((perms & std::filesystem::perms::owner_read) == std::filesystem::perms::owner_read) {45+ if ((perms & std::filesystem::perms::owner_read) == std::filesystem::perms::owner_read) {
46- return true;46+ return true;
47- }47+ }
48- return false;48+ return false;
49-}49+}
50- 50+ 
51-// HCCL DataType mapping51+// HCCL DataType mapping
52-std::map<at::ScalarType, HcclDataType> kScalarTypeToHcclDataType = {52+std::map<at::ScalarType, HcclDataType> kScalarTypeToHcclDataType = {
53- {at::kByte, HCCL_DATA_TYPE_UINT8},53+ {at::kByte, HCCL_DATA_TYPE_UINT8},
54- {at::kChar, HCCL_DATA_TYPE_INT8},54+ {at::kChar, HCCL_DATA_TYPE_INT8},
55- {at::kShort, HCCL_DATA_TYPE_INT16},55+ {at::kShort, HCCL_DATA_TYPE_INT16},
56- {at::kInt, HCCL_DATA_TYPE_INT32},56+ {at::kInt, HCCL_DATA_TYPE_INT32},
57- {at::kLong, HCCL_DATA_TYPE_INT64},57+ {at::kLong, HCCL_DATA_TYPE_INT64},
58- {at::kHalf, HCCL_DATA_TYPE_FP16},58+ {at::kHalf, HCCL_DATA_TYPE_FP16},
59- {at::kFloat, HCCL_DATA_TYPE_FP32},59+ {at::kFloat, HCCL_DATA_TYPE_FP32},
60- {at::ScalarType::UInt16, HCCL_DATA_TYPE_UINT16},60+ {at::ScalarType::UInt16, HCCL_DATA_TYPE_UINT16},
61- {at::ScalarType::UInt32, HCCL_DATA_TYPE_UINT32},61+ {at::ScalarType::UInt32, HCCL_DATA_TYPE_UINT32},
62- {at::ScalarType::UInt64, HCCL_DATA_TYPE_UINT64},62+ {at::ScalarType::UInt64, HCCL_DATA_TYPE_UINT64},
63- {at::kDouble, HCCL_DATA_TYPE_FP64},63+ {at::kDouble, HCCL_DATA_TYPE_FP64},
64- {at::kBool, HCCL_DATA_TYPE_UINT8},64+ {at::kBool, HCCL_DATA_TYPE_UINT8},
65- {at::kBFloat16, HCCL_DATA_TYPE_BFP16},65+ {at::kBFloat16, HCCL_DATA_TYPE_BFP16},
66- {at::ScalarType::Float8_e4m3fn, HCCL_DATA_TYPE_FP8E4M3},66+ {at::ScalarType::Float8_e4m3fn, HCCL_DATA_TYPE_FP8E4M3},
67- {at::ScalarType::Float8_e5m2, HCCL_DATA_TYPE_FP8E5M2},67+ {at::ScalarType::Float8_e5m2, HCCL_DATA_TYPE_FP8E5M2},
68-};68+};
69- 69+ 
70-std::map<HcclDataType, std::string> kHcclDataTypeToStringMap = {70+std::map<HcclDataType, std::string> kHcclDataTypeToStringMap = {
71- {HCCL_DATA_TYPE_UINT8, "at::kByte/at::kBool"},71+ {HCCL_DATA_TYPE_UINT8, "at::kByte/at::kBool"},
72- {HCCL_DATA_TYPE_INT8, "at::kChar"},72+ {HCCL_DATA_TYPE_INT8, "at::kChar"},
73- {HCCL_DATA_TYPE_INT16, "at::kShort"},73+ {HCCL_DATA_TYPE_INT16, "at::kShort"},
74- {HCCL_DATA_TYPE_INT32, "at::kInt"},74+ {HCCL_DATA_TYPE_INT32, "at::kInt"},
75- {HCCL_DATA_TYPE_UINT16, "at::ScalarType::UInt16"},75+ {HCCL_DATA_TYPE_UINT16, "at::ScalarType::UInt16"},
76- {HCCL_DATA_TYPE_UINT32, "at::ScalarType::UInt32"},76+ {HCCL_DATA_TYPE_UINT32, "at::ScalarType::UInt32"},
77- {HCCL_DATA_TYPE_UINT64, "at::ScalarType::UInt64"},77+ {HCCL_DATA_TYPE_UINT64, "at::ScalarType::UInt64"},
78- {HCCL_DATA_TYPE_INT64, "at::kLong"},78+ {HCCL_DATA_TYPE_INT64, "at::kLong"},
79- {HCCL_DATA_TYPE_FP16, "at::kHalf"},79+ {HCCL_DATA_TYPE_FP16, "at::kHalf"},
80- {HCCL_DATA_TYPE_FP32, "at::kFloat"},80+ {HCCL_DATA_TYPE_FP32, "at::kFloat"},
81- {HCCL_DATA_TYPE_FP64, "at::kDouble"},81+ {HCCL_DATA_TYPE_FP64, "at::kDouble"},
82- {HCCL_DATA_TYPE_BFP16, "at::kBFloat16"},82+ {HCCL_DATA_TYPE_BFP16, "at::kBFloat16"},
83- {HCCL_DATA_TYPE_FP8E4M3, "at::ScalarType::Float8_e4m3fn"},83+ {HCCL_DATA_TYPE_FP8E4M3, "at::ScalarType::Float8_e4m3fn"},
84- {HCCL_DATA_TYPE_FP8E5M2, "at::ScalarType::Float8_e5m2"},84+ {HCCL_DATA_TYPE_FP8E5M2, "at::ScalarType::Float8_e5m2"},
85-};85+};
86- 86+ 
87-// Helper function that gets the data type and issues error if not supported87+// Helper function that gets the data type and issues error if not supported
88-HcclDataType getHcclDataType(at::ScalarType type)88+HcclDataType getHcclDataType(at::ScalarType type)
89-{89+{
90- try {90+ try {
91- return kScalarTypeToHcclDataType.at(type);91+ return kScalarTypeToHcclDataType.at(type);
92- } catch (std::out_of_range& e) {92+ } catch (std::out_of_range& e) {
93- throw std::runtime_error("Unsupported data type for HCCL process group" + DIST_ERROR(ErrCode::NOT_SUPPORT));93+ throw std::runtime_error("Unsupported data type for HCCL process group" + DIST_ERROR(ErrCode::NOT_SUPPORT));
94- }94+ }
95-}95+}
96- 96+ 
97-std::string getHcclDataTypeSerialString(HcclDataType type)97+std::string getHcclDataTypeSerialString(HcclDataType type)
98-{98+{
99- const auto& iter = kHcclDataTypeToStringMap.find(type);99+ const auto& iter = kHcclDataTypeToStringMap.find(type);
100- if (iter != kHcclDataTypeToStringMap.cend()) {100+ if (iter != kHcclDataTypeToStringMap.cend()) {
101- return iter->second;101+ return iter->second;
102- } else {102+ } else {
103- TORCH_NPU_WARN_ONCE("Can not serialize undefined hccl data type.");103+ TORCH_NPU_WARN_ONCE("Can not serialize undefined hccl data type.");
104- return "";104+ return "";
105- }105+ }
106-}106+}
107- 107+ 
108-bool isSupportHcclCommName()108+bool isSupportHcclCommName()
109-{109+{
110- return at_npu::hccl::isHcclFeatureSupported(HcclCommConfigCapability::HCCL_COMM_CONFIG_COMM_NAME);110+ return at_npu::hccl::isHcclFeatureSupported(HcclCommConfigCapability::HCCL_COMM_CONFIG_COMM_NAME);
111-}111+}
112- 112+ 
113-HCCLComm::HCCLComm(HcclComm hcclComm) : hcclComm_(hcclComm), hcclAsyncErr_(HCCL_SUCCESS),113+HCCLComm::HCCLComm(HcclComm hcclComm) : hcclComm_(hcclComm), hcclAsyncErr_(HCCL_SUCCESS),
114- hcclCommType(0), p2pPeer(0) {}114+ hcclCommType(0), p2pPeer(0) {}
115- 115+
116-HCCLComm::~HCCLComm()116+HCCLComm::~HCCLComm()
117-{117+{
118- destroyHcclComm();118+ destroyHcclComm();
119-}119+}
120- 120+ 
121-std::shared_ptr<HCCLComm> HCCLComm::create(121+std::shared_ptr<HCCLComm> HCCLComm::create(
122- int numRanks,122+ int numRanks,
123- int rank,123+ int rank,
124- HcclRootInfo& rootInfo)124+ HcclRootInfo& rootInfo)
125-{125+{
126- auto comm = std::make_shared<HCCLComm>();126+ auto comm = std::make_shared<HCCLComm>();
127- HCCL_CHECK_ERROR(hcclCommInitRootInfo(numRanks, &rootInfo, rank, &(comm->hcclComm_)));127+ HCCL_CHECK_ERROR(hcclCommInitRootInfo(numRanks, &rootInfo, rank, &(comm->hcclComm_)));
128- c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},128+ c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},
129- c10_npu::ReleasePriority::PriorityMiddle);129+ c10_npu::ReleasePriority::PriorityMiddle);
130- return comm;130+ return comm;
131-}131+}
132- 132+ 
133-std::shared_ptr<HCCLComm> HCCLComm::create_config(133+std::shared_ptr<HCCLComm> HCCLComm::create_config(
134- int numRanks,134+ int numRanks,
135- int rank,135+ int rank,
136- HcclRootInfo& rootInfo,136+ HcclRootInfo& rootInfo,
137- HcclCommConfig* config)137+ HcclCommConfig* config)
138-{138+{
139- auto comm = std::make_shared<HCCLComm>();139+ auto comm = std::make_shared<HCCLComm>();
140- HCCL_CHECK_ERROR(hcclCommInitRootInfoConfig(numRanks, &rootInfo, rank, config, &(comm->hcclComm_)));140+ HCCL_CHECK_ERROR(hcclCommInitRootInfoConfig(numRanks, &rootInfo, rank, config, &(comm->hcclComm_)));
141- c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},141+ c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},
142- c10_npu::ReleasePriority::PriorityMiddle);142+ c10_npu::ReleasePriority::PriorityMiddle);
143- return comm;143+ return comm;
144-}144+}
145- 145+ 
146-std::shared_ptr<HCCLComm> HCCLComm::createGlobalHcclComm(146+std::shared_ptr<HCCLComm> HCCLComm::createGlobalHcclComm(
147- const char *clusterInfo,147+ const char *clusterInfo,
148- uint32_t rank,148+ uint32_t rank,
149- HcclCommConfig* config)149+ HcclCommConfig* config)
150-{150+{
151- auto comm = std::make_shared<HCCLComm>();151+ auto comm = std::make_shared<HCCLComm>();
152- if (hcclCommInitClusterInfoConfig(clusterInfo, rank, config, &(comm->hcclComm_)) != HCCL_SUCCESS) {152+ if (hcclCommInitClusterInfoConfig(clusterInfo, rank, config, &(comm->hcclComm_)) != HCCL_SUCCESS) {
153- return nullptr;153+ return nullptr;
154- }154+ }
155- c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},155+ c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {comm->destroyHcclComm();},
156- c10_npu::ReleasePriority::PriorityMiddle);156+ c10_npu::ReleasePriority::PriorityMiddle);
157- return comm;157+ return comm;
158-}158+}
159- 159+ 
160-std::shared_ptr<HCCLComm> HCCLComm::createSubHcclComm(160+std::shared_ptr<HCCLComm> HCCLComm::createSubHcclComm(
161- std::shared_ptr<HCCLComm> comm,161+ std::shared_ptr<HCCLComm> comm,
162- uint32_t rankNum,162+ uint32_t rankNum,
163- uint32_t *rankIds,163+ uint32_t *rankIds,
164- uint64_t subCommId,164+ uint64_t subCommId,
165- uint32_t subCommRankId,165+ uint32_t subCommRankId,
166- HcclCommConfig* config)166+ HcclCommConfig* config)
167-{167+{
168- auto subComm = std::make_shared<HCCLComm>();168+ auto subComm = std::make_shared<HCCLComm>();
169- if (hcclCreateSubCommConfig(&(comm->hcclComm_), rankNum, rankIds, subCommId, subCommRankId,169+ if (hcclCreateSubCommConfig(&(comm->hcclComm_), rankNum, rankIds, subCommId, subCommRankId,
170- config, &(subComm->hcclComm_)) != HCCL_SUCCESS) {170+ config, &(subComm->hcclComm_)) != HCCL_SUCCESS) {
171- return nullptr;171+ return nullptr;
172- }172+ }
173- c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {subComm->destroyHcclComm();},173+ c10_npu::NpuSysCtrl::GetInstance().RegisterReleaseFn([=]() ->void {subComm->destroyHcclComm();},
174- c10_npu::ReleasePriority::PriorityMiddle);174+ c10_npu::ReleasePriority::PriorityMiddle);
175- return subComm;175+ return subComm;
176-}176+}
177- 177+ 
178-// Move constructable178+// Move constructable
179-HCCLComm::HCCLComm(HCCLComm&& other)179+HCCLComm::HCCLComm(HCCLComm&& other)
180-{180+{
181- std::swap(hcclComm_, other.hcclComm_);181+ std::swap(hcclComm_, other.hcclComm_);
182- std::swap(hcclAsyncErr_, other.hcclAsyncErr_);182+ std::swap(hcclAsyncErr_, other.hcclAsyncErr_);
183- std::swap(hcclCommType, other.hcclCommType);183+ std::swap(hcclCommType, other.hcclCommType);
184- std::swap(p2pPeer, other.p2pPeer);184+ std::swap(p2pPeer, other.p2pPeer);
185-}185+}
186- 186+ 
187-// Move assignable187+// Move assignable
188-HCCLComm& HCCLComm::operator=(HCCLComm&& other)188+HCCLComm& HCCLComm::operator=(HCCLComm&& other)
189-{189+{
190- std::swap(hcclComm_, other.hcclComm_);190+ std::swap(hcclComm_, other.hcclComm_);
191- std::swap(hcclAsyncErr_, other.hcclAsyncErr_);191+ std::swap(hcclAsyncErr_, other.hcclAsyncErr_);
192- std::swap(hcclCommType, other.hcclCommType);192+ std::swap(hcclCommType, other.hcclCommType);
193- std::swap(p2pPeer, other.p2pPeer);193+ std::swap(p2pPeer, other.p2pPeer);
194- return *this;194+ return *this;
195-}195+}
196- 196+ 
197-void HCCLComm::destroyHcclComm()197+void HCCLComm::destroyHcclComm()
198-{198+{
199- std::unique_lock<std::mutex> lock(mutex_);199+ std::unique_lock<std::mutex> lock(mutex_);
200- if (hcclComm_) {200+ if (hcclComm_) {
201- hcclCommDestroy(hcclComm_);201+ hcclCommDestroy(hcclComm_);
202- hcclComm_ = nullptr;202+ hcclComm_ = nullptr;
203- }203+ }
204-}204+}
205- 205+ 
206-HcclResult HCCLComm::checkForHcclError()206+HcclResult HCCLComm::checkForHcclError()
207-{207+{
208- std::unique_lock<std::mutex> lock(mutex_);208+ std::unique_lock<std::mutex> lock(mutex_);
209-#ifdef ENABLE_HCCL_ERROR_CHECKING209+#ifdef ENABLE_HCCL_ERROR_CHECKING
210- if (hcclAsyncErr_ != HCCL_SUCCESS) {210+ if (hcclAsyncErr_ != HCCL_SUCCESS) {
211- return hcclAsyncErr_;211+ return hcclAsyncErr_;
212- }212+ }
213- if (hcclComm_ != nullptr) {213+ if (hcclComm_ != nullptr) {
214- HcclResult result = hcclGetCommAsyncError(hcclComm_, &hcclAsyncErr_);214+ HcclResult result = hcclGetCommAsyncError(hcclComm_, &hcclAsyncErr_);
215- if (result != HCCL_SUCCESS) {215+ if (result != HCCL_SUCCESS) {
216- std::string temp_str = std::string("Failed to get HCCL error code: ") + std::to_string(result);216+ std::string temp_str = std::string("Failed to get HCCL error code: ") + std::to_string(result);
217- const char* errmsg = temp_str.c_str();217+ const char* errmsg = temp_str.c_str();
218- ASCEND_LOGE("%s", errmsg);218+ ASCEND_LOGE("%s", errmsg);
219- LOG(ERROR) << c10::str(errmsg);219+ LOG(ERROR) << c10::str(errmsg);
220- return result; // return this error result instead of hcclAsyncErr_220+ return result; // return this error result instead of hcclAsyncErr_
221- }221+ }
222- }222+ }
223- return hcclAsyncErr_;223+ return hcclAsyncErr_;
224-#else224+#else
225- // Always return success, if error checks are disabled.225+ // Always return success, if error checks are disabled.
226- return HCCL_SUCCESS;226+ return HCCL_SUCCESS;
227-#endif227+#endif
228-}228+}
229- 229+ 
230-void DebugInfoWriter::write(const std::string &hcclTrace)230+void DebugInfoWriter::write(const std::string &hcclTrace)
231-{231+{
232- // Open a file for writing. The ios::binary flag is used to write data as232+ // Open a file for writing. The ios::binary flag is used to write data as
233- // binary.233+ // binary.
234- std::ofstream file(filename_, std::ios::binary);234+ std::ofstream file(filename_, std::ios::binary);
235- 235+ 
236- // Check if the file was opened successfully.236+ // Check if the file was opened successfully.
237- if (!file.is_open()) {237+ if (!file.is_open()) {
238- LOG(ERROR) << "Error opening file for writing HCCLPG debug info: "238+ LOG(ERROR) << "Error opening file for writing HCCLPG debug info: "
239- << filename_;239+ << filename_;
240- return;240+ return;
241- }241+ }
242- 242+ 
243- file.write(hcclTrace.data(), hcclTrace.size());243+ file.write(hcclTrace.data(), hcclTrace.size());
244- LOG(INFO) << "Finished writing HCCLPG debug info to " << filename_;244+ LOG(INFO) << "Finished writing HCCLPG debug info to " << filename_;
245-}245+}
246- 246+ 
247-DebugInfoWriter &DebugInfoWriter::getWriter(int rank)247+DebugInfoWriter &DebugInfoWriter::getWriter(int rank)
248-{248+{
249- if (writer_ == nullptr) {249+ if (writer_ == nullptr) {
250- std::string fileNamePrefix = c10d::getCvarString(250+ std::string fileNamePrefix = c10d::getCvarString(
251- {"TORCH_HCCL_DEBUG_INFO_TEMP_FILE"}, "/tmp/hccl_trace_rank_");251+ {"TORCH_HCCL_DEBUG_INFO_TEMP_FILE"}, "/tmp/hccl_trace_rank_");
252- // Using std::unique_ptr here to auto-delete the writer object252+ // Using std::unique_ptr here to auto-delete the writer object
253- // when the pointer itself is destroyed.253+ // when the pointer itself is destroyed.
254- std::unique_ptr<DebugInfoWriter> writerPtr(254+ std::unique_ptr<DebugInfoWriter> writerPtr(
255- new DebugInfoWriter(fileNamePrefix, rank));255+ new DebugInfoWriter(fileNamePrefix, rank));
256- DebugInfoWriter::registerWriter(std::move(writerPtr));256+ DebugInfoWriter::registerWriter(std::move(writerPtr));
257- }257+ }
258- return *writer_;258+ return *writer_;
259-}259+}
260- 260+ 
261-void DebugInfoWriter::registerWriter(std::unique_ptr<DebugInfoWriter> writer)261+void DebugInfoWriter::registerWriter(std::unique_ptr<DebugInfoWriter> writer)
262-{262+{
263- TORCH_CHECK_WITH(263+ TORCH_CHECK_WITH(
264- DistBackendError,264+ DistBackendError,
265- !hasWriterRegistered_.load(),265+ !hasWriterRegistered_.load(),
266- "debugInfoWriter already registered");266+ "debugInfoWriter already registered");
267- hasWriterRegistered_.store(true);267+ hasWriterRegistered_.store(true);
268- writer_ = std::move(writer);268+ writer_ = std::move(writer);
269-}269+}
270- 270+ 
271-std::unique_ptr<DebugInfoWriter> DebugInfoWriter::writer_ = nullptr;271+std::unique_ptr<DebugInfoWriter> DebugInfoWriter::writer_ = nullptr;
272-std::atomic<bool> DebugInfoWriter::hasWriterRegistered_(false);272+std::atomic<bool> DebugInfoWriter::hasWriterRegistered_(false);
273- 273+ 
274-struct HcclBufferNameKey {274+struct HcclBufferNameKey {
275- c10::DeviceIndex device_index;275+ c10::DeviceIndex device_index;
276- std::string name;276+ std::string name;
277- bool operator<(const HcclBufferNameKey& other) const277+ bool operator<(const HcclBufferNameKey& other) const
278- {278+ {
279- if (device_index != other.device_index) {279+ if (device_index != other.device_index) {
280- return device_index < other.device_index;280+ return device_index < other.device_index;
281- }281+ }
282- return name < other.name; // sort by name if device_index is the same282+ return name < other.name; // sort by name if device_index is the same
283- }283+ }
284-};284+};
285- 285+ 
286-struct HcclBufferNameStreamMap {286+struct HcclBufferNameStreamMap {
287- std::map<HcclBufferNameKey, c10_npu::NPUStream> map;287+ std::map<HcclBufferNameKey, c10_npu::NPUStream> map;
288- std::mutex mutex;288+ std::mutex mutex;
289-} g_BufferNameStreamMap = {};289+} g_BufferNameStreamMap = {};
290- 290+ 
291-c10::optional<c10_npu::NPUStream> getHcclStreamByBufferName(const std::string &name, c10::DeviceIndex device_index)291+c10::optional<c10_npu::NPUStream> getHcclStreamByBufferName(const std::string &name, c10::DeviceIndex device_index)
292-{292+{
293- std::unique_lock<std::mutex> lock(g_BufferNameStreamMap.mutex);293+ std::unique_lock<std::mutex> lock(g_BufferNameStreamMap.mutex);
294- auto &map = g_BufferNameStreamMap.map;294+ auto &map = g_BufferNameStreamMap.map;
295- auto it = map.find({device_index, name});295+ auto it = map.find({device_index, name});
296- if (it == map.end()) {296+ if (it == map.end()) {
297- return {};297+ return {};
298- }298+ }
299- return it->second;299+ return it->second;
300-}300+}
301- 301+ 
302-bool setHcclStreamByBufferName(const std::string &name, c10::DeviceIndex device_index, c10_npu::NPUStream steam)302+bool setHcclStreamByBufferName(const std::string &name, c10::DeviceIndex device_index, c10_npu::NPUStream steam)
303-{303+{
304- HcclBufferNameKey key = {device_index, name};304+ HcclBufferNameKey key = {device_index, name};
305- std::unique_lock<std::mutex> lock(g_BufferNameStreamMap.mutex);305+ std::unique_lock<std::mutex> lock(g_BufferNameStreamMap.mutex);
306- auto &map = g_BufferNameStreamMap.map;306+ auto &map = g_BufferNameStreamMap.map;
307- auto pair = map.insert({key, steam});307+ auto pair = map.insert({key, steam});
308- return pair.second;308+ return pair.second;
309-}309+}
310- 310+ 
311-} // namespace c10d_npu311+} // namespace c10d_npu
Mtorch_npu/csrc/distributed/HcclCompile.h+495-495
@@ -1,495 +1,495 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <c10/util/CallOnce.h>3+#include <c10/util/CallOnce.h>
4-#include "torch_npu/csrc/core/npu/NPUException.h"4+#include "torch_npu/csrc/core/npu/NPUException.h"
5-#include "torch_npu/csrc/core/npu/register/FunctionLoader.h"5+#include "torch_npu/csrc/core/npu/register/FunctionLoader.h"
6-#include "torch_npu/csrc/core/npu/NpuVariables.h"6+#include "torch_npu/csrc/core/npu/NpuVariables.h"
7- 7+ 
8-namespace c10d_npu {8+namespace c10d_npu {
9-#undef LOAD_FUNCTION9+#undef LOAD_FUNCTION
10-#define LOAD_FUNCTION(funcName) \10+#define LOAD_FUNCTION(funcName) \
11- REGISTER_FUNCTION(libhccl, funcName)11+ REGISTER_FUNCTION(libhccl, funcName)
12-#undef GET_FUNC12+#undef GET_FUNC
13-#define GET_FUNC(funcName) \13+#define GET_FUNC(funcName) \
14- GET_FUNCTION(libhccl, funcName)14+ GET_FUNCTION(libhccl, funcName)
15- 15+ 
16-REGISTER_LIBRARY(libhccl)16+REGISTER_LIBRARY(libhccl)
17-LOAD_FUNCTION(HcclAlltoAllV)17+LOAD_FUNCTION(HcclAlltoAllV)
18-LOAD_FUNCTION(HcclAllGatherV)18+LOAD_FUNCTION(HcclAllGatherV)
19-LOAD_FUNCTION(HcclReduceScatterV)19+LOAD_FUNCTION(HcclReduceScatterV)
20-LOAD_FUNCTION(HcclReduce)20+LOAD_FUNCTION(HcclReduce)
21-LOAD_FUNCTION(HcclGetCommAsyncError)21+LOAD_FUNCTION(HcclGetCommAsyncError)
22-LOAD_FUNCTION(HcclScatter)22+LOAD_FUNCTION(HcclScatter)
23-LOAD_FUNCTION(HcclBatchSendRecv)23+LOAD_FUNCTION(HcclBatchSendRecv)
24-LOAD_FUNCTION(HcclAlltoAll)24+LOAD_FUNCTION(HcclAlltoAll)
25-LOAD_FUNCTION(HcclCommInitRootInfoConfig)25+LOAD_FUNCTION(HcclCommInitRootInfoConfig)
26-LOAD_FUNCTION(HcclGetCommConfigCapability)26+LOAD_FUNCTION(HcclGetCommConfigCapability)
27-LOAD_FUNCTION(HcclCommInitClusterInfoConfig)27+LOAD_FUNCTION(HcclCommInitClusterInfoConfig)
28-LOAD_FUNCTION(HcclCreateSubCommConfig)28+LOAD_FUNCTION(HcclCreateSubCommConfig)
29-LOAD_FUNCTION(HcclCommWorkingDevNicSet)29+LOAD_FUNCTION(HcclCommWorkingDevNicSet)
30-LOAD_FUNCTION(HcclCommRegister)30+LOAD_FUNCTION(HcclCommRegister)
31-LOAD_FUNCTION(HcclCommDeregister)31+LOAD_FUNCTION(HcclCommDeregister)
32-LOAD_FUNCTION(HcclCommExchangeMem)32+LOAD_FUNCTION(HcclCommExchangeMem)
33-LOAD_FUNCTION(HcclGetRootInfo)33+LOAD_FUNCTION(HcclGetRootInfo)
34-LOAD_FUNCTION(HcclCommDestroy)34+LOAD_FUNCTION(HcclCommDestroy)
35-LOAD_FUNCTION(HcclSend)35+LOAD_FUNCTION(HcclSend)
36-LOAD_FUNCTION(HcclRecv)36+LOAD_FUNCTION(HcclRecv)
37-LOAD_FUNCTION(HcclAllReduce)37+LOAD_FUNCTION(HcclAllReduce)
38-LOAD_FUNCTION(HcclBroadcast)38+LOAD_FUNCTION(HcclBroadcast)
39-LOAD_FUNCTION(HcclAllGather)39+LOAD_FUNCTION(HcclAllGather)
40-LOAD_FUNCTION(HcclReduceScatter)40+LOAD_FUNCTION(HcclReduceScatter)
41-LOAD_FUNCTION(HcclCommInitAll)41+LOAD_FUNCTION(HcclCommInitAll)
42-LOAD_FUNCTION(HcclCommInitRootInfo)42+LOAD_FUNCTION(HcclCommInitRootInfo)
43- 43+ 
44-REGISTER_LIBRARY(libhcomm)44+REGISTER_LIBRARY(libhcomm)
45-REGISTER_FUNCTION(libhcomm, HcclGroupStart)45+REGISTER_FUNCTION(libhcomm, HcclGroupStart)
46-REGISTER_FUNCTION(libhcomm, HcclGroupEnd)46+REGISTER_FUNCTION(libhcomm, HcclGroupEnd)
47- 47+ 
48-extern HcclResult hcclGetRootInfo(HcclRootInfo *rootInfo)48+extern HcclResult hcclGetRootInfo(HcclRootInfo *rootInfo)
49-{49+{
50- using HcclGetRootInfoFunc = HcclResult(*)(HcclRootInfo *);50+ using HcclGetRootInfoFunc = HcclResult(*)(HcclRootInfo *);
51- static HcclGetRootInfoFunc func = nullptr;51+ static HcclGetRootInfoFunc func = nullptr;
52- if (func == nullptr) {52+ if (func == nullptr) {
53- func = (HcclGetRootInfoFunc)GET_FUNC(HcclGetRootInfo)53+ func = (HcclGetRootInfoFunc)GET_FUNC(HcclGetRootInfo)
54- }54+ }
55- TORCH_CHECK(func, "Failed to find function ", "HcclGetRootInfo", DIST_ERROR(ErrCode::NOT_FOUND));55+ TORCH_CHECK(func, "Failed to find function ", "HcclGetRootInfo", DIST_ERROR(ErrCode::NOT_FOUND));
56- auto ret = func(rootInfo);56+ auto ret = func(rootInfo);
57- return ret;57+ return ret;
58-}58+}
59- 59+ 
60-extern HcclResult hcclCommDestroy(HcclComm comm)60+extern HcclResult hcclCommDestroy(HcclComm comm)
61-{61+{
62- using HcclCommDestroyFunc = HcclResult(*)(HcclComm);62+ using HcclCommDestroyFunc = HcclResult(*)(HcclComm);
63- static HcclCommDestroyFunc func = nullptr;63+ static HcclCommDestroyFunc func = nullptr;
64- if (func == nullptr) {64+ if (func == nullptr) {
65- func = (HcclCommDestroyFunc)GET_FUNC(HcclCommDestroy)65+ func = (HcclCommDestroyFunc)GET_FUNC(HcclCommDestroy)
66- }66+ }
67- TORCH_CHECK(func, "Failed to find function ", "HcclCommDestroy", DIST_ERROR(ErrCode::NOT_FOUND));67+ TORCH_CHECK(func, "Failed to find function ", "HcclCommDestroy", DIST_ERROR(ErrCode::NOT_FOUND));
68- auto ret = func(comm);68+ auto ret = func(comm);
69- return ret;69+ return ret;
70-}70+}
71- 71+ 
72-extern HcclResult hcclSend(void *sendBuf, uint64_t count, HcclDataType dataType, uint32_t destRank,72+extern HcclResult hcclSend(void *sendBuf, uint64_t count, HcclDataType dataType, uint32_t destRank,
73- HcclComm comm, aclrtStream stream)73+ HcclComm comm, aclrtStream stream)
74-{74+{
75- using HcclSendFunc = HcclResult(*)(75+ using HcclSendFunc = HcclResult(*)(
76- void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);76+ void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);
77- static HcclSendFunc func = nullptr;77+ static HcclSendFunc func = nullptr;
78- if (func == nullptr) {78+ if (func == nullptr) {
79- func = (HcclSendFunc)GET_FUNC(HcclSend)79+ func = (HcclSendFunc)GET_FUNC(HcclSend)
80- }80+ }
81- TORCH_CHECK(func, "Failed to find function ", "HcclSend", DIST_ERROR(ErrCode::NOT_FOUND));81+ TORCH_CHECK(func, "Failed to find function ", "HcclSend", DIST_ERROR(ErrCode::NOT_FOUND));
82- auto ret = func(sendBuf, count, dataType, destRank, comm, stream);82+ auto ret = func(sendBuf, count, dataType, destRank, comm, stream);
83- return ret;83+ return ret;
84-}84+}
85- 85+ 
86-extern HcclResult hcclRecv(void *recvBuf, uint64_t count, HcclDataType dataType, uint32_t srcRank,86+extern HcclResult hcclRecv(void *recvBuf, uint64_t count, HcclDataType dataType, uint32_t srcRank,
87- HcclComm comm, aclrtStream stream)87+ HcclComm comm, aclrtStream stream)
88-{88+{
89- using HcclRecvFunc = HcclResult(*)(89+ using HcclRecvFunc = HcclResult(*)(
90- void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);90+ void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);
91- static HcclRecvFunc func = nullptr;91+ static HcclRecvFunc func = nullptr;
92- if (func == nullptr) {92+ if (func == nullptr) {
93- func = (HcclRecvFunc)GET_FUNC(HcclRecv)93+ func = (HcclRecvFunc)GET_FUNC(HcclRecv)
94- }94+ }
95- TORCH_CHECK(func, "Failed to find function ", "HcclRecv", DIST_ERROR(ErrCode::NOT_FOUND));95+ TORCH_CHECK(func, "Failed to find function ", "HcclRecv", DIST_ERROR(ErrCode::NOT_FOUND));
96- auto ret = func(recvBuf, count, dataType, srcRank, comm, stream);96+ auto ret = func(recvBuf, count, dataType, srcRank, comm, stream);
97- return ret;97+ return ret;
98-}98+}
99- 99+ 
100-extern HcclResult hcclCommInitAll(uint32_t ndev, int32_t *devices, HcclComm *comms)100+extern HcclResult hcclCommInitAll(uint32_t ndev, int32_t *devices, HcclComm *comms)
101-{101+{
102- using HcclCommInitAllFunc = HcclResult(*)(102+ using HcclCommInitAllFunc = HcclResult(*)(
103- uint32_t, int32_t *, HcclComm *);103+ uint32_t, int32_t *, HcclComm *);
104- static HcclCommInitAllFunc func = nullptr;104+ static HcclCommInitAllFunc func = nullptr;
105- if (func == nullptr) {105+ if (func == nullptr) {
106- func = (HcclCommInitAllFunc)GET_FUNC(HcclCommInitAll)106+ func = (HcclCommInitAllFunc)GET_FUNC(HcclCommInitAll)
107- }107+ }
108- TORCH_CHECK(func, "Failed to find function ", "HcclCommInitAll", DIST_ERROR(ErrCode::NOT_FOUND));108+ TORCH_CHECK(func, "Failed to find function ", "HcclCommInitAll", DIST_ERROR(ErrCode::NOT_FOUND));
109- auto ret = func(ndev, devices, comms);109+ auto ret = func(ndev, devices, comms);
110- return ret;110+ return ret;
111-}111+}
112- 112+ 
113-extern HcclResult hcclAllGather(void *sendBuf, void *recvBuf, uint64_t sendCount, HcclDataType dataType,113+extern HcclResult hcclAllGather(void *sendBuf, void *recvBuf, uint64_t sendCount, HcclDataType dataType,
114- HcclComm comm, aclrtStream stream)114+ HcclComm comm, aclrtStream stream)
115-{115+{
116- using HcclAllGatherFunc = HcclResult(*)(116+ using HcclAllGatherFunc = HcclResult(*)(
117- void *, void *, uint64_t, HcclDataType, HcclComm, aclrtStream);117+ void *, void *, uint64_t, HcclDataType, HcclComm, aclrtStream);
118- static HcclAllGatherFunc func = nullptr;118+ static HcclAllGatherFunc func = nullptr;
119- if (func == nullptr) {119+ if (func == nullptr) {
120- func = (HcclAllGatherFunc)GET_FUNC(HcclAllGather)120+ func = (HcclAllGatherFunc)GET_FUNC(HcclAllGather)
121- }121+ }
122- TORCH_CHECK(func, "Failed to find function ", "HcclAllGather", DIST_ERROR(ErrCode::NOT_FOUND));122+ TORCH_CHECK(func, "Failed to find function ", "HcclAllGather", DIST_ERROR(ErrCode::NOT_FOUND));
123- auto ret = func(sendBuf, recvBuf, sendCount, dataType, comm, stream);123+ auto ret = func(sendBuf, recvBuf, sendCount, dataType, comm, stream);
124- return ret;124+ return ret;
125-}125+}
126- 126+ 
127-extern HcclResult hcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType dataType,127+extern HcclResult hcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType dataType,
128- HcclReduceOp op, HcclComm comm, aclrtStream stream)128+ HcclReduceOp op, HcclComm comm, aclrtStream stream)
129-{129+{
130- using HcclAllReduceFunc = HcclResult(*)(130+ using HcclAllReduceFunc = HcclResult(*)(
131- void *, void *, uint64_t, HcclDataType, HcclReduceOp, HcclComm, aclrtStream);131+ void *, void *, uint64_t, HcclDataType, HcclReduceOp, HcclComm, aclrtStream);
132- static HcclAllReduceFunc func = nullptr;132+ static HcclAllReduceFunc func = nullptr;
133- if (func == nullptr) {133+ if (func == nullptr) {
134- func = (HcclAllReduceFunc)GET_FUNC(HcclAllReduce)134+ func = (HcclAllReduceFunc)GET_FUNC(HcclAllReduce)
135- }135+ }
136- TORCH_CHECK(func, "Failed to find function ", "HcclAllReduce", DIST_ERROR(ErrCode::NOT_FOUND));136+ TORCH_CHECK(func, "Failed to find function ", "HcclAllReduce", DIST_ERROR(ErrCode::NOT_FOUND));
137- auto ret = func(sendBuf, recvBuf, count, dataType, op, comm, stream);137+ auto ret = func(sendBuf, recvBuf, count, dataType, op, comm, stream);
138- return ret;138+ return ret;
139-}139+}
140- 140+ 
141-extern HcclResult hcclBroadcast(void *buf, uint64_t count, HcclDataType dataType, uint32_t root, HcclComm comm,141+extern HcclResult hcclBroadcast(void *buf, uint64_t count, HcclDataType dataType, uint32_t root, HcclComm comm,
142- aclrtStream stream)142+ aclrtStream stream)
143-{143+{
144- using HcclBroadcastFunc = HcclResult(*)(144+ using HcclBroadcastFunc = HcclResult(*)(
145- void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);145+ void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);
146- static HcclBroadcastFunc func = nullptr;146+ static HcclBroadcastFunc func = nullptr;
147- if (func == nullptr) {147+ if (func == nullptr) {
148- func = (HcclBroadcastFunc)GET_FUNC(HcclBroadcast)148+ func = (HcclBroadcastFunc)GET_FUNC(HcclBroadcast)
149- }149+ }
150- TORCH_CHECK(func, "Failed to find function ", "HcclBroadcast", DIST_ERROR(ErrCode::NOT_FOUND));150+ TORCH_CHECK(func, "Failed to find function ", "HcclBroadcast", DIST_ERROR(ErrCode::NOT_FOUND));
151- auto ret = func(buf, count, dataType, root, comm, stream);151+ auto ret = func(buf, count, dataType, root, comm, stream);
152- return ret;152+ return ret;
153-}153+}
154- 154+ 
155-extern HcclResult hcclCommInitRootInfo(uint32_t nRanks, const HcclRootInfo *rootInfo, uint32_t rank, HcclComm *comm)155+extern HcclResult hcclCommInitRootInfo(uint32_t nRanks, const HcclRootInfo *rootInfo, uint32_t rank, HcclComm *comm)
156-{156+{
157- using HcclCommInitRootInfoFunc = HcclResult(*)(157+ using HcclCommInitRootInfoFunc = HcclResult(*)(
158- uint32_t, const HcclRootInfo *, uint32_t, HcclComm *);158+ uint32_t, const HcclRootInfo *, uint32_t, HcclComm *);
159- static HcclCommInitRootInfoFunc func = nullptr;159+ static HcclCommInitRootInfoFunc func = nullptr;
160- if (func == nullptr) {160+ if (func == nullptr) {
161- func = (HcclCommInitRootInfoFunc)GET_FUNC(HcclCommInitRootInfo)161+ func = (HcclCommInitRootInfoFunc)GET_FUNC(HcclCommInitRootInfo)
162- }162+ }
163- TORCH_CHECK(func, "Failed to find function ", "HcclCommInitRootInfo", DIST_ERROR(ErrCode::NOT_FOUND));163+ TORCH_CHECK(func, "Failed to find function ", "HcclCommInitRootInfo", DIST_ERROR(ErrCode::NOT_FOUND));
164- auto ret = func(nRanks, rootInfo, rank, comm);164+ auto ret = func(nRanks, rootInfo, rank, comm);
165- return ret;165+ return ret;
166-}166+}
167- 167+ 
168-extern HcclResult hcclReduceScatter(void *sendBuf, void *recvBuf, uint64_t recvCount, HcclDataType dataType,168+extern HcclResult hcclReduceScatter(void *sendBuf, void *recvBuf, uint64_t recvCount, HcclDataType dataType,
169- HcclReduceOp op, HcclComm comm, aclrtStream stream)169+ HcclReduceOp op, HcclComm comm, aclrtStream stream)
170-{170+{
171- using HcclReduceScatterFunc = HcclResult(*)(171+ using HcclReduceScatterFunc = HcclResult(*)(
172- void *, void *, uint64_t, HcclDataType, HcclReduceOp, HcclComm, aclrtStream);172+ void *, void *, uint64_t, HcclDataType, HcclReduceOp, HcclComm, aclrtStream);
173- static HcclReduceScatterFunc func = nullptr;173+ static HcclReduceScatterFunc func = nullptr;
174- if (func == nullptr) {174+ if (func == nullptr) {
175- func = (HcclReduceScatterFunc)GET_FUNC(HcclReduceScatter);175+ func = (HcclReduceScatterFunc)GET_FUNC(HcclReduceScatter);
176- }176+ }
177- TORCH_CHECK(func, "Failed to find function ", "HcclReduceScatter", DIST_ERROR(ErrCode::NOT_FOUND));177+ TORCH_CHECK(func, "Failed to find function ", "HcclReduceScatter", DIST_ERROR(ErrCode::NOT_FOUND));
178- auto ret = func(sendBuf, recvBuf, recvCount, dataType, op, comm, stream);178+ auto ret = func(sendBuf, recvBuf, recvCount, dataType, op, comm, stream);
179- return ret;179+ return ret;
180-}180+}
181- 181+ 
182-extern HcclResult hcclAlltoAllV(const void *sendBuf, const void *sendCounts, const void *sdispls,182+extern HcclResult hcclAlltoAllV(const void *sendBuf, const void *sendCounts, const void *sdispls,
183- HcclDataType sendType, const void *recvBuf, const void *recvCounts, const void *rdispls,183+ HcclDataType sendType, const void *recvBuf, const void *recvCounts, const void *rdispls,
184- HcclDataType recvType, HcclComm comm, aclrtStream stream)184+ HcclDataType recvType, HcclComm comm, aclrtStream stream)
185-{185+{
186- using HcclAlltoAllVFunc = HcclResult(*)(186+ using HcclAlltoAllVFunc = HcclResult(*)(
187- const void *, const void *, const void *, HcclDataType,187+ const void *, const void *, const void *, HcclDataType,
188- const void *, const void *, const void *, HcclDataType,188+ const void *, const void *, const void *, HcclDataType,
189- HcclComm, aclrtStream);189+ HcclComm, aclrtStream);
190- static HcclAlltoAllVFunc func = nullptr;190+ static HcclAlltoAllVFunc func = nullptr;
191- if (func == nullptr) {191+ if (func == nullptr) {
192- func = (HcclAlltoAllVFunc)GET_FUNC(HcclAlltoAllV);192+ func = (HcclAlltoAllVFunc)GET_FUNC(HcclAlltoAllV);
193- }193+ }
194- TORCH_CHECK(func, "Failed to find function ", "HcclAlltoAllV", DIST_ERROR(ErrCode::NOT_FOUND));194+ TORCH_CHECK(func, "Failed to find function ", "HcclAlltoAllV", DIST_ERROR(ErrCode::NOT_FOUND));
195- auto ret = func(sendBuf, sendCounts, sdispls, sendType,195+ auto ret = func(sendBuf, sendCounts, sdispls, sendType,
196- recvBuf, recvCounts, rdispls, recvType, comm, stream);196+ recvBuf, recvCounts, rdispls, recvType, comm, stream);
197- return ret;197+ return ret;
198-}198+}
199- 199+ 
200-extern HcclResult hcclAllGatherV(const void *sendBuf, uint64_t sendCount,200+extern HcclResult hcclAllGatherV(const void *sendBuf, uint64_t sendCount,
201- const void *recvBuf, const void *recvCounts, const void *rdispls,201+ const void *recvBuf, const void *recvCounts, const void *rdispls,
202- HcclDataType dataType, HcclComm comm, aclrtStream stream)202+ HcclDataType dataType, HcclComm comm, aclrtStream stream)
203-{203+{
204- using HcclAllGatherVFunc = HcclResult(*)(204+ using HcclAllGatherVFunc = HcclResult(*)(
205- const void *, uint64_t,205+ const void *, uint64_t,
206- const void *, const void *, const void *,206+ const void *, const void *, const void *,
207- HcclDataType, HcclComm, aclrtStream);207+ HcclDataType, HcclComm, aclrtStream);
208- static HcclAllGatherVFunc func = nullptr;208+ static HcclAllGatherVFunc func = nullptr;
209- if (func == nullptr) {209+ if (func == nullptr) {
210- func = (HcclAllGatherVFunc)GET_FUNC(HcclAllGatherV);210+ func = (HcclAllGatherVFunc)GET_FUNC(HcclAllGatherV);
211- }211+ }
212- TORCH_CHECK(func, "Failed to find function ", "HcclAllGatherV", DIST_ERROR(ErrCode::NOT_FOUND));212+ TORCH_CHECK(func, "Failed to find function ", "HcclAllGatherV", DIST_ERROR(ErrCode::NOT_FOUND));
213- auto ret = func(sendBuf, sendCount, recvBuf, recvCounts, rdispls, dataType, comm, stream);213+ auto ret = func(sendBuf, sendCount, recvBuf, recvCounts, rdispls, dataType, comm, stream);
214- return ret;214+ return ret;
215-}215+}
216- 216+ 
217-extern HcclResult hcclReduceScatterV(const void *sendBuf, const void *sendCounts, const void *sdispls,217+extern HcclResult hcclReduceScatterV(const void *sendBuf, const void *sendCounts, const void *sdispls,
218- const void *recvBuf, uint64_t recvCount,218+ const void *recvBuf, uint64_t recvCount,
219- HcclDataType dataType, HcclReduceOp op, HcclComm comm, aclrtStream stream)219+ HcclDataType dataType, HcclReduceOp op, HcclComm comm, aclrtStream stream)
220-{220+{
221- using HcclReduceScatterVFunc = HcclResult(*)(221+ using HcclReduceScatterVFunc = HcclResult(*)(
222- const void *, const void *, const void *,222+ const void *, const void *, const void *,
223- const void *, uint64_t,223+ const void *, uint64_t,
224- HcclDataType, HcclReduceOp, HcclComm, aclrtStream);224+ HcclDataType, HcclReduceOp, HcclComm, aclrtStream);
225- static HcclReduceScatterVFunc func = nullptr;225+ static HcclReduceScatterVFunc func = nullptr;
226- if (func == nullptr) {226+ if (func == nullptr) {
227- func = (HcclReduceScatterVFunc)GET_FUNC(HcclReduceScatterV);227+ func = (HcclReduceScatterVFunc)GET_FUNC(HcclReduceScatterV);
228- }228+ }
229- TORCH_CHECK(func, "Failed to find function ", "HcclReduceScatterV", DIST_ERROR(ErrCode::NOT_FOUND));229+ TORCH_CHECK(func, "Failed to find function ", "HcclReduceScatterV", DIST_ERROR(ErrCode::NOT_FOUND));
230- auto ret = func(sendBuf, sendCounts, sdispls, recvBuf, recvCount, dataType, op, comm, stream);230+ auto ret = func(sendBuf, sendCounts, sdispls, recvBuf, recvCount, dataType, op, comm, stream);
231- return ret;231+ return ret;
232-}232+}
233- 233+ 
234-extern HcclResult hcclReduce(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType sendType,234+extern HcclResult hcclReduce(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType sendType,
235- HcclReduceOp op, uint32_t root, HcclComm comm, aclrtStream stream)235+ HcclReduceOp op, uint32_t root, HcclComm comm, aclrtStream stream)
236-{236+{
237- using HcclReduceVFunc = HcclResult(*)(237+ using HcclReduceVFunc = HcclResult(*)(
238- void *, void *, uint64_t, HcclDataType, HcclReduceOp, uint32_t, HcclComm, aclrtStream);238+ void *, void *, uint64_t, HcclDataType, HcclReduceOp, uint32_t, HcclComm, aclrtStream);
239- static HcclReduceVFunc func = nullptr;239+ static HcclReduceVFunc func = nullptr;
240- if (func == nullptr) {240+ if (func == nullptr) {
241- func = (HcclReduceVFunc)GET_FUNC(HcclReduce);241+ func = (HcclReduceVFunc)GET_FUNC(HcclReduce);
242- }242+ }
243- TORCH_CHECK(func, "Failed to find function ", "HcclReduce", DIST_ERROR(ErrCode::NOT_FOUND));243+ TORCH_CHECK(func, "Failed to find function ", "HcclReduce", DIST_ERROR(ErrCode::NOT_FOUND));
244- auto ret = func(sendBuf, recvBuf, count, sendType, op, root, comm, stream);244+ auto ret = func(sendBuf, recvBuf, count, sendType, op, root, comm, stream);
245- return ret;245+ return ret;
246-}246+}
247- 247+ 
248-HcclResult hcclGetCommAsyncError(HcclComm comm, HcclResult* asyncError)248+HcclResult hcclGetCommAsyncError(HcclComm comm, HcclResult* asyncError)
249-{249+{
250- using HcclGetCommAsyncErrorVFunc = HcclResult(*)(HcclComm, HcclResult*);250+ using HcclGetCommAsyncErrorVFunc = HcclResult(*)(HcclComm, HcclResult*);
251- static HcclGetCommAsyncErrorVFunc func = nullptr;251+ static HcclGetCommAsyncErrorVFunc func = nullptr;
252- if (func == nullptr) {252+ if (func == nullptr) {
253- func = (HcclGetCommAsyncErrorVFunc)GET_FUNC(HcclGetCommAsyncError);253+ func = (HcclGetCommAsyncErrorVFunc)GET_FUNC(HcclGetCommAsyncError);
254- }254+ }
255- TORCH_CHECK(func, "Failed to find function ", "HcclGetCommAsyncError", DIST_ERROR(ErrCode::NOT_FOUND));255+ TORCH_CHECK(func, "Failed to find function ", "HcclGetCommAsyncError", DIST_ERROR(ErrCode::NOT_FOUND));
256- auto ret = func(comm, asyncError);256+ auto ret = func(comm, asyncError);
257- return ret;257+ return ret;
258-}258+}
259- 259+ 
260-HcclResult hcclScatter(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType dataType, uint32_t root,260+HcclResult hcclScatter(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType dataType, uint32_t root,
261- HcclComm comm, aclrtStream stream)261+ HcclComm comm, aclrtStream stream)
262-{262+{
263- using HcclScatterVFunc = HcclResult(*)(void *, void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);263+ using HcclScatterVFunc = HcclResult(*)(void *, void *, uint64_t, HcclDataType, uint32_t, HcclComm, aclrtStream);
264- static HcclScatterVFunc func = nullptr;264+ static HcclScatterVFunc func = nullptr;
265- if (func == nullptr) {265+ if (func == nullptr) {
266- func = (HcclScatterVFunc)GET_FUNC(HcclScatter);266+ func = (HcclScatterVFunc)GET_FUNC(HcclScatter);
267- }267+ }
268- TORCH_CHECK(func, "Failed to find function ", "HcclScatter", DIST_ERROR(ErrCode::NOT_FOUND));268+ TORCH_CHECK(func, "Failed to find function ", "HcclScatter", DIST_ERROR(ErrCode::NOT_FOUND));
269- auto ret = func(sendBuf, recvBuf, count, dataType, root, comm, stream);269+ auto ret = func(sendBuf, recvBuf, count, dataType, root, comm, stream);
270- return ret;270+ return ret;
271-}271+}
272- 272+ 
273-HcclResult hcclBatchIsendIrecv(void* sendRecvInfo, uint32_t itemNum, HcclComm comm, aclrtStream stream)273+HcclResult hcclBatchIsendIrecv(void* sendRecvInfo, uint32_t itemNum, HcclComm comm, aclrtStream stream)
274-{274+{
275- using HcclBatchIsendIrecvVFunc = HcclResult(*)(275+ using HcclBatchIsendIrecvVFunc = HcclResult(*)(
276- void *, uint32_t, HcclComm, aclrtStream);276+ void *, uint32_t, HcclComm, aclrtStream);
277- static HcclBatchIsendIrecvVFunc func = nullptr;277+ static HcclBatchIsendIrecvVFunc func = nullptr;
278- if (func == nullptr) {278+ if (func == nullptr) {
279- func = (HcclBatchIsendIrecvVFunc)GET_FUNC(HcclBatchSendRecv);279+ func = (HcclBatchIsendIrecvVFunc)GET_FUNC(HcclBatchSendRecv);
280- }280+ }
281- TORCH_CHECK(func, "Failed to find function ", "HcclBatchSendRecv", DIST_ERROR(ErrCode::NOT_FOUND));281+ TORCH_CHECK(func, "Failed to find function ", "HcclBatchSendRecv", DIST_ERROR(ErrCode::NOT_FOUND));
282- auto ret = func(sendRecvInfo, itemNum, comm, stream);282+ auto ret = func(sendRecvInfo, itemNum, comm, stream);
283- return ret;283+ return ret;
284-}284+}
285- 285+ 
286-HcclResult hcclAlltoAll(const void *sendBuf, uint64_t sendCount, HcclDataType sendType,286+HcclResult hcclAlltoAll(const void *sendBuf, uint64_t sendCount, HcclDataType sendType,
287- const void *recvBuf, uint64_t recvCount, HcclDataType recvType,287+ const void *recvBuf, uint64_t recvCount, HcclDataType recvType,
288- HcclComm comm, aclrtStream stream)288+ HcclComm comm, aclrtStream stream)
289-{289+{
290- using HcclAlltoAllFunc = HcclResult(*)(290+ using HcclAlltoAllFunc = HcclResult(*)(
291- const void *, uint64_t, HcclDataType,291+ const void *, uint64_t, HcclDataType,
292- const void *, uint64_t, HcclDataType,292+ const void *, uint64_t, HcclDataType,
293- HcclComm, aclrtStream);293+ HcclComm, aclrtStream);
294- static HcclAlltoAllFunc func = nullptr;294+ static HcclAlltoAllFunc func = nullptr;
295- if (func == nullptr) {295+ if (func == nullptr) {
296- func = (HcclAlltoAllFunc)GET_FUNC(HcclAlltoAll);296+ func = (HcclAlltoAllFunc)GET_FUNC(HcclAlltoAll);
297- }297+ }
298- TORCH_CHECK(func, "Failed to find function ", "HcclAlltoAll", DIST_ERROR(ErrCode::NOT_FOUND));298+ TORCH_CHECK(func, "Failed to find function ", "HcclAlltoAll", DIST_ERROR(ErrCode::NOT_FOUND));
299- auto ret = func(sendBuf, sendCount, sendType,299+ auto ret = func(sendBuf, sendCount, sendType,
300- recvBuf, recvCount, recvType, comm, stream);300+ recvBuf, recvCount, recvType, comm, stream);
301- return ret;301+ return ret;
302-}302+}
303- 303+ 
304-bool hcclCommInitRootInfoConfigExist()304+bool hcclCommInitRootInfoConfigExist()
305-{305+{
306- static c10::once_flag flag;306+ static c10::once_flag flag;
307- static bool exist = false;307+ static bool exist = false;
308- c10::call_once(flag, [&]() {308+ c10::call_once(flag, [&]() {
309- auto func = GET_FUNC(HcclCommInitRootInfoConfig)309+ auto func = GET_FUNC(HcclCommInitRootInfoConfig)
310- if (func != nullptr) {310+ if (func != nullptr) {
311- exist = true;311+ exist = true;
312- }312+ }
313- });313+ });
314- return exist;314+ return exist;
315-}315+}
316- 316+ 
317-bool hcclAllGatherVExist()317+bool hcclAllGatherVExist()
318-{318+{
319- static c10::once_flag flag;319+ static c10::once_flag flag;
320- static bool exist = false;320+ static bool exist = false;
321- c10::call_once(flag, [&]() {321+ c10::call_once(flag, [&]() {
322- auto func = GET_FUNC(HcclAllGatherV)322+ auto func = GET_FUNC(HcclAllGatherV)
323- if (func != nullptr &&323+ if (func != nullptr &&
324- c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend310P1 &&324+ c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend310P1 &&
325- c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) {325+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) {
326- exist = true;326+ exist = true;
327- }327+ }
328- });328+ });
329- return exist;329+ return exist;
330-}330+}
331- 331+ 
332-bool hcclReduceScatterVExist()332+bool hcclReduceScatterVExist()
333-{333+{
334- static c10::once_flag flag;334+ static c10::once_flag flag;
335- static bool exist = false;335+ static bool exist = false;
336- c10::call_once(flag, [&]() {336+ c10::call_once(flag, [&]() {
337- auto func = GET_FUNC(HcclReduceScatterV)337+ auto func = GET_FUNC(HcclReduceScatterV)
338- if (func != nullptr &&338+ if (func != nullptr &&
339- ((c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend310P1 &&339+ ((c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend310P1 &&
340- c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) ||340+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) ||
341- c10_npu::GetSocVersion() == c10_npu::SocVersion::Ascend950)) {341+ c10_npu::GetSocVersion() == c10_npu::SocVersion::Ascend950)) {
342- exist = true;342+ exist = true;
343- }343+ }
344- });344+ });
345- return exist;345+ return exist;
346-}346+}
347- 347+ 
348-HcclResult hcclCommInitRootInfoConfig(uint32_t nRanks, const HcclRootInfo *rootInfo, uint32_t rank, HcclCommConfig* config, HcclComm *comm)348+HcclResult hcclCommInitRootInfoConfig(uint32_t nRanks, const HcclRootInfo *rootInfo, uint32_t rank, HcclCommConfig* config, HcclComm *comm)
349-{349+{
350- using HcclCommInitRootInfoConfigFunc = HcclResult(*)(350+ using HcclCommInitRootInfoConfigFunc = HcclResult(*)(
351- uint32_t, const HcclRootInfo *, uint32_t, HcclCommConfig*, HcclComm *);351+ uint32_t, const HcclRootInfo *, uint32_t, HcclCommConfig*, HcclComm *);
352- static HcclCommInitRootInfoConfigFunc func = nullptr;352+ static HcclCommInitRootInfoConfigFunc func = nullptr;
353- if (func == nullptr) {353+ if (func == nullptr) {
354- func = (HcclCommInitRootInfoConfigFunc)GET_FUNC(HcclCommInitRootInfoConfig)354+ func = (HcclCommInitRootInfoConfigFunc)GET_FUNC(HcclCommInitRootInfoConfig)
355- }355+ }
356- TORCH_CHECK(func, "Failed to find function ", "HcclCommInitRootInfoConfig", DIST_ERROR(ErrCode::NOT_FOUND));356+ TORCH_CHECK(func, "Failed to find function ", "HcclCommInitRootInfoConfig", DIST_ERROR(ErrCode::NOT_FOUND));
357- auto ret = func(nRanks, rootInfo, rank, config, comm);357+ auto ret = func(nRanks, rootInfo, rank, config, comm);
358- return ret;358+ return ret;
359-}359+}
360- 360+ 
361-bool isHcclFeatureSupported(HcclCommConfigCapability configParameter)361+bool isHcclFeatureSupported(HcclCommConfigCapability configParameter)
362-{362+{
363- using HcclGetCommConfigCapabilityFunc = uint32_t(*)();363+ using HcclGetCommConfigCapabilityFunc = uint32_t(*)();
364- static HcclGetCommConfigCapabilityFunc func = (HcclGetCommConfigCapabilityFunc) GET_FUNC(364+ static HcclGetCommConfigCapabilityFunc func = (HcclGetCommConfigCapabilityFunc) GET_FUNC(
365- HcclGetCommConfigCapability);365+ HcclGetCommConfigCapability);
366- if (func == nullptr) {366+ if (func == nullptr) {
367- return false;367+ return false;
368- }368+ }
369- return configParameter < func();369+ return configParameter < func();
370-}370+}
371- 371+ 
372-bool hcclCommInitClusterInfoConfigExist()372+bool hcclCommInitClusterInfoConfigExist()
373-{373+{
374- const static bool isClusterInitExist = []() -> bool {374+ const static bool isClusterInitExist = []() -> bool {
375- auto func = GET_FUNC(HcclCommInitClusterInfoConfig)375+ auto func = GET_FUNC(HcclCommInitClusterInfoConfig)
376- return func != nullptr;376+ return func != nullptr;
377- }();377+ }();
378- return isClusterInitExist;378+ return isClusterInitExist;
379-}379+}
380- 380+ 
381-HcclResult hcclCommInitClusterInfoConfig(const char *clusterInfo, uint32_t rank, HcclCommConfig *config, HcclComm *comm)381+HcclResult hcclCommInitClusterInfoConfig(const char *clusterInfo, uint32_t rank, HcclCommConfig *config, HcclComm *comm)
382-{382+{
383- using HcclCommInitClusterInfoConfigFunc = HcclResult(*)(const char *, uint32_t, HcclCommConfig *, HcclComm *);383+ using HcclCommInitClusterInfoConfigFunc = HcclResult(*)(const char *, uint32_t, HcclCommConfig *, HcclComm *);
384- static HcclCommInitClusterInfoConfigFunc func = nullptr;384+ static HcclCommInitClusterInfoConfigFunc func = nullptr;
385- if (func == nullptr) {385+ if (func == nullptr) {
386- func = (HcclCommInitClusterInfoConfigFunc)GET_FUNC(HcclCommInitClusterInfoConfig)386+ func = (HcclCommInitClusterInfoConfigFunc)GET_FUNC(HcclCommInitClusterInfoConfig)
387- }387+ }
388- TORCH_CHECK(func, "Failed to find function ", "HcclCommInitClusterInfoConfig", DIST_ERROR(ErrCode::NOT_FOUND));388+ TORCH_CHECK(func, "Failed to find function ", "HcclCommInitClusterInfoConfig", DIST_ERROR(ErrCode::NOT_FOUND));
389- auto ret = func(clusterInfo, rank, config, comm);389+ auto ret = func(clusterInfo, rank, config, comm);
390- return ret;390+ return ret;
391-}391+}
392- 392+ 
393-bool hcclCreateSubCommConfigExist()393+bool hcclCreateSubCommConfigExist()
394-{394+{
395- const static bool isCreateSubCommExist = []() -> bool {395+ const static bool isCreateSubCommExist = []() -> bool {
396- auto func = GET_FUNC(HcclCreateSubCommConfig)396+ auto func = GET_FUNC(HcclCreateSubCommConfig)
397- return func != nullptr;397+ return func != nullptr;
398- }();398+ }();
399- return isCreateSubCommExist;399+ return isCreateSubCommExist;
400-}400+}
401- 401+ 
402-HcclResult hcclCreateSubCommConfig(HcclComm *comm, uint32_t rankNum, uint32_t *rankIds, uint64_t subCommId, uint32_t subCommRankId,402+HcclResult hcclCreateSubCommConfig(HcclComm *comm, uint32_t rankNum, uint32_t *rankIds, uint64_t subCommId, uint32_t subCommRankId,
403- HcclCommConfig* config, HcclComm *subComm)403+ HcclCommConfig* config, HcclComm *subComm)
404-{404+{
405- using HcclCreateSubCommConfigFunc = HcclResult(*)(HcclComm *, uint32_t, uint32_t *, uint64_t, uint32_t, HcclCommConfig *, HcclComm *);405+ using HcclCreateSubCommConfigFunc = HcclResult(*)(HcclComm *, uint32_t, uint32_t *, uint64_t, uint32_t, HcclCommConfig *, HcclComm *);
406- static HcclCreateSubCommConfigFunc func = nullptr;406+ static HcclCreateSubCommConfigFunc func = nullptr;
407- if (func == nullptr) {407+ if (func == nullptr) {
408- func = (HcclCreateSubCommConfigFunc)GET_FUNC(HcclCreateSubCommConfig)408+ func = (HcclCreateSubCommConfigFunc)GET_FUNC(HcclCreateSubCommConfig)
409- }409+ }
410- TORCH_CHECK(func, "Failed to find function ", "HcclCreateSubCommConfig", DIST_ERROR(ErrCode::NOT_FOUND));410+ TORCH_CHECK(func, "Failed to find function ", "HcclCreateSubCommConfig", DIST_ERROR(ErrCode::NOT_FOUND));
411- auto ret = func(comm, rankNum, rankIds, subCommId, subCommRankId, config, subComm);411+ auto ret = func(comm, rankNum, rankIds, subCommId, subCommRankId, config, subComm);
412- return ret;412+ return ret;
413-}413+}
414- 414+ 
415-bool hcclCommWorkingDevNicSetExist()415+bool hcclCommWorkingDevNicSetExist()
416-{416+{
417- const static bool isHcclCommWorkingDevNicSetExist = []() -> bool {417+ const static bool isHcclCommWorkingDevNicSetExist = []() -> bool {
418- auto func = GET_FUNC(HcclCommWorkingDevNicSet)418+ auto func = GET_FUNC(HcclCommWorkingDevNicSet)
419- return func != nullptr;419+ return func != nullptr;
420- }();420+ }();
421- return isHcclCommWorkingDevNicSetExist;421+ return isHcclCommWorkingDevNicSetExist;
422-}422+}
423- 423+ 
424-HcclResult hcclCommWorkingDevNicSet(HcclComm comm, uint32_t *ranks, bool *useBackup, uint32_t nRanks)424+HcclResult hcclCommWorkingDevNicSet(HcclComm comm, uint32_t *ranks, bool *useBackup, uint32_t nRanks)
425-{425+{
426- using HcclCommWorkingDevNicSetFunc = HcclResult(*)(HcclComm, uint32_t *, bool *, uint32_t);426+ using HcclCommWorkingDevNicSetFunc = HcclResult(*)(HcclComm, uint32_t *, bool *, uint32_t);
427- static HcclCommWorkingDevNicSetFunc func = nullptr;427+ static HcclCommWorkingDevNicSetFunc func = nullptr;
428- if (func == nullptr) {428+ if (func == nullptr) {
429- func = (HcclCommWorkingDevNicSetFunc)GET_FUNC(HcclCommWorkingDevNicSet)429+ func = (HcclCommWorkingDevNicSetFunc)GET_FUNC(HcclCommWorkingDevNicSet)
430- }430+ }
431- TORCH_CHECK(func, "Failed to find function ", "HcclCommWorkingDevNicSet", DIST_ERROR(ErrCode::NOT_FOUND));431+ TORCH_CHECK(func, "Failed to find function ", "HcclCommWorkingDevNicSet", DIST_ERROR(ErrCode::NOT_FOUND));
432- auto ret = func(comm, ranks, useBackup, nRanks);432+ auto ret = func(comm, ranks, useBackup, nRanks);
433- return ret;433+ return ret;
434-}434+}
435- 435+ 
436-HcclResult hcclCommRegister(HcclComm comm, void *addr, uint64_t size, void **handle, uint32_t flag)436+HcclResult hcclCommRegister(HcclComm comm, void *addr, uint64_t size, void **handle, uint32_t flag)
437-{437+{
438- using HcclCommRegisterFunc = HcclResult(*)(HcclComm, void *, uint64_t, void **, uint32_t);438+ using HcclCommRegisterFunc = HcclResult(*)(HcclComm, void *, uint64_t, void **, uint32_t);
439- static HcclCommRegisterFunc func = nullptr;439+ static HcclCommRegisterFunc func = nullptr;
440- if (func == nullptr) {440+ if (func == nullptr) {
441- func = (HcclCommRegisterFunc)GET_FUNC(HcclCommRegister)441+ func = (HcclCommRegisterFunc)GET_FUNC(HcclCommRegister)
442- }442+ }
443- TORCH_CHECK(func, "Failed to find function ", "HcclCommRegister", DIST_ERROR(ErrCode::NOT_FOUND));443+ TORCH_CHECK(func, "Failed to find function ", "HcclCommRegister", DIST_ERROR(ErrCode::NOT_FOUND));
444- auto ret = func(comm, addr, size, handle, flag);444+ auto ret = func(comm, addr, size, handle, flag);
445- return ret;445+ return ret;
446-}446+}
447- 447+ 
448-HcclResult hcclCommDeregister(HcclComm comm, void *handle)448+HcclResult hcclCommDeregister(HcclComm comm, void *handle)
449-{449+{
450- using HcclCommDeregisterFunc = HcclResult(*)(HcclComm, void *);450+ using HcclCommDeregisterFunc = HcclResult(*)(HcclComm, void *);
451- static HcclCommDeregisterFunc func = nullptr;451+ static HcclCommDeregisterFunc func = nullptr;
452- if (func == nullptr) {452+ if (func == nullptr) {
453- func = (HcclCommDeregisterFunc)GET_FUNC(HcclCommDeregister)453+ func = (HcclCommDeregisterFunc)GET_FUNC(HcclCommDeregister)
454- }454+ }
455- TORCH_CHECK(func, "Failed to find function ", "HcclCommDeregister", DIST_ERROR(ErrCode::NOT_FOUND));455+ TORCH_CHECK(func, "Failed to find function ", "HcclCommDeregister", DIST_ERROR(ErrCode::NOT_FOUND));
456- auto ret = func(comm, handle);456+ auto ret = func(comm, handle);
457- return ret;457+ return ret;
458-}458+}
459- 459+ 
460-HcclResult hcclCommExchangeMem(HcclComm comm, void *windowHandle, uint32_t *peerRanks, uint32_t peerRankNum)460+HcclResult hcclCommExchangeMem(HcclComm comm, void *windowHandle, uint32_t *peerRanks, uint32_t peerRankNum)
461-{461+{
462- using HcclCommExchangeMemFunc = HcclResult(*)(HcclComm, void *, uint32_t *, uint32_t);462+ using HcclCommExchangeMemFunc = HcclResult(*)(HcclComm, void *, uint32_t *, uint32_t);
463- static HcclCommExchangeMemFunc func = nullptr;463+ static HcclCommExchangeMemFunc func = nullptr;
464- if (func == nullptr) {464+ if (func == nullptr) {
465- func = (HcclCommExchangeMemFunc)GET_FUNC(HcclCommExchangeMem)465+ func = (HcclCommExchangeMemFunc)GET_FUNC(HcclCommExchangeMem)
466- }466+ }
467- TORCH_CHECK(func, "Failed to find function ", "HcclCommExchangeMem", DIST_ERROR(ErrCode::NOT_FOUND));467+ TORCH_CHECK(func, "Failed to find function ", "HcclCommExchangeMem", DIST_ERROR(ErrCode::NOT_FOUND));
468- auto ret = func(comm, windowHandle, peerRanks, peerRankNum);468+ auto ret = func(comm, windowHandle, peerRanks, peerRankNum);
469- return ret;469+ return ret;
470-}470+}
471- 471+ 
472-HcclResult hcclGroupStart()472+HcclResult hcclGroupStart()
473-{473+{
474- using hcclGroupStartFunc = HcclResult(*)();474+ using hcclGroupStartFunc = HcclResult(*)();
475- static hcclGroupStartFunc func = nullptr;475+ static hcclGroupStartFunc func = nullptr;
476- if (func == nullptr) {476+ if (func == nullptr) {
477- func = (hcclGroupStartFunc)GET_FUNCTION(libhcomm, HcclGroupStart)477+ func = (hcclGroupStartFunc)GET_FUNCTION(libhcomm, HcclGroupStart)
478- }478+ }
479- TORCH_CHECK(func, "Failed to find function ", "HcclGroupStart", DIST_ERROR(ErrCode::NOT_FOUND));479+ TORCH_CHECK(func, "Failed to find function ", "HcclGroupStart", DIST_ERROR(ErrCode::NOT_FOUND));
480- auto ret = func();480+ auto ret = func();
481- return ret;481+ return ret;
482-}482+}
483- 483+ 
484-HcclResult hcclGroupEnd()484+HcclResult hcclGroupEnd()
485-{485+{
486- using hcclGroupEndFunc = HcclResult(*)();486+ using hcclGroupEndFunc = HcclResult(*)();
487- static hcclGroupEndFunc func = nullptr;487+ static hcclGroupEndFunc func = nullptr;
488- if (func == nullptr) {488+ if (func == nullptr) {
489- func = (hcclGroupEndFunc)GET_FUNCTION(libhcomm, HcclGroupEnd)489+ func = (hcclGroupEndFunc)GET_FUNCTION(libhcomm, HcclGroupEnd)
490- }490+ }
491- TORCH_CHECK(func, "Failed to find function ", "HcclGroupEnd", DIST_ERROR(ErrCode::NOT_FOUND));491+ TORCH_CHECK(func, "Failed to find function ", "HcclGroupEnd", DIST_ERROR(ErrCode::NOT_FOUND));
492- auto ret = func();492+ auto ret = func();
493- return ret;493+ return ret;
494-}494+}
495-} // namespace c10d_npu495+} // namespace c10d_npu
Mtorch_npu/csrc/distributed/Init.cpp+614-614
@@ -1,614 +1,614 @@
1-#include <deque>1+#include <deque>
2- 2+ 
3-#include <torch/custom_class.h>3+#include <torch/custom_class.h>
4-#include <torch/csrc/python_headers.h>4+#include <torch/csrc/python_headers.h>
5-#include <c10/util/intrusive_ptr.h>5+#include <c10/util/intrusive_ptr.h>
6-#include <c10/util/irange.h>6+#include <c10/util/irange.h>
7-#include <c10d/ProcessGroup.hpp>7+#include <c10d/ProcessGroup.hpp>
8-#include <c10d/comm.hpp>8+#include <c10d/comm.hpp>
9-#include <c10d/Work.hpp>9+#include <c10d/Work.hpp>
10-#include <pybind11/chrono.h>10+#include <pybind11/chrono.h>
11- 11+ 
12-#include <torch/csrc/Exceptions.h>12+#include <torch/csrc/Exceptions.h>
13-#include <ATen/core/functional.h>13+#include <ATen/core/functional.h>
14-#include <torch/csrc/jit/python/pybind_utils.h>14+#include <torch/csrc/jit/python/pybind_utils.h>
15-#include <torch/csrc/utils/object_ptr.h>15+#include <torch/csrc/utils/object_ptr.h>
16-#include <torch/csrc/utils/pybind.h>16+#include <torch/csrc/utils/pybind.h>
17-#include <torch/csrc/utils/tensor_flatten.h>17+#include <torch/csrc/utils/tensor_flatten.h>
18-#include <torch/csrc/distributed/c10d/python_comm_hook.h>18+#include <torch/csrc/distributed/c10d/python_comm_hook.h>
19- 19+ 
20-#include "torch_npu/csrc/distributed/rpc/init.h"20+#include "torch_npu/csrc/distributed/rpc/init.h"
21-#include "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"21+#include "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"
22-#include "torch_npu/csrc/distributed/ProcessGroupLCCL.hpp"22+#include "torch_npu/csrc/distributed/ProcessGroupLCCL.hpp"
23-#include "torch_npu/csrc/distributed/reducer.hpp"23+#include "torch_npu/csrc/distributed/reducer.hpp"
24-#include "torch_npu/csrc/distributed/ParallelTcpStore.hpp"24+#include "torch_npu/csrc/distributed/ParallelTcpStore.hpp"
25-#include "torch_npu/csrc/aten/NPUNativeFunctions.h"25+#include "torch_npu/csrc/aten/NPUNativeFunctions.h"
26-#include "torch_npu/csrc/aten/CustomFunctions.h"26+#include "torch_npu/csrc/aten/CustomFunctions.h"
27-#include "torch_npu/csrc/core/NPUBridge.h"27+#include "torch_npu/csrc/core/NPUBridge.h"
28-#include "torch_npu/csrc/distributed/Init.h"28+#include "torch_npu/csrc/distributed/Init.h"
29- 29+ 
30- 30+ 
31-namespace {31+namespace {
32- 32+ 
33-// Wrapper to ensure GIL is released before destructing ProcessGroupGloo33+// Wrapper to ensure GIL is released before destructing ProcessGroupGloo
34-template <typename T>34+template <typename T>
35-class IntrusivePtrNoGilDestructor {35+class IntrusivePtrNoGilDestructor {
36- c10::intrusive_ptr<T> impl_;36+ c10::intrusive_ptr<T> impl_;
37- 37+ 
38-public:38+public:
39- IntrusivePtrNoGilDestructor() = default;39+ IntrusivePtrNoGilDestructor() = default;
40- IntrusivePtrNoGilDestructor(const IntrusivePtrNoGilDestructor&) = default;40+ IntrusivePtrNoGilDestructor(const IntrusivePtrNoGilDestructor&) = default;
41- IntrusivePtrNoGilDestructor(IntrusivePtrNoGilDestructor&&) = default;41+ IntrusivePtrNoGilDestructor(IntrusivePtrNoGilDestructor&&) = default;
42- IntrusivePtrNoGilDestructor& operator=(const IntrusivePtrNoGilDestructor&) =42+ IntrusivePtrNoGilDestructor& operator=(const IntrusivePtrNoGilDestructor&) =
43- default;43+ default;
44- IntrusivePtrNoGilDestructor& operator=(IntrusivePtrNoGilDestructor&&) =44+ IntrusivePtrNoGilDestructor& operator=(IntrusivePtrNoGilDestructor&&) =
45- default;45+ default;
46- IntrusivePtrNoGilDestructor(c10::intrusive_ptr<T> impl)46+ IntrusivePtrNoGilDestructor(c10::intrusive_ptr<T> impl)
47- : impl_(std::move(impl)) {}47+ : impl_(std::move(impl)) {}
48- explicit IntrusivePtrNoGilDestructor(T* impl)48+ explicit IntrusivePtrNoGilDestructor(T* impl)
49- : impl_(c10::intrusive_ptr<T>::unsafe_steal_from_new(impl)) {}49+ : impl_(c10::intrusive_ptr<T>::unsafe_steal_from_new(impl)) {}
50- ~IntrusivePtrNoGilDestructor() {50+ ~IntrusivePtrNoGilDestructor() {
51- if (impl_) {51+ if (impl_) {
52- if (PyGILState_Check() != 0) {52+ if (PyGILState_Check() != 0) {
53- pybind11::gil_scoped_release release;53+ pybind11::gil_scoped_release release;
54- impl_.reset();54+ impl_.reset();
55- } else {55+ } else {
56- impl_.reset();56+ impl_.reset();
57- }57+ }
58- }58+ }
59- }59+ }
60- T& operator*() const noexcept {60+ T& operator*() const noexcept {
61- return *impl_;61+ return *impl_;
62- }62+ }
63- T* operator->() const noexcept {63+ T* operator->() const noexcept {
64- return impl_.get();64+ return impl_.get();
65- }65+ }
66- C10_NODISCARD T* get() const noexcept {66+ C10_NODISCARD T* get() const noexcept {
67- return impl_.get();67+ return impl_.get();
68- }68+ }
69- void reset() noexcept {69+ void reset() noexcept {
70- impl_.reset();70+ impl_.reset();
71- }71+ }
72- operator bool() const noexcept {72+ operator bool() const noexcept {
73- return impl_;73+ return impl_;
74- }74+ }
75-};75+};
76- 76+ 
77-} // anonymous namespace77+} // anonymous namespace
78- 78+ 
79-PYBIND11_DECLARE_HOLDER_TYPE(T, IntrusivePtrNoGilDestructor<T>, true);79+PYBIND11_DECLARE_HOLDER_TYPE(T, IntrusivePtrNoGilDestructor<T>, true);
80- 80+ 
81- 81+ 
82-namespace torch_npu {82+namespace torch_npu {
83-namespace distributed {83+namespace distributed {
84- 84+ 
85-template <typename T>85+template <typename T>
86-using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>;86+using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>;
87- 87+ 
88-template <typename T>88+template <typename T>
89-using intrusive_ptr_class_ = py::class_<T, c10::intrusive_ptr<T>>;89+using intrusive_ptr_class_ = py::class_<T, c10::intrusive_ptr<T>>;
90- 90+ 
91-template <typename T>91+template <typename T>
92-using intrusive_ptr_no_gil_destructor_class_ =92+using intrusive_ptr_no_gil_destructor_class_ =
93- py::class_<T, IntrusivePtrNoGilDestructor<T>>;93+ py::class_<T, IntrusivePtrNoGilDestructor<T>>;
94- 94+ 
95- 95+ 
96-class BroadcastWork {96+class BroadcastWork {
97-public:97+public:
98- inline std::vector<at::Tensor> cast_tensors(at::TensorList tensors) const98+ inline std::vector<at::Tensor> cast_tensors(at::TensorList tensors) const
99- {99+ {
100- static auto cast_back_to_ori_format = [](const at::Tensor &t) {100+ static auto cast_back_to_ori_format = [](const at::Tensor &t) {
101- return at_npu::native::custom_ops::npu_format_cast(t, torch_npu::NPUBridge::GetNpuStorageImpl(t)->npu_desc_.origin_format_);101+ return at_npu::native::custom_ops::npu_format_cast(t, torch_npu::NPUBridge::GetNpuStorageImpl(t)->npu_desc_.origin_format_);
102- };102+ };
103- return c10::fmap(tensors, cast_back_to_ori_format);103+ return c10::fmap(tensors, cast_back_to_ori_format);
104- }104+ }
105- 105+ 
106- BroadcastWork(106+ BroadcastWork(
107- const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,107+ const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,
108- std::vector<at::Tensor> bucket_tensors,108+ std::vector<at::Tensor> bucket_tensors,
109- int root_rank = 0)109+ int root_rank = 0)
110- : bucket_tensors_(std::move(bucket_tensors)),110+ : bucket_tensors_(std::move(bucket_tensors)),
111- cast_tensors_(cast_tensors(bucket_tensors_)),111+ cast_tensors_(cast_tensors(bucket_tensors_)),
112- flat_tensor_({torch::utils::flatten_dense_tensors(cast_tensors_)}) {112+ flat_tensor_({torch::utils::flatten_dense_tensors(cast_tensors_)}) {
113- c10d::BroadcastOptions broadcastOptions;113+ c10d::BroadcastOptions broadcastOptions;
114- broadcastOptions.rootRank = root_rank;114+ broadcastOptions.rootRank = root_rank;
115- work_ = process_group->broadcast(flat_tensor_, broadcastOptions);115+ work_ = process_group->broadcast(flat_tensor_, broadcastOptions);
116- }116+ }
117- 117+ 
118- void finish()118+ void finish()
119- {119+ {
120- work_->wait();120+ work_->wait();
121- auto output_tensors = torch::utils::unflatten_dense_tensors(121+ auto output_tensors = torch::utils::unflatten_dense_tensors(
122- flat_tensor_.front(), cast_tensors_);122+ flat_tensor_.front(), cast_tensors_);
123- TORCH_INTERNAL_ASSERT(output_tensors.size() == bucket_tensors_.size(), DIST_ERROR(ErrCode::PARAM));123+ TORCH_INTERNAL_ASSERT(output_tensors.size() == bucket_tensors_.size(), DIST_ERROR(ErrCode::PARAM));
124- for (const auto i : c10::irange(output_tensors.size())) {124+ for (const auto i : c10::irange(output_tensors.size())) {
125- bucket_tensors_[i].copy_(output_tensors[i], true);125+ bucket_tensors_[i].copy_(output_tensors[i], true);
126- }126+ }
127- }127+ }
128- 128+ 
129-protected:129+protected:
130- // The list of tensors to broadcast. They are guaranteed to be130+ // The list of tensors to broadcast. They are guaranteed to be
131- // placed on the same device and have the same dtype.131+ // placed on the same device and have the same dtype.
132- std::vector<at::Tensor> bucket_tensors_;132+ std::vector<at::Tensor> bucket_tensors_;
133- // Some tensors with format, such as FRACTAL_Z, 5HD, may be padded to133+ // Some tensors with format, such as FRACTAL_Z, 5HD, may be padded to
134- // keep alignment with 16*16 cube kernel which will modify storage as134+ // keep alignment with 16*16 cube kernel which will modify storage as
135- // input tensor for cat operation during flatten to a buffer tensor.135+ // input tensor for cat operation during flatten to a buffer tensor.
136- // So, it needs to cast all bucket tensors to tensors with format HCHW136+ // So, it needs to cast all bucket tensors to tensors with format HCHW
137- std::vector<at::Tensor> cast_tensors_;137+ std::vector<at::Tensor> cast_tensors_;
138- // The vector with a single flattened tensor containing the contents138+ // The vector with a single flattened tensor containing the contents
139- // of the tensors in bucket_tensors_. It must be stored in a vector139+ // of the tensors in bucket_tensors_. It must be stored in a vector
140- // because c10d::ProcessGroup::broadcast takes a vector argument.140+ // because c10d::ProcessGroup::broadcast takes a vector argument.
141- std::vector<at::Tensor> flat_tensor_;141+ std::vector<at::Tensor> flat_tensor_;
142- 142+ 
143-private:143+private:
144- 144+ 
145- // The broadcast work that is kicked off upon construction.145+ // The broadcast work that is kicked off upon construction.
146- c10::intrusive_ptr<c10d::Work> work_;146+ c10::intrusive_ptr<c10d::Work> work_;
147-};147+};
148- 148+ 
149-// Broadcast many tensors to all processes in the process group.149+// Broadcast many tensors to all processes in the process group.
150-void broadcast_coalesced(150+void broadcast_coalesced(
151- c10::intrusive_ptr<c10d::ProcessGroup> process_group,151+ c10::intrusive_ptr<c10d::ProcessGroup> process_group,
152- at::TensorList tensors,152+ at::TensorList tensors,
153- size_t buffer_size,153+ size_t buffer_size,
154- int rank)154+ int rank)
155-{155+{
156- // Coalesce tensors into buckets taking into account the maximum buffer size.156+ // Coalesce tensors into buckets taking into account the maximum buffer size.
157- // This routine is multi-device aware, so the tensors can be split across157+ // This routine is multi-device aware, so the tensors can be split across
158- // multiple devices and can contain a mix of CPU and CUDA tensors.158+ // multiple devices and can contain a mix of CPU and CUDA tensors.
159- std::vector<std::vector<size_t>> buckets;159+ std::vector<std::vector<size_t>> buckets;
160- std::tie(buckets, std::ignore) =160+ std::tie(buckets, std::ignore) =
161- c10d_npu::compute_bucket_assignment_by_size(tensors.vec(), {buffer_size});161+ c10d_npu::compute_bucket_assignment_by_size(tensors.vec(), {buffer_size});
162- 162+ 
163- // Returns tensor at specified index in input tensor list.163+ // Returns tensor at specified index in input tensor list.
164- const auto lookup = [&tensors](size_t index) { return tensors[index]; };164+ const auto lookup = [&tensors](size_t index) { return tensors[index]; };
165- 165+ 
166- // We maintain a maximum of 2 in flight broadcast operations to avoid166+ // We maintain a maximum of 2 in flight broadcast operations to avoid
167- // allocating too much memory (in case the specified tensors are very large).167+ // allocating too much memory (in case the specified tensors are very large).
168- std::deque<BroadcastWork> in_flight;168+ std::deque<BroadcastWork> in_flight;
169- constexpr auto max_in_flight = 2;169+ constexpr auto max_in_flight = 2;
170- for (const auto& bucket : buckets) {170+ for (const auto& bucket : buckets) {
171- if (in_flight.size() >= max_in_flight) {171+ if (in_flight.size() >= max_in_flight) {
172- in_flight.front().finish();172+ in_flight.front().finish();
173- in_flight.pop_front();173+ in_flight.pop_front();
174- }174+ }
175- in_flight.emplace_back(process_group, c10::fmap(bucket, lookup), rank);175+ in_flight.emplace_back(process_group, c10::fmap(bucket, lookup), rank);
176- }176+ }
177- 177+ 
178- while (!in_flight.empty()) {178+ while (!in_flight.empty()) {
179- in_flight.front().finish();179+ in_flight.front().finish();
180- in_flight.pop_front();180+ in_flight.pop_front();
181- }181+ }
182-}182+}
183- 183+ 
184-// Called from DDP's Python API to create a c10d Python comm hook object.184+// Called from DDP's Python API to create a c10d Python comm hook object.
185-// The input state and callable comm_hook are Python objects. It later calls185+// The input state and callable comm_hook are Python objects. It later calls
186-// register_comm_hook function of the reducer input to register the hook.186+// register_comm_hook function of the reducer input to register the hook.
187-void _register_comm_hook(187+void _register_comm_hook(
188- c10d_npu::Reducer& reducer,188+ c10d_npu::Reducer& reducer,
189- py::object state,189+ py::object state,
190- py::object comm_hook)190+ py::object comm_hook)
191-{191+{
192- reducer.register_comm_hook(std::make_unique<::c10d::PythonCommHook>(192+ reducer.register_comm_hook(std::make_unique<::c10d::PythonCommHook>(
193- std::move(state), std::move(comm_hook)));193+ std::move(state), std::move(comm_hook)));
194-}194+}
195- 195+ 
196-// Called from DDP's Python API to create a c10d C++ comm hook.196+// Called from DDP's Python API to create a c10d C++ comm hook.
197-// The input is an enum hook type. It later calls register_builtin_comm_hook197+// The input is an enum hook type. It later calls register_builtin_comm_hook
198-// function of the reducer input to set the hook type.198+// function of the reducer input to set the hook type.
199-void _register_builtin_comm_hook(199+void _register_builtin_comm_hook(
200- c10d_npu::Reducer& reducer,200+ c10d_npu::Reducer& reducer,
201- ::c10d::BuiltinCommHookType comm_hook_type)201+ ::c10d::BuiltinCommHookType comm_hook_type)
202-{202+{
203- reducer.register_builtin_comm_hook(comm_hook_type);203+ reducer.register_builtin_comm_hook(comm_hook_type);
204-}204+}
205- 205+ 
206-PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs)206+PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs)
207-{207+{
208- auto torch_npu_C_module = THPObjectPtr(PyImport_ImportModule("torch_npu._C"));208+ auto torch_npu_C_module = THPObjectPtr(PyImport_ImportModule("torch_npu._C"));
209- if (!torch_npu_C_module) {209+ if (!torch_npu_C_module) {
210- throw python_error();210+ throw python_error();
211- }211+ }
212- auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();212+ auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();
213- 213+
214- auto m =214+ auto m =
215- torch_npu_C_m.def_submodule("_distributed_c10d", "distributed c10d bindings");215+ torch_npu_C_m.def_submodule("_distributed_c10d", "distributed c10d bindings");
216- auto module = py::handle(m).cast<py::module>();216+ auto module = py::handle(m).cast<py::module>();
217- 217+ 
218- module.def("_compute_bucket_assignment_by_size",218+ module.def("_compute_bucket_assignment_by_size",
219- [](const std::vector<at::Tensor>& tensors,219+ [](const std::vector<at::Tensor>& tensors,
220- const std::vector<size_t>& bucket_size_limits,220+ const std::vector<size_t>& bucket_size_limits,
221- const std::vector<bool>& expect_sparse_gradient,221+ const std::vector<bool>& expect_sparse_gradient,
222- const std::vector<int64_t>& tensor_indices,222+ const std::vector<int64_t>& tensor_indices,
223- const c10::optional<std::shared_ptr<::c10d::Logger>>& logger) {223+ const c10::optional<std::shared_ptr<::c10d::Logger>>& logger) {
224- if (logger.has_value()) {224+ if (logger.has_value()) {
225- std::weak_ptr<::c10d::Logger> logger_weakref = logger.value();225+ std::weak_ptr<::c10d::Logger> logger_weakref = logger.value();
226- return ::c10d_npu::compute_bucket_assignment_by_size(tensors, bucket_size_limits, expect_sparse_gradient, tensor_indices, {logger_weakref});226+ return ::c10d_npu::compute_bucket_assignment_by_size(tensors, bucket_size_limits, expect_sparse_gradient, tensor_indices, {logger_weakref});
227- } else {227+ } else {
228- return ::c10d_npu::compute_bucket_assignment_by_size(tensors, bucket_size_limits, expect_sparse_gradient, tensor_indices, {});228+ return ::c10d_npu::compute_bucket_assignment_by_size(tensors, bucket_size_limits, expect_sparse_gradient, tensor_indices, {});
229- }229+ }
230- },230+ },
231- py::arg("tensors"),231+ py::arg("tensors"),
232- py::arg("bucket_size"),232+ py::arg("bucket_size"),
233- py::arg("expect_sparse_gradient") = std::vector<bool>(),233+ py::arg("expect_sparse_gradient") = std::vector<bool>(),
234- py::arg("tensor_indices") = std::vector<int64_t>(),234+ py::arg("tensor_indices") = std::vector<int64_t>(),
235- py::arg("logger") = c10::optional<std::shared_ptr<::c10d::Logger>>{},235+ py::arg("logger") = c10::optional<std::shared_ptr<::c10d::Logger>>{},
236- py::call_guard<py::gil_scoped_release>());236+ py::call_guard<py::gil_scoped_release>());
237- 237+ 
238- module.def("_verify_params_across_processes",238+ module.def("_verify_params_across_processes",
239- [](const c10::intrusive_ptr<::c10d::ProcessGroup>& process_group,239+ [](const c10::intrusive_ptr<::c10d::ProcessGroup>& process_group,
240- const std::vector<at::Tensor>& params,240+ const std::vector<at::Tensor>& params,
241- const c10::optional<std::shared_ptr<::c10d::Logger>>& logger) {241+ const c10::optional<std::shared_ptr<::c10d::Logger>>& logger) {
242- if (logger.has_value()) {242+ if (logger.has_value()) {
243- std::weak_ptr<::c10d::Logger> logger_weakref = logger.value();243+ std::weak_ptr<::c10d::Logger> logger_weakref = logger.value();
244- c10d_npu::verify_params_across_processes(process_group, params, {logger_weakref});244+ c10d_npu::verify_params_across_processes(process_group, params, {logger_weakref});
245- } else {245+ } else {
246- c10d_npu::verify_params_across_processes(process_group, params, {});246+ c10d_npu::verify_params_across_processes(process_group, params, {});
247- }247+ }
248- },248+ },
249- py::arg("process_group"),249+ py::arg("process_group"),
250- py::arg("params"),250+ py::arg("params"),
251- py::arg("logger") = c10::optional<std::shared_ptr<::c10d::Logger>>{},251+ py::arg("logger") = c10::optional<std::shared_ptr<::c10d::Logger>>{},
252- py::call_guard<py::gil_scoped_release>());252+ py::call_guard<py::gil_scoped_release>());
253- 253+ 
254- module254+ module
255- .def("_register_comm_hook",255+ .def("_register_comm_hook",
256- &_register_comm_hook,256+ &_register_comm_hook,
257- py::arg("reducer"),257+ py::arg("reducer"),
258- py::arg("state"),258+ py::arg("state"),
259- py::arg("comm_hook"),259+ py::arg("comm_hook"),
260- py::call_guard<py::gil_scoped_release>())260+ py::call_guard<py::gil_scoped_release>())
261- .def("_register_builtin_comm_hook",261+ .def("_register_builtin_comm_hook",
262- &_register_builtin_comm_hook,262+ &_register_builtin_comm_hook,
263- py::arg("reducer"),263+ py::arg("reducer"),
264- py::arg("comm_hook_type"));264+ py::arg("comm_hook_type"));
265- 265+ 
266- module.def("_broadcast_coalesced",266+ module.def("_broadcast_coalesced",
267- // Define a lambda such that the pybind11 prototype can take a std::vector267+ // Define a lambda such that the pybind11 prototype can take a std::vector
268- // for the tensor list argument, but still pass it to the underlying268+ // for the tensor list argument, but still pass it to the underlying
269- // function as a c10::ArrayRef.269+ // function as a c10::ArrayRef.
270- [](c10::intrusive_ptr<::c10d::ProcessGroup> process_group,270+ [](c10::intrusive_ptr<::c10d::ProcessGroup> process_group,
271- std::vector<at::Tensor> tensors, // NOLINT271+ std::vector<at::Tensor> tensors, // NOLINT
272- size_t buffer_size,272+ size_t buffer_size,
273- int rank) {273+ int rank) {
274- torch_npu::distributed::broadcast_coalesced(274+ torch_npu::distributed::broadcast_coalesced(
275- std::move(process_group), tensors, buffer_size, rank);275+ std::move(process_group), tensors, buffer_size, rank);
276- },276+ },
277- py::arg("process_group"),277+ py::arg("process_group"),
278- py::arg("tensors"),278+ py::arg("tensors"),
279- py::arg("buffer_size"),279+ py::arg("buffer_size"),
280- // The source of truth rank to broadcast the tensors from.280+ // The source of truth rank to broadcast the tensors from.
281- py::arg("src") = 0,281+ py::arg("src") = 0,
282- py::call_guard<py::gil_scoped_release>());282+ py::call_guard<py::gil_scoped_release>());
283- 283+ 
284- module.def("_is_support_hccl_comm_name", &c10d_npu::isSupportHcclCommName);284+ module.def("_is_support_hccl_comm_name", &c10d_npu::isSupportHcclCommName);
285- 285+ 
286- shared_ptr_class_<c10d_npu::Reducer>(module, "Reducer")286+ shared_ptr_class_<c10d_npu::Reducer>(module, "Reducer")
287- .def(py::init<287+ .def(py::init<
288- std::vector<at::Tensor>,288+ std::vector<at::Tensor>,
289- std::vector<std::vector<size_t>>,289+ std::vector<std::vector<size_t>>,
290- std::vector<size_t>,290+ std::vector<size_t>,
291- c10::intrusive_ptr<::c10d::ProcessGroup>,291+ c10::intrusive_ptr<::c10d::ProcessGroup>,
292- std::vector<bool>,292+ std::vector<bool>,
293- int64_t,293+ int64_t,
294- bool,294+ bool,
295- bool,295+ bool,
296- std::unordered_map<size_t, std::string>,296+ std::unordered_map<size_t, std::string>,
297- int64_t>(),297+ int64_t>(),
298- py::arg("params"),298+ py::arg("params"),
299- py::arg("bucket_indices"),299+ py::arg("bucket_indices"),
300- py::arg("per_bucket_size_limits"),300+ py::arg("per_bucket_size_limits"),
301- py::arg("process_group"),301+ py::arg("process_group"),
302- py::arg("expect_sparse_gradients") = std::vector<bool>(),302+ py::arg("expect_sparse_gradients") = std::vector<bool>(),
303- py::arg("bucket_bytes_cap") = ::c10d::kDefaultBucketBytesCap,303+ py::arg("bucket_bytes_cap") = ::c10d::kDefaultBucketBytesCap,
304- py::arg("find_unused_parameters") = false,304+ py::arg("find_unused_parameters") = false,
305- py::arg("gradient_as_bucket_view") = false,305+ py::arg("gradient_as_bucket_view") = false,
306- py::arg("param_to_name_mapping") =306+ py::arg("param_to_name_mapping") =
307- std::unordered_map<size_t, std::string>(),307+ std::unordered_map<size_t, std::string>(),
308- py::arg("first_bucket_bytes_cap") = ::c10d::kDefaultFirstBucketBytes,308+ py::arg("first_bucket_bytes_cap") = ::c10d::kDefaultFirstBucketBytes,
309- py::call_guard<py::gil_scoped_release>())309+ py::call_guard<py::gil_scoped_release>())
310- .def("prepare_for_forward",310+ .def("prepare_for_forward",
311- &c10d_npu::Reducer::prepare_for_forward,311+ &c10d_npu::Reducer::prepare_for_forward,
312- py::call_guard<py::gil_scoped_release>())312+ py::call_guard<py::gil_scoped_release>())
313- .def("prepare_for_backward",313+ .def("prepare_for_backward",
314- &c10d_npu::Reducer::prepare_for_backward,314+ &c10d_npu::Reducer::prepare_for_backward,
315- py::call_guard<py::gil_scoped_release>())315+ py::call_guard<py::gil_scoped_release>())
316- .def("prepare_for_backward",316+ .def("prepare_for_backward",
317- [](c10d_npu::Reducer& reducer, const at::Tensor& output)317+ [](c10d_npu::Reducer& reducer, const at::Tensor& output)
318- -> void { reducer.prepare_for_backward({output}); },318+ -> void { reducer.prepare_for_backward({output}); },
319- py::call_guard<py::gil_scoped_release>())319+ py::call_guard<py::gil_scoped_release>())
320- .def("get_backward_stats", &c10d_npu::Reducer::get_backward_stats)320+ .def("get_backward_stats", &c10d_npu::Reducer::get_backward_stats)
321- .def("_install_post_backward_futures", [](::c10d_npu::Reducer& reducer, const std::vector<std::shared_ptr<torch::jit::PythonFutureWrapper>>& futs) {321+ .def("_install_post_backward_futures", [](::c10d_npu::Reducer& reducer, const std::vector<std::shared_ptr<torch::jit::PythonFutureWrapper>>& futs) {
322- c10::List<c10::intrusive_ptr<c10::ivalue::Future>> futures(c10::FutureType::create(c10::TensorType::get()));322+ c10::List<c10::intrusive_ptr<c10::ivalue::Future>> futures(c10::FutureType::create(c10::TensorType::get()));
323- for (const auto &fut : futs) {323+ for (const auto &fut : futs) {
324- futures.push_back(fut->fut);324+ futures.push_back(fut->fut);
325- }325+ }
326- reducer.install_futures(std::move(futures));326+ reducer.install_futures(std::move(futures));
327- },327+ },
328- py::call_guard<py::gil_scoped_release>())328+ py::call_guard<py::gil_scoped_release>())
329- .def("_rebuild_buckets",329+ .def("_rebuild_buckets",
330- &::c10d_npu::Reducer::rebuild_buckets,330+ &::c10d_npu::Reducer::rebuild_buckets,
331- py::call_guard<py::gil_scoped_release>())331+ py::call_guard<py::gil_scoped_release>())
332- .def("_get_zeros_like_grad_buckets",332+ .def("_get_zeros_like_grad_buckets",
333- [](::c10d_npu::Reducer& reducer) {333+ [](::c10d_npu::Reducer& reducer) {
334- return reducer.get_grad_buckets(true);334+ return reducer.get_grad_buckets(true);
335- },335+ },
336- py::call_guard<py::gil_scoped_release>())336+ py::call_guard<py::gil_scoped_release>())
337- .def("_push_all_rebuilt_params",337+ .def("_push_all_rebuilt_params",
338- &::c10d_npu::Reducer::push_rebuilt_params_for_all_indices,338+ &::c10d_npu::Reducer::push_rebuilt_params_for_all_indices,
339- py::call_guard<py::gil_scoped_release>())339+ py::call_guard<py::gil_scoped_release>())
340- .def("_set_forward_pass_work_handle",340+ .def("_set_forward_pass_work_handle",
341- &::c10d_npu::Reducer::set_forward_pass_work_handle,341+ &::c10d_npu::Reducer::set_forward_pass_work_handle,
342- py::call_guard<py::gil_scoped_release>())342+ py::call_guard<py::gil_scoped_release>())
343- .def("_get_local_used_map",343+ .def("_get_local_used_map",
344- &::c10d_npu::Reducer::get_local_used_map_on_device)344+ &::c10d_npu::Reducer::get_local_used_map_on_device)
345- .def("_set_ddp_runtime_logging_sample_rate",345+ .def("_set_ddp_runtime_logging_sample_rate",
346- &::c10d_npu::Reducer::set_ddp_runtime_logging_sample_rate,346+ &::c10d_npu::Reducer::set_ddp_runtime_logging_sample_rate,
347- py::arg("sample_rate"),347+ py::arg("sample_rate"),
348- py::call_guard<py::gil_scoped_release>())348+ py::call_guard<py::gil_scoped_release>())
349- .def("_set_static_graph",349+ .def("_set_static_graph",
350- &::c10d_npu::Reducer::set_static_graph,350+ &::c10d_npu::Reducer::set_static_graph,
351- py::call_guard<py::gil_scoped_release>())351+ py::call_guard<py::gil_scoped_release>())
352- .def("_ddp_graph_static",352+ .def("_ddp_graph_static",
353- &::c10d_npu::Reducer::ddp_graph_static,353+ &::c10d_npu::Reducer::ddp_graph_static,
354- py::call_guard<py::gil_scoped_release>())354+ py::call_guard<py::gil_scoped_release>())
355- .def("_delay_all_reduce",355+ .def("_delay_all_reduce",
356- &::c10d_npu::Reducer::delay_all_reduce,356+ &::c10d_npu::Reducer::delay_all_reduce,
357- py::call_guard<py::gil_scoped_release>())357+ py::call_guard<py::gil_scoped_release>())
358- .def("_run_comm_hook",358+ .def("_run_comm_hook",
359- [](::c10d_npu::Reducer& reducer, ::c10d::GradBucket& bucket)359+ [](::c10d_npu::Reducer& reducer, ::c10d::GradBucket& bucket)
360- -> std::shared_ptr<torch::jit::PythonFutureWrapper> {360+ -> std::shared_ptr<torch::jit::PythonFutureWrapper> {
361- c10::intrusive_ptr<c10::ivalue::Future> fut =361+ c10::intrusive_ptr<c10::ivalue::Future> fut =
362- reducer.run_comm_hook(bucket);362+ reducer.run_comm_hook(bucket);
363- return std::make_shared<torch::jit::PythonFutureWrapper>(fut);363+ return std::make_shared<torch::jit::PythonFutureWrapper>(fut);
364- },364+ },
365- py::call_guard<py::gil_scoped_release>())365+ py::call_guard<py::gil_scoped_release>())
366- .def("set_logger",366+ .def("set_logger",
367- [](::c10d_npu::Reducer& reducer,367+ [](::c10d_npu::Reducer& reducer,
368- const std::shared_ptr<::c10d::Logger> logger) {368+ const std::shared_ptr<::c10d::Logger> logger) {
369- std::weak_ptr<::c10d::Logger> logger_weakref = logger;369+ std::weak_ptr<::c10d::Logger> logger_weakref = logger;
370- reducer.set_logger(logger_weakref);370+ reducer.set_logger(logger_weakref);
371- });371+ });
372- 372+ 
373- py::module_ dist = py::module_::import("torch._C._distributed_c10d");373+ py::module_ dist = py::module_::import("torch._C._distributed_c10d");
374- auto processGroupHCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupHCCL>(374+ auto processGroupHCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupHCCL>(
375- module, "ProcessGroupHCCL", dist.attr("Backend"))375+ module, "ProcessGroupHCCL", dist.attr("Backend"))
376- .def(py::init<const c10::intrusive_ptr<::c10d::Store>&,376+ .def(py::init<const c10::intrusive_ptr<::c10d::Store>&,
377- int,377+ int,
378- int,378+ int,
379- c10::intrusive_ptr<::c10d_npu::ProcessGroupHCCL::Options>>(),379+ c10::intrusive_ptr<::c10d_npu::ProcessGroupHCCL::Options>>(),
380- py::call_guard<py::gil_scoped_release>())380+ py::call_guard<py::gil_scoped_release>())
381- .def(py::init([](const c10::intrusive_ptr<::c10d::Store>& store,381+ .def(py::init([](const c10::intrusive_ptr<::c10d::Store>& store,
382- int rank,382+ int rank,
383- int size,383+ int size,
384- const std::chrono::milliseconds& timeout) {384+ const std::chrono::milliseconds& timeout) {
385- auto options = ::c10d_npu::ProcessGroupHCCL::Options::create();385+ auto options = ::c10d_npu::ProcessGroupHCCL::Options::create();
386- options->is_high_priority_stream = false;386+ options->is_high_priority_stream = false;
387- options->timeout = timeout;387+ options->timeout = timeout;
388- return c10::make_intrusive<::c10d_npu::ProcessGroupHCCL>(388+ return c10::make_intrusive<::c10d_npu::ProcessGroupHCCL>(
389- store, rank, size, options);389+ store, rank, size, options);
390- }),390+ }),
391- py::arg("store"),391+ py::arg("store"),
392- py::arg("rank"),392+ py::arg("rank"),
393- py::arg("size"),393+ py::arg("size"),
394- py::arg("timeout") = kProcessGroupDefaultTimeout,394+ py::arg("timeout") = kProcessGroupDefaultTimeout,
395- py::call_guard<py::gil_scoped_release>())395+ py::call_guard<py::gil_scoped_release>())
396- .def("get_hccl_comm", &::c10d_npu::ProcessGroupHCCL::getHcclComm)396+ .def("get_hccl_comm", &::c10d_npu::ProcessGroupHCCL::getHcclComm)
397- .def("_set_hccl_comm_name", &::c10d_npu::ProcessGroupHCCL::setHcclCommName)397+ .def("_set_hccl_comm_name", &::c10d_npu::ProcessGroupHCCL::setHcclCommName)
398- .def("resume_hccl_comm", &::c10d_npu::ProcessGroupHCCL::resumeHcclComm)398+ .def("resume_hccl_comm", &::c10d_npu::ProcessGroupHCCL::resumeHcclComm)
399- .def("_set_switch_nic_comm",399+ .def("_set_switch_nic_comm",
400- &::c10d_npu::ProcessGroupHCCL::setSwitchNicComm,400+ &::c10d_npu::ProcessGroupHCCL::setSwitchNicComm,
401- py::arg("rankid"),401+ py::arg("rankid"),
402- py::arg("nRanks"),402+ py::arg("nRanks"),
403- py::arg("ranks") = std::vector<uint32_t>{},403+ py::arg("ranks") = std::vector<uint32_t>{},
404- py::arg("useBackup") = std::vector<bool>{})404+ py::arg("useBackup") = std::vector<bool>{})
405- .def("abort_hccl_comm", &::c10d_npu::ProcessGroupHCCL::abortAndClearHcclComm)405+ .def("abort_hccl_comm", &::c10d_npu::ProcessGroupHCCL::abortAndClearHcclComm)
406- .def("_delete_tcpstore_key", &::c10d_npu::ProcessGroupHCCL::deleteTCPStoreKey)406+ .def("_delete_tcpstore_key", &::c10d_npu::ProcessGroupHCCL::deleteTCPStoreKey)
407- .def("set_watchdog_status", &::c10d_npu::ProcessGroupHCCL::setWatchdogStatus)407+ .def("set_watchdog_status", &::c10d_npu::ProcessGroupHCCL::setWatchdogStatus)
408- .def("clear_workmeta_list", &::c10d_npu::ProcessGroupHCCL::clearWorkMetaList)408+ .def("clear_workmeta_list", &::c10d_npu::ProcessGroupHCCL::clearWorkMetaList)
409- .def("get_hccl_comm_name",409+ .def("get_hccl_comm_name",
410- [](::c10d_npu::ProcessGroupHCCL &pg, int rankid, py::args args, py::kwargs kwargs)410+ [](::c10d_npu::ProcessGroupHCCL &pg, int rankid, py::args args, py::kwargs kwargs)
411- -> std::string {411+ -> std::string {
412- bool init_comm = true;412+ bool init_comm = true;
413- if (kwargs.contains("init_comm")) {413+ if (kwargs.contains("init_comm")) {
414- init_comm = py::cast<bool>(kwargs["init_comm"]);414+ init_comm = py::cast<bool>(kwargs["init_comm"]);
415- }415+ }
416- return pg.getHcclCommName(rankid, init_comm);416+ return pg.getHcclCommName(rankid, init_comm);
417- })417+ })
418- .def("_get_stream_id", &::c10d_npu::ProcessGroupHCCL::getStreamId,418+ .def("_get_stream_id", &::c10d_npu::ProcessGroupHCCL::getStreamId,
419- py::arg("p2p") = false,419+ py::arg("p2p") = false,
420- py::arg("peer") = -1)420+ py::arg("peer") = -1)
421- .def("get_coll_stream_id", &::c10d_npu::ProcessGroupHCCL::getCollNpuStreamId,421+ .def("get_coll_stream_id", &::c10d_npu::ProcessGroupHCCL::getCollNpuStreamId,
422- py::arg("device"))422+ py::arg("device"))
423- .def("get_p2p_stream_id", &::c10d_npu::ProcessGroupHCCL::getP2PStreamId,423+ .def("get_p2p_stream_id", &::c10d_npu::ProcessGroupHCCL::getP2PStreamId,
424- py::arg("device"),424+ py::arg("device"),
425- py::arg("peer"),425+ py::arg("peer"),
426- py::arg("is_batched"))426+ py::arg("is_batched"))
427- .def("_window_register_and_exchange", &::c10d_npu::ProcessGroupHCCL::windowRegisterAndExchange,427+ .def("_window_register_and_exchange", &::c10d_npu::ProcessGroupHCCL::windowRegisterAndExchange,
428- py::arg("window_size"),428+ py::arg("window_size"),
429- py::arg("peer_ranks"))429+ py::arg("peer_ranks"))
430- .def("_get_window_mem", &::c10d_npu::ProcessGroupHCCL::getWindowMem)430+ .def("_get_window_mem", &::c10d_npu::ProcessGroupHCCL::getWindowMem)
431- .def_property_readonly("options", &::c10d_npu::ProcessGroupHCCL::getOptions)431+ .def_property_readonly("options", &::c10d_npu::ProcessGroupHCCL::getOptions)
432- .def("batch_isend_irecv",432+ .def("batch_isend_irecv",
433- [](::c10d_npu::ProcessGroupHCCL &pg, std::vector<std::string> &op_type,433+ [](::c10d_npu::ProcessGroupHCCL &pg, std::vector<std::string> &op_type,
434- std::vector<at::Tensor> &tensors,434+ std::vector<at::Tensor> &tensors,
435- std::vector<int64_t> remote_rank_list)435+ std::vector<int64_t> remote_rank_list)
436- -> c10::intrusive_ptr<c10d::Work> {436+ -> c10::intrusive_ptr<c10d::Work> {
437- return pg.batch_isend_irecv(op_type, tensors, remote_rank_list);437+ return pg.batch_isend_irecv(op_type, tensors, remote_rank_list);
438- },438+ },
439- py::call_guard<py::gil_scoped_release>())439+ py::call_guard<py::gil_scoped_release>())
440- .def("reduce_scatter_tensor_uneven",440+ .def("reduce_scatter_tensor_uneven",
441- &::c10d_npu::ProcessGroupHCCL::_reduce_scatter_base_uneven,441+ &::c10d_npu::ProcessGroupHCCL::_reduce_scatter_base_uneven,
442- py::arg("output"),442+ py::arg("output"),
443- py::arg("input"),443+ py::arg("input"),
444- py::arg("input_split_sizes") = std::vector<int64_t>{},444+ py::arg("input_split_sizes") = std::vector<int64_t>{},
445- py::arg("opts") = ::c10d::ReduceScatterOptions(),445+ py::arg("opts") = ::c10d::ReduceScatterOptions(),
446- py::call_guard<py::gil_scoped_release>())446+ py::call_guard<py::gil_scoped_release>())
447- .def("all_gather_into_tensor_uneven",447+ .def("all_gather_into_tensor_uneven",
448- &::c10d_npu::ProcessGroupHCCL::_allgather_base_uneven,448+ &::c10d_npu::ProcessGroupHCCL::_allgather_base_uneven,
449- py::arg("output"),449+ py::arg("output"),
450- py::arg("input"),450+ py::arg("input"),
451- py::arg("output_split_sizes") = std::vector<int64_t>{},451+ py::arg("output_split_sizes") = std::vector<int64_t>{},
452- py::arg("opts") = ::c10d::AllgatherOptions(),452+ py::arg("opts") = ::c10d::AllgatherOptions(),
453- py::call_guard<py::gil_scoped_release>())453+ py::call_guard<py::gil_scoped_release>())
454- .def("_set_default_timeout",454+ .def("_set_default_timeout",
455- &::c10d_npu::ProcessGroupHCCL::setTimeout,455+ &::c10d_npu::ProcessGroupHCCL::setTimeout,
456- py::arg("timeout"),456+ py::arg("timeout"),
457- py::call_guard<py::gil_scoped_release>())457+ py::call_guard<py::gil_scoped_release>())
458- .def("_add_ephemeral_timeout",458+ .def("_add_ephemeral_timeout",
459- &::c10d_npu::ProcessGroupHCCL::addEphemeralTimeout,459+ &::c10d_npu::ProcessGroupHCCL::addEphemeralTimeout,
460- py::arg("timeout"),460+ py::arg("timeout"),
461- py::call_guard<py::gil_scoped_release>());461+ py::call_guard<py::gil_scoped_release>());
462- 462+ 
463- intrusive_ptr_class_<::c10d_npu::ProcessGroupHCCL::Options>(463+ intrusive_ptr_class_<::c10d_npu::ProcessGroupHCCL::Options>(
464- processGroupHCCL,464+ processGroupHCCL,
465- "Options",465+ "Options",
466- dist.attr("Backend").attr("Options"))466+ dist.attr("Backend").attr("Options"))
467- .def(py::init<>())467+ .def(py::init<>())
468- .def_readwrite("op_timeout", &::c10d_npu::ProcessGroupHCCL::Options::opTimeout)468+ .def_readwrite("op_timeout", &::c10d_npu::ProcessGroupHCCL::Options::opTimeout)
469- .def_readwrite("is_high_priority_stream",469+ .def_readwrite("is_high_priority_stream",
470- &::c10d_npu::ProcessGroupHCCL::Options::is_high_priority_stream)470+ &::c10d_npu::ProcessGroupHCCL::Options::is_high_priority_stream)
471- .def_readwrite("global_ranks_in_group",471+ .def_readwrite("global_ranks_in_group",
472- &::c10d_npu::ProcessGroupHCCL::Options::global_ranks_in_group)472+ &::c10d_npu::ProcessGroupHCCL::Options::global_ranks_in_group)
473- .def_readwrite("hccl_config", &::c10d_npu::ProcessGroupHCCL::Options::hccl_config)473+ .def_readwrite("hccl_config", &::c10d_npu::ProcessGroupHCCL::Options::hccl_config)
474- .def_readwrite("group_id",474+ .def_readwrite("group_id",
475- &::c10d_npu::ProcessGroupHCCL::Options::group_id);475+ &::c10d_npu::ProcessGroupHCCL::Options::group_id);
476- 476+
477- // bind for ProcessGroupLCCL477+ // bind for ProcessGroupLCCL
478- auto processGroupLCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupLCCL>(478+ auto processGroupLCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupLCCL>(
479- module, "ProcessGroupLCCL", dist.attr("Backend"))479+ module, "ProcessGroupLCCL", dist.attr("Backend"))
480- .def(py::init<const c10::intrusive_ptr<::c10d::Store>&, int, int>(),480+ .def(py::init<const c10::intrusive_ptr<::c10d::Store>&, int, int>(),
481- py::call_guard<py::gil_scoped_release>());481+ py::call_guard<py::gil_scoped_release>());
482- 482+ 
483- auto cDist = py::module_::import("torch._C._distributed_c10d");483+ auto cDist = py::module_::import("torch._C._distributed_c10d");
484- auto parallelStore = intrusive_ptr_no_gil_destructor_class_<::c10d::ParallelTcpStore>(484+ auto parallelStore = intrusive_ptr_no_gil_destructor_class_<::c10d::ParallelTcpStore>(
485- module, "ParallelStore", cDist.attr("Store"), R"(485+ module, "ParallelStore", cDist.attr("Store"), R"(
486-A TCP-Parallel-Epoll-based distributed key-value store implementation. The server store holds486+A TCP-Parallel-Epoll-based distributed key-value store implementation. The server store holds
487-the data, while the client stores can connect to the server store over TCP and487+the data, while the client stores can connect to the server store over TCP and
488-perform actions such as :meth:`~torch.distributed.store.set` to insert a key-value488+perform actions such as :meth:`~torch.distributed.store.set` to insert a key-value
489-pair, :meth:`~torch.distributed.store.get` to retrieve a key-value pair, etc. There489+pair, :meth:`~torch.distributed.store.get` to retrieve a key-value pair, etc. There
490-should always be one server store initialized because the client store(s) will wait for490+should always be one server store initialized because the client store(s) will wait for
491-the server to establish a connection.491+the server to establish a connection.
492- 492+ 
493-Arguments:493+Arguments:
494- host_name (str): The hostname or IP Address the server store should run on.494+ host_name (str): The hostname or IP Address the server store should run on.
495- port (int): The port on which the server store should listen for incoming requests.495+ port (int): The port on which the server store should listen for incoming requests.
496- world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates a non-fixed number of store users).496+ world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates a non-fixed number of store users).
497- agentRun(bool): The client(worker), agentRun is False. The agent(proxy), agentRun is True.497+ agentRun(bool): The client(worker), agentRun is False. The agent(proxy), agentRun is True.
498- agentPid(int): Generally, a single `torch_run` is launched on a node. If multiple `torch_run` are launched on a node, the agentPid refers to the process ID (PID) of each `torch_run`.498+ agentPid(int): Generally, a single `torch_run` is launched on a node. If multiple `torch_run` are launched on a node, the agentPid refers to the process ID (PID) of each `torch_run`.
499- The pid transmit to the worker through environment variable and is used for local socket communication.499+ The pid transmit to the worker through environment variable and is used for local socket communication.
500- is_master (bool, optional): True when initializing the server store and False for client stores. Default is False.500+ is_master (bool, optional): True when initializing the server store and False for client stores. Default is False.
501- enableTiered(bool, optional): parallel tcpstore tiered optimization, if True, The agent adds a proxy role, the worker on the node connects to the proxy via Unix Domain Socket.501+ enableTiered(bool, optional): parallel tcpstore tiered optimization, if True, The agent adds a proxy role, the worker on the node connects to the proxy via Unix Domain Socket.
502- and the proxy connects to the server via TCP, completing the establishment and communication of the connection., Default is False.502+ and the proxy connects to the server via TCP, completing the establishment and communication of the connection., Default is False.
503- timeout (timedelta, optional): Timeout used by the store during initialization and for methods such as :meth:`~torch.distributed.store.get` and :meth:`~torch.distributed.store.wait`. Default is timedelta(seconds=300)503+ timeout (timedelta, optional): Timeout used by the store during initialization and for methods such as :meth:`~torch.distributed.store.get` and :meth:`~torch.distributed.store.wait`. Default is timedelta(seconds=300)
504- wait_for_worker (bool, optional): Whether to wait for all the workers to connect with the server store. This is only applicable when world_size is a fixed value. Default is True.504+ wait_for_worker (bool, optional): Whether to wait for all the workers to connect with the server store. This is only applicable when world_size is a fixed value. Default is True.
505- 505+ 
506---enable_tiered_parallel_tcpstore = "false":506+--enable_tiered_parallel_tcpstore = "false":
507-Example::507+Example::
508- >>> import torch_npu.distributed as dist508+ >>> import torch_npu.distributed as dist
509- >>> from datetime import timedelta509+ >>> from datetime import timedelta
510- >>> # Run on process 1 (server)510+ >>> # Run on process 1 (server)
511- >>> server_store = dist.ParallelStore("127.0.0.1", 1234, 2, True, 100, True, timedelta(seconds=30))511+ >>> server_store = dist.ParallelStore("127.0.0.1", 1234, 2, True, 100, True, timedelta(seconds=30))
512- >>> # Run on process 2 (client)512+ >>> # Run on process 2 (client)
513- >>> client_store = dist.ParallelStore("127.0.0.1", 1234, 2, False, 100, False)513+ >>> client_store = dist.ParallelStore("127.0.0.1", 1234, 2, False, 100, False)
514- >>> # Use any of the store methods from either the client or server after initialization514+ >>> # Use any of the store methods from either the client or server after initialization
515- >>> server_store.set("first_key", "first_value")515+ >>> server_store.set("first_key", "first_value")
516- >>> client_store.get("first_key")516+ >>> client_store.get("first_key")
517- 517+ 
518---enable_tiered_parallel_tcpstore = "true":518+--enable_tiered_parallel_tcpstore = "true":
519-Example::519+Example::
520- >>> import torch_npu.distributed as dist520+ >>> import torch_npu.distributed as dist
521- >>> from datetime import timedelta521+ >>> from datetime import timedelta
522- >>> # Run on process 1 (server proxy)522+ >>> # Run on process 1 (server proxy)
523- >>> server_store = dist.ParallelStore("127.0.0.1", 1234, 2, True, 100, True, True, timedelta(seconds=30))523+ >>> server_store = dist.ParallelStore("127.0.0.1", 1234, 2, True, 100, True, True, timedelta(seconds=30))
524- >>> # Run on process 2 (client)524+ >>> # Run on process 2 (client)
525- >>> client_store = dist.ParallelStore("127.0.0.1", 1234, 2, False, 100, False, True)525+ >>> client_store = dist.ParallelStore("127.0.0.1", 1234, 2, False, 100, False, True)
526- >>> # Use any of the store methods from either the client or server and proxy after initialization526+ >>> # Use any of the store methods from either the client or server and proxy after initialization
527- >>> server_store.set("first_key", "first_value")527+ >>> server_store.set("first_key", "first_value")
528- >>> client_store.get("first_key")528+ >>> client_store.get("first_key")
529- )")529+ )")
530- 530+ 
531- .def(py::init([](const std::string &host,531+ .def(py::init([](const std::string &host,
532- uint16_t port,532+ uint16_t port,
533- int worldSize,533+ int worldSize,
534- bool agentRun,534+ bool agentRun,
535- uint32_t agentPid,535+ uint32_t agentPid,
536- bool isServer,536+ bool isServer,
537- bool enableTiered,537+ bool enableTiered,
538- std::chrono::milliseconds timeout,538+ std::chrono::milliseconds timeout,
539- bool waitWorkers,539+ bool waitWorkers,
540- bool multiTenant) {540+ bool multiTenant) {
541- c10::optional<std::size_t> numWorkers = c10::nullopt;541+ c10::optional<std::size_t> numWorkers = c10::nullopt;
542- if (worldSize > -1) {542+ if (worldSize > -1) {
543- numWorkers = static_cast<std::size_t>(worldSize);543+ numWorkers = static_cast<std::size_t>(worldSize);
544- }544+ }
545- ::c10d::TCPStoreOptions opts{ port, isServer, numWorkers, waitWorkers, timeout, multiTenant };545+ ::c10d::TCPStoreOptions opts{ port, isServer, numWorkers, waitWorkers, timeout, multiTenant };
546- return c10::make_intrusive <::c10d::ParallelTcpStore>(host, agentRun, agentPid, enableTiered, opts);546+ return c10::make_intrusive <::c10d::ParallelTcpStore>(host, agentRun, agentPid, enableTiered, opts);
547- }),547+ }),
548- py::arg("host") = "127.0.0.1",548+ py::arg("host") = "127.0.0.1",
549- py::arg("port") = 29500,549+ py::arg("port") = 29500,
550- py::arg("world_size") = -1,550+ py::arg("world_size") = -1,
551- py::arg("agent_run") = false,551+ py::arg("agent_run") = false,
552- py::arg("agent_pid") = -1,552+ py::arg("agent_pid") = -1,
553- py::arg("is_server") = false,553+ py::arg("is_server") = false,
554- py::arg("enable_tiered") = false,554+ py::arg("enable_tiered") = false,
555- py::arg("timeout") = std::chrono::milliseconds(300000),555+ py::arg("timeout") = std::chrono::milliseconds(300000),
556- py::arg("wait_workers") = true,556+ py::arg("wait_workers") = true,
557- py::arg("multi_tenant") = false);557+ py::arg("multi_tenant") = false);
558- 558+ 
559- module.def("_dump_hccl_trace_json",559+ module.def("_dump_hccl_trace_json",
560- [](std::optional<bool> includeCollectives,560+ [](std::optional<bool> includeCollectives,
561- std::optional<bool> onlyActive) {561+ std::optional<bool> onlyActive) {
562- return py::bytes(::c10d_npu::dump_hccl_trace_json(562+ return py::bytes(::c10d_npu::dump_hccl_trace_json(
563- includeCollectives.value_or(true), onlyActive.value_or(false)));563+ includeCollectives.value_or(true), onlyActive.value_or(false)));
564- },564+ },
565- py::arg("includeCollectives") = std::optional<bool>(),565+ py::arg("includeCollectives") = std::optional<bool>(),
566- py::arg("onlyActive") = std::optional<bool>(),566+ py::arg("onlyActive") = std::optional<bool>(),
567- R"(567+ R"(
568- Arguments:568+ Arguments:
569- includeCollectives(bool, optional): Whether to include collective work traces. Default is True.569+ includeCollectives(bool, optional): Whether to include collective work traces. Default is True.
570- onlyActive (bool, optional): Whether to only include active collective work traces. Default is False.570+ onlyActive (bool, optional): Whether to only include active collective work traces. Default is False.
571- Returns:571+ Returns:
572- Stringified json work traces.572+ Stringified json work traces.
573- Default settings return everything - i.e. contains HCCL comm dumps and collective traces.573+ Default settings return everything - i.e. contains HCCL comm dumps and collective traces.
574- )");574+ )");
575- module.def("_dump_hccl_trace",575+ module.def("_dump_hccl_trace",
576- [](std::optional<bool> includeCollectives,576+ [](std::optional<bool> includeCollectives,
577- std::optional<bool> includeStackTraces,577+ std::optional<bool> includeStackTraces,
578- std::optional<bool> onlyActive) {578+ std::optional<bool> onlyActive) {
579- return py::bytes(::c10d_npu::dump_hccl_trace(579+ return py::bytes(::c10d_npu::dump_hccl_trace(
580- includeCollectives.value_or(true),580+ includeCollectives.value_or(true),
581- includeStackTraces.value_or(true),581+ includeStackTraces.value_or(true),
582- onlyActive.value_or(false)));582+ onlyActive.value_or(false)));
583- },583+ },
584- py::arg("includeCollectives") = std::optional<bool>(),584+ py::arg("includeCollectives") = std::optional<bool>(),
585- py::arg("includeStackTraces") = std::optional<bool>(),585+ py::arg("includeStackTraces") = std::optional<bool>(),
586- py::arg("onlyActive") = std::optional<bool>(),586+ py::arg("onlyActive") = std::optional<bool>(),
587- R"(587+ R"(
588- Arguments:588+ Arguments:
589- includeCollectives(bool, optional): Whether to include collective work traces. Default is True.589+ includeCollectives(bool, optional): Whether to include collective work traces. Default is True.
590- includeStackTraces(bool, optional): Whether to include stacktraces in the collective work traces. Default is True.590+ includeStackTraces(bool, optional): Whether to include stacktraces in the collective work traces. Default is True.
591- onlyActive (bool, optional): Whether to only include active collective work traces. Default is False.591+ onlyActive (bool, optional): Whether to only include active collective work traces. Default is False.
592- Returns:592+ Returns:
593- Stringified pickle work traces.593+ Stringified pickle work traces.
594- Default settings return everything - i.e. contains HCCL comm dumps and collective traces.594+ Default settings return everything - i.e. contains HCCL comm dumps and collective traces.
595- )");595+ )");
596- 596+ 
597- Py_RETURN_TRUE;597+ Py_RETURN_TRUE;
598-}598+}
599- 599+ 
600-// c10d methods on torch._C600+// c10d methods on torch._C
601-static PyMethodDef methods[] = { // NOLINT601+static PyMethodDef methods[] = { // NOLINT
602- {"_c10d_npu_init", c10d_npu_init, METH_NOARGS, nullptr},602+ {"_c10d_npu_init", c10d_npu_init, METH_NOARGS, nullptr},
603-#ifdef USE_RPC_FRAMEWORK603+#ifdef USE_RPC_FRAMEWORK
604- {"_rpc_npu_init", rpc::rpc_npu_init, METH_NOARGS, nullptr},604+ {"_rpc_npu_init", rpc::rpc_npu_init, METH_NOARGS, nullptr},
605-#endif605+#endif
606- {nullptr, nullptr, 0, nullptr}};606+ {nullptr, nullptr, 0, nullptr}};
607- 607+ 
608-PyMethodDef* python_functions()608+PyMethodDef* python_functions()
609-{609+{
610- return methods;610+ return methods;
611-}611+}
612- 612+ 
613-} // namespace distributed613+} // namespace distributed
614-} // namespace torch_npu614+} // namespace torch_npu
Mtorch_npu/csrc/distributed/Init.h+11-11
@@ -1,12 +1,12 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <torch/csrc/python_headers.h>3+#include <torch/csrc/python_headers.h>
4-#include "torch_npu/csrc/core/npu/NPUMacros.h"4+#include "torch_npu/csrc/core/npu/NPUMacros.h"
5- 5+ 
6-namespace torch_npu {6+namespace torch_npu {
7-namespace distributed {7+namespace distributed {
8- 8+ 
9-TORCH_NPU_API PyMethodDef* python_functions();9+TORCH_NPU_API PyMethodDef* python_functions();
10- 10+ 
11-} // namespace distributed11+} // namespace distributed
12} // namespace torch_npu12} // namespace torch_npu
Mtorch_npu/csrc/distributed/default_comm_hooks.cpp+57-57
@@ -1,57 +1,57 @@
1-#include <c10d/default_comm_hooks.hpp>1+#include <c10d/default_comm_hooks.hpp>
2-#include <c10/core/ScalarType.h>2+#include <c10/core/ScalarType.h>
3-#include <c10/util/Exception.h>3+#include <c10/util/Exception.h>
4- 4+ 
5-#include <c10d/ProcessGroup.hpp>5+#include <c10d/ProcessGroup.hpp>
6-#include <c10d/comm.hpp>6+#include <c10d/comm.hpp>
7-#include <torch/torch.h>7+#include <torch/torch.h>
8- 8+ 
9-#include "torch_npu/csrc/core/npu/NPUException.h"9+#include "torch_npu/csrc/core/npu/NPUException.h"
10- 10+ 
11-namespace c10d {11+namespace c10d {
12- 12+ 
13-c10::intrusive_ptr<c10::ivalue::Future> AllReduceCommHook::runHook(13+c10::intrusive_ptr<c10::ivalue::Future> AllReduceCommHook::runHook(
14- GradBucket& bucket)14+ GradBucket& bucket)
15-{15+{
16- std::vector<at::Tensor> tensors = {bucket.getBufferRef()};16+ std::vector<at::Tensor> tensors = {bucket.getBufferRef()};
17- // Apply the division first to avoid overflow, especially for FP16.17+ // Apply the division first to avoid overflow, especially for FP16.
18- tensors[0] /= state_->getSize();18+ tensors[0] /= state_->getSize();
19- return state_->allreduce(tensors)->getFuture();19+ return state_->allreduce(tensors)->getFuture();
20-}20+}
21- 21+ 
22-c10::intrusive_ptr<c10::ivalue::Future> FP16CompressCommHook::runHook(22+c10::intrusive_ptr<c10::ivalue::Future> FP16CompressCommHook::runHook(
23- GradBucket& bucket)23+ GradBucket& bucket)
24-{24+{
25- auto compressed_tensor = bucket.getBufferRef().to(torch::kFloat16);25+ auto compressed_tensor = bucket.getBufferRef().to(torch::kFloat16);
26- // Apply the division first to avoid overflow.26+ // Apply the division first to avoid overflow.
27- compressed_tensor /= state_->getSize();27+ compressed_tensor /= state_->getSize();
28- std::vector<at::Tensor> tensors = {compressed_tensor};28+ std::vector<at::Tensor> tensors = {compressed_tensor};
29- 29+ 
30- auto allreduce_fut = state_->allreduce(tensors)->getFuture();30+ auto allreduce_fut = state_->allreduce(tensors)->getFuture();
31- auto decompressed_tensor = bucket.getBufferRef();31+ auto decompressed_tensor = bucket.getBufferRef();
32- auto decompress = [decompressed_tensor](c10::ivalue::Future& allreduce_fut) {32+ auto decompress = [decompressed_tensor](c10::ivalue::Future& allreduce_fut) {
33- auto result = allreduce_fut.value();33+ auto result = allreduce_fut.value();
34- TORCH_INTERNAL_ASSERT(34+ TORCH_INTERNAL_ASSERT(
35- result.isTensorList(),35+ result.isTensorList(),
36- "ProcessGroup::allreduce should return TensorList", DIST_ERROR(ErrCode::INTERNAL));36+ "ProcessGroup::allreduce should return TensorList", DIST_ERROR(ErrCode::INTERNAL));
37- 37+ 
38- auto reduce_tensor = result.toTensorVector()[0];38+ auto reduce_tensor = result.toTensorVector()[0];
39- TORCH_INTERNAL_ASSERT_DEBUG_ONLY(39+ TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
40- reduce_tensor.scalar_type() == at::ScalarType::Half,40+ reduce_tensor.scalar_type() == at::ScalarType::Half,
41- "Expected reduced tensor to be fp16 in FP16CompressHook, but got type ",41+ "Expected reduced tensor to be fp16 in FP16CompressHook, but got type ",
42- reduce_tensor.scalar_type(), DIST_ERROR(ErrCode::TYPE)42+ reduce_tensor.scalar_type(), DIST_ERROR(ErrCode::TYPE)
43- );43+ );
44- decompressed_tensor.copy_(reduce_tensor);44+ decompressed_tensor.copy_(reduce_tensor);
45- return c10::IValue(decompressed_tensor);45+ return c10::IValue(decompressed_tensor);
46- };46+ };
47- 47+ 
48- return allreduce_fut->then(decompress, allreduce_fut->elementType());48+ return allreduce_fut->then(decompress, allreduce_fut->elementType());
49-}49+}
50- 50+ 
51-c10::intrusive_ptr<c10::ivalue::Future> _AllReduceBySumCommHook::runHook(GradBucket& bucket)51+c10::intrusive_ptr<c10::ivalue::Future> _AllReduceBySumCommHook::runHook(GradBucket& bucket)
52-{52+{
53- std::vector<at::Tensor> tensors = {bucket.getBufferRef()};53+ std::vector<at::Tensor> tensors = {bucket.getBufferRef()};
54- return state_->allreduce(tensors)->getFuture();54+ return state_->allreduce(tensors)->getFuture();
55-}55+}
56- 56+ 
57-} // namespace c10d57+} // namespace c10d
Mtorch_npu/csrc/distributed/default_comm_hooks.hpp+48-48
@@ -1,48 +1,48 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <c10d/ProcessGroup.hpp>3+#include <c10d/ProcessGroup.hpp>
4-#include <c10d/comm.hpp>4+#include <c10d/comm.hpp>
5- 5+ 
6-namespace c10d {6+namespace c10d {
7- 7+ 
8-enum class BuiltinCommHookType {8+enum class BuiltinCommHookType {
9- ALLREDUCE = 1,9+ ALLREDUCE = 1,
10- FP16_COMPRESS = 2,10+ FP16_COMPRESS = 2,
11-};11+};
12- 12+ 
13-class AllReduceCommHook : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {13+class AllReduceCommHook : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
14-public:14+public:
15- explicit AllReduceCommHook(c10::intrusive_ptr<ProcessGroup> state)15+ explicit AllReduceCommHook(c10::intrusive_ptr<ProcessGroup> state)
16- : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}16+ : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
17- 17+ 
18- ~AllReduceCommHook() override = default;18+ ~AllReduceCommHook() override = default;
19- 19+ 
20- c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;20+ c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
21-};21+};
22- 22+ 
23-class FP16CompressCommHook : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {23+class FP16CompressCommHook : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
24-public:24+public:
25- explicit FP16CompressCommHook(c10::intrusive_ptr<ProcessGroup> state)25+ explicit FP16CompressCommHook(c10::intrusive_ptr<ProcessGroup> state)
26- : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}26+ : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
27- 27+ 
28- ~FP16CompressCommHook() override = default;28+ ~FP16CompressCommHook() override = default;
29- 29+ 
30- c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;30+ c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
31-};31+};
32- 32+ 
33-// Almost same as AllReduceCommHook, but without division inside the hook.33+// Almost same as AllReduceCommHook, but without division inside the hook.
34-// This enables the optimization of fusing copy and division and saves one scan34+// This enables the optimization of fusing copy and division and saves one scan
35-// over all the input parameters, when no communication hook is provided by the user.35+// over all the input parameters, when no communication hook is provided by the user.
36-// Only used internally and not released as a public built-in communication hook.36+// Only used internally and not released as a public built-in communication hook.
37-class _AllReduceBySumCommHook37+class _AllReduceBySumCommHook
38- : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {38+ : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
39-public:39+public:
40- explicit _AllReduceBySumCommHook(c10::intrusive_ptr<ProcessGroup> state)40+ explicit _AllReduceBySumCommHook(c10::intrusive_ptr<ProcessGroup> state)
41- : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}41+ : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
42- 42+ 
43- ~_AllReduceBySumCommHook() override = default;43+ ~_AllReduceBySumCommHook() override = default;
44- 44+ 
45- c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;45+ c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
46-};46+};
47- 47+ 
48-} // namespace c10d48+} // namespace c10d
Mtorch_npu/csrc/distributed/reducer.cpp+2161-2161
Mtorch_npu/csrc/distributed/reducer.hpp+622-622
@@ -1,622 +1,622 @@
1-#pragma once1+#pragma once
2- 2+ 
3-#include <atomic>3+#include <atomic>
4-#include <memory>4+#include <memory>
5-#include <mutex>5+#include <mutex>
6-#include <tuple>6+#include <tuple>
7-#include <unordered_map>7+#include <unordered_map>
8-#include <vector>8+#include <vector>
9- 9+ 
10-#include <ATen/core/ivalue_inl.h>10+#include <ATen/core/ivalue_inl.h>
11-#include <c10/macros/Macros.h>11+#include <c10/macros/Macros.h>
12-#include <c10/util/ApproximateClock.h>12+#include <c10/util/ApproximateClock.h>
13-#include <c10/util/intrusive_ptr.h>13+#include <c10/util/intrusive_ptr.h>
14-#include <c10d/ProcessGroup.hpp>14+#include <c10d/ProcessGroup.hpp>
15-#include <c10d/Utils.hpp>15+#include <c10d/Utils.hpp>
16-#include <c10d/Work.hpp>16+#include <c10d/Work.hpp>
17-#include <c10d/comm.hpp>17+#include <c10d/comm.hpp>
18-#include <c10d/default_comm_hooks.hpp>18+#include <c10d/default_comm_hooks.hpp>
19-#include <torch/csrc/autograd/function.h>19+#include <torch/csrc/autograd/function.h>
20-#include <torch/csrc/autograd/profiler.h>20+#include <torch/csrc/autograd/profiler.h>
21-#include <torch/csrc/autograd/variable.h>21+#include <torch/csrc/autograd/variable.h>
22-#ifndef _WIN3222+#ifndef _WIN32
23-#include <torch/csrc/distributed/autograd/context/context.h>23+#include <torch/csrc/distributed/autograd/context/context.h>
24-#endif24+#endif
25-#include <c10d/logger.hpp>25+#include <c10d/logger.hpp>
26-#include <c10d/debug.h>26+#include <c10d/debug.h>
27- 27+ 
28-namespace c10d_npu {28+namespace c10d_npu {
29- 29+ 
30-constexpr int kDefaultFirstBucketBytes = int(1024 * 1024);30+constexpr int kDefaultFirstBucketBytes = int(1024 * 1024);
31-constexpr int kDefaultBucketBytesCap = int(25 * 1024 * 1024);31+constexpr int kDefaultBucketBytesCap = int(25 * 1024 * 1024);
32-// Collect runtime stats once for every kDDPRuntimeLoggingSampleRate iterations.32+// Collect runtime stats once for every kDDPRuntimeLoggingSampleRate iterations.
33-constexpr int kDDPRuntimeLoggingSampleRate = 100;33+constexpr int kDDPRuntimeLoggingSampleRate = 100;
34-constexpr int kUnsetTime = -1;34+constexpr int kUnsetTime = -1;
35- 35+ 
36-inline int64_t current_time_in_nanos()36+inline int64_t current_time_in_nanos()
37-{37+{
38- return c10::getTime();38+ return c10::getTime();
39-}39+}
40- 40+ 
41-// Forward declaration41+// Forward declaration
42-class Logger;42+class Logger;
43- 43+ 
44-class TORCH_API Timer {44+class TORCH_API Timer {
45-private:45+private:
46- // The timestamp of forward call start time in each iteration.46+ // The timestamp of forward call start time in each iteration.
47- int64_t forward_start_time = kUnsetTime;47+ int64_t forward_start_time = kUnsetTime;
48- // The timestamp of backward computation start and end time in each48+ // The timestamp of backward computation start and end time in each
49- // iteration.49+ // iteration.
50- int64_t backward_compute_start_time = kUnsetTime;50+ int64_t backward_compute_start_time = kUnsetTime;
51- int64_t backward_compute_end_time = kUnsetTime;51+ int64_t backward_compute_end_time = kUnsetTime;
52- // The timestamp of first communication call start time in each iteration.52+ // The timestamp of first communication call start time in each iteration.
53- int64_t backward_comm_start_time = kUnsetTime;53+ int64_t backward_comm_start_time = kUnsetTime;
54- // The timestamp of last communication call end time in each iteration.54+ // The timestamp of last communication call end time in each iteration.
55- int64_t backward_comm_end_time = kUnsetTime;55+ int64_t backward_comm_end_time = kUnsetTime;
56-public:56+public:
57- enum class Event {57+ enum class Event {
58- kForwardStart,58+ kForwardStart,
59- kBackwardComputeStart,59+ kBackwardComputeStart,
60- kBackwardComputeEnd,60+ kBackwardComputeEnd,
61- kBackwardCommStart,61+ kBackwardCommStart,
62- kBackwardCommEnd,62+ kBackwardCommEnd,
63- };63+ };
64- 64+ 
65- // Record the current event, i.e., mark it as having occurred now. Default65+ // Record the current event, i.e., mark it as having occurred now. Default
66- // CPU implementation.66+ // CPU implementation.
67- virtual void record(Event event) {67+ virtual void record(Event event) {
68- getTimeRef(event) = current_time_in_nanos();68+ getTimeRef(event) = current_time_in_nanos();
69- }69+ }
70- 70+ 
71- // Return the difference between when two events occurred, in nanoseconds.71+ // Return the difference between when two events occurred, in nanoseconds.
72- // Or nullopt if one of them hasn't been recorded.72+ // Or nullopt if one of them hasn't been recorded.
73- virtual c10::optional<int64_t> measureDifference(Event start, Event end) = 0;73+ virtual c10::optional<int64_t> measureDifference(Event start, Event end) = 0;
74- 74+ 
75- virtual ~Timer() = default;75+ virtual ~Timer() = default;
76- 76+ 
77- // Return host-side timestamp, or nullopt if it has not yet been recorded.77+ // Return host-side timestamp, or nullopt if it has not yet been recorded.
78- c10::optional<int64_t> getTimestamp(Event event) {78+ c10::optional<int64_t> getTimestamp(Event event) {
79- auto time = getTimeRef(event);79+ auto time = getTimeRef(event);
80- if (time == kUnsetTime) {80+ if (time == kUnsetTime) {
81- return c10::nullopt;81+ return c10::nullopt;
82- } else {82+ } else {
83- return time;83+ return time;
84- }84+ }
85- }85+ }
86- 86+ 
87- // Return host-side time member variable corresponding to the given event.87+ // Return host-side time member variable corresponding to the given event.
88- int64_t& getTimeRef(Event event) {88+ int64_t& getTimeRef(Event event) {
89- switch (event) {89+ switch (event) {
90- case Event::kForwardStart:90+ case Event::kForwardStart:
91- return forward_start_time;91+ return forward_start_time;
92- case Event::kBackwardComputeStart:92+ case Event::kBackwardComputeStart:
93- return backward_compute_start_time;93+ return backward_compute_start_time;
94- case Event::kBackwardComputeEnd:94+ case Event::kBackwardComputeEnd:
95- return backward_compute_end_time;95+ return backward_compute_end_time;
96- case Event::kBackwardCommStart:96+ case Event::kBackwardCommStart:
97- return backward_comm_start_time;97+ return backward_comm_start_time;
98- case Event::kBackwardCommEnd:98+ case Event::kBackwardCommEnd:
99- return backward_comm_end_time;99+ return backward_comm_end_time;
100- default:100+ default:
101- TORCH_INTERNAL_ASSERT(false);101+ TORCH_INTERNAL_ASSERT(false);
102- }102+ }
103- }103+ }
104-};104+};
105- 105+ 
106-// Local accumulator type for a single bucket.106+// Local accumulator type for a single bucket.
107-struct BucketAccumulator {107+struct BucketAccumulator {
108- std::vector<size_t> indices;108+ std::vector<size_t> indices;
109- size_t size = 0;109+ size_t size = 0;
110- size_t size_limit = 0;110+ size_t size_limit = 0;
111-};111+};
112- 112+ 
113-C10_DECLARE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device);113+C10_DECLARE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device);
114- 114+ 
115-class Reducer {115+class Reducer {
116-public:116+public:
117- // The constructor takes a list of variables for every model replica.117+ // The constructor takes a list of variables for every model replica.
118- // The bucket assignment for this reducer is specified as a list of118+ // The bucket assignment for this reducer is specified as a list of
119- // buckets, each of which is specified as a list of indices into the119+ // buckets, each of which is specified as a list of indices into the
120- // variables list for **a single replica** (i.e. `variables[0]`).120+ // variables list for **a single replica** (i.e. `variables[0]`).
121- explicit Reducer(121+ explicit Reducer(
122- std::vector<at::Tensor> params,122+ std::vector<at::Tensor> params,
123- std::vector<std::vector<size_t>> bucket_indices,123+ std::vector<std::vector<size_t>> bucket_indices,
124- std::vector<size_t> per_bucket_size_limits,124+ std::vector<size_t> per_bucket_size_limits,
125- c10::intrusive_ptr<c10d::ProcessGroup> process_group,125+ c10::intrusive_ptr<c10d::ProcessGroup> process_group,
126- std::vector<bool> expect_sparse_gradients,126+ std::vector<bool> expect_sparse_gradients,
127- int64_t bucket_bytes_cap,127+ int64_t bucket_bytes_cap,
128- bool find_unused_parameters,128+ bool find_unused_parameters,
129- bool gradient_as_bucket_view,129+ bool gradient_as_bucket_view,
130- std::unordered_map<size_t, std::string> paramNames,130+ std::unordered_map<size_t, std::string> paramNames,
131- int64_t first_bucket_bytes_cap);131+ int64_t first_bucket_bytes_cap);
132- 132+ 
133- ~Reducer() noexcept(false);133+ ~Reducer() noexcept(false);
134- 134+ 
135- // To (re-)initialize bucket assignment, pass a list of buckets, each135+ // To (re-)initialize bucket assignment, pass a list of buckets, each
136- // of which is specified by a list of indices in the variables list.136+ // of which is specified by a list of indices in the variables list.
137- // This function performs validation that the variables within a bucket137+ // This function performs validation that the variables within a bucket
138- // all live on the same device and have the same dimensionality.138+ // all live on the same device and have the same dimensionality.
139- void initialize_buckets(139+ void initialize_buckets(
140- std::vector<std::vector<size_t>> bucket_indices,140+ std::vector<std::vector<size_t>> bucket_indices,
141- std::vector<size_t> per_bucket_sizes);141+ std::vector<size_t> per_bucket_sizes);
142- 142+ 
143- // This function is called when the forward function has produced an output,143+ // This function is called when the forward function has produced an output,
144- // and the user wishes to reduce gradients in the backwards pass.144+ // and the user wishes to reduce gradients in the backwards pass.
145- // If they don't, and wish to accumulate gradients before reducing them,145+ // If they don't, and wish to accumulate gradients before reducing them,
146- // a call to this function can simply be omitted.146+ // a call to this function can simply be omitted.
147- void prepare_for_backward(const std::vector<at::Tensor>& outputs);147+ void prepare_for_backward(const std::vector<at::Tensor>& outputs);
148- 148+ 
149- // Called at the begginning of forward() inside DistributedDataParallel,149+ // Called at the begginning of forward() inside DistributedDataParallel,
150- // right now it caputures the starting time of forward in each iteration.150+ // right now it caputures the starting time of forward in each iteration.
151- void prepare_for_forward();151+ void prepare_for_forward();
152- 152+ 
153- // Returns the relative time in nanoseconds when gradients were ready,153+ // Returns the relative time in nanoseconds when gradients were ready,
154- // with respect to the time `prepare_for_backward` was called. The154+ // with respect to the time `prepare_for_backward` was called. The
155- // vector is for parameters for a single model replica.155+ // vector is for parameters for a single model replica.
156- std::vector<int64_t> get_backward_stats() const156+ std::vector<int64_t> get_backward_stats() const
157- {157+ {
158- return backward_stats_;158+ return backward_stats_;
159- }159+ }
160- 160+ 
161- // Registers a hook to the reducer. The hook is `CommHookInterface`161+ // Registers a hook to the reducer. The hook is `CommHookInterface`
162- // type to allow both Python and CPP hooks. This function can only162+ // type to allow both Python and CPP hooks. This function can only
163- // be called once before calling backward.163+ // be called once before calling backward.
164- // Cannot combine with the call of `register_builtin_comm_hook`.164+ // Cannot combine with the call of `register_builtin_comm_hook`.
165- void register_comm_hook(std::unique_ptr<c10d::CommHookInterface> iface);165+ void register_comm_hook(std::unique_ptr<c10d::CommHookInterface> iface);
166- 166+ 
167- // Registers a built-in C++ comm hook to the reducer. This function can only167+ // Registers a built-in C++ comm hook to the reducer. This function can only
168- // be called once before calling backward.168+ // be called once before calling backward.
169- // Cannot combine with the call of `register_comm_hook`.169+ // Cannot combine with the call of `register_comm_hook`.
170- void register_builtin_comm_hook(c10d::BuiltinCommHookType comm_hook_type);170+ void register_builtin_comm_hook(c10d::BuiltinCommHookType comm_hook_type);
171- 171+ 
172- // Runs allreduce or installed communication hook given GradBucket instance.172+ // Runs allreduce or installed communication hook given GradBucket instance.
173- c10::intrusive_ptr<c10::ivalue::Future> run_comm_hook(173+ c10::intrusive_ptr<c10::ivalue::Future> run_comm_hook(
174- c10d::GradBucket& grad_bucket);174+ c10d::GradBucket& grad_bucket);
175- 175+ 
176- // Runs default allreduce hook.176+ // Runs default allreduce hook.
177- c10::intrusive_ptr<c10::ivalue::Future> run_allreduce_hook(177+ c10::intrusive_ptr<c10::ivalue::Future> run_allreduce_hook(
178- c10d::GradBucket& grad_bucket);178+ c10d::GradBucket& grad_bucket);
179- 179+ 
180- // Returns gradient buckets in sequential order of buckets_. This is the order180+ // Returns gradient buckets in sequential order of buckets_. This is the order
181- // in which buckets are reduced across processes. If return_zero_tensors=true,181+ // in which buckets are reduced across processes. If return_zero_tensors=true,
182- // will return zero tensors of the same shape instead of the true tensors.182+ // will return zero tensors of the same shape instead of the true tensors.
183- std::vector<c10d::GradBucket> get_grad_buckets(183+ std::vector<c10d::GradBucket> get_grad_buckets(
184- bool return_zero_tensors = true) const;184+ bool return_zero_tensors = true) const;
185- 185+ 
186- // Rebuild buckets based on rebuilt_params_ and rebuilt_param_indices_186+ // Rebuild buckets based on rebuilt_params_ and rebuilt_param_indices_
187- // according to when tensors received grads in the backward pass.187+ // according to when tensors received grads in the backward pass.
188- bool rebuild_buckets();188+ bool rebuild_buckets();
189- 189+ 
190- // Install futures that should be awaited at end of backwards. Currently these190+ // Install futures that should be awaited at end of backwards. Currently these
191- // are only used by user-defined custom buffer reduction hooks, but can be generalized191+ // are only used by user-defined custom buffer reduction hooks, but can be generalized
192- // to any user-originating futures that need to be awaited.192+ // to any user-originating futures that need to be awaited.
193- void install_futures(c10::List<c10::intrusive_ptr<c10::ivalue::Future>> futs);193+ void install_futures(c10::List<c10::intrusive_ptr<c10::ivalue::Future>> futs);
194- 194+ 
195- // Returns true if we should rebuild buckets, else false. We only rebuild195+ // Returns true if we should rebuild buckets, else false. We only rebuild
196- // buckets once after the first iteration and never rebuild them if196+ // buckets once after the first iteration and never rebuild them if
197- // find_unused_parameters_.197+ // find_unused_parameters_.
198- inline bool should_rebuild_buckets() const198+ inline bool should_rebuild_buckets() const
199- {199+ {
200- return (static_graph_ || !find_unused_parameters_) && !has_rebuilt_bucket_;200+ return (static_graph_ || !find_unused_parameters_) && !has_rebuilt_bucket_;
201- }201+ }
202- 202+ 
203- // Pushes all parameters to be rebuilt.203+ // Pushes all parameters to be rebuilt.
204- void push_rebuilt_params_for_all_indices();204+ void push_rebuilt_params_for_all_indices();
205- 205+ 
206- // Creates and sets ForwardPassWorkHandle given a ProcessGroup::Work and the206+ // Creates and sets ForwardPassWorkHandle given a ProcessGroup::Work and the
207- // corresponding tensor being reduced.207+ // corresponding tensor being reduced.
208- void set_forward_pass_work_handle(208+ void set_forward_pass_work_handle(
209- c10::intrusive_ptr<c10d::Work> forwardPassWorkHandle,209+ c10::intrusive_ptr<c10d::Work> forwardPassWorkHandle,
210- bool useStaticWorldSize);210+ bool useStaticWorldSize);
211- 211+ 
212- // Retrieve on-device tensors used to track locally unused parameters. It is212+ // Retrieve on-device tensors used to track locally unused parameters. It is
213- // a tensor where index i = 1 if the Variable with that index has been used.213+ // a tensor where index i = 1 if the Variable with that index has been used.
214- at::Tensor get_local_used_map_on_device() const;214+ at::Tensor get_local_used_map_on_device() const;
215- 215+ 
216- // An function for users to set sample_rate of collecting216+ // An function for users to set sample_rate of collecting
217- // runtime stats. The time stats will be recorded for the217+ // runtime stats. The time stats will be recorded for the
218- // first 10 iterations, after 10 iteratons time stats will be218+ // first 10 iterations, after 10 iteratons time stats will be
219- // recorded once every "sample_rate" training iterations.219+ // recorded once every "sample_rate" training iterations.
220- void set_ddp_runtime_logging_sample_rate(int sample_rate);220+ void set_ddp_runtime_logging_sample_rate(int sample_rate);
221- 221+ 
222- // Specify the training graph is static.222+ // Specify the training graph is static.
223- void set_static_graph();223+ void set_static_graph();
224- 224+ 
225- // Delay all reduce to be after all gradients' calculation is complete.225+ // Delay all reduce to be after all gradients' calculation is complete.
226- void delay_all_reduce();226+ void delay_all_reduce();
227- 227+ 
228- // Weak reference to associated DDP logger. The reference is weak to avoid228+ // Weak reference to associated DDP logger. The reference is weak to avoid
229- // refcycle between reducer and logger.229+ // refcycle between reducer and logger.
230- void set_logger(std::weak_ptr<c10d::Logger> logger);230+ void set_logger(std::weak_ptr<c10d::Logger> logger);
231- 231+ 
232- // When graph is not explicitly set by user as static and has unused232+ // When graph is not explicitly set by user as static and has unused
233- // parameters, this will return whether the graph has been static until the233+ // parameters, this will return whether the graph has been static until the
234- // current iteration, which means unused params set has not changed.234+ // current iteration, which means unused params set has not changed.
235- bool ddp_graph_static();235+ bool ddp_graph_static();
236- 236+ 
237-protected:237+protected:
238- // Forward declaration.238+ // Forward declaration.
239- struct Bucket;239+ struct Bucket;
240- 240+ 
241- void push_rebuilt_params(const size_t& index);241+ void push_rebuilt_params(const size_t& index);
242- 242+ 
243- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)243+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
244- mutable std::mutex mutex_;244+ mutable std::mutex mutex_;
245- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)245+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
246- const std::vector<at::Tensor> params_;246+ const std::vector<at::Tensor> params_;
247- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)247+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
248- const c10::intrusive_ptr<::c10d::ProcessGroup> process_group_;248+ const c10::intrusive_ptr<::c10d::ProcessGroup> process_group_;
249- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)249+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
250- std::vector<bool> expect_sparse_gradients_;250+ std::vector<bool> expect_sparse_gradients_;
251- 251+ 
252- std::vector<std::shared_ptr<torch::autograd::Node>>252+ std::vector<std::shared_ptr<torch::autograd::Node>>
253- grad_accumulators_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)253+ grad_accumulators_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)
254- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)254+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
255- std::unordered_map<torch::autograd::Node*, size_t> gradAccToVariableMap_;255+ std::unordered_map<torch::autograd::Node*, size_t> gradAccToVariableMap_;
256- std::vector<std::pair<uintptr_t, std::shared_ptr<torch::autograd::Node>>>256+ std::vector<std::pair<uintptr_t, std::shared_ptr<torch::autograd::Node>>>
257- hooks_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)257+ hooks_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)
258- 258+ 
259- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)259+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
260- bool expect_autograd_hooks_;260+ bool expect_autograd_hooks_;
261- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)261+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
262- bool require_finalize_;262+ bool require_finalize_;
263- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)263+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
264- size_t next_bucket_;264+ size_t next_bucket_;
265- 265+ 
266- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)266+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
267- bool has_marked_unused_parameters_;267+ bool has_marked_unused_parameters_;
268- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)268+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
269- const bool find_unused_parameters_;269+ const bool find_unused_parameters_;
270- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)270+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
271- const bool gradient_as_bucket_view_;271+ const bool gradient_as_bucket_view_;
272- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)272+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
273- std::vector<size_t> unused_parameters_;273+ std::vector<size_t> unused_parameters_;
274- // Previous iteration's unused params, used for checking if unused parameters274+ // Previous iteration's unused params, used for checking if unused parameters
275- // change between iterations. Only filled during the first backwards call.275+ // change between iterations. Only filled during the first backwards call.
276- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)276+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
277- std::vector<size_t> prev_iteration_unused_parameters_;277+ std::vector<size_t> prev_iteration_unused_parameters_;
278- // Whether graph is static or not. When user does not explicitly set static278+ // Whether graph is static or not. When user does not explicitly set static
279- // graph, the only possible dynamism is set of unused parameters changing279+ // graph, the only possible dynamism is set of unused parameters changing
280- // between iterations which is tracked by this flag.280+ // between iterations which is tracked by this flag.
281- // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)281+ // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
282- bool ddp_graph_static_{true};282+ bool ddp_graph_static_{true};
283- // Locally used parameter maps indicating if parameters are used locally283+ // Locally used parameter maps indicating if parameters are used locally
284- // during the current iteration or no_sync session if no_sync is on.284+ // during the current iteration or no_sync session if no_sync is on.
285- // Each map is a one-dim int32 tensor of number of parameters. These tensors285+ // Each map is a one-dim int32 tensor of number of parameters. These tensors
286- // are marked in autograd_hook to indicate the corresponding param has been286+ // are marked in autograd_hook to indicate the corresponding param has been
287- // used, and get allreduced in the end of backward step of current iteration287+ // used, and get allreduced in the end of backward step of current iteration
288- // or no_sync session for figuring out the globally unused parameters.288+ // or no_sync session for figuring out the globally unused parameters.
289- //289+ //
290- // local_used_map_: CPU tensor for bookkeeping locally used params290+ // local_used_map_: CPU tensor for bookkeeping locally used params
291- // local_used_map_dev_: dev tensor for reducing globally unused params291+ // local_used_map_dev_: dev tensor for reducing globally unused params
292- at::Tensor local_used_map_;292+ at::Tensor local_used_map_;
293- at::Tensor local_used_map_dev_;293+ at::Tensor local_used_map_dev_;
294- // Indicate that reduction is done and D2H copy is done as well.294+ // Indicate that reduction is done and D2H copy is done as well.
295- bool local_used_map_reduced_;295+ bool local_used_map_reduced_;
296- 296+ 
297- // Weak pointer to associated DDP logger.297+ // Weak pointer to associated DDP logger.
298- std::weak_ptr<c10d::Logger> logger_;298+ std::weak_ptr<c10d::Logger> logger_;
299- // List of futures installed by Reducer::install_futures that should be awaited299+ // List of futures installed by Reducer::install_futures that should be awaited
300- // at the end of backwards pass.300+ // at the end of backwards pass.
301- c10::optional<c10::List<c10::intrusive_ptr<c10::ivalue::Future>>> installed_futures_{c10::nullopt};301+ c10::optional<c10::List<c10::intrusive_ptr<c10::ivalue::Future>>> installed_futures_{c10::nullopt};
302- 302+ 
303- // Work handle for allreduce on local_used_map_303+ // Work handle for allreduce on local_used_map_
304- c10::intrusive_ptr<c10d::Work> local_used_work_;304+ c10::intrusive_ptr<c10d::Work> local_used_work_;
305- 305+ 
306- void mark_variable_ready_dense(size_t variable_index);306+ void mark_variable_ready_dense(size_t variable_index);
307- 307+ 
308- void mark_variable_ready_sparse(size_t variable_index);308+ void mark_variable_ready_sparse(size_t variable_index);
309- 309+ 
310- void mark_variable_ready(size_t variable_index);310+ void mark_variable_ready(size_t variable_index);
311- 311+ 
312- void autograd_hook(size_t index);312+ void autograd_hook(size_t index);
313- 313+ 
314- void mark_bucket_ready(size_t bucket_index);314+ void mark_bucket_ready(size_t bucket_index);
315- 315+ 
316- void finalize_bucket_dense(Bucket& replica);316+ void finalize_bucket_dense(Bucket& replica);
317- 317+ 
318- void finalize_backward();318+ void finalize_backward();
319- 319+ 
320- // Returns list of model parameters corresponding to the given bucket.320+ // Returns list of model parameters corresponding to the given bucket.
321- // bucket_index is a key to cache after buckets are rebuilt, after which this321+ // bucket_index is a key to cache after buckets are rebuilt, after which this
322- // mapping never changes.322+ // mapping never changes.
323- std::vector<at::Tensor> get_variables_for_bucket(323+ std::vector<at::Tensor> get_variables_for_bucket(
324- size_t bucket_index, const Bucket& bucket) const;324+ size_t bucket_index, const Bucket& bucket) const;
325- 325+ 
326- // Asserts that the reduction for the previous iteration has finished before326+ // Asserts that the reduction for the previous iteration has finished before
327- // rebuilding buckets or kicking off the next one.327+ // rebuilding buckets or kicking off the next one.
328- void ensure_prior_reduction_finished();328+ void ensure_prior_reduction_finished();
329- 329+ 
330- // Broadcast rebuilt buckets from rank 0 to other ranks before initializing330+ // Broadcast rebuilt buckets from rank 0 to other ranks before initializing
331- // the buckets331+ // the buckets
332- void sync_bucket_indices(std::vector<std::vector<size_t>>& bucket_indices);332+ void sync_bucket_indices(std::vector<std::vector<size_t>>& bucket_indices);
333- 333+ 
334- // We'd like to use DistAutogradContext::GradCallback here but dist autograd334+ // We'd like to use DistAutogradContext::GradCallback here but dist autograd
335- // doesn't exist under Windows. So we just directly use the concrete type but335+ // doesn't exist under Windows. So we just directly use the concrete type but
336- // to preserve and enforce our original intent we do a static assert when dist336+ // to preserve and enforce our original intent we do a static assert when dist
337- // autograd is available.337+ // autograd is available.
338- using GradCallback = std::function<bool(at::Tensor&)>;338+ using GradCallback = std::function<bool(at::Tensor&)>;
339-#ifndef _WIN32339+#ifndef _WIN32
340- static_assert(340+ static_assert(
341- std::is_same<341+ std::is_same<
342- GradCallback,342+ GradCallback,
343- torch::distributed::autograd::DistAutogradContext::GradCallback>::343+ torch::distributed::autograd::DistAutogradContext::GradCallback>::
344- value,344+ value,
345- "");345+ "");
346-#endif346+#endif
347- void runGradCallbackForVariable(at::Tensor& variable, GradCallback&& cb);347+ void runGradCallbackForVariable(at::Tensor& variable, GradCallback&& cb);
348- 348+ 
349- // A bucket replica represents [1..N] gradients to be reduced,349+ // A bucket replica represents [1..N] gradients to be reduced,
350- // with the same dtype, on the same device.350+ // with the same dtype, on the same device.
351- //351+ //
352- // Batching gradients together before reducing them can result in lower352+ // Batching gradients together before reducing them can result in lower
353- // overhead and/or faster time to completion. Only gradients of the same type353+ // overhead and/or faster time to completion. Only gradients of the same type
354- // and on the same device can be batched. The tensor that represents the354+ // and on the same device can be batched. The tensor that represents the
355- // flattened gradient uses the same type and is placed on the same device.355+ // flattened gradient uses the same type and is placed on the same device.
356- // Buckets are filled as the gradients they hold are computed (triggered by356+ // Buckets are filled as the gradients they hold are computed (triggered by
357- // autograd hooks). Buckets are reduced in a predetermined order that is357+ // autograd hooks). Buckets are reduced in a predetermined order that is
358- // identical across processes.358+ // identical across processes.
359- struct BucketReplica {359+ struct BucketReplica {
360- // Flattened (1 dimensional) contents of bucket.360+ // Flattened (1 dimensional) contents of bucket.
361- at::Tensor contents;361+ at::Tensor contents;
362- 362+ 
363- // Views into contents for each grad. Each view will be created with363+ // Views into contents for each grad. Each view will be created with
364- // layout (sizes + strides) matching the grad's expected layout364+ // layout (sizes + strides) matching the grad's expected layout
365- // ("Gradient Layout Contract" in torch/csrc/autograd/AccumulateGrad.h).365+ // ("Gradient Layout Contract" in torch/csrc/autograd/AccumulateGrad.h).
366- // `bucket_views_in[i].copy_(grad)` and366+ // `bucket_views_in[i].copy_(grad)` and
367- // `grad.copy_(bucket_views_out[i])`367+ // `grad.copy_(bucket_views_out[i])`
368- // provide convenient ways to move grad data in/out of contents.368+ // provide convenient ways to move grad data in/out of contents.
369- // The reason we keep two states for bucket_views is that if DDP369+ // The reason we keep two states for bucket_views is that if DDP
370- // communication hook was registered, `bucket_views_out` could be370+ // communication hook was registered, `bucket_views_out` could be
371- // re-initialized with the value of hook's `future_work`. We still need to371+ // re-initialized with the value of hook's `future_work`. We still need to
372- // keep a separate view reference to replica's original contents for372+ // keep a separate view reference to replica's original contents for
373- // `bucket_views_in[i].copy_(grad)` call.373+ // `bucket_views_in[i].copy_(grad)` call.
374- std::vector<at::Tensor> bucket_views_in;374+ std::vector<at::Tensor> bucket_views_in;
375- std::vector<at::Tensor> bucket_views_out;375+ std::vector<at::Tensor> bucket_views_out;
376- 376+ 
377- // Variables that contribute to this bucket replica. Use refcounted value377+ // Variables that contribute to this bucket replica. Use refcounted value
378- // here so that we can easily unflatten the bucket contents into the378+ // here so that we can easily unflatten the bucket contents into the
379- // participating variables after reduction has completed.379+ // participating variables after reduction has completed.
380- std::vector<at::Tensor> variables;380+ std::vector<at::Tensor> variables;
381- 381+ 
382- // Per-variable offset/length into the flat bucket contents tensor and grad382+ // Per-variable offset/length into the flat bucket contents tensor and grad
383- // bucket.383+ // bucket.
384- std::vector<size_t> offsets;384+ std::vector<size_t> offsets;
385- std::vector<size_t> lengths;385+ std::vector<size_t> lengths;
386- 386+ 
387- // Per-variable sizes into the grad bucekt.387+ // Per-variable sizes into the grad bucekt.
388- std::vector<c10::IntArrayRef> sizes_vec;388+ std::vector<c10::IntArrayRef> sizes_vec;
389- 389+ 
390- // Number of tensors to be added before this bucket is complete.390+ // Number of tensors to be added before this bucket is complete.
391- // This is reset to `variables.size()` every iteration.391+ // This is reset to `variables.size()` every iteration.
392- size_t pending;392+ size_t pending;
393- };393+ };
394- 394+ 
395- // This function is called inside `initialize_buckets`, it initializes both395+ // This function is called inside `initialize_buckets`, it initializes both
396- // bucket_views_in and bucket_views_out into the contents tensor for each396+ // bucket_views_in and bucket_views_out into the contents tensor for each
397- // variable's grad. Views serve as entry points to copy_ each grad's data397+ // variable's grad. Views serve as entry points to copy_ each grad's data
398- // in/out of the flat contents tensor.398+ // in/out of the flat contents tensor.
399- void initialize_bucket_views(BucketReplica& replica, at::Tensor& contents);399+ void initialize_bucket_views(BucketReplica& replica, at::Tensor& contents);
400- 400+ 
401- // This function is called inside `finalize_backward`, it happens only if401+ // This function is called inside `finalize_backward`, it happens only if
402- // DDP communication hook was registered to recreate just bucket_views_out402+ // DDP communication hook was registered to recreate just bucket_views_out
403- // with the result of `future_work`.403+ // with the result of `future_work`.
404- void populate_bucket_views_out(BucketReplica& replica, at::Tensor& tensor) const;404+ void populate_bucket_views_out(BucketReplica& replica, at::Tensor& tensor) const;
405- 405+ 
406- // If gradient_as_bucket_view_ is false, after allreduce buckets,406+ // If gradient_as_bucket_view_ is false, after allreduce buckets,
407- // copy bucket results back to grads.407+ // copy bucket results back to grads.
408- void copy_bucket_to_grad(408+ void copy_bucket_to_grad(
409- at::Tensor& variable,409+ at::Tensor& variable,
410- Reducer::BucketReplica& replica,410+ Reducer::BucketReplica& replica,
411- size_t intra_bucket_index,411+ size_t intra_bucket_index,
412- bool global_unused);412+ bool global_unused);
413- // Check layout of grad and bucket_view before copying the grad to bucket.413+ // Check layout of grad and bucket_view before copying the grad to bucket.
414- void check_grad_layout(const at::Tensor& grad, const at::Tensor& bucket_view);414+ void check_grad_layout(const at::Tensor& grad, const at::Tensor& bucket_view);
415- // If gradient_as_bucket_view_ is false, before allreduce buckets,415+ // If gradient_as_bucket_view_ is false, before allreduce buckets,
416- // copy grads to buckets.416+ // copy grads to buckets.
417- void copy_grad_to_bucket(const at::Tensor& grad, at::Tensor& bucket_view);417+ void copy_grad_to_bucket(const at::Tensor& grad, at::Tensor& bucket_view);
418- // A bucket holds N bucket replicas (1 per model replica).418+ // A bucket holds N bucket replicas (1 per model replica).
419- //419+ //
420- // If every bucket in this struct is ready, the reduction can be kicked off.420+ // If every bucket in this struct is ready, the reduction can be kicked off.
421- // One bucket per replica. Reduction is kicked off when every bucket is ready.421+ // One bucket per replica. Reduction is kicked off when every bucket is ready.
422- //422+ //
423- struct Bucket {423+ struct Bucket {
424- std::vector<BucketReplica> replicas;424+ std::vector<BucketReplica> replicas;
425- 425+ 
426- // Global indices of participating variables in the bucket426+ // Global indices of participating variables in the bucket
427- std::vector<size_t> variable_indices;427+ std::vector<size_t> variable_indices;
428- 428+ 
429- // Number of replicas to be marked done before this bucket is ready.429+ // Number of replicas to be marked done before this bucket is ready.
430- size_t pending;430+ size_t pending;
431- 431+ 
432- // Keep work handle around when this set of buckets is being reduced.432+ // Keep work handle around when this set of buckets is being reduced.
433- c10::intrusive_ptr<c10d::Work> work;433+ c10::intrusive_ptr<c10d::Work> work;
434- 434+ 
435- // Keep future work handle around DDP comm hook.435+ // Keep future work handle around DDP comm hook.
436- // If no hook is registered, a temporary vanilla allreduce hook will be436+ // If no hook is registered, a temporary vanilla allreduce hook will be
437- // used.437+ // used.
438- c10::intrusive_ptr<at::ivalue::Future> future_work;438+ c10::intrusive_ptr<at::ivalue::Future> future_work;
439- 439+ 
440- // If this bucket should expect a single sparse gradient.440+ // If this bucket should expect a single sparse gradient.
441- // Implies: replicas[i].variables.size() == 1.441+ // Implies: replicas[i].variables.size() == 1.
442- bool expect_sparse_gradient = false;442+ bool expect_sparse_gradient = false;
443- // "Limit" of cumulative parameter sizes that this bucket manages. It is443+ // "Limit" of cumulative parameter sizes that this bucket manages. It is
444- // actually a soft limit because we don't shard parameters across buckets444+ // actually a soft limit because we don't shard parameters across buckets
445- // so a single parameter may push it over the cap.445+ // so a single parameter may push it over the cap.
446- size_t bucket_size_limit;446+ size_t bucket_size_limit;
447- };447+ };
448- 448+ 
449- std::vector<Bucket> buckets_;449+ std::vector<Bucket> buckets_;
450- 450+ 
451- // A variable locator locates a particular variable in the bucket451+ // A variable locator locates a particular variable in the bucket
452- // structure. The `bucket_index` field points to the bucket in the `buckets_`452+ // structure. The `bucket_index` field points to the bucket in the `buckets_`
453- // vector. The `intra_bucket_index` field points to the index of the variable453+ // vector. The `intra_bucket_index` field points to the index of the variable
454- // in any of the vector fields in the bucket replica.454+ // in any of the vector fields in the bucket replica.
455- struct VariableLocator {455+ struct VariableLocator {
456- // Index into the `buckets_` variable.456+ // Index into the `buckets_` variable.
457- size_t bucket_index;457+ size_t bucket_index;
458- // Index of parameter in single bucket replica.458+ // Index of parameter in single bucket replica.
459- size_t intra_bucket_index;459+ size_t intra_bucket_index;
460- 460+ 
461- VariableLocator() = default;461+ VariableLocator() = default;
462- 462+ 
463- VariableLocator(size_t bucket_index_, size_t intra_bucket_index_)463+ VariableLocator(size_t bucket_index_, size_t intra_bucket_index_)
464- {464+ {
465- bucket_index = bucket_index_;465+ bucket_index = bucket_index_;
466- intra_bucket_index = intra_bucket_index_;466+ intra_bucket_index = intra_bucket_index_;
467- }467+ }
468- };468+ };
469- 469+ 
470- // Map the index of a variable to its location in the bucket structure.470+ // Map the index of a variable to its location in the bucket structure.
471- std::vector<VariableLocator> variable_locators_;471+ std::vector<VariableLocator> variable_locators_;
472- 472+ 
473- // track the number of iterations to synchronize grads in training so far.473+ // track the number of iterations to synchronize grads in training so far.
474- long num_iterations_;474+ long num_iterations_;
475- // track the number of buckets that have been ready for475+ // track the number of buckets that have been ready for
476- // communication calls like allReduce or communication hooks.476+ // communication calls like allReduce or communication hooks.
477- int num_buckets_ready_;477+ int num_buckets_ready_;
478- 478+ 
479- // Timing information.479+ // Timing information.
480- int64_t backward_compute_start_time_ = -1;480+ int64_t backward_compute_start_time_ = -1;
481- std::unique_ptr<Timer> timer_;481+ std::unique_ptr<Timer> timer_;
482- 482+ 
483- // We collect the relative timestamp of every gradient being ready483+ // We collect the relative timestamp of every gradient being ready
484- // when executing autograd. This can be used to derive a timeline of484+ // when executing autograd. This can be used to derive a timeline of
485- // the point in time buckets were ready, or ideal bucket assignment/ordering.485+ // the point in time buckets were ready, or ideal bucket assignment/ordering.
486- std::vector<int64_t> backward_stats_;486+ std::vector<int64_t> backward_stats_;
487- 487+ 
488- bool should_collect_runtime_stats();488+ bool should_collect_runtime_stats();
489- void record_forward_compute_start_time();489+ void record_forward_compute_start_time();
490- void record_backward_compute_start_time();490+ void record_backward_compute_start_time();
491- void record_backward_compute_end_time();491+ void record_backward_compute_end_time();
492- void record_backward_comm_start_time();492+ void record_backward_comm_start_time();
493- void record_backward_comm_end_time();493+ void record_backward_comm_end_time();
494- 494+ 
495- int get_ddp_runtime_logging_sample_rate() const;495+ int get_ddp_runtime_logging_sample_rate() const;
496- int ddp_runtime_logging_sample_rate_ = kDDPRuntimeLoggingSampleRate;496+ int ddp_runtime_logging_sample_rate_ = kDDPRuntimeLoggingSampleRate;
497- 497+ 
498- bool is_multi_device_module_ = false;498+ bool is_multi_device_module_ = false;
499- 499+ 
500- // Following variables are to help build dynamic bucket order500+ // Following variables are to help build dynamic bucket order
501- bool has_rebuilt_bucket_;501+ bool has_rebuilt_bucket_;
502- std::vector<at::Tensor> rebuilt_params_;502+ std::vector<at::Tensor> rebuilt_params_;
503- std::vector<int64_t> rebuilt_param_indices_;503+ std::vector<int64_t> rebuilt_param_indices_;
504- const int64_t bucket_bytes_cap_;504+ const int64_t bucket_bytes_cap_;
505- 505+ 
506-#ifndef _WIN32506+#ifndef _WIN32
507- struct RpcContext {507+ struct RpcContext {
508- using ContextPtr = torch::distributed::autograd::ContextPtr;508+ using ContextPtr = torch::distributed::autograd::ContextPtr;
509- // The shared_ptr is to hold the context instance.509+ // The shared_ptr is to hold the context instance.
510- ContextPtr context_ptr_holder;510+ ContextPtr context_ptr_holder;
511- std::atomic<ContextPtr::element_type*> context_ptr{nullptr};511+ std::atomic<ContextPtr::element_type*> context_ptr{nullptr};
512- 512+ 
513- void set(ContextPtr&& new_context_ptr);513+ void set(ContextPtr&& new_context_ptr);
514- };514+ };
515- RpcContext rpc_context_;515+ RpcContext rpc_context_;
516-#endif516+#endif
517- 517+ 
518- // A struct containing work handle and tensor for allreduce scheduled in518+ // A struct containing work handle and tensor for allreduce scheduled in
519- // forward pass, if applicable.519+ // forward pass, if applicable.
520- struct ForwardPassAllreduceWork {520+ struct ForwardPassAllreduceWork {
521- c10::intrusive_ptr<c10d::Work> workHandle;521+ c10::intrusive_ptr<c10d::Work> workHandle;
522- at::Tensor resultTensor;522+ at::Tensor resultTensor;
523- // whether we should divide by the initial world_size or the no. of523+ // whether we should divide by the initial world_size or the no. of
524- // remaining DDP ranks.524+ // remaining DDP ranks.
525- bool useStaticWorldSize = false;525+ bool useStaticWorldSize = false;
526- };526+ };
527- 527+ 
528- // Handle for the currently scheduled allreduce in the forward pass, if528+ // Handle for the currently scheduled allreduce in the forward pass, if
529- // applicable.529+ // applicable.
530- ForwardPassAllreduceWork forwardPassWorkHandle_;530+ ForwardPassAllreduceWork forwardPassWorkHandle_;
531- 531+ 
532- // Division factor for reduction of gradients.532+ // Division factor for reduction of gradients.
533- // Equal to the process group size, with an exception of handling uneven533+ // Equal to the process group size, with an exception of handling uneven
534- // input.534+ // input.
535- int div_factor_;535+ int div_factor_;
536- 536+ 
537- bool static_graph_;537+ bool static_graph_;
538- 538+ 
539- // Key: size_t (index), Value: the number of times that a variable's539+ // Key: size_t (index), Value: the number of times that a variable's
540- // autograd_hook() should be triggered before marking this variable's grad as540+ // autograd_hook() should be triggered before marking this variable's grad as
541- // ready for communication. Map will not change after 1st iteration.541+ // ready for communication. Map will not change after 1st iteration.
542- std::unordered_map<size_t, int> numGradHooksTriggeredMap_;542+ std::unordered_map<size_t, int> numGradHooksTriggeredMap_;
543- // Key: size_t (index), Value: the number of times that a variable's543+ // Key: size_t (index), Value: the number of times that a variable's
544- // autograd_hook() are left to be triggered before marking this variable's544+ // autograd_hook() are left to be triggered before marking this variable's
545- // grad as ready for communication. Map will change after 1st iteration to545+ // grad as ready for communication. Map will change after 1st iteration to
546- // track a grad is ready for communication or not.546+ // track a grad is ready for communication or not.
547- std::unordered_map<size_t, int> numGradHooksTriggeredMapPerIteration_;547+ std::unordered_map<size_t, int> numGradHooksTriggeredMapPerIteration_;
548- 548+ 
549-private:549+private:
550- // reset counting for buckets before backward starts550+ // reset counting for buckets before backward starts
551- void reset_bucket_counting();551+ void reset_bucket_counting();
552- // search unused parameters beore backward starts552+ // search unused parameters beore backward starts
553- void search_unused_parameters(553+ void search_unused_parameters(
554- const std::vector<torch::autograd::Variable>& outputs);554+ const std::vector<torch::autograd::Variable>& outputs);
555- void set_divide_factor();555+ void set_divide_factor();
556- // kick off all reduce for the ready bucket556+ // kick off all reduce for the ready bucket
557- void all_reduce_bucket(Bucket& bucket);557+ void all_reduce_bucket(Bucket& bucket);
558- // kick off all reduce to local used map, it can help find global unused558+ // kick off all reduce to local used map, it can help find global unused
559- // parameters559+ // parameters
560- void all_reduce_local_used_map();560+ void all_reduce_local_used_map();
561- // initialize locally used parameter maps561+ // initialize locally used parameter maps
562- void initialize_local_used_map();562+ void initialize_local_used_map();
563- // get current cuda stream563+ // get current cuda stream
564- const c10::Stream get_current_stream();564+ const c10::Stream get_current_stream();
565- bool dynamic_graph_find_unused() const;565+ bool dynamic_graph_find_unused() const;
566- bool static_graph_first_iteration() const;566+ bool static_graph_first_iteration() const;
567- bool static_graph_after_first_iteration() const;567+ bool static_graph_after_first_iteration() const;
568- 568+ 
569- // comm_hook_ is used to access the DDP communication hook if registered.569+ // comm_hook_ is used to access the DDP communication hook if registered.
570- std::unique_ptr<c10d::CommHookInterface> comm_hook_;570+ std::unique_ptr<c10d::CommHookInterface> comm_hook_;
571- // Debug level setting. It is parsed once when Reducer is constructed, and571+ // Debug level setting. It is parsed once when Reducer is constructed, and
572- // remains the same across a single invocation of DDP training.572+ // remains the same across a single invocation of DDP training.
573- c10d::DebugLevel ddp_debug_level_;573+ c10d::DebugLevel ddp_debug_level_;
574- // Mapping of variable index to fully qualified name of model to notify users574+ // Mapping of variable index to fully qualified name of model to notify users
575- // about errors when certain parameters do not get gradient.575+ // about errors when certain parameters do not get gradient.
576- std::unordered_map<size_t, std::string> param_names_;576+ std::unordered_map<size_t, std::string> param_names_;
577- // Variable indices stored sequentially in order of when the gradient is ready577+ // Variable indices stored sequentially in order of when the gradient is ready
578- // for the current backwards pass.578+ // for the current backwards pass.
579- std::vector<int> grad_ready_order_indices_;579+ std::vector<int> grad_ready_order_indices_;
580- // Bytes capacity of first bucket, can be configured by user580+ // Bytes capacity of first bucket, can be configured by user
581- int64_t first_bucket_bytes_cap_;581+ int64_t first_bucket_bytes_cap_;
582- // Per iteration set of parameter indices that have been marked ready.582+ // Per iteration set of parameter indices that have been marked ready.
583- std::unordered_set<size_t> perIterationReadyParams_;583+ std::unordered_set<size_t> perIterationReadyParams_;
584- // Retrieves parameter names that have not been marked as ready as part of584+ // Retrieves parameter names that have not been marked as ready as part of
585- // previous iteration.585+ // previous iteration.
586- std::vector<std::string> getUnmarkedParamsForIteration();586+ std::vector<std::string> getUnmarkedParamsForIteration();
587- // Retrives parameter indices that have not been marked as ready as part of587+ // Retrives parameter indices that have not been marked as ready as part of
588- // previous iteration.588+ // previous iteration.
589- std::vector<size_t> getUnmarkedParamIndicesForIteration();589+ std::vector<size_t> getUnmarkedParamIndicesForIteration();
590- // Raises appropriate error if mark_variable_ready is called on the same590+ // Raises appropriate error if mark_variable_ready is called on the same
591- // variable twice, which is unexpected.591+ // variable twice, which is unexpected.
592- void checkAndRaiseMarkedTwiceError(size_t curVariableIndex);592+ void checkAndRaiseMarkedTwiceError(size_t curVariableIndex);
593- // Retrieves parameter corresponding to the given VariableIndex.593+ // Retrieves parameter corresponding to the given VariableIndex.
594- at::Tensor& get_param_from_index(size_t index);594+ at::Tensor& get_param_from_index(size_t index);
595- 595+ 
596- // Cached bucket index to model parameter mapping. Populated after buckets596+ // Cached bucket index to model parameter mapping. Populated after buckets
597- // are rebuilt after which this mapping is static.597+ // are rebuilt after which this mapping is static.
598- mutable std::unordered_map<size_t, std::vector<at::Tensor>> cached_variables_for_bucket_;598+ mutable std::unordered_map<size_t, std::vector<at::Tensor>> cached_variables_for_bucket_;
599- 599+ 
600- friend class Logger;600+ friend class Logger;
601-};601+};
602- 602+ 
603-// This is equivalent to take_tensors but returns indices into the603+// This is equivalent to take_tensors but returns indices into the
604-// tensor list argument for bucket assignment. Also, it is aware604+// tensor list argument for bucket assignment. Also, it is aware
605-// of device placement and will not allow buckets to span devices.605+// of device placement and will not allow buckets to span devices.
606-// The index of tensors[i] assigned to bucket is tensor_indices[i],606+// The index of tensors[i] assigned to bucket is tensor_indices[i],
607-// when tensor_indices is empty, the index of tensors[i] assigned to607+// when tensor_indices is empty, the index of tensors[i] assigned to
608-// bucket is i.608+// bucket is i.
609-std::tuple<std::vector<std::vector<size_t>>, std::vector<size_t>> compute_bucket_assignment_by_size(609+std::tuple<std::vector<std::vector<size_t>>, std::vector<size_t>> compute_bucket_assignment_by_size(
610- const std::vector<at::Tensor>& tensors,610+ const std::vector<at::Tensor>& tensors,
611- const std::vector<size_t>& bucket_size,611+ const std::vector<size_t>& bucket_size,
612- const std::vector<bool>& expect_sparse_gradient = {},612+ const std::vector<bool>& expect_sparse_gradient = {},
613- const std::vector<int64_t>& tensor_indices = {},613+ const std::vector<int64_t>& tensor_indices = {},
614- const c10::optional<std::weak_ptr<c10d::Logger>>& logger = {});614+ const c10::optional<std::weak_ptr<c10d::Logger>>& logger = {});
615- 615+ 
616-// Verify models across all processes are the same as model on rank 0 with616+// Verify models across all processes are the same as model on rank 0 with
617-// respect to no. of params and matching dtype/size/layout.617+// respect to no. of params and matching dtype/size/layout.
618-void verify_params_across_processes(618+void verify_params_across_processes(
619- const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,619+ const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,
620- const std::vector<at::Tensor>& params,620+ const std::vector<at::Tensor>& params,
621- const c10::optional<std::weak_ptr<c10d::Logger>>& logger);621+ const c10::optional<std::weak_ptr<c10d::Logger>>& logger);
622-} // namespace c10d_npu622+} // namespace c10d_npu
Mtorch_npu/csrc/framework/CMakeLists.txt+11-11
@@ -1,12 +1,12 @@
1-FILE(GLOB _FRAMEWORK_SRCS1+FILE(GLOB _FRAMEWORK_SRCS
2- *.cpp2+ *.cpp
3- aoe/*.cpp3+ aoe/*.cpp
4- autograd/*.cpp4+ autograd/*.cpp
5- contiguous/*.cpp5+ contiguous/*.cpp
6- interface/*.cpp6+ interface/*.cpp
7- utils/*.cpp)7+ utils/*.cpp)
8- 8+ 
9-LIST(APPEND FRAMEWORK_SRCS ${_FRAMEWORK_SRCS})9+LIST(APPEND FRAMEWORK_SRCS ${_FRAMEWORK_SRCS})
10- 10+ 
11-# Pass to parent11+# Pass to parent
12set(FRAMEWORK_SRCS ${FRAMEWORK_SRCS} PARENT_SCOPE)12set(FRAMEWORK_SRCS ${FRAMEWORK_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/framework/StorageDescHelper.cpp+0-1
@@ -311,4 +311,3 @@ int64_t StorageDescHelper::GetValidMemorySize(const at::Tensor &tensor)
311 311 
312} // namespace native312} // namespace native
313} // namespace at_npu313} // namespace at_npu
314- 
Mtorch_npu/csrc/framework/contiguous/ReshapeOpt.cpp+0-1
@@ -87,4 +87,3 @@ bool CanUseMemcpyForOtherFormat(const at::Tensor &tensor)
87 87 
88} // namespace native88} // namespace native
89} // namespace at_npu89} // namespace at_npu
90- 
Mtorch_npu/csrc/framework/interface/LibAscendHal.cpp+0-1
@@ -67,4 +67,3 @@ bool isSyscntEnable()
67 67 
68} // namespace native68} // namespace native
69} // namespace at_npu69} // namespace at_npu
70- 
Mtorch_npu/csrc/framework/interface/LibAscendHal.h+0-1
@@ -8,4 +8,3 @@ bool isSyscntEnable();
8} // namespace native8} // namespace native
9} // namespace torchat_npu_npu9} // namespace torchat_npu_npu
10#endif10#endif
11- 
Mtorch_npu/csrc/npu/CMakeLists.txt+5-5
@@ -1,6 +1,6 @@
1-FILE(GLOB _NPU_SRCS *.cpp)1+FILE(GLOB _NPU_SRCS *.cpp)
2- 2+ 
3-LIST(APPEND NPU_SRCS ${_NPU_SRCS})3+LIST(APPEND NPU_SRCS ${_NPU_SRCS})
4- 4+ 
5-# Pass to parent5+# Pass to parent
6set(NPU_SRCS ${NPU_SRCS} PARENT_SCOPE)6set(NPU_SRCS ${NPU_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/profiler/CMakeLists.txt+5-5
@@ -1,6 +1,6 @@
1-FILE(GLOB _PROF_SRCS *.cpp unwind/*.cpp python/*.cpp)1+FILE(GLOB _PROF_SRCS *.cpp unwind/*.cpp python/*.cpp)
2- 2+ 
3-LIST(APPEND PROF_SRCS ${_PROF_SRCS})3+LIST(APPEND PROF_SRCS ${_PROF_SRCS})
4- 4+ 
5-# Pass to parent5+# Pass to parent
6set(PROF_SRCS ${PROF_SRCS} PARENT_SCOPE)6set(PROF_SRCS ${PROF_SRCS} PARENT_SCOPE)
Mtorch_npu/csrc/utils/TensorType.cpp+347-347
@@ -1,347 +1,347 @@
1-#include <c10/core/DeviceType.h>1+#include <c10/core/DeviceType.h>
2-#include <torch/csrc/autograd/utils/wrap_outputs.h>2+#include <torch/csrc/autograd/utils/wrap_outputs.h>
3- 3+ 
4-#include "torch_npu/csrc/utils/TensorType.h"4+#include "torch_npu/csrc/utils/TensorType.h"
5-#include "torch_npu/csrc/utils/LazyInit.h"5+#include "torch_npu/csrc/utils/LazyInit.h"
6- 6+ 
7-namespace torch_npu {7+namespace torch_npu {
8-namespace utils {8+namespace utils {
9-using namespace at;9+using namespace at;
10-using namespace torch::autograd;10+using namespace torch::autograd;
11- 11+ 
12-std::vector<std::pair<Backend, ScalarType>> all_declared_types_npu()12+std::vector<std::pair<Backend, ScalarType>> all_declared_types_npu()
13-{13+{
14- std::vector<std::pair<Backend, ScalarType>> ret;14+ std::vector<std::pair<Backend, ScalarType>> ret;
15- // can't easily iterate over enum classes, does not support BFloat16 now15+ // can't easily iterate over enum classes, does not support BFloat16 now
16- std::vector<Backend> backends = { c10::Backend::PrivateUse1 };16+ std::vector<Backend> backends = { c10::Backend::PrivateUse1 };
17- std::vector<ScalarType> scalar_types = { ScalarType::Byte, ScalarType::Char, ScalarType::Double,17+ std::vector<ScalarType> scalar_types = { ScalarType::Byte, ScalarType::Char, ScalarType::Double,
18- ScalarType::Float, ScalarType::Int, ScalarType::Long,18+ ScalarType::Float, ScalarType::Int, ScalarType::Long,
19- ScalarType::Short, ScalarType::Half, ScalarType::Bool,19+ ScalarType::Short, ScalarType::Half, ScalarType::Bool,
20- ScalarType::BFloat16 };20+ ScalarType::BFloat16 };
21- 21+ 
22- for (auto &backend : backends) {22+ for (auto &backend : backends) {
23- for (auto &scalar_type : scalar_types) {23+ for (auto &scalar_type : scalar_types) {
24- ret.emplace_back(std::make_pair(backend, scalar_type));24+ ret.emplace_back(std::make_pair(backend, scalar_type));
25- }25+ }
26- }26+ }
27- 27+ 
28- return ret;28+ return ret;
29-}29+}
30- 30+ 
31-struct PyTensorType {31+struct PyTensorType {
32- PyTypeObject py_type;32+ PyTypeObject py_type;
33- THPDtype *dtype;33+ THPDtype *dtype;
34- THPLayout *layout;34+ THPLayout *layout;
35- bool is_npu;35+ bool is_npu;
36- char name[64];36+ char name[64];
37- int backend;37+ int backend;
38- int scalar_type;38+ int scalar_type;
39- 39+ 
40- Backend get_backend() const40+ Backend get_backend() const
41- {41+ {
42- return static_cast<Backend>(backend);42+ return static_cast<Backend>(backend);
43- }43+ }
44- 44+ 
45- DispatchKey get_dispatch_key() const45+ DispatchKey get_dispatch_key() const
46- {46+ {
47- return backendToDispatchKey(static_cast<Backend>(backend));47+ return backendToDispatchKey(static_cast<Backend>(backend));
48- }48+ }
49- 49+ 
50- ScalarType get_scalar_type() const50+ ScalarType get_scalar_type() const
51- {51+ {
52- return static_cast<ScalarType>(scalar_type);52+ return static_cast<ScalarType>(scalar_type);
53- }53+ }
54-};54+};
55- 55+ 
56-static_assert(std::is_standard_layout<PyTensorType>::value, "PyTensorType must be standard layout");56+static_assert(std::is_standard_layout<PyTensorType>::value, "PyTensorType must be standard layout");
57- 57+ 
58-static void py_bind_tensor_types(const std::vector<PyTensorType> &tensor_types);58+static void py_bind_tensor_types(const std::vector<PyTensorType> &tensor_types);
59- 59+ 
60-static PyObject *Tensor_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)60+static PyObject *Tensor_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
61-{61+{
62- HANDLE_TH_ERRORS62+ HANDLE_TH_ERRORS
63- auto &tensor_type = *((PyTensorType *)type);63+ auto &tensor_type = *((PyTensorType *)type);
64- if (tensor_type.is_npu) {64+ if (tensor_type.is_npu) {
65- TORCH_NPU_WARN_ONCE("Warning: The torch.npu.*DtypeTensor constructors are no longer recommended. "65+ TORCH_NPU_WARN_ONCE("Warning: The torch.npu.*DtypeTensor constructors are no longer recommended. "
66- "It's best to use methods such as torch.tensor(data, dtype=*, device='npu') "66+ "It's best to use methods such as torch.tensor(data, dtype=*, device='npu') "
67- "to create tensors.");67+ "to create tensors.");
68- }68+ }
69- TORCH_CHECK_TYPE(!tensor_type.is_npu || c10_npu::device_count() != 0, "type ", tensor_type.name,69+ TORCH_CHECK_TYPE(!tensor_type.is_npu || c10_npu::device_count() != 0, "type ", tensor_type.name,
70- " not available. Torch not compiled with npu enabled.", PTA_ERROR(ErrCode::TYPE))70+ " not available. Torch not compiled with npu enabled.", PTA_ERROR(ErrCode::TYPE))
71- torch_npu::utils::npu_lazy_init();71+ torch_npu::utils::npu_lazy_init();
72- return THPVariable_Wrap(72+ return THPVariable_Wrap(
73- torch::utils::legacy_tensor_ctor(tensor_type.get_dispatch_key(), tensor_type.get_scalar_type(), args, kwargs));73+ torch::utils::legacy_tensor_ctor(tensor_type.get_dispatch_key(), tensor_type.get_scalar_type(), args, kwargs));
74- END_HANDLE_TH_ERRORS74+ END_HANDLE_TH_ERRORS
75-}75+}
76- 76+ 
77-static PyObject *Tensor_instancecheck(PyObject *_self, PyObject *arg)77+static PyObject *Tensor_instancecheck(PyObject *_self, PyObject *arg)
78-{78+{
79- HANDLE_TH_ERRORS79+ HANDLE_TH_ERRORS
80- auto self = (PyTensorType *)_self;80+ auto self = (PyTensorType *)_self;
81- if (THPVariable_Check(arg)) {81+ if (THPVariable_Check(arg)) {
82- const auto &var = THPVariable_Unpack(arg);82+ const auto &var = THPVariable_Unpack(arg);
83- 83+ 
84- if (legacyExtractDispatchKey(var.key_set()) == self->get_dispatch_key() &&84+ if (legacyExtractDispatchKey(var.key_set()) == self->get_dispatch_key() &&
85- var.scalar_type() == static_cast<ScalarType>(self->scalar_type)) {85+ var.scalar_type() == static_cast<ScalarType>(self->scalar_type)) {
86- Py_RETURN_TRUE;86+ Py_RETURN_TRUE;
87- }87+ }
88- }88+ }
89- Py_RETURN_FALSE;89+ Py_RETURN_FALSE;
90- END_HANDLE_TH_ERRORS90+ END_HANDLE_TH_ERRORS
91-}91+}
92- 92+ 
93-PyObject *Tensor_dtype(PyTensorType *self, void *unused)93+PyObject *Tensor_dtype(PyTensorType *self, void *unused)
94-{94+{
95- return torch::autograd::utils::wrap(self->dtype);95+ return torch::autograd::utils::wrap(self->dtype);
96-}96+}
97- 97+ 
98-PyObject *Tensor_layout(PyTensorType *self, void *unused)98+PyObject *Tensor_layout(PyTensorType *self, void *unused)
99-{99+{
100- return torch::autograd::utils::wrap(self->layout);100+ return torch::autograd::utils::wrap(self->layout);
101-}101+}
102- 102+ 
103-PyObject *Tensor_is_npu(PyTensorType *self, void *unused)103+PyObject *Tensor_is_npu(PyTensorType *self, void *unused)
104-{104+{
105- if (self->is_npu) {105+ if (self->is_npu) {
106- Py_RETURN_TRUE;106+ Py_RETURN_TRUE;
107- } else {107+ } else {
108- Py_RETURN_FALSE;108+ Py_RETURN_FALSE;
109- }109+ }
110-}110+}
111- 111+ 
112-PyObject *Tensor_is_sparse(PyTensorType *self, void *unused)112+PyObject *Tensor_is_sparse(PyTensorType *self, void *unused)
113-{113+{
114- if (self->layout->layout == at::Layout::Strided) {114+ if (self->layout->layout == at::Layout::Strided) {
115- Py_RETURN_FALSE;115+ Py_RETURN_FALSE;
116- } else {116+ } else {
117- Py_RETURN_TRUE;117+ Py_RETURN_TRUE;
118- }118+ }
119-}119+}
120- 120+ 
121-static struct PyMethodDef metaclass_methods[] = {121+static struct PyMethodDef metaclass_methods[] = {
122- {"__instancecheck__", Tensor_instancecheck, METH_O, nullptr},122+ {"__instancecheck__", Tensor_instancecheck, METH_O, nullptr},
123- {nullptr}123+ {nullptr}
124-};124+};
125- 125+ 
126-using getter = PyObject *(*)(PyObject *, void *);126+using getter = PyObject *(*)(PyObject *, void *);
127- 127+ 
128-static struct PyGetSetDef metaclass_properties[] = {128+static struct PyGetSetDef metaclass_properties[] = {
129- {"dtype", (getter)Tensor_dtype, nullptr, nullptr, nullptr},129+ {"dtype", (getter)Tensor_dtype, nullptr, nullptr, nullptr},
130- {"layout", (getter)Tensor_layout, nullptr, nullptr, nullptr},130+ {"layout", (getter)Tensor_layout, nullptr, nullptr, nullptr},
131- {"is_npu", (getter)Tensor_is_npu, nullptr, nullptr, nullptr},131+ {"is_npu", (getter)Tensor_is_npu, nullptr, nullptr, nullptr},
132- {"is_sparse", (getter)Tensor_is_sparse, nullptr, nullptr, nullptr},132+ {"is_sparse", (getter)Tensor_is_sparse, nullptr, nullptr, nullptr},
133- {nullptr}133+ {nullptr}
134-};134+};
135- 135+ 
136-static PyTypeObject metaclass = {136+static PyTypeObject metaclass = {
137- PyVarObject_HEAD_INIT(nullptr, 0) "torch.tensortype", /* tp_name */137+ PyVarObject_HEAD_INIT(nullptr, 0) "torch.tensortype", /* tp_name */
138- sizeof(PyTypeObject) /* tp_basicsize */138+ sizeof(PyTypeObject) /* tp_basicsize */
139-};139+};
140- 140+ 
141-static void py_initialize_metaclass(PyTypeObject &metaclass)141+static void py_initialize_metaclass(PyTypeObject &metaclass)
142-{142+{
143- metaclass.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;143+ metaclass.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
144- metaclass.tp_methods = metaclass_methods;144+ metaclass.tp_methods = metaclass_methods;
145- metaclass.tp_getset = metaclass_properties;145+ metaclass.tp_getset = metaclass_properties;
146- metaclass.tp_base = &PyType_Type;146+ metaclass.tp_base = &PyType_Type;
147- if (PyType_Ready(&metaclass) < 0) {147+ if (PyType_Ready(&metaclass) < 0) {
148- throw python_error();148+ throw python_error();
149- }149+ }
150-}150+}
151- 151+ 
152-static PyTypeObject tensor_type_prototype = {152+static PyTypeObject tensor_type_prototype = {
153- PyVarObject_HEAD_INIT(&metaclass, 0) nullptr, /* tp_name */153+ PyVarObject_HEAD_INIT(&metaclass, 0) nullptr, /* tp_name */
154- sizeof(PyTensorType) /* tp_basicsize */154+ sizeof(PyTensorType) /* tp_basicsize */
155-};155+};
156- 156+ 
157-static void py_initialize_tensor_type(PyTypeObject &type, const char *name, PyObject *tp_dict)157+static void py_initialize_tensor_type(PyTypeObject &type, const char *name, PyObject *tp_dict)
158-{158+{
159- // NOTE: we don't use the typical static declaration of PyTypeObject because159+ // NOTE: we don't use the typical static declaration of PyTypeObject because
160- // we need to initialize as many types as there are VariableType instances.160+ // we need to initialize as many types as there are VariableType instances.
161- // We copy the basic object fields from a prototype definition and initialize161+ // We copy the basic object fields from a prototype definition and initialize
162- // the remaining fields below.162+ // the remaining fields below.
163- memcpy(&type, &tensor_type_prototype, sizeof(PyTypeObject));163+ memcpy(&type, &tensor_type_prototype, sizeof(PyTypeObject));
164- // Subclassing from torch.<ScalarType>Tensor isn't supported.164+ // Subclassing from torch.<ScalarType>Tensor isn't supported.
165- // (Py_TPFLAGS_BASETYPE omitted). Subclassing torch.Tensor still allowed.165+ // (Py_TPFLAGS_BASETYPE omitted). Subclassing torch.Tensor still allowed.
166- type.tp_flags = Py_TPFLAGS_DEFAULT;166+ type.tp_flags = Py_TPFLAGS_DEFAULT;
167- type.tp_name = name;167+ type.tp_name = name;
168- type.tp_new = Tensor_new;168+ type.tp_new = Tensor_new;
169- if (PyType_Ready(&type) < 0) {169+ if (PyType_Ready(&type) < 0) {
170- throw python_error();170+ throw python_error();
171- }171+ }
172- if (PyDict_Merge(type.tp_dict, tp_dict, 0) < 0) {172+ if (PyDict_Merge(type.tp_dict, tp_dict, 0) < 0) {
173- throw python_error();173+ throw python_error();
174- }174+ }
175-}175+}
176- 176+ 
177-static std::string get_module(Backend backend)177+static std::string get_module(Backend backend)
178-{178+{
179- switch (backend) {179+ switch (backend) {
180- case Backend::CPU:180+ case Backend::CPU:
181- return "torch";181+ return "torch";
182- case Backend::CUDA:182+ case Backend::CUDA:
183- return "torch.cuda";183+ return "torch.cuda";
184- case Backend::SparseCPU:184+ case Backend::SparseCPU:
185- return "torch.sparse";185+ return "torch.sparse";
186- case Backend::SparseCUDA:186+ case Backend::SparseCUDA:
187- return "torch.cuda.sparse";187+ return "torch.cuda.sparse";
188- case Backend::PrivateUse1:188+ case Backend::PrivateUse1:
189- return "torch." + c10::get_privateuse1_backend();189+ return "torch." + c10::get_privateuse1_backend();
190- default:190+ default:
191- AT_ERROR("invalid backend: ", c10::toString(backend));191+ AT_ERROR("invalid backend: ", c10::toString(backend));
192- }192+ }
193-}193+}
194- 194+ 
195-static std::string get_name(Backend backend, ScalarType scalarType)195+static std::string get_name(Backend backend, ScalarType scalarType)
196-{196+{
197- std::ostringstream ss;197+ std::ostringstream ss;
198- ss << get_module(backend) << "." << toString(scalarType) << "Tensor";198+ ss << get_module(backend) << "." << toString(scalarType) << "Tensor";
199- return ss.str();199+ return ss.str();
200-}200+}
201- 201+ 
202-static void set_type(PyTensorType &type_obj, Backend backend, ScalarType scalarType)202+static void set_type(PyTensorType &type_obj, Backend backend, ScalarType scalarType)
203-{203+{
204- // This field is lazily initialized from backend and scalar_type204+ // This field is lazily initialized from backend and scalar_type
205- type_obj.backend = static_cast<int>(backend);205+ type_obj.backend = static_cast<int>(backend);
206- type_obj.scalar_type = static_cast<int>(scalarType);206+ type_obj.scalar_type = static_cast<int>(scalarType);
207- type_obj.layout = torch::getTHPLayout(c10::layout_from_backend(backend));207+ type_obj.layout = torch::getTHPLayout(c10::layout_from_backend(backend));
208- type_obj.dtype = torch::getTHPDtype(scalarType);208+ type_obj.dtype = torch::getTHPDtype(scalarType);
209- type_obj.is_npu = (backend == c10::Backend::PrivateUse1);209+ type_obj.is_npu = (backend == c10::Backend::PrivateUse1);
210-}210+}
211- 211+ 
212-static void set_name(PyTensorType &type_obj, const std::string &name)212+static void set_name(PyTensorType &type_obj, const std::string &name)
213-{213+{
214- size_t n = sizeof(type_obj.name);214+ size_t n = sizeof(type_obj.name);
215- strncpy(type_obj.name, name.c_str(), n);215+ strncpy(type_obj.name, name.c_str(), n);
216- type_obj.name[n - 1] = '\0';216+ type_obj.name[n - 1] = '\0';
217-}217+}
218- 218+ 
219-static THPObjectPtr get_tensor_dict()219+static THPObjectPtr get_tensor_dict()
220-{220+{
221- auto torch = THPObjectPtr(PyImport_ImportModule("torch"));221+ auto torch = THPObjectPtr(PyImport_ImportModule("torch"));
222- if (!torch) {222+ if (!torch) {
223- throw python_error();223+ throw python_error();
224- }224+ }
225- 225+ 
226- auto tensor_class = THPObjectPtr(PyObject_GetAttrString(torch, "Tensor"));226+ auto tensor_class = THPObjectPtr(PyObject_GetAttrString(torch, "Tensor"));
227- if (!tensor_class) {227+ if (!tensor_class) {
228- throw python_error();228+ throw python_error();
229- }229+ }
230- 230+ 
231- auto tensor_type = (PyTypeObject *)tensor_class.get();231+ auto tensor_type = (PyTypeObject *)tensor_class.get();
232- TORCH_CHECK(tensor_type->tp_base, "missing base type for Tensor", PTA_ERROR(ErrCode::TYPE));232+ TORCH_CHECK(tensor_type->tp_base, "missing base type for Tensor", PTA_ERROR(ErrCode::TYPE));
233- 233+ 
234- auto res = THPObjectPtr(PyDict_New());234+ auto res = THPObjectPtr(PyDict_New());
235- if (!res) {235+ if (!res) {
236- throw python_error();236+ throw python_error();
237- }237+ }
238- 238+ 
239- if (PyDict_Merge(res.get(), tensor_type->tp_dict, 0) < 0) {239+ if (PyDict_Merge(res.get(), tensor_type->tp_dict, 0) < 0) {
240- throw python_error();240+ throw python_error();
241- }241+ }
242- if (PyDict_Merge(res.get(), tensor_type->tp_base->tp_dict, 0) < 0) {242+ if (PyDict_Merge(res.get(), tensor_type->tp_base->tp_dict, 0) < 0) {
243- throw python_error();243+ throw python_error();
244- }244+ }
245- 245+ 
246- return res;246+ return res;
247-}247+}
248- 248+ 
249-static std::vector<PyTensorType> tensor_types;249+static std::vector<PyTensorType> tensor_types;
250- 250+ 
251-static void initialize_npu_aten_types(std::vector<PyTensorType> &tensor_types)251+static void initialize_npu_aten_types(std::vector<PyTensorType> &tensor_types)
252-{252+{
253- // only initialize npu types253+ // only initialize npu types
254- auto declared_types = all_declared_types_npu();254+ auto declared_types = all_declared_types_npu();
255- tensor_types.resize(declared_types.size());255+ tensor_types.resize(declared_types.size());
256- 256+ 
257- for (size_t i = 0, end = declared_types.size(); i != end; i++) {257+ for (size_t i = 0, end = declared_types.size(); i != end; i++) {
258- auto &tensor_type = tensor_types[i];258+ auto &tensor_type = tensor_types[i];
259- Backend backend = declared_types[i].first;259+ Backend backend = declared_types[i].first;
260- ScalarType scalar_type = declared_types[i].second;260+ ScalarType scalar_type = declared_types[i].second;
261- set_type(tensor_type, backend, scalar_type);261+ set_type(tensor_type, backend, scalar_type);
262- set_name(tensor_type, get_name(backend, scalar_type));262+ set_name(tensor_type, get_name(backend, scalar_type));
263- }263+ }
264-}264+}
265- 265+ 
266-void _initialize_python_bindings()266+void _initialize_python_bindings()
267-{267+{
268- // Initialize the at::Type* pointers, name, and properties of the PyTensorType268+ // Initialize the at::Type* pointers, name, and properties of the PyTensorType
269- // vector. After this call, the vector must not be resized.269+ // vector. After this call, the vector must not be resized.
270- initialize_npu_aten_types(tensor_types);270+ initialize_npu_aten_types(tensor_types);
271- 271+ 
272- // Initialize the Python metaclass for the torch.FloatTensor, etc. types.272+ // Initialize the Python metaclass for the torch.FloatTensor, etc. types.
273- // The metaclass handles __instancecheck__ checks and binds the dtype property273+ // The metaclass handles __instancecheck__ checks and binds the dtype property
274- // on the type objects.274+ // on the type objects.
275- py_initialize_metaclass(metaclass);275+ py_initialize_metaclass(metaclass);
276- 276+ 
277- // Get the tp_dict of the Variable class. We copy function definitions277+ // Get the tp_dict of the Variable class. We copy function definitions
278- // onto each Tensor type object so that they can be accessed via e.g.278+ // onto each Tensor type object so that they can be accessed via e.g.
279- // `torch.npu.FloatTensor.add`.279+ // `torch.npu.FloatTensor.add`.
280- auto tensor_dict = get_tensor_dict();280+ auto tensor_dict = get_tensor_dict();
281- 281+ 
282- // Initialize each Python type object torch.npu.FloatTensor, torch.npu.DoubleTensor, etc.282+ // Initialize each Python type object torch.npu.FloatTensor, torch.npu.DoubleTensor, etc.
283- for (auto &tensor_type : tensor_types) {283+ for (auto &tensor_type : tensor_types) {
284- py_initialize_tensor_type(tensor_type.py_type, tensor_type.name, tensor_dict.get());284+ py_initialize_tensor_type(tensor_type.py_type, tensor_type.name, tensor_dict.get());
285- }285+ }
286- 286+ 
287- // Add the type objects to their corresponding modules. e.g. torch.npu.FloatTensor287+ // Add the type objects to their corresponding modules. e.g. torch.npu.FloatTensor
288- // is added to the `torch_npu` module as `FloatTensor`. Also add all the type288+ // is added to the `torch_npu` module as `FloatTensor`. Also add all the type
289- // objects to the set torch_npu._tensor_classes.289+ // objects to the set torch_npu._tensor_classes.
290- py_bind_tensor_types(tensor_types);290+ py_bind_tensor_types(tensor_types);
291-}291+}
292- 292+ 
293-static void py_bind_tensor_types(const std::vector<PyTensorType> &tensor_types)293+static void py_bind_tensor_types(const std::vector<PyTensorType> &tensor_types)
294-{294+{
295- auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));295+ auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));
296- if (!torch_module) {296+ if (!torch_module) {
297- throw python_error();297+ throw python_error();
298- }298+ }
299- 299+ 
300- auto tensor_classes = THPObjectPtr(PyObject_GetAttrString(torch_module.get(), "_tensor_classes"));300+ auto tensor_classes = THPObjectPtr(PyObject_GetAttrString(torch_module.get(), "_tensor_classes"));
301- if (!tensor_classes) {301+ if (!tensor_classes) {
302- throw python_error();302+ throw python_error();
303- }303+ }
304- 304+ 
305- for (auto &tensor_type : tensor_types) {305+ for (auto &tensor_type : tensor_types) {
306- auto name = std::string(tensor_type.name);306+ auto name = std::string(tensor_type.name);
307- auto idx = name.rfind('.');307+ auto idx = name.rfind('.');
308- auto type_name = name.substr(idx + 1);308+ auto type_name = name.substr(idx + 1);
309- auto module_name = name.substr(0, idx);309+ auto module_name = name.substr(0, idx);
310- 310+ 
311- auto module_obj = THPObjectPtr(PyImport_ImportModule(module_name.c_str()));311+ auto module_obj = THPObjectPtr(PyImport_ImportModule(module_name.c_str()));
312- if (!module_obj) {312+ if (!module_obj) {
313- throw python_error();313+ throw python_error();
314- }314+ }
315- 315+ 
316- PyObject *type_obj = (PyObject *)&tensor_type;316+ PyObject *type_obj = (PyObject *)&tensor_type;
317- Py_INCREF(type_obj);317+ Py_INCREF(type_obj);
318- if (PyModule_AddObject(module_obj.get(), type_name.c_str(), type_obj) < 0) {318+ if (PyModule_AddObject(module_obj.get(), type_name.c_str(), type_obj) < 0) {
319- throw python_error();319+ throw python_error();
320- }320+ }
321- if (PySet_Add(tensor_classes.get(), type_obj) < 0) {321+ if (PySet_Add(tensor_classes.get(), type_obj) < 0) {
322- throw python_error();322+ throw python_error();
323- }323+ }
324- }324+ }
325-}325+}
326- 326+ 
327-// Callback for python part. Used for additional initialization of python classes327+// Callback for python part. Used for additional initialization of python classes
328-static PyObject *THPModule_initExtension(PyObject *_unused, PyObject *noargs)328+static PyObject *THPModule_initExtension(PyObject *_unused, PyObject *noargs)
329-{329+{
330- HANDLE_TH_ERRORS330+ HANDLE_TH_ERRORS
331- _initialize_python_bindings();331+ _initialize_python_bindings();
332- Py_RETURN_NONE;332+ Py_RETURN_NONE;
333- END_HANDLE_TH_ERRORS333+ END_HANDLE_TH_ERRORS
334-}334+}
335- 335+ 
336-// autograd methods on torch._C336+// autograd methods on torch._C
337-static PyMethodDef TorchNpuExtensionMethods[] = {337+static PyMethodDef TorchNpuExtensionMethods[] = {
338- {"_initExtension", (PyCFunction)THPModule_initExtension, METH_NOARGS, nullptr},338+ {"_initExtension", (PyCFunction)THPModule_initExtension, METH_NOARGS, nullptr},
339- {nullptr, nullptr, 0, nullptr}339+ {nullptr, nullptr, 0, nullptr}
340-};340+};
341- 341+ 
342-PyMethodDef *npu_extension_functions()342+PyMethodDef *npu_extension_functions()
343-{343+{
344- return TorchNpuExtensionMethods;344+ return TorchNpuExtensionMethods;
345-}345+}
346-}346+}
347-}347+}
Mtorch_npu/csrc/utils/TensorType.h+16-16
@@ -1,16 +1,16 @@
1-#include <torch/csrc/utils/tensor_new.h>1+#include <torch/csrc/utils/tensor_new.h>
2- 2+ 
3-#include "torch_npu/csrc/core/npu/NPUMacros.h"3+#include "torch_npu/csrc/core/npu/NPUMacros.h"
4-#include "torch_npu/csrc/core/npu/NPUFunctions.h"4+#include "torch_npu/csrc/core/npu/NPUFunctions.h"
5- 5+ 
6-namespace torch_npu {6+namespace torch_npu {
7-namespace utils {7+namespace utils {
8- 8+ 
9-// Initializes the Python tensor type objects: torch.npu.FloatTensor,9+// Initializes the Python tensor type objects: torch.npu.FloatTensor,
10-// torch.npu.DoubleTensor, etc. and binds them in their containing modules.10+// torch.npu.DoubleTensor, etc. and binds them in their containing modules.
11-void _initialize_python_bindings();11+void _initialize_python_bindings();
12- 12+ 
13-TORCH_NPU_API PyMethodDef* npu_extension_functions();13+TORCH_NPU_API PyMethodDef* npu_extension_functions();
14- 14+ 
15-}15+}
16-}16+}
Mtorch_npu/distributed/__init__.py+30-30
@@ -1,30 +1,30 @@
1-__all__ = [1+__all__ = [
2- "is_hccl_available", "reinit_process_group", "reduce_scatter_tensor_uneven", "all_gather_into_tensor_uneven"2+ "is_hccl_available", "reinit_process_group", "reduce_scatter_tensor_uneven", "all_gather_into_tensor_uneven"
3-]3+]
4- 4+ 
5-from torch.distributed import _make_nccl_premul_sum as _make_hccl_premul_sum5+from torch.distributed import _make_nccl_premul_sum as _make_hccl_premul_sum
6- 6+ 
7-import torch_npu7+import torch_npu
8- 8+ 
9- 9+ 
10-def is_available():10+def is_available():
11- """11+ """
12- Returns ``True`` if the distributed package is available. Otherwise,12+ Returns ``True`` if the distributed package is available. Otherwise,
13- ``torch.distributed`` does not expose any other APIs. Currently,13+ ``torch.distributed`` does not expose any other APIs. Currently,
14- ``torch.distributed`` is available on Linux, MacOS and Windows. Set14+ ``torch.distributed`` is available on Linux, MacOS and Windows. Set
15- ``USE_DISTRIBUTED=1`` to enable it when building PyTorch from source.15+ ``USE_DISTRIBUTED=1`` to enable it when building PyTorch from source.
16- Currently, the default value is ``USE_DISTRIBUTED=1`` for Linux and Windows,16+ Currently, the default value is ``USE_DISTRIBUTED=1`` for Linux and Windows,
17- ``USE_DISTRIBUTED=0`` for MacOS.17+ ``USE_DISTRIBUTED=0`` for MacOS.
18- """18+ """
19- return hasattr(torch_npu._C, "_c10d_npu_init")19+ return hasattr(torch_npu._C, "_c10d_npu_init")
20- 20+ 
21- 21+ 
22-from torch_npu._C._distributed_c10d import (22+from torch_npu._C._distributed_c10d import (
23- ParallelStore,23+ ParallelStore,
24- _verify_params_across_processes,24+ _verify_params_across_processes,
25- _is_support_hccl_comm_name,25+ _is_support_hccl_comm_name,
26-)26+)
27- 27+ 
28- 28+ 
29-from torch_npu.distributed import tensor, nn29+from torch_npu.distributed import tensor, nn
30-from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven30+from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven
Mtorch_npu/dynamo/trace_rule.py+97-97
@@ -1,98 +1,98 @@
1-import torch1+import torch
2-from torch._dynamo.variables import TorchInGraphFunctionVariable2+from torch._dynamo.variables import TorchInGraphFunctionVariable
3-from torch._dynamo.trace_rules import manual_torch_name_rule_map, SkipFunctionVariable3+from torch._dynamo.trace_rules import manual_torch_name_rule_map, SkipFunctionVariable
4-import torch._dynamo.variables.torch as torch_module4+import torch._dynamo.variables.torch as torch_module
5-from torch._dynamo.utils import common_constant_types5+from torch._dynamo.utils import common_constant_types
6-import torch_npu6+import torch_npu
7- 7+ 
8-__all__ = []8+__all__ = []
9- 9+ 
10-torch_non_c_binding_in_graph_functions_npu = dict.fromkeys(10+torch_non_c_binding_in_graph_functions_npu = dict.fromkeys(
11- [11+ [
12- "torch.npu.current_stream",12+ "torch.npu.current_stream",
13- "torch.npu.default_stream",13+ "torch.npu.default_stream",
14- "torch.npu.stream",14+ "torch.npu.stream",
15- "torch.npu.set_stream",15+ "torch.npu.set_stream",
16- "torch.npu.current_device",16+ "torch.npu.current_device",
17- "torch.npu.get_device_capability",17+ "torch.npu.get_device_capability",
18- "torch.npu.get_device_properties",18+ "torch.npu.get_device_properties",
19- "torch.npu.graphs.graph_pool_handle",19+ "torch.npu.graphs.graph_pool_handle",
20- "torch.npu.ipc_collect",20+ "torch.npu.ipc_collect",
21- "torch.npu.is_available",21+ "torch.npu.is_available",
22- "torch.npu.memory._dump_snapshot",22+ "torch.npu.memory._dump_snapshot",
23- "torch.npu.memory._free_mutex",23+ "torch.npu.memory._free_mutex",
24- "torch.npu.memory._record_memory_history_impl",24+ "torch.npu.memory._record_memory_history_impl",
25- "torch.npu.memory._set_allocator_settings",25+ "torch.npu.memory._set_allocator_settings",
26- "torch.npu.memory.empty_cache",26+ "torch.npu.memory.empty_cache",
27- "torch.npu.mem_get_info",27+ "torch.npu.mem_get_info",
28- "torch.npu.memory.reset_accumulated_host_memory_stats",28+ "torch.npu.memory.reset_accumulated_host_memory_stats",
29- "torch.npu.memory.reset_accumulated_memory_stats",29+ "torch.npu.memory.reset_accumulated_memory_stats",
30- "torch.npu.memory.reset_max_memory_allocated",30+ "torch.npu.memory.reset_max_memory_allocated",
31- "torch.npu.memory.reset_max_memory_cached",31+ "torch.npu.memory.reset_max_memory_cached",
32- "torch.npu.memory.reset_peak_host_memory_stats",32+ "torch.npu.memory.reset_peak_host_memory_stats",
33- "torch.npu.memory.reset_peak_memory_stats",33+ "torch.npu.memory.reset_peak_memory_stats",
34- "torch.npu.memory.get_per_process_memory_fraction",34+ "torch.npu.memory.get_per_process_memory_fraction",
35- "torch.npu.memory.set_per_process_memory_fraction",35+ "torch.npu.memory.set_per_process_memory_fraction",
36- "torch.npu.random.manual_seed_all",36+ "torch.npu.random.manual_seed_all",
37- "torch.npu.random.manual_seed",37+ "torch.npu.random.manual_seed",
38- "torch.npu.random.seed_all",38+ "torch.npu.random.seed_all",
39- "torch.npu.random.seed",39+ "torch.npu.random.seed",
40- "torch.npu.set_sync_debug_mode",40+ "torch.npu.set_sync_debug_mode",
41- "torch.npu._set_rng_state_offset",41+ "torch.npu._set_rng_state_offset",
42- "torch.npu._get_generator",42+ "torch.npu._get_generator",
43- "torch.npu._memory_viz._frames_fmt",43+ "torch.npu._memory_viz._frames_fmt",
44- "torch.npu._memory_viz._frame_fmt",44+ "torch.npu._memory_viz._frame_fmt",
45- "torch.npu.amp.autocast_mode.custom_bwd",45+ "torch.npu.amp.autocast_mode.custom_bwd",
46- "torch.npu.amp.autocast_mode.custom_fwd",46+ "torch.npu.amp.autocast_mode.custom_fwd",
47- "torch.npu._get_current_allocator",47+ "torch.npu._get_current_allocator",
48- "torch.npu.is_bf16_supported",48+ "torch.npu.is_bf16_supported",
49- "torch.npu.memory._get_current_allocator",49+ "torch.npu.memory._get_current_allocator",
50- ],50+ ],
51- TorchInGraphFunctionVariable,51+ TorchInGraphFunctionVariable,
52-)52+)
53- 53+ 
54-torch_c_binding_in_graph_functions_npu = dict.fromkeys(54+torch_c_binding_in_graph_functions_npu = dict.fromkeys(
55- [55+ [
56- "torch_npu._C._npu_changeCurrentAllocator",56+ "torch_npu._C._npu_changeCurrentAllocator",
57- "torch_npu._C._npu_npuCachingAllocator_set_allocator_settings",57+ "torch_npu._C._npu_npuCachingAllocator_set_allocator_settings",
58- "torch_npu._C._npu_emptyCache",58+ "torch_npu._C._npu_emptyCache",
59- "torch_npu._C._npu_getAllocator",59+ "torch_npu._C._npu_getAllocator",
60- "torch_npu._C._npu_getCheckpointState",60+ "torch_npu._C._npu_getCheckpointState",
61- "torch_npu._C._npu_getCurrentStream",61+ "torch_npu._C._npu_getCurrentStream",
62- "torch_npu._C._npu_getDefaultStream",62+ "torch_npu._C._npu_getDefaultStream",
63- "torch_npu._C._npu_init",63+ "torch_npu._C._npu_init",
64- "torch_npu._C._npu_ipc_collect",64+ "torch_npu._C._npu_ipc_collect",
65- "torch_npu._C._npu_resetAccumulatedHostMemoryStats",65+ "torch_npu._C._npu_resetAccumulatedHostMemoryStats",
66- "torch_npu._C._npu_resetPeakHostMemoryStats",66+ "torch_npu._C._npu_resetPeakHostMemoryStats",
67- "torch_npu._C._npu_resetPeakMemoryStats",67+ "torch_npu._C._npu_resetPeakMemoryStats",
68- "torch_npu._C._npu_set_sync_debug_mode",68+ "torch_npu._C._npu_set_sync_debug_mode",
69- "torch_npu._C._npu_setDevice",69+ "torch_npu._C._npu_setDevice",
70- "torch_npu._C._npu_getMemoryFraction",70+ "torch_npu._C._npu_getMemoryFraction",
71- "torch_npu._C._npu_setMemoryFraction",71+ "torch_npu._C._npu_setMemoryFraction",
72- "torch_npu._C._npu_synchronize",72+ "torch_npu._C._npu_synchronize",
73- "torch_npu._C._npu_resetAccumulatedMemoryStats",73+ "torch_npu._C._npu_resetAccumulatedMemoryStats",
74- "torch_npu._C._npu_hasPrimaryContext",74+ "torch_npu._C._npu_hasPrimaryContext",
75- "torch_npu._C._npu_setStream",75+ "torch_npu._C._npu_setStream",
76- ],76+ ],
77- TorchInGraphFunctionVariable,77+ TorchInGraphFunctionVariable,
78-)78+)
79- 79+ 
80-skip_functions_npu = dict.fromkeys(80+skip_functions_npu = dict.fromkeys(
81- [81+ [
82- "torch_npu.npu.utils.synchronize",82+ "torch_npu.npu.utils.synchronize",
83- "torch.npu.set_device",83+ "torch.npu.set_device",
84- ],84+ ],
85- SkipFunctionVariable85+ SkipFunctionVariable
86-)86+)
87- 87+ 
88- 88+ 
89-def _patch_npu_trace_rules():89+def _patch_npu_trace_rules():
90- torch._dynamo.trace_rules.clear_lru_cache()90+ torch._dynamo.trace_rules.clear_lru_cache()
91- torch._dynamo.trace_rules.torch_name_rule_map.append(torch_non_c_binding_in_graph_functions_npu)91+ torch._dynamo.trace_rules.torch_name_rule_map.append(torch_non_c_binding_in_graph_functions_npu)
92- torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu)92+ torch._dynamo.trace_rules.torch_name_rule_map.append(torch_c_binding_in_graph_functions_npu)
93- torch._dynamo.trace_rules.torch_name_rule_map.append(skip_functions_npu)93+ torch._dynamo.trace_rules.torch_name_rule_map.append(skip_functions_npu)
94- torch_module.constant_fold_functions[torch.npu.current_device] = True94+ torch_module.constant_fold_functions[torch.npu.current_device] = True
95- torch_module.constant_fold_functions[torch.npu.get_device_properties] = True95+ torch_module.constant_fold_functions[torch.npu.get_device_properties] = True
96- torch_module.constant_fold_functions_need_guards[torch.npu.current_device] = True96+ torch_module.constant_fold_functions_need_guards[torch.npu.current_device] = True
97- torch_module.constant_fold_functions[torch.npu.is_available] = True97+ torch_module.constant_fold_functions[torch.npu.is_available] = True
98 common_constant_types.add(torch_npu._C._NPUDeviceProperties)98 common_constant_types.add(torch_npu._C._NPUDeviceProperties)
Mtorch_npu/npu/amp/__init__.py+5-5
@@ -1,6 +1,6 @@
1-__all__ = [1+__all__ = [
2- "autocast", "GradScaler", "custom_fwd", "custom_bwd"2+ "autocast", "GradScaler", "custom_fwd", "custom_bwd"
3-]3+]
4- 4+ 
5-from .autocast_mode import autocast, custom_fwd, custom_bwd # noqa: F4015+from .autocast_mode import autocast, custom_fwd, custom_bwd # noqa: F401
O
OopenLiBingCI5月17日

此条代码评论区间+2至+5

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。

likedislike
6from .grad_scaler import GradScaler # noqa: F4016from .grad_scaler import GradScaler # noqa: F401
Mtorch_npu/npu/amp/autocast_mode.py+97-97
@@ -1,98 +1,98 @@
1-__all__ = ["autocast", "custom_fwd", "custom_bwd"]1+__all__ = ["autocast", "custom_fwd", "custom_bwd"]
2- 2+ 
3- 3+ 
4-import functools4+import functools
5-import collections5+import collections
6-from typing import Any6+from typing import Any
7-from typing_extensions import deprecated7+from typing_extensions import deprecated
8- 8+ 
9-try:9+try:
10- import numpy as np10+ import numpy as np
11- 11+ 
12- HAS_NUMPY = True12+ HAS_NUMPY = True
13-except ModuleNotFoundError:13+except ModuleNotFoundError:
14- np = None # type: ignore[assignment]14+ np = None # type: ignore[assignment]
15- 15+ 
16-import torch16+import torch
17-import torch_npu17+import torch_npu
O
OopenLiBingCI5月17日

此条代码评论区间+12至+17

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
18-from torch_npu.utils._error_code import ErrCode, pta_error18+from torch_npu.utils._error_code import ErrCode, pta_error
19- 19+ 
20- 20+ 
21-class autocast(torch.amp.autocast_mode.autocast):21+class autocast(torch.amp.autocast_mode.autocast):
22- r"""22+ r"""
23- See :class:`torch.autocast`.23+ See :class:`torch.autocast`.
24- ``torch.npu.amp.autocast(args...)`` is equivalent to ``torch.autocast("npu", args...)``24+ ``torch.npu.amp.autocast(args...)`` is equivalent to ``torch.autocast("npu", args...)``
25- """25+ """
26- 26+ 
27- def __init__(self, enabled: bool = True, dtype: torch.dtype = torch.float16, cache_enabled: bool = True):27+ def __init__(self, enabled: bool = True, dtype: torch.dtype = torch.float16, cache_enabled: bool = True):
28- if torch._jit_internal.is_scripting():28+ if torch._jit_internal.is_scripting():
29- self._enabled = enabled29+ self._enabled = enabled
30- self.device = "npu"30+ self.device = "npu"
31- self.fast_dtype = dtype31+ self.fast_dtype = dtype
32- return32+ return
33- super().__init__("npu", enabled=enabled, dtype=dtype, cache_enabled=cache_enabled)33+ super().__init__("npu", enabled=enabled, dtype=dtype, cache_enabled=cache_enabled)
34- 34+ 
35- def __enter__(self):35+ def __enter__(self):
36- if torch._jit_internal.is_scripting():36+ if torch._jit_internal.is_scripting():
37- return self37+ return self
38- return super().__enter__()38+ return super().__enter__()
39- 39+ 
40- def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): # type: ignore[override]40+ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): # type: ignore[override]
41- if torch._jit_internal.is_scripting():41+ if torch._jit_internal.is_scripting():
42- return None42+ return None
O
OopenLiBingCI5月17日

此条代码评论区间+37至+42

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
43- return super().__exit__(exc_type, exc_val, exc_tb)43+ return super().__exit__(exc_type, exc_val, exc_tb)
44- 44+ 
45- def __call__(self, func):45+ def __call__(self, func):
46- if torch._jit_internal.is_scripting():46+ if torch._jit_internal.is_scripting():
47- return func47+ return func
48- return super().__call__(func)48+ return super().__call__(func)
49- 49+ 
50- 50+ 
51-# Casts Tensors and containers of Tensors. Special-cases passthroughs for strings and np.ndarrays, which51+# Casts Tensors and containers of Tensors. Special-cases passthroughs for strings and np.ndarrays, which
52-# may be falsely detected as "Iterables."52+# may be falsely detected as "Iterables."
53-def _cast(value, dtype):53+def _cast(value, dtype):
54- if isinstance(value, torch.Tensor):54+ if isinstance(value, torch.Tensor):
55- is_eligible = (value.is_floating_point() and value.device.type == 'npu' and (value.dtype is not torch.float64))55+ is_eligible = (value.is_floating_point() and value.device.type == 'npu' and (value.dtype is not torch.float64))
56- return value.to(dtype) if is_eligible else value56+ return value.to(dtype) if is_eligible else value
57- elif isinstance(value, (str, bytes)):57+ elif isinstance(value, (str, bytes)):
58- return value58+ return value
59- elif HAS_NUMPY and isinstance(value, np.ndarray):59+ elif HAS_NUMPY and isinstance(value, np.ndarray):
60- return value60+ return value
61- elif isinstance(value, collections.abc.Mapping):61+ elif isinstance(value, collections.abc.Mapping):
62- return {_cast(k, dtype): _cast(v, dtype) for k, v in value.items()}62+ return {_cast(k, dtype): _cast(v, dtype) for k, v in value.items()}
63- elif isinstance(value, collections.abc.Iterable):63+ elif isinstance(value, collections.abc.Iterable):
64- iterable = map(lambda v: _cast(v, dtype), value)64+ iterable = map(lambda v: _cast(v, dtype), value)
65- if isinstance(value, list) or isinstance(value, tuple):65+ if isinstance(value, list) or isinstance(value, tuple):
66- return type(value)(iterable)66+ return type(value)(iterable)
67- else:67+ else:
68- return iterable68+ return iterable
69- else:69+ else:
70- return value70+ return value
71- 71+ 
72- 72+ 
73-@deprecated(73+@deprecated(
74- "`torch_npu.npu.amp.custom_fwd(args...)` is deprecated. "74+ "`torch_npu.npu.amp.custom_fwd(args...)` is deprecated. "
75- "Please use `torch.amp.custom_fwd(args..., device_type='npu')` instead.",75+ "Please use `torch.amp.custom_fwd(args..., device_type='npu')` instead.",
76- category=FutureWarning,76+ category=FutureWarning,
77-)77+)
78-def custom_fwd(fwd=None, *, cast_inputs=None):78+def custom_fwd(fwd=None, *, cast_inputs=None):
79- """79+ """
80- ``torch_npu.npu.amp.custom_fwd(args...)`` is deprecated. Please use80+ ``torch_npu.npu.amp.custom_fwd(args...)`` is deprecated. Please use
81- ``torch.amp.custom_fwd(args..., device_type='npu')`` instead.81+ ``torch.amp.custom_fwd(args..., device_type='npu')`` instead.
82- """82+ """
83- return functools.partial(torch.amp.custom_fwd, device_type="npu")(83+ return functools.partial(torch.amp.custom_fwd, device_type="npu")(
84- fwd=fwd, cast_inputs=cast_inputs84+ fwd=fwd, cast_inputs=cast_inputs
85- )85+ )
86- 86+ 
87- 87+ 
88-@deprecated(88+@deprecated(
89- "`torch_npu.npu.amp.custom_bwd(args...)` is deprecated. "89+ "`torch_npu.npu.amp.custom_bwd(args...)` is deprecated. "
90- "Please use `torch.amp.custom_bwd(args..., device_type='npu')` instead.",90+ "Please use `torch.amp.custom_bwd(args..., device_type='npu')` instead.",
91- category=FutureWarning,91+ category=FutureWarning,
92-)92+)
93-def custom_bwd(bwd):93+def custom_bwd(bwd):
94- """94+ """
95- ``torch_npu.npu.amp.custom_bwd(args...)`` is deprecated. Please use95+ ``torch_npu.npu.amp.custom_bwd(args...)`` is deprecated. Please use
96- ``torch.amp.custom_bwd(args..., device_type='npu')`` instead.96+ ``torch.amp.custom_bwd(args..., device_type='npu')`` instead.
97- """97+ """
98 return functools.partial(torch.amp.custom_bwd, device_type="npu")(bwd)98 return functools.partial(torch.amp.custom_bwd, device_type="npu")(bwd)
Mtorch_npu/npu/amp/common.py+5-5
@@ -1,5 +1,5 @@
1-import torch_npu1+import torch_npu
2- 2+ 
3- 3+ 
4-def amp_definitely_not_available():4+def amp_definitely_not_available():
5- return not torch_npu.npu.is_available()5+ return not torch_npu.npu.is_available()
Mtorch_npu/npu/amp/grad_scaler.py+505-505
@@ -1,505 +1,505 @@
1-import warnings1+import warnings
2-from collections import defaultdict2+from collections import defaultdict
3-import collections.abc as container_abcs3+import collections.abc as container_abcs
4-from typing import List4+from typing import List
5- 5+ 
6-import torch6+import torch
7-import torch.distributed as dist7+import torch.distributed as dist
8-from torch.amp.grad_scaler import _MultiDeviceReplicator, OptState, _refresh_per_optimizer_state8+from torch.amp.grad_scaler import _MultiDeviceReplicator, OptState, _refresh_per_optimizer_state
9-from torch.amp.grad_scaler import GradScaler as BaseGradScaler9+from torch.amp.grad_scaler import GradScaler as BaseGradScaler
10-import torch_npu10+import torch_npu
11-from torch_npu.utils._error_code import ErrCode, pta_error11+from torch_npu.utils._error_code import ErrCode, pta_error
12-from .common import amp_definitely_not_available12+from .common import amp_definitely_not_available
13- 13+ 
14- 14+ 
15-class _NpuMultiDeviceReplicator(_MultiDeviceReplicator):15+class _NpuMultiDeviceReplicator(_MultiDeviceReplicator):
16- """16+ """
17- Lazily serves copies of a tensor to requested devices. Copies are cached per-device.17+ Lazily serves copies of a tensor to requested devices. Copies are cached per-device.
18- """18+ """
19- 19+ 
20- def __init__(self, master_tensor: torch.Tensor) -> None:20+ def __init__(self, master_tensor: torch.Tensor) -> None:
21- if not master_tensor.is_npu:21+ if not master_tensor.is_npu:
22- raise ValueError("Device type of master_tensor should be npu." + pta_error(ErrCode.VALUE))22+ raise ValueError("Device type of master_tensor should be npu." + pta_error(ErrCode.VALUE))
23- self.master = master_tensor23+ self.master = master_tensor
24- self._per_device_tensors = {}24+ self._per_device_tensors = {}
25- 25+ 
26- 26+ 
27-class GradScaler(BaseGradScaler):27+class GradScaler(BaseGradScaler):
28- """28+ """
29- An instance ``scaler`` of :class:`GradScaler` helps perform the steps of gradient scaling29+ An instance ``scaler`` of :class:`GradScaler` helps perform the steps of gradient scaling
30- conveniently.30+ conveniently.
31- 31+ 
32- * ``scaler.scale(loss)`` multiplies a given loss by ``scaler``'s current scale factor.32+ * ``scaler.scale(loss)`` multiplies a given loss by ``scaler``'s current scale factor.
33- * ``scaler.step(optimizer)`` safely unscales gradients and calls ``optimizer.step()``.33+ * ``scaler.step(optimizer)`` safely unscales gradients and calls ``optimizer.step()``.
34- * ``scaler.update()`` updates ``scaler``'s scale factor.34+ * ``scaler.update()`` updates ``scaler``'s scale factor.
35- 35+ 
36- Example::36+ Example::
37- 37+ 
38- # Creates a GradScaler once at the beginning of training.38+ # Creates a GradScaler once at the beginning of training.
39- scaler = GradScaler()39+ scaler = GradScaler()
40- 40+ 
41- for epoch in epochs:41+ for epoch in epochs:
42- for input, target in data:42+ for input, target in data:
43- optimizer.zero_grad()43+ optimizer.zero_grad()
44- output = model(input)44+ output = model(input)
45- loss = loss_fn(output, target)45+ loss = loss_fn(output, target)
46- 46+ 
47- # Scales loss. Calls backward() on scaled loss to create scaled gradients.47+ # Scales loss. Calls backward() on scaled loss to create scaled gradients.
48- scaler.scale(loss).backward()48+ scaler.scale(loss).backward()
49- 49+ 
50- # scaler.step() first unscales gradients of the optimizer's params.50+ # scaler.step() first unscales gradients of the optimizer's params.
51- # If gradients don't contain infs/NaNs, optimizer.step() is then called,51+ # If gradients don't contain infs/NaNs, optimizer.step() is then called,
52- # otherwise, optimizer.step() is skipped.52+ # otherwise, optimizer.step() is skipped.
53- scaler.step(optimizer)53+ scaler.step(optimizer)
54- 54+ 
55- # Updates the scale for next iteration.55+ # Updates the scale for next iteration.
56- scaler.update()56+ scaler.update()
57- 57+ 
58- See the :ref:`Automatic Mixed Precision examples<amp-examples>` for usage58+ See the :ref:`Automatic Mixed Precision examples<amp-examples>` for usage
59- (along with autocasting) in more complex cases like gradient clipping, gradient accumulation, gradient penalty,59+ (along with autocasting) in more complex cases like gradient clipping, gradient accumulation, gradient penalty,
60- and multiple losses/optimizers.60+ and multiple losses/optimizers.
61- 61+ 
62- ``scaler`` dynamically estimates the scale factor each iteration. To minimize gradient underflow,62+ ``scaler`` dynamically estimates the scale factor each iteration. To minimize gradient underflow,
63- a large scale factor should be used. However, ``float16`` values can "overflow" (become inf or NaN) if63+ a large scale factor should be used. However, ``float16`` values can "overflow" (become inf or NaN) if
64- the scale factor is too large. Therefore, the optimal scale factor is the largest factor that can be used64+ the scale factor is too large. Therefore, the optimal scale factor is the largest factor that can be used
65- without incurring inf or NaN gradient values.65+ without incurring inf or NaN gradient values.
66- ``scaler`` approximates the optimal scale factor over time by checking the gradients for infs and NaNs during every66+ ``scaler`` approximates the optimal scale factor over time by checking the gradients for infs and NaNs during every
67- ``scaler.step(optimizer)`` (or optional separate ``scaler.unscale_(optimizer)``, see :meth:`unscale_`).67+ ``scaler.step(optimizer)`` (or optional separate ``scaler.unscale_(optimizer)``, see :meth:`unscale_`).
68- 68+ 
69- * If infs/NaNs are found, ``scaler.step(optimizer)`` skips the underlying ``optimizer.step()`` (so the params69+ * If infs/NaNs are found, ``scaler.step(optimizer)`` skips the underlying ``optimizer.step()`` (so the params
70- themselves remain uncorrupted) and ``update()`` multiplies the scale by ``backoff_factor``.70+ themselves remain uncorrupted) and ``update()`` multiplies the scale by ``backoff_factor``.
71- 71+ 
72- * If no infs/NaNs are found, ``scaler.step(optimizer)`` runs the underlying ``optimizer.step()`` as usual.72+ * If no infs/NaNs are found, ``scaler.step(optimizer)`` runs the underlying ``optimizer.step()`` as usual.
73- If ``growth_interval`` unskipped iterations occur consecutively, ``update()`` multiplies the scale by73+ If ``growth_interval`` unskipped iterations occur consecutively, ``update()`` multiplies the scale by
74- ``growth_factor``.74+ ``growth_factor``.
75- 75+ 
76- The scale factor often causes infs/NaNs to appear in gradients for the first few iterations as its76+ The scale factor often causes infs/NaNs to appear in gradients for the first few iterations as its
77- value calibrates. ``scaler.step`` will skip the underlying ``optimizer.step()`` for these77+ value calibrates. ``scaler.step`` will skip the underlying ``optimizer.step()`` for these
78- iterations. After that, step skipping should occur rarely (once every few hundred or thousand iterations).78+ iterations. After that, step skipping should occur rarely (once every few hundred or thousand iterations).
79- 79+ 
80- Args:80+ Args:
81- init_scale (float, optional, default=2.**16): Initial scale factor.81+ init_scale (float, optional, default=2.**16): Initial scale factor.
82- growth_factor (float, optional, default=2.0): Factor by which the scale is multiplied during82+ growth_factor (float, optional, default=2.0): Factor by which the scale is multiplied during
83- :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.83+ :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.
84- backoff_factor (float, optional, default=0.5): Factor by which the scale is multiplied during84+ backoff_factor (float, optional, default=0.5): Factor by which the scale is multiplied during
85- :meth:`update` if inf/NaN gradients occur in an iteration.85+ :meth:`update` if inf/NaN gradients occur in an iteration.
86- growth_interval (int, optional, default=2000): Number of consecutive iterations without inf/NaN gradients86+ growth_interval (int, optional, default=2000): Number of consecutive iterations without inf/NaN gradients
87- that must occur for the scale to be multiplied by ``growth_factor``.87+ that must occur for the scale to be multiplied by ``growth_factor``.
88- dynamic (bool, optional, default=True): If ``False``, use static loss scale.88+ dynamic (bool, optional, default=True): If ``False``, use static loss scale.
89- enabled (bool, optional, default=True): If ``False``, disables gradient scaling. :meth:`step` simply89+ enabled (bool, optional, default=True): If ``False``, disables gradient scaling. :meth:`step` simply
90- invokes the underlying ``optimizer.step()``, and other methods become no-ops.90+ invokes the underlying ``optimizer.step()``, and other methods become no-ops.
91- """91+ """
92- 92+ 
93- def __init__(self,93+ def __init__(self,
94- init_scale=2. ** 16,94+ init_scale=2. ** 16,
95- growth_factor=2.0,95+ growth_factor=2.0,
96- backoff_factor=0.5,96+ backoff_factor=0.5,
97- growth_interval=2000,97+ growth_interval=2000,
98- dynamic=True,98+ dynamic=True,
99- enabled=True):99+ enabled=True):
100- if enabled and amp_definitely_not_available():100+ if enabled and amp_definitely_not_available():
101- warnings.warn("torch_npu.amp.GradScaler is enabled, but NPU is not available. Disabling.")101+ warnings.warn("torch_npu.amp.GradScaler is enabled, but NPU is not available. Disabling.")
102- self._enabled = False102+ self._enabled = False
103- else:103+ else:
104- self._enabled = enabled104+ self._enabled = enabled
105- 105+ 
106- if self._enabled:106+ if self._enabled:
107- if growth_factor <= 1.0:107+ if growth_factor <= 1.0:
108- raise ValueError("The growth factor must be > 1.0." + pta_error(ErrCode.VALUE))108+ raise ValueError("The growth factor must be > 1.0." + pta_error(ErrCode.VALUE))
109- if backoff_factor >= 1.0:109+ if backoff_factor >= 1.0:
110- raise ValueError("The backoff factor must be < 1.0." + pta_error(ErrCode.VALUE))110+ raise ValueError("The backoff factor must be < 1.0." + pta_error(ErrCode.VALUE))
111- 111+ 
112- self._init_scale = init_scale112+ self._init_scale = init_scale
113- # self._scale will be lazily initialized during the first call to scale()113+ # self._scale will be lazily initialized during the first call to scale()
114- self._scale = None114+ self._scale = None
115- self._growth_factor = growth_factor115+ self._growth_factor = growth_factor
116- self._backoff_factor = backoff_factor116+ self._backoff_factor = backoff_factor
117- self._growth_interval = growth_interval117+ self._growth_interval = growth_interval
118- self._dynamic = dynamic118+ self._dynamic = dynamic
119- self._init_growth_tracker = 0119+ self._init_growth_tracker = 0
120- # self._growth_tracker will be lazily initialized during the first call to scale()120+ # self._growth_tracker will be lazily initialized during the first call to scale()
121- self._growth_tracker = None121+ self._growth_tracker = None
122- self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)122+ self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
123- self._has_overflow = False123+ self._has_overflow = False
124- self._clear_overflow_flag = False124+ self._clear_overflow_flag = False
125- self._dist_initialized = False125+ self._dist_initialized = False
126- self._dist_overflow_count = None126+ self._dist_overflow_count = None
127- 127+ 
128- def _lazy_init_scale_growth_tracker(self, dev):128+ def _lazy_init_scale_growth_tracker(self, dev):
129- if self._growth_tracker is not None:129+ if self._growth_tracker is not None:
130- raise RuntimeError("_growth_tracker initialized before _scale" + pta_error(ErrCode.VALUE))130+ raise RuntimeError("_growth_tracker initialized before _scale" + pta_error(ErrCode.VALUE))
131- 131+ 
132- self._scale = torch.full((), self._init_scale, dtype=torch.float32).pin_memory().to(dev, non_blocking=True)132+ self._scale = torch.full((), self._init_scale, dtype=torch.float32).pin_memory().to(dev, non_blocking=True)
133- self._growth_tracker = torch.full((), self._init_growth_tracker, dtype=torch.int32)133+ self._growth_tracker = torch.full((), self._init_growth_tracker, dtype=torch.int32)
134- self._growth_tracker = self._growth_tracker.pin_memory().to(dev, non_blocking=True)134+ self._growth_tracker = self._growth_tracker.pin_memory().to(dev, non_blocking=True)
135- 135+ 
136- def _lazy_init_dist_flag_and_dist_overflow_count(self):136+ def _lazy_init_dist_flag_and_dist_overflow_count(self):
137- if self._dist_overflow_count is not None:137+ if self._dist_overflow_count is not None:
138- raise RuntimeError("_dist_overflow_count initialized before _scale" + pta_error(ErrCode.VALUE))138+ raise RuntimeError("_dist_overflow_count initialized before _scale" + pta_error(ErrCode.VALUE))
139- try:139+ try:
140- if dist.is_initialized():140+ if dist.is_initialized():
141- self._dist_initialized = True141+ self._dist_initialized = True
142- except AttributeError as err:142+ except AttributeError as err:
143- print("torch.distributed has no attribute is_initialized")143+ print("torch.distributed has no attribute is_initialized")
144- 144+ 
145- self._dist_overflow_count = torch.Tensor([0.]).to('npu')145+ self._dist_overflow_count = torch.Tensor([0.]).to('npu')
146- 146+ 
147- def scale(self, outputs):147+ def scale(self, outputs):
148- """148+ """
149- Multiplies ('scales') a tensor or list of tensors by the scale factor.149+ Multiplies ('scales') a tensor or list of tensors by the scale factor.
150- 150+ 
151- Returns scaled outputs. If this instance of :class:`GradScaler` is not enabled, outputs are returned151+ Returns scaled outputs. If this instance of :class:`GradScaler` is not enabled, outputs are returned
152- unmodified.152+ unmodified.
153- 153+ 
154- Args:154+ Args:
155- outputs (Tensor or iterable of Tensors): Outputs to scale.155+ outputs (Tensor or iterable of Tensors): Outputs to scale.
156- """156+ """
157- if not self._enabled:157+ if not self._enabled:
158- return outputs158+ return outputs
159- 159+ 
160- if self._dist_overflow_count is None:160+ if self._dist_overflow_count is None:
161- self._lazy_init_dist_flag_and_dist_overflow_count()161+ self._lazy_init_dist_flag_and_dist_overflow_count()
162- if self._dist_overflow_count is None:162+ if self._dist_overflow_count is None:
163- raise RuntimeError("_dist_overflow_count is None." + pta_error(ErrCode.VALUE))163+ raise RuntimeError("_dist_overflow_count is None." + pta_error(ErrCode.VALUE))
164- 164+ 
165- if self._dynamic and not self._clear_overflow_flag:165+ if self._dynamic and not self._clear_overflow_flag:
166- if not torch_npu.npu.utils.is_support_inf_nan():166+ if not torch_npu.npu.utils.is_support_inf_nan():
167- GradScaler.clear_npu_overflow_flag()167+ GradScaler.clear_npu_overflow_flag()
168- self._clear_overflow_flag = True168+ self._clear_overflow_flag = True
169- 169+ 
170- # Short-circuit for the common case.170+ # Short-circuit for the common case.
171- if isinstance(outputs, torch.Tensor):171+ if isinstance(outputs, torch.Tensor):
172- if not outputs.is_npu:172+ if not outputs.is_npu:
173- raise ValueError("Device type of outputs should be npu." + pta_error(ErrCode.VALUE))173+ raise ValueError("Device type of outputs should be npu." + pta_error(ErrCode.VALUE))
174- if self._scale is None:174+ if self._scale is None:
175- self._lazy_init_scale_growth_tracker(outputs.device)175+ self._lazy_init_scale_growth_tracker(outputs.device)
176- if self._scale is None:176+ if self._scale is None:
177- raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))177+ raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))
178- return outputs * self._scale.to(device=outputs.device, non_blocking=True)178+ return outputs * self._scale.to(device=outputs.device, non_blocking=True)
179- 179+ 
180- # Invoke the more complex machinery only if we're treating multiple outputs.180+ # Invoke the more complex machinery only if we're treating multiple outputs.
181- stash: List[_NpuMultiDeviceReplicator] = [] # holds a reference that can be overwritten by apply_scale181+ stash: List[_NpuMultiDeviceReplicator] = [] # holds a reference that can be overwritten by apply_scale
182- 182+ 
183- def apply_scale(val):183+ def apply_scale(val):
184- if isinstance(val, torch.Tensor):184+ if isinstance(val, torch.Tensor):
185- if not val.is_npu:185+ if not val.is_npu:
186- raise ValueError("Device type of val should be npu." + pta_error(ErrCode.VALUE))186+ raise ValueError("Device type of val should be npu." + pta_error(ErrCode.VALUE))
187- if len(stash) == 0:187+ if len(stash) == 0:
188- if self._scale is None:188+ if self._scale is None:
189- self._lazy_init_scale_growth_tracker(val.device)189+ self._lazy_init_scale_growth_tracker(val.device)
190- if self._scale is None:190+ if self._scale is None:
191- raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))191+ raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))
192- stash.append(_NpuMultiDeviceReplicator(self._scale))192+ stash.append(_NpuMultiDeviceReplicator(self._scale))
193- return val * stash[0].get(val.device)193+ return val * stash[0].get(val.device)
194- elif isinstance(val, container_abcs.Iterable):194+ elif isinstance(val, container_abcs.Iterable):
195- iterable = map(apply_scale, val)195+ iterable = map(apply_scale, val)
196- if isinstance(val, list) or isinstance(val, tuple):196+ if isinstance(val, list) or isinstance(val, tuple):
197- return type(val)(iterable)197+ return type(val)(iterable)
198- else:198+ else:
199- return iterable199+ return iterable
200- else:200+ else:
201- raise TypeError("outputs must be a Tensor or an iterable of Tensors" + pta_error(ErrCode.TYPE))201+ raise TypeError("outputs must be a Tensor or an iterable of Tensors" + pta_error(ErrCode.TYPE))
202- 202+ 
203- return apply_scale(outputs)203+ return apply_scale(outputs)
204- 204+ 
205- def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16):205+ def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16):
206- per_device_found_inf = _NpuMultiDeviceReplicator(found_inf)206+ per_device_found_inf = _NpuMultiDeviceReplicator(found_inf)
207- per_device_inv_scale = _NpuMultiDeviceReplicator(inv_scale)207+ per_device_inv_scale = _NpuMultiDeviceReplicator(inv_scale)
208- 208+ 
209- # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.209+ # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
210- # There could be hundreds of grads, so we'd like to iterate through them just once.210+ # There could be hundreds of grads, so we'd like to iterate through them just once.
211- # However, we don't know their devices or dtypes in advance.211+ # However, we don't know their devices or dtypes in advance.
212- 212+ 
213- # Google says mypy struggles with defaultdicts type annotations.213+ # Google says mypy struggles with defaultdicts type annotations.
214- per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))214+ per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))
215- with torch.no_grad():215+ with torch.no_grad():
216- if hasattr(optimizer, 'is_fused_optimizer'):216+ if hasattr(optimizer, 'is_fused_optimizer'):
217- if not optimizer.is_params_grads_combined:217+ if not optimizer.is_params_grads_combined:
218- optimizer._maybe_init_combined_params_and_grads()218+ optimizer._maybe_init_combined_params_and_grads()
219- 219+ 
220- device = found_inf.device220+ device = found_inf.device
221- for grads_combined_one_dtype in optimizer.grads_all_group_combined:221+ for grads_combined_one_dtype in optimizer.grads_all_group_combined:
222- if grads_combined_one_dtype is None:222+ if grads_combined_one_dtype is None:
223- continue223+ continue
224- if self._dynamic:224+ if self._dynamic:
225- torch._amp_foreach_non_finite_check_and_unscale_(225+ torch._amp_foreach_non_finite_check_and_unscale_(
226- [grads_combined_one_dtype],226+ [grads_combined_one_dtype],
227- per_device_found_inf.get(device),227+ per_device_found_inf.get(device),
228- per_device_inv_scale.get(device))228+ per_device_inv_scale.get(device))
229- if per_device_found_inf.get(device).item() > 0:229+ if per_device_found_inf.get(device).item() > 0:
230- self._has_overflow = True230+ self._has_overflow = True
231- else:231+ else:
232- grads_combined_one_dtype.mul_(232+ grads_combined_one_dtype.mul_(
233- per_device_inv_scale.get(device))233+ per_device_inv_scale.get(device))
234- else:234+ else:
235- for group in optimizer.param_groups:235+ for group in optimizer.param_groups:
236- for param in group["params"]:236+ for param in group["params"]:
237- if param.grad is None:237+ if param.grad is None:
238- continue238+ continue
239- if (not allow_fp16) and param.grad.dtype == torch.float16:239+ if (not allow_fp16) and param.grad.dtype == torch.float16:
240- raise TypeError("Attempting to unscale FP16 gradients." + pta_error(ErrCode.TYPE))240+ raise TypeError("Attempting to unscale FP16 gradients." + pta_error(ErrCode.TYPE))
241- if param.grad.is_sparse:241+ if param.grad.is_sparse:
242- # is_coalesced() == False means the sparse grad has values with duplicate indices.242+ # is_coalesced() == False means the sparse grad has values with duplicate indices.
243- # coalesce() deduplicates indices and adds all values that have the same index.243+ # coalesce() deduplicates indices and adds all values that have the same index.
244- # For scaled fp16 values, there's a good chance coalescing will cause overflow,244+ # For scaled fp16 values, there's a good chance coalescing will cause overflow,
245- # so we should check the coalesced _values().245+ # so we should check the coalesced _values().
246- if param.grad.dtype == torch.float16:246+ if param.grad.dtype == torch.float16:
247- param.grad = param.grad.coalesce()247+ param.grad = param.grad.coalesce()
248- to_unscale = param.grad._values()248+ to_unscale = param.grad._values()
249- else:249+ else:
250- to_unscale = param.grad250+ to_unscale = param.grad
251- 251+ 
252- per_device_and_dtype_grads[to_unscale.device][to_unscale.dtype].append(to_unscale)252+ per_device_and_dtype_grads[to_unscale.device][to_unscale.dtype].append(to_unscale)
253- 253+ 
254- for device, per_dtype_grads in per_device_and_dtype_grads.items():254+ for device, per_dtype_grads in per_device_and_dtype_grads.items():
255- for grads in per_dtype_grads.values():255+ for grads in per_dtype_grads.values():
256- if self._dynamic:256+ if self._dynamic:
257- torch._amp_foreach_non_finite_check_and_unscale_(grads,257+ torch._amp_foreach_non_finite_check_and_unscale_(grads,
258- per_device_found_inf.get(device),258+ per_device_found_inf.get(device),
259- per_device_inv_scale.get(device))259+ per_device_inv_scale.get(device))
260- if per_device_found_inf.get(device).item() > 0:260+ if per_device_found_inf.get(device).item() > 0:
261- self._has_overflow = True261+ self._has_overflow = True
262- else:262+ else:
263- for grad in grads:263+ for grad in grads:
264- grad.mul_(per_device_inv_scale.get(device))264+ grad.mul_(per_device_inv_scale.get(device))
265- 265+ 
266- self._sync_dist_overflow_count()266+ self._sync_dist_overflow_count()
267- if self._has_overflow:267+ if self._has_overflow:
268- per_device_found_inf.get(found_inf.device).add_(1)268+ per_device_found_inf.get(found_inf.device).add_(1)
269- else:269+ else:
270- per_device_found_inf.get(found_inf.device)270+ per_device_found_inf.get(found_inf.device)
271- 271+ 
272- return per_device_found_inf._per_device_tensors272+ return per_device_found_inf._per_device_tensors
273- 273+ 
274- def unscale_(self, optimizer):274+ def unscale_(self, optimizer):
275- """275+ """
276- Divides ("unscales") the optimizer's gradient tensors by the scale factor.276+ Divides ("unscales") the optimizer's gradient tensors by the scale factor.
277- 277+ 
278- :meth:`unscale_` is optional, serving cases where you need to278+ :meth:`unscale_` is optional, serving cases where you need to
279- :ref:`modify or inspect gradients<working-with-unscaled-gradients>`279+ :ref:`modify or inspect gradients<working-with-unscaled-gradients>`
280- between the backward pass(es) and :meth:`step`.280+ between the backward pass(es) and :meth:`step`.
281- If :meth:`unscale_` is not called explicitly, gradients will be unscaled automatically during :meth:`step`.281+ If :meth:`unscale_` is not called explicitly, gradients will be unscaled automatically during :meth:`step`.
282- 282+ 
283- Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients::283+ Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients::
284- 284+ 
285- ...285+ ...
286- scaler.scale(loss).backward()286+ scaler.scale(loss).backward()
287- scaler.unscale_(optimizer)287+ scaler.unscale_(optimizer)
288- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)288+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
289- scaler.step(optimizer)289+ scaler.step(optimizer)
290- scaler.update()290+ scaler.update()
291- 291+ 
292- Args:292+ Args:
293- optimizer (torch.optim.Optimizer): Optimizer that owns the gradients to be unscaled.293+ optimizer (torch.optim.Optimizer): Optimizer that owns the gradients to be unscaled.
294- 294+ 
295- .. note::295+ .. note::
296- :meth:`unscale_` does not incur a CPU-NPU sync.296+ :meth:`unscale_` does not incur a CPU-NPU sync.
297- 297+ 
298- .. warning::298+ .. warning::
299- :meth:`unscale_` should only be called once per optimizer per :meth:`step` call,299+ :meth:`unscale_` should only be called once per optimizer per :meth:`step` call,
300- and only after all gradients for that optimizer's assigned parameters have been accumulated.300+ and only after all gradients for that optimizer's assigned parameters have been accumulated.
301- Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError.301+ Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError.
302- 302+ 
303- .. warning::303+ .. warning::
304- :meth:`unscale_` may unscale sparse gradients out of place, replacing the ``.grad`` attribute.304+ :meth:`unscale_` may unscale sparse gradients out of place, replacing the ``.grad`` attribute.
305- """305+ """
306- if not self._enabled:306+ if not self._enabled:
307- return307+ return
308- 308+ 
309- self._check_scale_growth_tracker("unscale_")309+ self._check_scale_growth_tracker("unscale_")
310- 310+ 
311- optimizer_state = self._per_optimizer_states[id(optimizer)]311+ optimizer_state = self._per_optimizer_states[id(optimizer)]
312- 312+ 
313- if optimizer_state["stage"] is OptState.UNSCALED:313+ if optimizer_state["stage"] is OptState.UNSCALED:
314- raise RuntimeError("unscale_() has already been called on this optimizer since the last update()." +314+ raise RuntimeError("unscale_() has already been called on this optimizer since the last update()." +
315- pta_error(ErrCode.INTERNAL))315+ pta_error(ErrCode.INTERNAL))
316- elif optimizer_state["stage"] is OptState.STEPPED:316+ elif optimizer_state["stage"] is OptState.STEPPED:
317- raise RuntimeError("unscale_() is being called after step()." +317+ raise RuntimeError("unscale_() is being called after step()." +
318- pta_error(ErrCode.INTERNAL))318+ pta_error(ErrCode.INTERNAL))
319- 319+ 
320- # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.320+ # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
321- if self._scale is None:321+ if self._scale is None:
322- raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))322+ raise RuntimeError("_scale is None." + pta_error(ErrCode.VALUE))
323- inv_scale = self._scale.float().reciprocal()323+ inv_scale = self._scale.float().reciprocal()
324- found_inf = torch.full((), 0.0, dtype=torch.float32).pin_memory().to(self._scale.device, non_blocking=True)324+ found_inf = torch.full((), 0.0, dtype=torch.float32).pin_memory().to(self._scale.device, non_blocking=True)
325- 325+ 
326- optimizer_state["found_inf_per_device"] = self._unscale_grads_(optimizer, inv_scale, found_inf, False)326+ optimizer_state["found_inf_per_device"] = self._unscale_grads_(optimizer, inv_scale, found_inf, False)
327- optimizer_state["stage"] = OptState.UNSCALED327+ optimizer_state["stage"] = OptState.UNSCALED
328- 328+ 
329- def _maybe_opt_step(self, optimizer, optimizer_state, *args, **kwargs):329+ def _maybe_opt_step(self, optimizer, optimizer_state, *args, **kwargs):
330- retval = None330+ retval = None
331- if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()) and not self._has_overflow:331+ if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()) and not self._has_overflow:
332- retval = optimizer.step(*args, **kwargs)332+ retval = optimizer.step(*args, **kwargs)
333- else:333+ else:
334- print("Gradient overflow. Skipping step")334+ print("Gradient overflow. Skipping step")
335- return retval335+ return retval
336- 336+ 
337- def step(self, optimizer, *args, **kwargs):337+ def step(self, optimizer, *args, **kwargs):
338- """338+ """
339- :meth:`step` carries out the following two operations:339+ :meth:`step` carries out the following two operations:
340- 340+ 
341- 1. Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``341+ 1. Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
342- earlier in the iteration). As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.342+ earlier in the iteration). As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
343- 2. If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled343+ 2. If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
344- gradients. Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.344+ gradients. Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.
345- 345+ 
346- ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.346+ ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.
347- 347+ 
348- Returns the return value of ``optimizer.step(*args, **kwargs)``.348+ Returns the return value of ``optimizer.step(*args, **kwargs)``.
349- 349+ 
350- Args:350+ Args:
351- optimizer (torch.optim.Optimizer): Optimizer that applies the gradients.351+ optimizer (torch.optim.Optimizer): Optimizer that applies the gradients.
352- args: Any arguments.352+ args: Any arguments.
353- kwargs: Any keyword arguments.353+ kwargs: Any keyword arguments.
354- 354+ 
355- .. warning::355+ .. warning::
356- Closure use is not currently supported.356+ Closure use is not currently supported.
357- """357+ """
358- if (not self._enabled):358+ if (not self._enabled):
359- return optimizer.step(*args, **kwargs)359+ return optimizer.step(*args, **kwargs)
360- 360+ 
361- if "closure" in kwargs:361+ if "closure" in kwargs:
362- raise RuntimeError("Closure use is not currently supported if GradScaler is enabled." +362+ raise RuntimeError("Closure use is not currently supported if GradScaler is enabled." +
363- pta_error(ErrCode.NOT_SUPPORT))363+ pta_error(ErrCode.NOT_SUPPORT))
364- 364+ 
365- self._check_scale_growth_tracker("step")365+ self._check_scale_growth_tracker("step")
366- 366+ 
367- optimizer_state = self._per_optimizer_states[id(optimizer)]367+ optimizer_state = self._per_optimizer_states[id(optimizer)]
368- 368+ 
369- if optimizer_state["stage"] is OptState.STEPPED:369+ if optimizer_state["stage"] is OptState.STEPPED:
370- raise RuntimeError("step() has already been called since the last update()." +370+ raise RuntimeError("step() has already been called since the last update()." +
371- pta_error(ErrCode.INTERNAL))371+ pta_error(ErrCode.INTERNAL))
372- 372+ 
373- retval = None373+ retval = None
374- 374+ 
375- if (hasattr(optimizer, "_step_supports_amp_scaling") and optimizer._step_supports_amp_scaling):375+ if (hasattr(optimizer, "_step_supports_amp_scaling") and optimizer._step_supports_amp_scaling):
376- # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.376+ # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
377- # The contract with custom optimizers is that their step() should accept an additional,377+ # The contract with custom optimizers is that their step() should accept an additional,
378- # optional grad_scaler kwarg. We append self to the kwargs so the custom optimizer has full information:378+ # optional grad_scaler kwarg. We append self to the kwargs so the custom optimizer has full information:
379- # it can query its own state, invoke unscale_ on itself, etc379+ # it can query its own state, invoke unscale_ on itself, etc
380- retval = optimizer.step(*args, **dict(kwargs, grad_scaler=self))380+ retval = optimizer.step(*args, **dict(kwargs, grad_scaler=self))
381- optimizer_state["stage"] = OptState.STEPPED381+ optimizer_state["stage"] = OptState.STEPPED
382- return retval382+ return retval
383- 383+ 
384- if optimizer_state["stage"] is OptState.READY:384+ if optimizer_state["stage"] is OptState.READY:
385- self.unscale_(optimizer)385+ self.unscale_(optimizer)
386- 386+ 
387- if len(optimizer_state["found_inf_per_device"]) <= 0:387+ if len(optimizer_state["found_inf_per_device"]) <= 0:
388- raise RuntimeError("No inf checks were recorded for this optimizer." + pta_error(ErrCode.INTERNAL))388+ raise RuntimeError("No inf checks were recorded for this optimizer." + pta_error(ErrCode.INTERNAL))
389- 389+ 
390- if self._dynamic:390+ if self._dynamic:
391- retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)391+ retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
392- optimizer_state["stage"] = OptState.STEPPED392+ optimizer_state["stage"] = OptState.STEPPED
393- return retval393+ return retval
394- 394+ 
395- retval = optimizer.step(*args, **kwargs)395+ retval = optimizer.step(*args, **kwargs)
396- optimizer_state["stage"] = OptState.STEPPED396+ optimizer_state["stage"] = OptState.STEPPED
397- 397+ 
398- return retval398+ return retval
399- 399+ 
400- def update(self, new_scale=None):400+ def update(self, new_scale=None):
401- """401+ """
402- Updates the scale factor.402+ Updates the scale factor.
403- 403+ 
404- If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``404+ If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
405- to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,405+ to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
406- the scale is multiplied by ``growth_factor`` to increase it.406+ the scale is multiplied by ``growth_factor`` to increase it.
407- 407+ 
408- Passing ``new_scale`` sets the scale directly.408+ Passing ``new_scale`` sets the scale directly.
409- 409+ 
410- Args:410+ Args:
411- new_scale (float or :class:`torch.npu.FloatTensor`, optional, default=None): New scale factor.411+ new_scale (float or :class:`torch.npu.FloatTensor`, optional, default=None): New scale factor.
412- 412+ 
413- .. warning::413+ .. warning::
414- :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has414+ :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
415- been invoked for all optimizers used this iteration.415+ been invoked for all optimizers used this iteration.
416- """416+ """
417- if not self._enabled:417+ if not self._enabled:
418- return418+ return
419- 419+ 
420- _scale, _ = self._check_scale_growth_tracker("update")420+ _scale, _ = self._check_scale_growth_tracker("update")
421- 421+ 
422- if new_scale is not None:422+ if new_scale is not None:
423- # Accept a new user-defined scale.423+ # Accept a new user-defined scale.
424- if isinstance(new_scale, float):424+ if isinstance(new_scale, float):
425- self._scale = torch.full((), new_scale, dtype=torch.float32)425+ self._scale = torch.full((), new_scale, dtype=torch.float32)
426- self._scale = self._scale.pin_memory().to(_scale.device, non_blocking=True)426+ self._scale = self._scale.pin_memory().to(_scale.device, non_blocking=True)
427- else:427+ else:
428- reason = "new_scale should be a float or a 1-element torch.npu.FloatTensor with requires_grad=False."428+ reason = "new_scale should be a float or a 1-element torch.npu.FloatTensor with requires_grad=False."
429- if not isinstance(new_scale, torch.npu.FloatTensor): # type: ignore[attr-defined]429+ if not isinstance(new_scale, torch.npu.FloatTensor): # type: ignore[attr-defined]
430- raise TypeError(reason + pta_error(ErrCode.TYPE))430+ raise TypeError(reason + pta_error(ErrCode.TYPE))
431- if new_scale.numel() != 1:431+ if new_scale.numel() != 1:
O
OopenLiBingCI5月17日

此条代码评论区间+427至+431

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
432- raise ValueError(reason + pta_error(ErrCode.VALUE))432+ raise ValueError(reason + pta_error(ErrCode.VALUE))
433- if new_scale.requires_grad:433+ if new_scale.requires_grad:
434- raise ValueError(reason + pta_error(ErrCode.VALUE))434+ raise ValueError(reason + pta_error(ErrCode.VALUE))
435- self._scale = new_scale435+ self._scale = new_scale
436- elif self._dynamic:436+ elif self._dynamic:
437- self._npu_update_scale()437+ self._npu_update_scale()
438- 438+ 
439- # To prepare for next iteration, clear the data collected from optimizers this iteration.439+ # To prepare for next iteration, clear the data collected from optimizers this iteration.
440- self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)440+ self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
441- self._has_overflow = False441+ self._has_overflow = False
442- self._clear_overflow_flag = False442+ self._clear_overflow_flag = False
443- 443+ 
444- def state_dict(self):444+ def state_dict(self):
445- state = super(GradScaler, self).state_dict()445+ state = super(GradScaler, self).state_dict()
446- if self._enabled:446+ if self._enabled:
447- state["dynamic"] = self._dynamic447+ state["dynamic"] = self._dynamic
448- return state448+ return state
449- 449+ 
450- def load_state_dict(self, state_dict):450+ def load_state_dict(self, state_dict):
451- if not self._enabled:451+ if not self._enabled:
452- return452+ return
453- 453+ 
454- if len(state_dict) == 0:454+ if len(state_dict) == 0:
455- raise RuntimeError("The source state dict is empty, possibly because it was saved "455+ raise RuntimeError("The source state dict is empty, possibly because it was saved "
456- "from a disabled instance of GradScaler." + pta_error(ErrCode.VALUE))456+ "from a disabled instance of GradScaler." + pta_error(ErrCode.VALUE))
457- 457+ 
458- super(GradScaler, self).load_state_dict(state_dict)458+ super(GradScaler, self).load_state_dict(state_dict)
459- self._dynamic = state_dict["dynamic"]459+ self._dynamic = state_dict["dynamic"]
460- 460+ 
461- @staticmethod461+ @staticmethod
462- def get_npu_overflow_flag():462+ def get_npu_overflow_flag():
463- float_status = torch.zeros(8).pin_memory().to('npu', non_blocking=True)463+ float_status = torch.zeros(8).pin_memory().to('npu', non_blocking=True)
464- result = torch_npu.npu_get_float_status(float_status)464+ result = torch_npu.npu_get_float_status(float_status)
465- if (result.cpu()[0] != 0):465+ if (result.cpu()[0] != 0):
466- return True466+ return True
467- else:467+ else:
468- return False468+ return False
469- 469+ 
470- @staticmethod470+ @staticmethod
471- def clear_npu_overflow_flag():471+ def clear_npu_overflow_flag():
472- float_status = torch.zeros(8).pin_memory().to('npu', non_blocking=True)472+ float_status = torch.zeros(8).pin_memory().to('npu', non_blocking=True)
473- result = torch_npu.npu_clear_float_status(float_status)473+ result = torch_npu.npu_clear_float_status(float_status)
474- 474+ 
475- def _sync_dist_overflow_count(self):475+ def _sync_dist_overflow_count(self):
476- if torch_npu.npu.utils.is_support_inf_nan():476+ if torch_npu.npu.utils.is_support_inf_nan():
477- return477+ return
478- if self._dynamic and self._dist_initialized:478+ if self._dynamic and self._dist_initialized:
479- if self._has_overflow:479+ if self._has_overflow:
480- self._dist_overflow_count.add_(1)480+ self._dist_overflow_count.add_(1)
481- dist.all_reduce(self._dist_overflow_count)481+ dist.all_reduce(self._dist_overflow_count)
482- self._dist_overflow_count.zero_()482+ self._dist_overflow_count.zero_()
483- else:483+ else:
484- dist.all_reduce(self._dist_overflow_count)484+ dist.all_reduce(self._dist_overflow_count)
485- if self._dist_overflow_count.item() != 0:485+ if self._dist_overflow_count.item() != 0:
486- self._has_overflow = True486+ self._has_overflow = True
487- self._dist_overflow_count.zero_()487+ self._dist_overflow_count.zero_()
488- 488+ 
489- def _npu_update_scale(self):489+ def _npu_update_scale(self):
490- if self._has_overflow:490+ if self._has_overflow:
491- self._scale.mul_(self._backoff_factor)491+ self._scale.mul_(self._backoff_factor)
492- self._growth_tracker.zero_()492+ self._growth_tracker.zero_()
493- print(("Loss scaler reducing loss scale "493+ print(("Loss scaler reducing loss scale "
494- "to {}").format(self._scale.item()))494+ "to {}").format(self._scale.item()))
495- else:495+ else:
496- # Entering this branch means we just carried out a successful step,496+ # Entering this branch means we just carried out a successful step,
497- # so growth_tracker is incremented before comparing to growth_interval.497+ # so growth_tracker is incremented before comparing to growth_interval.
498- self._growth_tracker.add_(1)498+ self._growth_tracker.add_(1)
499- if self._growth_tracker.item() == self._growth_interval:499+ if self._growth_tracker.item() == self._growth_interval:
500- new_scale = self._scale * self._growth_factor500+ new_scale = self._scale * self._growth_factor
501- if not torch.isinf(new_scale):501+ if not torch.isinf(new_scale):
502- self._scale = new_scale502+ self._scale = new_scale
503- self._growth_tracker.zero_()503+ self._growth_tracker.zero_()
504- print(("Loss scaler increasing loss scale "504+ print(("Loss scaler increasing loss scale "
505- "to {}").format(self._scale.item()))505+ "to {}").format(self._scale.item()))
Mtorch_npu/testing/common_distributed.py+109-109
@@ -1,109 +1,109 @@
1-import unittest1+import unittest
2-from functools import wraps2+from functools import wraps
3-from typing import (3+from typing import (
4- Tuple,4+ Tuple,
5- Dict,5+ Dict,
6- Any,6+ Any,
7-)7+)
8-from collections import namedtuple8+from collections import namedtuple
9-import sys9+import sys
10-import os10+import os
11-from contextlib import contextmanager11+from contextlib import contextmanager
12-import torch12+import torch
13-import torch.distributed as dist13+import torch.distributed as dist
14-import torch_npu14+import torch_npu
15- 15+ 
16-TestSkip = namedtuple('TestSkip', 'exit_code, message')16+TestSkip = namedtuple('TestSkip', 'exit_code, message')
17-TEST_SKIPS = {17+TEST_SKIPS = {
18- "multi-npu": TestSkip(75, "Multi-NPU condition not satisfied"),18+ "multi-npu": TestSkip(75, "Multi-NPU condition not satisfied"),
19- "multi-npu-1": TestSkip(75, "Need at least 1 ASCEND devices"),19+ "multi-npu-1": TestSkip(75, "Need at least 1 ASCEND devices"),
20- "multi-npu-2": TestSkip(75, "Need at least 2 ASCEND devices"),20+ "multi-npu-2": TestSkip(75, "Need at least 2 ASCEND devices"),
21- "multi-npu-3": TestSkip(75, "Need at least 3 ASCEND devices"),21+ "multi-npu-3": TestSkip(75, "Need at least 3 ASCEND devices"),
22- "multi-npu-4": TestSkip(75, "Need at least 4 ASCEND devices"),22+ "multi-npu-4": TestSkip(75, "Need at least 4 ASCEND devices"),
23- "multi-npu-5": TestSkip(75, "Need at least 5 ASCEND devices"),23+ "multi-npu-5": TestSkip(75, "Need at least 5 ASCEND devices"),
24- "multi-npu-6": TestSkip(75, "Need at least 6 ASCEND devices"),24+ "multi-npu-6": TestSkip(75, "Need at least 6 ASCEND devices"),
25- "multi-npu-7": TestSkip(75, "Need at least 7 ASCEND devices"),25+ "multi-npu-7": TestSkip(75, "Need at least 7 ASCEND devices"),
26- "multi-npu-8": TestSkip(75, "Need at least 8 ASCEND devices"),26+ "multi-npu-8": TestSkip(75, "Need at least 8 ASCEND devices"),
27- "hccl":TestSkip(76, "c10d not compiled with HCCL support"),27+ "hccl":TestSkip(76, "c10d not compiled with HCCL support"),
28- "known_issues":TestSkip(77, "Test skipped due to known issues"),28+ "known_issues":TestSkip(77, "Test skipped due to known issues"),
29-}29+}
30- 30+ 
31- 31+ 
32-def skipIfUnsupportMultiNPU(npu_number_needed):32+def skipIfUnsupportMultiNPU(npu_number_needed):
33- def skip_dec(func):33+ def skip_dec(func):
34- @wraps(func)34+ @wraps(func)
35- def wrapper(self, *args, **kwargs):35+ def wrapper(self, *args, **kwargs):
36- if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:36+ if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:
37- raise unittest.SkipTest(f"Multi-NPU {npu_number_needed} condition not satisfied")37+ raise unittest.SkipTest(f"Multi-NPU {npu_number_needed} condition not satisfied")
38- return func(self, *args, **kwargs)38+ return func(self, *args, **kwargs)
39- return wrapper39+ return wrapper
40- return skip_dec40+ return skip_dec
41- 41+ 
42- 42+ 
43-def with_comms(func):43+def with_comms(func):
44- if func is None:44+ if func is None:
45- raise RuntimeError("Test function is None.")45+ raise RuntimeError("Test function is None.")
46- 46+ 
47- def get_device_type(self):47+ def get_device_type(self):
48- if torch.npu.is_available() and torch.npu.device_count() >= self.world_size:48+ if torch.npu.is_available() and torch.npu.device_count() >= self.world_size:
49- return "npu"49+ return "npu"
50- return "cpu"50+ return "cpu"
51- 51+ 
52- @wraps(func) # pyre-ignore[6]52+ @wraps(func) # pyre-ignore[6]
53- def wrapper(53+ def wrapper(
54- self, *args: Tuple[object], **kwargs: Dict[str, Any] # type: ignore[misc]54+ self, *args: Tuple[object], **kwargs: Dict[str, Any] # type: ignore[misc]
55- ) -> None:55+ ) -> None:
56- 56+ 
57- pg_backend = (57+ pg_backend = (
O
OopenLiBingCI5月17日

此条代码评论区间+52至+57

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
58- "hccl" if get_device_type(self) == "npu" else "gloo"58+ "hccl" if get_device_type(self) == "npu" else "gloo"
59- )59+ )
60- if pg_backend == "hccl" and torch.npu.device_count() < self.world_size:60+ if pg_backend == "hccl" and torch.npu.device_count() < self.world_size:
61- raise RuntimeError(TEST_SKIPS[f"multi-npu-{self.world_size}"].message)61+ raise RuntimeError(TEST_SKIPS[f"multi-npu-{self.world_size}"].message)
62- 62+ 
63- init_pg(backend=pg_backend, world_size=self.world_size, rank=self.rank, file_name=self.file_name)63+ init_pg(backend=pg_backend, world_size=self.world_size, rank=self.rank, file_name=self.file_name)
64- 64+ 
65- torch.npu.manual_seed(0)65+ torch.npu.manual_seed(0)
66- torch.npu.initial_seed()66+ torch.npu.initial_seed()
67- func(self, *args, **kwargs) # type: ignore[misc]67+ func(self, *args, **kwargs) # type: ignore[misc]
68- self.destroy_pg()68+ self.destroy_pg()
69- 69+ 
70- return wrapper70+ return wrapper
O
OopenLiBingCI5月17日

此条代码评论区间+65至+70

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。

likedislike
71- 71+ 
72- 72+ 
73-def init_pg(backend: str = "hccl", world_size=1, rank=0, file_name="file://") -> None:73+def init_pg(backend: str = "hccl", world_size=1, rank=0, file_name="file://") -> None:
74- if backend == "hccl" and torch.npu.device_count() < world_size:74+ if backend == "hccl" and torch.npu.device_count() < world_size:
75- raise RuntimeError(TEST_SKIPS[f"multi-npu-{world_size}"].message)75+ raise RuntimeError(TEST_SKIPS[f"multi-npu-{world_size}"].message)
76- 76+ 
77- if backend not in ["hccl", "gloo"]:77+ if backend not in ["hccl", "gloo"]:
78- raise RuntimeError(f"Backend {backend} not supported!")78+ raise RuntimeError(f"Backend {backend} not supported!")
79- 79+ 
80- dist.init_process_group(80+ dist.init_process_group(
81- backend=backend,81+ backend=backend,
82- world_size=world_size,82+ world_size=world_size,
83- rank=rank, # pyre-ignore[16]83+ rank=rank, # pyre-ignore[16]
84- init_method=f"file://{file_name}", # pyre-ignore[16]84+ init_method=f"file://{file_name}", # pyre-ignore[16]
85- )85+ )
86- 86+ 
87- # set device for hccl pg for collectives87+ # set device for hccl pg for collectives
88- if backend == "hccl":88+ if backend == "hccl":
89- torch.npu.set_device(rank)89+ torch.npu.set_device(rank)
90- 90+ 
91- 91+ 
92-@contextmanager92+@contextmanager
93-def _dynamo_dist_per_rank_init(rank, world_size, init_pg_=True):93+def _dynamo_dist_per_rank_init(rank, world_size, init_pg_=True):
94- # To avoid multiple inheritance from _dynamo.test_case.TestCase and MultiProcessTestCase,94+ # To avoid multiple inheritance from _dynamo.test_case.TestCase and MultiProcessTestCase,
95- # Just manually implement the most important part of the dynamo behavior to reset/clear.95+ # Just manually implement the most important part of the dynamo behavior to reset/clear.
96- torch_npu.npu.set_device(rank)96+ torch_npu.npu.set_device(rank)
97- os.environ['MASTER_ADDR'] = 'localhost'97+ os.environ['MASTER_ADDR'] = 'localhost'
98- os.environ['MASTER_PORT'] = '6789'98+ os.environ['MASTER_PORT'] = '6789'
99- if init_pg_:99+ if init_pg_:
100- dist.init_process_group(backend="hccl", rank=rank, world_size=world_size)100+ dist.init_process_group(backend="hccl", rank=rank, world_size=world_size)
101- torch._dynamo.reset()101+ torch._dynamo.reset()
102- torch._dynamo.utils.counters.clear()102+ torch._dynamo.utils.counters.clear()
103- try:103+ try:
104- yield104+ yield
105- finally:105+ finally:
106- torch._dynamo.reset()106+ torch._dynamo.reset()
107- torch._dynamo.utils.counters.clear()107+ torch._dynamo.utils.counters.clear()
108- if init_pg_:108+ if init_pg_:
109- dist.destroy_process_group()109+ dist.destroy_process_group()
Mtorch_npu/testing/decorator.py+142-142
@@ -1,142 +1,142 @@
1-from functools import wraps, partialmethod1+from functools import wraps, partialmethod
2- 2+ 
3-import os3+import os
4-import inspect4+import inspect
5-import itertools5+import itertools
6-import torch6+import torch
7- 7+ 
8- 8+ 
9-def feed_data(func, new_name, *args, **kwargs):9+def feed_data(func, new_name, *args, **kwargs):
10- """10+ """
11- This internal method decorator feeds the test data item to the test.11+ This internal method decorator feeds the test data item to the test.
12- """12+ """
13- @wraps(func)13+ @wraps(func)
14- def wrapper(self):14+ def wrapper(self):
15- return func(self, *args, **kwargs)15+ return func(self, *args, **kwargs)
16- wrapper.__name__ = new_name16+ wrapper.__name__ = new_name
17- wrapper.__wrapped__ = func17+ wrapper.__wrapped__ = func
18- return wrapper18+ return wrapper
19- 19+ 
20- 20+ 
21-def instantiate_tests(arg=None, **kwargs):21+def instantiate_tests(arg=None, **kwargs):
22- 22+ 
23- def wrapper(cls):23+ def wrapper(cls):
24- def gen_testcase(cls, func, name, key_list, func_args, value):24+ def gen_testcase(cls, func, name, key_list, func_args, value):
25- new_kwargs = dict(device="npu") if "device" in func_args else {}25+ new_kwargs = dict(device="npu") if "device" in func_args else {}
26- test_name = name26+ test_name = name
27- for k, v in zip(key_list, value):27+ for k, v in zip(key_list, value):
28- func_key = None28+ func_key = None
29- if k == "format":29+ if k == "format":
30- test_name += ("_" + str(v))30+ test_name += ("_" + str(v))
31- elif k == "dtype":31+ elif k == "dtype":
32- test_name += ("_" + str(v).split('.')[1])32+ test_name += ("_" + str(v).split('.')[1])
33- for _func_key in func_args:33+ for _func_key in func_args:
34- if k in _func_key:34+ if k in _func_key:
35- if func_key is not None:35+ if func_key is not None:
36- raise RuntimeError(f"Multiple matches for {k}")36+ raise RuntimeError(f"Multiple matches for {k}")
37- func_key = _func_key37+ func_key = _func_key
38- new_kwargs[func_key] = v38+ new_kwargs[func_key] = v
39- setattr(cls, test_name, feed_data(func, test_name, **new_kwargs))39+ setattr(cls, test_name, feed_data(func, test_name, **new_kwargs))
40- 40+ 
41- for name, func in list(cls.__dict__.items()):41+ for name, func in list(cls.__dict__.items()):
42- data = {}42+ data = {}
43- if hasattr(func, "dtypes"):43+ if hasattr(func, "dtypes"):
44- data['dtype'] = func.dtypes44+ data['dtype'] = func.dtypes
45- if hasattr(func, "formats"):45+ if hasattr(func, "formats"):
46- data['format'] = func.formats46+ data['format'] = func.formats
47- 47+ 
48- key_list = data.keys()48+ key_list = data.keys()
49- if not key_list:49+ if not key_list:
50- continue50+ continue
51- 51+ 
52- func_args = inspect.getfullargspec(func).args52+ func_args = inspect.getfullargspec(func).args
53- value_list = [data.get(key) for key in key_list]53+ value_list = [data.get(key) for key in key_list]
54- for value in itertools.product(*value_list):54+ for value in itertools.product(*value_list):
55- gen_testcase(cls, func, name, key_list, func_args, value)55+ gen_testcase(cls, func, name, key_list, func_args, value)
56- 56+ 
57- delattr(cls, name)57+ delattr(cls, name)
58- return cls58+ return cls
59- 59+ 
60- return wrapper(arg)60+ return wrapper(arg)
61- 61+ 
62- 62+ 
63-def gen_ops_testcase(cls, func, name, keys, value, op_info):63+def gen_ops_testcase(cls, func, name, keys, value, op_info):
64- new_kwargs = {}64+ new_kwargs = {}
65- test_name = f'{func.__name__}_{name}'65+ test_name = f'{func.__name__}_{name}'
66- 66+ 
67- for k, v in zip(keys, value):67+ for k, v in zip(keys, value):
68- if k == "npu_format":68+ if k == "npu_format":
69- test_name += ("_" + str(v))69+ test_name += ("_" + str(v))
70- elif k == "dtype":70+ elif k == "dtype":
71- test_name += ("_" + str(v).split('.')[1])71+ test_name += ("_" + str(v).split('.')[1])
72- new_kwargs[k] = v72+ new_kwargs[k] = v
73- 73+ 
74- new_kwargs['op'] = op_info74+ new_kwargs['op'] = op_info
75- new_func = partialmethod(func, **new_kwargs)75+ new_func = partialmethod(func, **new_kwargs)
76- 76+ 
77- setattr(cls, test_name, new_func)77+ setattr(cls, test_name, new_func)
78- for decorator in op_info.get_decorators(cls.__name__, func.__name__, 'cpu', value[0], {}):78+ for decorator in op_info.get_decorators(cls.__name__, func.__name__, 'cpu', value[0], {}):
79- setattr(cls, test_name, decorator(new_func))79+ setattr(cls, test_name, decorator(new_func))
80- 80+ 
81- 81+ 
82-def gen_op_input(testcase, func, op_info):82+def gen_op_input(testcase, func, op_info):
83- data = {83+ data = {
84- 'dtype': func.dtypes if hasattr(func, "dtypes") else op_info.dtypesIfNPU, 84+ 'dtype': func.dtypes if hasattr(func, "dtypes") else op_info.dtypesIfNPU,
85- 'npu_format': func.formats if hasattr(func, "formats") else op_info.formats85+ 'npu_format': func.formats if hasattr(func, "formats") else op_info.formats
86- }86+ }
87- 87+ 
88- if 'test_variant_consistency_eager' in testcase:88+ if 'test_variant_consistency_eager' in testcase:
89- if torch.float32 in op_info.dtypesIfNPU:89+ if torch.float32 in op_info.dtypesIfNPU:
90- data['dtype'] = {torch.float32}90+ data['dtype'] = {torch.float32}
91- else:91+ else:
92- data['dtype'] = {list(op_info.dtypesIfNPU)[-1]}92+ data['dtype'] = {list(op_info.dtypesIfNPU)[-1]}
93- 93+ 
94- return data94+ return data
95- 95+ 
96- 96+ 
97-def instantiate_ops_tests(op_db):97+def instantiate_ops_tests(op_db):
98- 98+ 
99- def wrapper(cls):99+ def wrapper(cls):
100- testcases = [x for x in dir(cls) if x.startswith('test_')]100+ testcases = [x for x in dir(cls) if x.startswith('test_')]
101- for testcase in testcases: 101+ for testcase in testcases:
102- if hasattr(cls, testcase):102+ if hasattr(cls, testcase):
103- func = getattr(cls, testcase)103+ func = getattr(cls, testcase)
104- for op_info in op_db:104+ for op_info in op_db:
105- data = gen_op_input(testcase, func, op_info)105+ data = gen_op_input(testcase, func, op_info)
106- keys = data.keys()106+ keys = data.keys()
107- values = [data.get(key) for key in keys]107+ values = [data.get(key) for key in keys]
108- 108+ 
109- for value in itertools.product(*values):109+ for value in itertools.product(*values):
110- gen_ops_testcase(cls, func, op_info.name, keys, value, op_info)110+ gen_ops_testcase(cls, func, op_info.name, keys, value, op_info)
111- 111+ 
112- delattr(cls, testcase)112+ delattr(cls, testcase)
113- 113+ 
114- return cls114+ return cls
115- 115+
116- return wrapper116+ return wrapper
117- 117+ 
118- 118+ 
119-class Dtypes(object):119+class Dtypes(object):
120- 120+ 
121- def __init__(self, *args):121+ def __init__(self, *args):
122- if (args is None or len(args) == 0):122+ if (args is None or len(args) == 0):
123- raise RuntimeError("No dtypes given")123+ raise RuntimeError("No dtypes given")
124- if not all(isinstance(arg, torch.dtype) for arg in args):124+ if not all(isinstance(arg, torch.dtype) for arg in args):
125- raise RuntimeError("Unknown dtype in {0}".format(str(args)))125+ raise RuntimeError("Unknown dtype in {0}".format(str(args)))
126- self.args = args126+ self.args = args
127- 127+ 
128- def __call__(self, fn):128+ def __call__(self, fn):
129- fn.dtypes = self.args129+ fn.dtypes = self.args
130- return fn130+ return fn
131- 131+ 
132- 132+ 
133-class Formats(object):133+class Formats(object):
134- 134+ 
135- def __init__(self, *args):135+ def __init__(self, *args):
136- if args is None or len(args) == 0:136+ if args is None or len(args) == 0:
137- raise RuntimeError("No formats given")137+ raise RuntimeError("No formats given")
138- self.args = args138+ self.args = args
139- 139+ 
140- def __call__(self, fn):140+ def __call__(self, fn):
141- fn.formats = self.args141+ fn.formats = self.args
142- return fn142+ return fn
Mtorch_npu/utils/__init__.py+0-1
@@ -23,4 +23,3 @@ from .combine_tensors import (
23)23)
24from .flops_count import _FlopsCounter as FlopsCounter24from .flops_count import _FlopsCounter as FlopsCounter
25from .serialization import save_async25from .serialization import save_async
26- 
Mtorch_npu/utils/cpp_extension.py+60-60
@@ -1,60 +1,60 @@
1-import os1+import os
2-import setuptools2+import setuptools
3- 3+ 
4-import torch4+import torch
5-import torch.utils.cpp_extension as TorchExtension5+import torch.utils.cpp_extension as TorchExtension
6- 6+ 
7-import torch_npu7+import torch_npu
8- 8+ 
9-PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))9+PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
10- 10+ 
11- 11+ 
12-def NpuExtension(name, sources, *args, **kwargs):12+def NpuExtension(name, sources, *args, **kwargs):
13- r'''13+ r'''
14- Creates a :class:`setuptools.Extension` for C++.14+ Creates a :class:`setuptools.Extension` for C++.
15- 15+ 
16- Convenience method that creates a :class:`setuptools.Extension` with the16+ Convenience method that creates a :class:`setuptools.Extension` with the
17- bare minimum (but often sufficient) arguments to build a C++ extension.17+ bare minimum (but often sufficient) arguments to build a C++ extension.
18- 18+ 
19- All arguments are forwarded to the :class:`setuptools.Extension`19+ All arguments are forwarded to the :class:`setuptools.Extension`
20- constructor.20+ constructor.
21- 21+ 
22- Example:22+ Example:
23- >>> from setuptools import setup23+ >>> from setuptools import setup
24- >>> from torch_npu.utils.cpp_extension import NpuExtension24+ >>> from torch_npu.utils.cpp_extension import NpuExtension
25- >>> setup(25+ >>> setup(
26- name='extension',26+ name='extension',
27- ext_modules=[27+ ext_modules=[
28- NpuExtension(28+ NpuExtension(
29- name='extension',29+ name='extension',
30- sources=['extension.cpp'],30+ sources=['extension.cpp'],
31- extra_compile_args=['-g']),31+ extra_compile_args=['-g']),
32- ],32+ ],
33- cmdclass={33+ cmdclass={
34- 'build_ext': BuildExtension34+ 'build_ext': BuildExtension
35- })35+ })
36- '''36+ '''
37- 37+ 
38- torch_npu_dir = PYTORCH_NPU_INSTALL_PATH38+ torch_npu_dir = PYTORCH_NPU_INSTALL_PATH
39- include_dirs = kwargs.get('include_dirs', [])39+ include_dirs = kwargs.get('include_dirs', [])
40- include_dirs.append(os.path.join(torch_npu_dir, 'include'))40+ include_dirs.append(os.path.join(torch_npu_dir, 'include'))
41- include_dirs.append(os.path.join(torch_npu_dir, 'include', 'third_party', 'acl', 'inc'))41+ include_dirs.append(os.path.join(torch_npu_dir, 'include', 'third_party', 'acl', 'inc'))
42- include_dirs.append(os.path.join(torch_npu_dir, 'include', 'third_party', 'hccl', 'inc'))42+ include_dirs.append(os.path.join(torch_npu_dir, 'include', 'third_party', 'hccl', 'inc'))
43- include_dirs += TorchExtension.include_paths()43+ include_dirs += TorchExtension.include_paths()
44- kwargs['include_dirs'] = include_dirs44+ kwargs['include_dirs'] = include_dirs
45- 45+ 
46- library_dirs = kwargs.get('library_dirs', [])46+ library_dirs = kwargs.get('library_dirs', [])
47- library_dirs.append(os.path.join(torch_npu_dir, 'lib'))47+ library_dirs.append(os.path.join(torch_npu_dir, 'lib'))
48- library_dirs += TorchExtension.library_paths()48+ library_dirs += TorchExtension.library_paths()
49- kwargs['library_dirs'] = library_dirs49+ kwargs['library_dirs'] = library_dirs
50- 50+ 
51- libraries = kwargs.get('libraries', [])51+ libraries = kwargs.get('libraries', [])
52- libraries.append('c10')52+ libraries.append('c10')
53- libraries.append('torch')53+ libraries.append('torch')
54- libraries.append('torch_cpu')54+ libraries.append('torch_cpu')
55- libraries.append('torch_python')55+ libraries.append('torch_python')
56- libraries.append('torch_npu')56+ libraries.append('torch_npu')
57- kwargs['libraries'] = libraries57+ kwargs['libraries'] = libraries
58- 58+ 
59- kwargs['language'] = 'c++'59+ kwargs['language'] = 'c++'
60- return setuptools.Extension(name, sources, *args, **kwargs)60+ return setuptools.Extension(name, sources, *args, **kwargs)
Mtorch_npu/utils/flops_count.py+0-1
@@ -38,4 +38,3 @@ class FlopsCounter(_FlopsCounter):
38 super().__init__()38 super().__init__()
39 warnings.warn("torch_npu.utils.flops_count.FlopsCounter() will be deprecated. "39 warnings.warn("torch_npu.utils.flops_count.FlopsCounter() will be deprecated. "
40 "If necessary, please use torch_npu.utils.FlopsCounter().", FutureWarning)40 "If necessary, please use torch_npu.utils.FlopsCounter().", FutureWarning)
41- 
Mtorch_npu/utils/syncbatchnorm.py+103-103
@@ -1,103 +1,103 @@
1-import torch1+import torch
2-import torch.distributed as dist2+import torch.distributed as dist
3-from torch.autograd.function import Function3+from torch.autograd.function import Function
4- 4+ 
5-import torch_npu5+import torch_npu
6-from torch_npu.utils._error_code import ErrCode, ops_error6+from torch_npu.utils._error_code import ErrCode, ops_error
7- 7+ 
8- 8+ 
9-__all__ = ["SyncBatchNorm"]9+__all__ = ["SyncBatchNorm"]
10- 10+ 
11- 11+ 
12-class SyncBatchNorm(Function):12+class SyncBatchNorm(Function):
13- 13+ 
14- @staticmethod14+ @staticmethod
15- def forward(self, input_tensor, weight, bias, running_mean, running_var, eps, momentum, process_group, world_size):15+ def forward(self, input_tensor, weight, bias, running_mean, running_var, eps, momentum, process_group, world_size):
16- input_tensor = input_tensor.contiguous()16+ input_tensor = input_tensor.contiguous()
17- input_shape = input_tensor.shape17+ input_shape = input_tensor.shape
18- input_tensor_ = input_tensor.reshape(input_shape[0], input_shape[1], 1, -1)18+ input_tensor_ = input_tensor.reshape(input_shape[0], input_shape[1], 1, -1)
19- # calculate sum/sum_square for input.19+ # calculate sum/sum_square for input.
20- sum_val, sum_square_val = torch_npu.batch_norm_reduce(input_tensor_, eps)20+ sum_val, sum_square_val = torch_npu.batch_norm_reduce(input_tensor_, eps)
21- 21+ 
22- count = torch.full((1,),22+ count = torch.full((1,),
23- input_tensor.numel() // input_tensor.size(1),23+ input_tensor.numel() // input_tensor.size(1),
24- dtype=sum_val.dtype,24+ dtype=sum_val.dtype,
25- device=sum_val.device)25+ device=sum_val.device)
26- 26+ 
27- num_channels = input_tensor.shape[1]27+ num_channels = input_tensor.shape[1]
28- # C, C, 1 -> (2C + 1)28+ # C, C, 1 -> (2C + 1)
29- combined = torch.cat([sum_val, sum_square_val, count], dim=0)29+ combined = torch.cat([sum_val, sum_square_val, count], dim=0)
30- # world_size * (2C + 1)30+ # world_size * (2C + 1)
31- combined_list = [torch.empty_like(combined) for k in range(world_size)]31+ combined_list = [torch.empty_like(combined) for k in range(world_size)]
32- dist.all_gather(combined_list, combined, process_group, async_op=False)32+ dist.all_gather(combined_list, combined, process_group, async_op=False)
33- combined = torch.stack(combined_list, dim=0)33+ combined = torch.stack(combined_list, dim=0)
34- # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 134+ # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
35- sum_all, square_sum_all, count_all = torch.split(combined, num_channels, dim=1)35+ sum_all, square_sum_all, count_all = torch.split(combined, num_channels, dim=1)
36- 36+ 
37- size = count_all.view(-1).sum()37+ size = count_all.view(-1).sum()
38- if size == 1:38+ if size == 1:
39- raise ValueError('Expected more than 1 value per channel when training, got input size {}'.format(size) +39+ raise ValueError('Expected more than 1 value per channel when training, got input size {}'.format(size) +
40- ops_error(ErrCode.VALUE))40+ ops_error(ErrCode.VALUE))
41- 41+ 
42- # calculate global mean & invstd42+ # calculate global mean & invstd
43- mean, invstd = torch_npu.batch_norm_gather_stats_update(input_tensor,43+ mean, invstd = torch_npu.batch_norm_gather_stats_update(input_tensor,
44- sum_all,44+ sum_all,
45- square_sum_all,45+ square_sum_all,
46- running_mean,46+ running_mean,
47- running_var,47+ running_var,
48- momentum,48+ momentum,
49- eps,49+ eps,
50- count_all.view(-1))50+ count_all.view(-1))
51- 51+ 
52- self.save_for_backward(input_tensor, weight, mean, invstd, count_all)52+ self.save_for_backward(input_tensor, weight, mean, invstd, count_all)
53- self.process_group = process_group53+ self.process_group = process_group
54- 54+ 
55- # apply element-wise normalization55+ # apply element-wise normalization
56- out = torch.batch_norm_elemt(input_tensor, weight, bias, mean, invstd, eps)56+ out = torch.batch_norm_elemt(input_tensor, weight, bias, mean, invstd, eps)
57- return out57+ return out
58- 58+ 
59- @staticmethod59+ @staticmethod
60- def backward(self, grad_output):60+ def backward(self, grad_output):
61- if not grad_output.is_contiguous(memory_format=torch.channels_last):61+ if not grad_output.is_contiguous(memory_format=torch.channels_last):
62- grad_output = grad_output.contiguous()62+ grad_output = grad_output.contiguous()
63- saved_input, weight, mean, invstd, count_tensor = self.saved_tensors63+ saved_input, weight, mean, invstd, count_tensor = self.saved_tensors
64- grad_input = grad_weight = grad_bias = None64+ grad_input = grad_weight = grad_bias = None
65- process_group = self.process_group65+ process_group = self.process_group
66- 66+ 
67- # calculate local stats as well as grad_weight / grad_bias67+ # calculate local stats as well as grad_weight / grad_bias
68- sum_dy, sum_dy_xmu, grad_weight, grad_bias = torch.batch_norm_backward_reduce(grad_output,68+ sum_dy, sum_dy_xmu, grad_weight, grad_bias = torch.batch_norm_backward_reduce(grad_output,
69- saved_input,69+ saved_input,
70- mean,70+ mean,
71- invstd,71+ invstd,
72- weight,72+ weight,
73- self.needs_input_grad[0],73+ self.needs_input_grad[0],
74- self.needs_input_grad[1],74+ self.needs_input_grad[1],
75- self.needs_input_grad[2])75+ self.needs_input_grad[2])
76- 76+ 
77- if self.needs_input_grad[0]:77+ if self.needs_input_grad[0]:
78- # synchronizing stats used to calculate input gradient.78+ # synchronizing stats used to calculate input gradient.
79- num_channels = sum_dy.shape[0]79+ num_channels = sum_dy.shape[0]
80- combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)80+ combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)
81- torch.distributed.all_reduce(81+ torch.distributed.all_reduce(
82- combined, torch.distributed.ReduceOp.SUM, process_group, async_op=False)82+ combined, torch.distributed.ReduceOp.SUM, process_group, async_op=False)
83- sum_dy, sum_dy_xmu = torch.split(combined, num_channels)83+ sum_dy, sum_dy_xmu = torch.split(combined, num_channels)
84- 84+ 
85- # backward pass for gradient calculation85+ # backward pass for gradient calculation
86- grad_input = torch.batch_norm_backward_elemt(grad_output,86+ grad_input = torch.batch_norm_backward_elemt(grad_output,
87- saved_input,87+ saved_input,
88- mean,88+ mean,
89- invstd,89+ invstd,
90- weight,90+ weight,
91- sum_dy,91+ sum_dy,
92- sum_dy_xmu,92+ sum_dy_xmu,
93- count_tensor)93+ count_tensor)
94- 94+ 
95- # synchronizing of grad_weight / grad_bias is not needed as distributed95+ # synchronizing of grad_weight / grad_bias is not needed as distributed
96- # training would handle all reduce.96+ # training would handle all reduce.
97- if weight is None or not self.needs_input_grad[1]:97+ if weight is None or not self.needs_input_grad[1]:
98- grad_weight = None98+ grad_weight = None
99- 99+ 
100- if weight is None or not self.needs_input_grad[2]:100+ if weight is None or not self.needs_input_grad[2]:
101- grad_bias = None101+ grad_bias = None
102- 102+ 
103- return grad_input, grad_weight, grad_bias, None, None, None, None, None, None103+ return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
Mtorch_npu/utils/tensor_methods.py+86-86
@@ -1,86 +1,86 @@
1-from functools import wraps1+from functools import wraps
2- 2+ 
3-import torch3+import torch
4- 4+ 
5-import torch_npu5+import torch_npu
6-from torch_npu.utils._error_code import ErrCode, pta_error6+from torch_npu.utils._error_code import ErrCode, pta_error
7- 7+ 
8- 8+ 
9-__all__ = []9+__all__ = []
10- 10+ 
11- 11+ 
12-def _npu(self, *args, **kwargs):12+def _npu(self, *args, **kwargs):
13- return torch_npu._C.npu(self, *args, **kwargs)13+ return torch_npu._C.npu(self, *args, **kwargs)
14- 14+ 
15- 15+ 
16-@property16+@property
17-def _is_npu(self):17+def _is_npu(self):
18- return torch_npu._C.is_npu(self)18+ return torch_npu._C.is_npu(self)
19- 19+ 
20- 20+ 
21-class _NPUTensortypeCache(object):21+class _NPUTensortypeCache(object):
22- init = False22+ init = False
23- tensortype_list = []23+ tensortype_list = []
24- tensortype_dict = {}24+ tensortype_dict = {}
25- 25+ 
26- @classmethod26+ @classmethod
27- def tensortype_list_dict_init(cls):27+ def tensortype_list_dict_init(cls):
28- if not cls.init:28+ if not cls.init:
29- cls.tensortype_list += [29+ cls.tensortype_list += [
30- torch_npu.npu.BoolTensor,30+ torch_npu.npu.BoolTensor,
31- torch_npu.npu.ByteTensor,31+ torch_npu.npu.ByteTensor,
32- torch_npu.npu.CharTensor,32+ torch_npu.npu.CharTensor,
33- torch_npu.npu.DoubleTensor,33+ torch_npu.npu.DoubleTensor,
34- torch_npu.npu.FloatTensor,34+ torch_npu.npu.FloatTensor,
35- torch_npu.npu.HalfTensor,35+ torch_npu.npu.HalfTensor,
36- torch_npu.npu.IntTensor,36+ torch_npu.npu.IntTensor,
37- torch_npu.npu.LongTensor,37+ torch_npu.npu.LongTensor,
38- torch_npu.npu.ShortTensor,38+ torch_npu.npu.ShortTensor,
39- torch_npu.npu.BFloat16Tensor,39+ torch_npu.npu.BFloat16Tensor,
40- ]40+ ]
41- 41+ 
42- cls.tensortype_str_list = [42+ cls.tensortype_str_list = [
43- "torch_npu.npu.BoolTensor",43+ "torch_npu.npu.BoolTensor",
44- "torch_npu.npu.ByteTensor",44+ "torch_npu.npu.ByteTensor",
45- "torch_npu.npu.CharTensor",45+ "torch_npu.npu.CharTensor",
46- "torch_npu.npu.DoubleTensor",46+ "torch_npu.npu.DoubleTensor",
47- "torch_npu.npu.FloatTensor",47+ "torch_npu.npu.FloatTensor",
48- "torch_npu.npu.HalfTensor",48+ "torch_npu.npu.HalfTensor",
49- "torch_npu.npu.IntTensor",49+ "torch_npu.npu.IntTensor",
50- "torch_npu.npu.LongTensor",50+ "torch_npu.npu.LongTensor",
51- "torch_npu.npu.ShortTensor",51+ "torch_npu.npu.ShortTensor",
52- "torch_npu.npu.BFloat16Tensor",52+ "torch_npu.npu.BFloat16Tensor",
53- ]53+ ]
54- 54+ 
55- for tensortype, tensortype_str in zip(cls.tensortype_list, cls.tensortype_str_list):55+ for tensortype, tensortype_str in zip(cls.tensortype_list, cls.tensortype_str_list):
56- cls.tensortype_dict[tensortype_str] = tensortype56+ cls.tensortype_dict[tensortype_str] = tensortype
57- cls.tensortype_dict[tensortype_str.replace('torch_npu.', 'torch.')] = tensortype57+ cls.tensortype_dict[tensortype_str.replace('torch_npu.', 'torch.')] = tensortype
58- 58+ 
59- cls.init = True59+ cls.init = True
60- 60+ 
61- @classmethod61+ @classmethod
62- def get_tensortype_list(cls):62+ def get_tensortype_list(cls):
63- return cls.tensortype_list63+ return cls.tensortype_list
64- 64+ 
65- @classmethod65+ @classmethod
66- def get_tensortype_dict(cls):66+ def get_tensortype_dict(cls):
67- return cls.tensortype_dict67+ return cls.tensortype_dict
68- 68+ 
69- 69+ 
70-def _npu_type(self, dtype=None, non_blocking=False, **kwargs):70+def _npu_type(self, dtype=None, non_blocking=False, **kwargs):
71- if dtype is None:71+ if dtype is None:
72- return self.type_raw(dtype, non_blocking, **kwargs)72+ return self.type_raw(dtype, non_blocking, **kwargs)
73- 73+
74- _NPUTensortypeCache.tensortype_list_dict_init()74+ _NPUTensortypeCache.tensortype_list_dict_init()
75- if isinstance(dtype, str) and dtype in _NPUTensortypeCache.get_tensortype_dict():75+ if isinstance(dtype, str) and dtype in _NPUTensortypeCache.get_tensortype_dict():
76- tensortype_class = _NPUTensortypeCache.get_tensortype_dict()[dtype]76+ tensortype_class = _NPUTensortypeCache.get_tensortype_dict()[dtype]
77- return self.to(dtype=tensortype_class.dtype, device='npu', non_blocking=non_blocking)77+ return self.to(dtype=tensortype_class.dtype, device='npu', non_blocking=non_blocking)
78- elif dtype in _NPUTensortypeCache.get_tensortype_list():78+ elif dtype in _NPUTensortypeCache.get_tensortype_list():
79- return self.to(dtype=dtype.dtype, device='npu', non_blocking=non_blocking)79+ return self.to(dtype=dtype.dtype, device='npu', non_blocking=non_blocking)
80- else:80+ else:
81- return self.type_raw(dtype, non_blocking, **kwargs)81+ return self.type_raw(dtype, non_blocking, **kwargs)
82- 82+ 
83- 83+ 
84-def _add_tensor_methods():84+def _add_tensor_methods():
85- torch.Tensor.type_raw = torch.Tensor.type85+ torch.Tensor.type_raw = torch.Tensor.type
86- torch.Tensor.type = _npu_type86+ torch.Tensor.type = _npu_type