From 0fd4053b16f264abdbc47d22ecfee30fefdd6590 Mon Sep 17 00:00:00 2001 From: luyang Date: Wed, 11 Jun 2025 16:59:55 +0800 Subject: [PATCH 1/5] adaptation_of_mlu_device_based_on_refactor_xccl_primitive --- cmake/caches/ci/cpu.cmake | 1 + cmake/caches/cn/cpu.cmake | 1 + cmake/oneflow.cmake | 5 +++++ oneflow/api/python/flags.cpp | 8 ++++++++ oneflow/core/auto_parallel/boxing_collector.cpp | 4 ++-- oneflow/core/common/auto_registration_factory.h | 3 ++- oneflow/core/functional/impl/nn_functor.cpp | 2 +- .../boxing/collective_boxing_sub_task_graph_builder.cpp | 2 ++ .../boxing/hierarchical_sub_task_graph_builder_impl.cpp | 2 +- oneflow/core/job/resource_desc.cpp | 2 ++ oneflow/core/job/runtime.cpp | 4 ++-- oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp | 4 ++-- oneflow/core/job_rewriter/job_completer.cpp | 4 ++-- .../job_rewriter/nccl_logical_chain_strict_order_pass.cpp | 4 ++-- oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp | 4 ++-- oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp | 4 ++-- oneflow/user/kernels/eager_nccl_s2s_kernel.cu | 2 +- oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_fusion_kernel.cpp | 7 ++++--- oneflow/user/kernels/nccl_logical_kernels.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp | 7 ++++--- 21 files changed, 50 insertions(+), 28 deletions(-) diff --git a/cmake/caches/ci/cpu.cmake b/cmake/caches/ci/cpu.cmake index 95303189483..48cd357159d 100644 --- a/cmake/caches/ci/cpu.cmake +++ b/cmake/caches/ci/cpu.cmake @@ -1,5 +1,6 @@ set(BUILD_CUDA NO CACHE BOOL "") set(BUILD_NPU NO CACHE BOOL "") +set(BUILD_MLU NO CACHE BOOL "") set(BUILD_GIT_VERSION YES CACHE BOOL "") set(BUILD_TESTING YES CACHE BOOL "") set(WITH_ONEDNN YES CACHE BOOL "") diff --git a/cmake/caches/cn/cpu.cmake b/cmake/caches/cn/cpu.cmake index b7fae9fece0..8ba7d965dac 100644 --- a/cmake/caches/cn/cpu.cmake +++ b/cmake/caches/cn/cpu.cmake @@ -1,5 +1,6 @@ set(BUILD_CUDA NO CACHE BOOL "") set(BUILD_NPU NO CACHE BOOL "") +set(BUILD_MLU NO CACHE BOOL "") set(BUILD_SHARED_LIBS YES CACHE BOOL "") set(THIRD_PARTY_MIRROR aliyun CACHE STRING "") set(PIP_INDEX_MIRROR "https://pypi.tuna.tsinghua.edu.cn/simple" CACHE STRING "") diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index bfbf87232e8..45ec5dcbf54 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -358,6 +358,11 @@ if(BUILD_NPU) endif() message(STATUS "BUILD_NPU: ${BUILD_NPU}") +if(BUILD_MLU) + add_definitions(-DWITH_MLU) +endif() +message(STATUS "BUILD_MLU: ${BUILD_MLU}") + if(BUILD_CUDA AND WITH_CUTLASS) if(CUDA_VERSION VERSION_GREATER_EQUAL "10.1") add_definitions(-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1) diff --git a/oneflow/api/python/flags.cpp b/oneflow/api/python/flags.cpp index 1f4f70a3f85..eb3e8c838eb 100644 --- a/oneflow/api/python/flags.cpp +++ b/oneflow/api/python/flags.cpp @@ -37,6 +37,14 @@ ONEFLOW_API_PYBIND11_MODULE("flags", m) { #endif // WITH_NPU }); + m.def("with_mlu", []() { +#ifdef WITH_MLU + return true; +#else + return false; +#endif // WITH_MLU + }); + m.def("cuda_version", []() { #ifdef WITH_CUDA return CUDA_VERSION; diff --git a/oneflow/core/auto_parallel/boxing_collector.cpp b/oneflow/core/auto_parallel/boxing_collector.cpp index fca8f5f734c..5fb8ba08518 100644 --- a/oneflow/core/auto_parallel/boxing_collector.cpp +++ b/oneflow/core/auto_parallel/boxing_collector.cpp @@ -581,7 +581,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const return Maybe::Ok(); } -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) // Use a general basic communication if no P in the consumer if (((Singleton::Get()->nccl_use_compute_stream() && producer_parallel_desc == consumer_parallel_desc) @@ -600,7 +600,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const // Otherwise, one-step transfer return Maybe::Ok(); } -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) if (JUST(ComputeLazyCopyCostBetweenNdSbp(sbp_producer, sbp_consumer, logical_blob_desc, producer_parallel_desc, consumer_parallel_desc, diff --git a/oneflow/core/common/auto_registration_factory.h b/oneflow/core/common/auto_registration_factory.h index 680b0e5125c..7af6b8d263a 100644 --- a/oneflow/core/common/auto_registration_factory.h +++ b/oneflow/core/common/auto_registration_factory.h @@ -68,7 +68,8 @@ struct AutoRegistrationFactory { bool has_creators() const { return creators_.get() != nullptr; } const HashMap& creators() const { - CHECK(has_creators()) << "Unregistered key type: " << typeid(Key).name(); + CHECK(has_creators()) << "Unregistered key type: " << typeid(Key).name() + << "Base type name:" << typeid(Base).name(); return *creators_.get(); } diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index 51f42367c07..f2d1a8a029b 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -2238,7 +2238,7 @@ class SparseSoftmaxCrossEntropyFunctor { const std::shared_ptr& label) const { if (!(logits->is_global() && label->is_global())) { return false; } // npu-implementation not support ms version yet -#ifdef WITH_NPU +#if defined(WITH_NPU) || defined(WITH_MLU) return false; #endif diff --git a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp index 4e5742460e0..ed19e1f2669 100644 --- a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp +++ b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp @@ -37,6 +37,8 @@ CollectiveBoxingSubTskGphBuilder::CollectiveBoxingSubTskGphBuilder() { builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kCUDA)); #elif defined(WITH_NPU) builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kNPU)); +#elif defined(WITH_MLU) + builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kMLU)); #else LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; #endif diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index ddd7d879ca5..6362b2f2343 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -102,7 +102,7 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (in_parallel_desc.device_type() == out_parallel_desc.device_type() && in_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index 1447279605a..815459c993c 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -75,6 +75,8 @@ bool ResourceDesc::nccl_use_compute_stream() const { return resource_.nccl_use_compute_stream(); #elif defined(WITH_NPU) return resource_.nccl_use_compute_stream(); +#elif defined(WITH_MLU) + return resource_.nccl_use_compute_stream(); #else return false; #endif diff --git a/oneflow/core/job/runtime.cpp b/oneflow/core/job/runtime.cpp index a87ac28d6ed..1f23fedcbeb 100644 --- a/oneflow/core/job/runtime.cpp +++ b/oneflow/core/job/runtime.cpp @@ -70,7 +70,7 @@ Runtime::Runtime( Singleton::Get()->AddPlan(plan); collective_boxing_scheduler_plan_token_ = Singleton::Get()->AddPlan(plan); -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) const auto& vaild_ccl_comm_mgr_device_types = EagerCclCommMgrBuilder::Get().vaild_ccl_comm_mgr_device_types(); if (!vaild_ccl_comm_mgr_device_types.empty() && !Singleton::Get()) { @@ -78,7 +78,7 @@ Runtime::Runtime( EagerCclCommMgrBuilder::Get().NewCclCommMgr(vaild_ccl_comm_mgr_device_types.front())); } Singleton::Get()->CreateCommFromPlan(plan); -#endif // defined(WITH_CUDA) || defined(WITH_NPU) +#endif // defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) } std::vector source_tasks; source_tasks.reserve(plan.task().size()); diff --git a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp index d02f6316092..ed96998c2d1 100644 --- a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp +++ b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp @@ -16,7 +16,7 @@ limitations under the License. #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/common/util.h" #include "oneflow/core/job/nd_sbp_util.h" -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/framework/framework.h" #include "oneflow/core/framework/nd_sbp.h" #include "oneflow/core/framework/instructions_builder.h" @@ -883,4 +883,4 @@ REGISTER_JOB_PASS("InsertNcclLogicalOpPass", InsertNcclLogicalOpPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) diff --git a/oneflow/core/job_rewriter/job_completer.cpp b/oneflow/core/job_rewriter/job_completer.cpp index ae061f8d83e..f6fe5a639af 100644 --- a/oneflow/core/job_rewriter/job_completer.cpp +++ b/oneflow/core/job_rewriter/job_completer.cpp @@ -153,7 +153,7 @@ Maybe JobCompleter::Complete(Job* job) { compile_tc->Count("[GraphCompile]" + job_name + " SystemOpFillJobNamePass", 1, true); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) if (Singleton::Get()->nccl_use_compute_stream()) { // NOTE(chengcheng): this pass need as last pass for insert correct op with nccl boxing. JUST(JobPass4Name("InsertNcclLogicalOpPass")(job, &job_pass_ctx)); @@ -169,7 +169,7 @@ Maybe JobCompleter::Complete(Job* job) { JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); } -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU JUST(JobPass4Name("LogicalChainPass")(job, &job_pass_ctx)); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); diff --git a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp index 33392b17a1c..d366602ebed 100644 --- a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -210,4 +210,4 @@ REGISTER_JOB_PASS("NcclLogicalChainStrictOrderPass", NcclLogicalChainStrictOrder } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp index 05e255493d3..0087b7446cd 100644 --- a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -293,4 +293,4 @@ REGISTER_JOB_PASS("NcclLogicalOpFusionPass", NcclLogicalOpFusionPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp index af81336ea70..a4111f61ae9 100644 --- a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp +++ b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp @@ -22,7 +22,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -254,4 +254,4 @@ REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxin } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu index 93070a08020..cc1e96a8d60 100644 --- a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu +++ b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu @@ -23,7 +23,7 @@ limitations under the License. #include "oneflow/core/ep/cuda/cuda_stream.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp index dded34c8eb2..a149319942c 100644 --- a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp @@ -25,7 +25,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -554,4 +554,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduc } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp index aba3546a3d8..ac22b2fa2f2 100644 --- a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "collective_communication/include/all_to_all.h" #include "collective_communication/include/reduce_scatter.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -703,7 +703,8 @@ size_t InferNcclLogicalFusionKernelTmpBufferSize(user_op::InferContext* ctx) { REGISTER_USER_KERNEL("_nccl_logical_fusion") .SetCreateFn() .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) - || (user_op::HobDeviceType() == DeviceType::kNPU)) + || (user_op::HobDeviceType() == DeviceType::kNPU) + || (user_op::HobDeviceType() == DeviceType::kMLU)) .SetInferTmpSizeFn(InferNcclLogicalFusionKernelTmpBufferSize); // TODO: SetIsMatchedHob support multi devices(not including cpu) @@ -711,4 +712,4 @@ REGISTER_USER_KERNEL("_nccl_logical_fusion") } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_kernels.cpp b/oneflow/user/kernels/nccl_logical_kernels.cpp index dde7be74c28..b721f74d5e4 100644 --- a/oneflow/user/kernels/nccl_logical_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_kernels.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/broadcast.h" #include "oneflow/user/kernels/collective_communication/include/reduce.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -640,4 +640,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_s2s"); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp index dadbd0f9096..21a4e98a1fd 100644 --- a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp @@ -30,7 +30,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -288,9 +288,10 @@ size_t InferTmpBufferSize(user_op::InferContext* ctx) { REGISTER_USER_KERNEL("_nccl_logical_send_recv") .SetCreateFn() .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) - || (user_op::HobDeviceType() == DeviceType::kNPU)) + || (user_op::HobDeviceType() == DeviceType::kNPU) + || (user_op::HobDeviceType() == DeviceType::kMLU)) .SetInferTmpSizeFn(InferTmpBufferSize); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU +#endif // WITH_CUDA || WITH_NPU || WITH_MLU From dfcf553fe7c6984b83ee21f5a9b5a9a20eee158e Mon Sep 17 00:00:00 2001 From: luyang Date: Tue, 1 Jul 2025 17:48:19 +0800 Subject: [PATCH 2/5] support num_accumulation_steps --- oneflow/user/kernels/pack_kernel.cpp | 4 +++- oneflow/user/kernels/unpack_kernel.cpp | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/oneflow/user/kernels/pack_kernel.cpp b/oneflow/user/kernels/pack_kernel.cpp index ea342c27029..c61923a27bb 100644 --- a/oneflow/user/kernels/pack_kernel.cpp +++ b/oneflow/user/kernels/pack_kernel.cpp @@ -69,7 +69,9 @@ class PackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_PACK_KERNEL, DEVICE_TYPE_SEQ) - +#if defined(WITH_MLU) +REGISTER_PACK_KERNEL(DeviceType::kMLU) +#endif #undef REGISTER_PACK_KERNEL } // namespace diff --git a/oneflow/user/kernels/unpack_kernel.cpp b/oneflow/user/kernels/unpack_kernel.cpp index 35b18165a44..482248c3add 100644 --- a/oneflow/user/kernels/unpack_kernel.cpp +++ b/oneflow/user/kernels/unpack_kernel.cpp @@ -62,7 +62,9 @@ class UnpackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_UNPACK_KERNEL, DEVICE_TYPE_SEQ) - +#if defined(WITH_MLU) +REGISTER_UNPACK_KERNEL(DeviceType::kMLU) +#endif #undef REGISTER_UNPACK_KERNEL } // namespace From 8dcc612a2730972e9da9124f8ccca0fbf9a1bd3f Mon Sep 17 00:00:00 2001 From: luyang Date: Tue, 1 Jul 2025 20:57:41 +0800 Subject: [PATCH 3/5] add definition 'WITH_DEVICES' --- cmake/oneflow.cmake | 24 ++++++++++++++++++- oneflow/api/python/flags.cpp | 8 +++++++ .../core/auto_parallel/boxing_collector.cpp | 4 ++-- oneflow/core/functional/impl/nn_functor.cpp | 2 +- ...llective_boxing_sub_task_graph_builder.cpp | 12 +++++++--- ...erarchical_sub_task_graph_builder_impl.cpp | 2 +- oneflow/core/job/resource_desc.cpp | 10 +++++++- oneflow/core/job/runtime.cpp | 4 ++-- .../insert_nccl_logical_op_pass.cpp | 4 ++-- oneflow/core/job_rewriter/job_completer.cpp | 4 ++-- .../nccl_logical_chain_strict_order_pass.cpp | 4 ++-- .../nccl_logical_op_fusion_pass.cpp | 4 ++-- .../kernel/nccl_send_recv_boxing_kernel.cpp | 4 ++-- .../kernels/nccl_logical_2d_sbp_kernels.cpp | 4 ++-- .../kernels/nccl_logical_fusion_kernel.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_kernels.cpp | 4 ++-- .../kernels/nccl_logical_send_recv_kernel.cpp | 4 ++-- oneflow/user/kernels/pack_kernel.cpp | 9 ++++++- oneflow/user/kernels/unpack_kernel.cpp | 10 +++++++- 19 files changed, 90 insertions(+), 31 deletions(-) diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index 45ec5dcbf54..19e30d66dda 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -353,15 +353,37 @@ if(BUILD_CUDA) PROPERTIES COMPILE_FLAGS "-DCUDA_REAL_ARCHS=\"${CUDA_REAL_ARCHS}\"") endif() +# init flag of devices +set(DEVICES_ENABLED FALSE) + if(BUILD_NPU) add_definitions(-DWITH_NPU) + set(DEVICES_ENABLED TRUE) endif() -message(STATUS "BUILD_NPU: ${BUILD_NPU}") +message(STATUS "NPU support enabled.") if(BUILD_MLU) add_definitions(-DWITH_MLU) + set(DEVICES_ENABLED TRUE) endif() +message(STATUS "MLU support enabled.") + +if(DEVICES_ENABLED) + # check WITH_DEVICES defined or not + get_directory_property(EXISTING_DEFS COMPILE_DEFINITIONS) + + if(NOT "WITH_DEVICES" IN_LIST EXISTING_DEFS) + add_definitions(-DWITH_DEVICES) + message(STATUS "Added generic device support definition") + else() + message(STATUS "Generic device support already defined") + endif() +endif() + +# show all devices status message(STATUS "BUILD_MLU: ${BUILD_MLU}") +message(STATUS "BUILD_NPU: ${BUILD_NPU}") +message(STATUS "Generic device support: ${DEVICE_ENABLED}") if(BUILD_CUDA AND WITH_CUTLASS) if(CUDA_VERSION VERSION_GREATER_EQUAL "10.1") diff --git a/oneflow/api/python/flags.cpp b/oneflow/api/python/flags.cpp index eb3e8c838eb..212fad21c21 100644 --- a/oneflow/api/python/flags.cpp +++ b/oneflow/api/python/flags.cpp @@ -29,6 +29,14 @@ ONEFLOW_API_PYBIND11_MODULE("flags", m) { #endif // WITH_CUDA }); + m.def("with_devices", []() { +#ifdef WITH_DEVICES + return true; +#else + return false; +#endif // WITH_DEVICES + }); + m.def("with_npu", []() { #ifdef WITH_NPU return true; diff --git a/oneflow/core/auto_parallel/boxing_collector.cpp b/oneflow/core/auto_parallel/boxing_collector.cpp index 5fb8ba08518..fdf1f6adb35 100644 --- a/oneflow/core/auto_parallel/boxing_collector.cpp +++ b/oneflow/core/auto_parallel/boxing_collector.cpp @@ -581,7 +581,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const return Maybe::Ok(); } -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) // Use a general basic communication if no P in the consumer if (((Singleton::Get()->nccl_use_compute_stream() && producer_parallel_desc == consumer_parallel_desc) @@ -600,7 +600,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const // Otherwise, one-step transfer return Maybe::Ok(); } -#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) +#endif // WITH_CUDA || WITH_DEVICES if (JUST(ComputeLazyCopyCostBetweenNdSbp(sbp_producer, sbp_consumer, logical_blob_desc, producer_parallel_desc, consumer_parallel_desc, diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index f2d1a8a029b..d6833da44b4 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -2238,7 +2238,7 @@ class SparseSoftmaxCrossEntropyFunctor { const std::shared_ptr& label) const { if (!(logits->is_global() && label->is_global())) { return false; } // npu-implementation not support ms version yet -#if defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_DEVICES) return false; #endif diff --git a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp index ed19e1f2669..ad2557a3ebf 100644 --- a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp +++ b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp @@ -35,12 +35,18 @@ CollectiveBoxingSubTskGphBuilder::CollectiveBoxingSubTskGphBuilder() { if (collective_boxing_conf.nccl_enable_all_to_all()) { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kCUDA)); -#elif defined(WITH_NPU) +#else + LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; +#endif + +#if defined(WITH_DEVICES) +#if defined(WITH_NPU) builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kNPU)); #elif defined(WITH_MLU) builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kMLU)); -#else - LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; +#elif defined(WITH_XPU) + builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kXPU)); +#endif #endif } chain_builder_.reset(new ChainSubTskGphBuilder(builders)); diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index 6362b2f2343..20333a207fe 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -102,7 +102,7 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (in_parallel_desc.device_type() == out_parallel_desc.device_type() && in_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index 815459c993c..e18ad03afaa 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -73,13 +73,21 @@ CollectiveBoxingConf ResourceDesc::collective_boxing_conf() const { bool ResourceDesc::nccl_use_compute_stream() const { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 return resource_.nccl_use_compute_stream(); -#elif defined(WITH_NPU) +#else + return false; +#endif + +#if defined(WITH_DEVICES) +#if defined(WITH_NPU) return resource_.nccl_use_compute_stream(); #elif defined(WITH_MLU) return resource_.nccl_use_compute_stream(); +#elif defined(WITH_XPU) + return resource_.nccl_use_compute_stream(); #else return false; #endif +#endif } void ResourceDesc::DumpCudnnConf(const JobConfigProto& job_conf) { diff --git a/oneflow/core/job/runtime.cpp b/oneflow/core/job/runtime.cpp index 1f23fedcbeb..aaf677b7a76 100644 --- a/oneflow/core/job/runtime.cpp +++ b/oneflow/core/job/runtime.cpp @@ -70,7 +70,7 @@ Runtime::Runtime( Singleton::Get()->AddPlan(plan); collective_boxing_scheduler_plan_token_ = Singleton::Get()->AddPlan(plan); -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) const auto& vaild_ccl_comm_mgr_device_types = EagerCclCommMgrBuilder::Get().vaild_ccl_comm_mgr_device_types(); if (!vaild_ccl_comm_mgr_device_types.empty() && !Singleton::Get()) { @@ -78,7 +78,7 @@ Runtime::Runtime( EagerCclCommMgrBuilder::Get().NewCclCommMgr(vaild_ccl_comm_mgr_device_types.front())); } Singleton::Get()->CreateCommFromPlan(plan); -#endif // defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#endif // defined(WITH_CUDA) || WITH_DEVICES } std::vector source_tasks; source_tasks.reserve(plan.task().size()); diff --git a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp index ed96998c2d1..69fda0fb1f9 100644 --- a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp +++ b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp @@ -16,7 +16,7 @@ limitations under the License. #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/common/util.h" #include "oneflow/core/job/nd_sbp_util.h" -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) #include "oneflow/core/framework/framework.h" #include "oneflow/core/framework/nd_sbp.h" #include "oneflow/core/framework/instructions_builder.h" @@ -883,4 +883,4 @@ REGISTER_JOB_PASS("InsertNcclLogicalOpPass", InsertNcclLogicalOpPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/core/job_rewriter/job_completer.cpp b/oneflow/core/job_rewriter/job_completer.cpp index f6fe5a639af..fbf9ba27bc3 100644 --- a/oneflow/core/job_rewriter/job_completer.cpp +++ b/oneflow/core/job_rewriter/job_completer.cpp @@ -153,7 +153,7 @@ Maybe JobCompleter::Complete(Job* job) { compile_tc->Count("[GraphCompile]" + job_name + " SystemOpFillJobNamePass", 1, true); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) if (Singleton::Get()->nccl_use_compute_stream()) { // NOTE(chengcheng): this pass need as last pass for insert correct op with nccl boxing. JUST(JobPass4Name("InsertNcclLogicalOpPass")(job, &job_pass_ctx)); @@ -169,7 +169,7 @@ Maybe JobCompleter::Complete(Job* job) { JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); } -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES JUST(JobPass4Name("LogicalChainPass")(job, &job_pass_ctx)); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); diff --git a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp index d366602ebed..788f6a230cc 100644 --- a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -210,4 +210,4 @@ REGISTER_JOB_PASS("NcclLogicalChainStrictOrderPass", NcclLogicalChainStrictOrder } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp index 0087b7446cd..99b5bf51b79 100644 --- a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) +#if defined(WITH_CUDA) || defined(WITH_DEVICES) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -293,4 +293,4 @@ REGISTER_JOB_PASS("NcclLogicalOpFusionPass", NcclLogicalOpFusionPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp index a4111f61ae9..7fc99a5a302 100644 --- a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp +++ b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp @@ -22,7 +22,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) namespace oneflow { @@ -254,4 +254,4 @@ REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxin } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp index a149319942c..ae80efa66ab 100644 --- a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp @@ -25,7 +25,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) namespace oneflow { @@ -554,4 +554,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduc } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp index ac22b2fa2f2..10fed2a6649 100644 --- a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "collective_communication/include/all_to_all.h" #include "collective_communication/include/reduce_scatter.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) namespace oneflow { @@ -712,4 +712,4 @@ REGISTER_USER_KERNEL("_nccl_logical_fusion") } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/user/kernels/nccl_logical_kernels.cpp b/oneflow/user/kernels/nccl_logical_kernels.cpp index b721f74d5e4..e4fc78f8dfd 100644 --- a/oneflow/user/kernels/nccl_logical_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_kernels.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/broadcast.h" #include "oneflow/user/kernels/collective_communication/include/reduce.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) namespace oneflow { @@ -640,4 +640,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_s2s"); } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp index 21a4e98a1fd..e8666e6f524 100644 --- a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp @@ -30,7 +30,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) namespace oneflow { @@ -294,4 +294,4 @@ REGISTER_USER_KERNEL("_nccl_logical_send_recv") } // namespace oneflow -#endif // WITH_CUDA || WITH_NPU || WITH_MLU +#endif // WITH_CUDA || WITH_DEVICES diff --git a/oneflow/user/kernels/pack_kernel.cpp b/oneflow/user/kernels/pack_kernel.cpp index c61923a27bb..03ff31dcf62 100644 --- a/oneflow/user/kernels/pack_kernel.cpp +++ b/oneflow/user/kernels/pack_kernel.cpp @@ -69,9 +69,16 @@ class PackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_PACK_KERNEL, DEVICE_TYPE_SEQ) -#if defined(WITH_MLU) +#if defined(WITH_DEVICES) +#if defined(WITH_NPU) +REGISTER_PACK_KERNEL(DeviceType::kNPU) +#elif defined(WITH_MLU) REGISTER_PACK_KERNEL(DeviceType::kMLU) +#elif defined(WITH_XPU) +REGISTER_PACK_KERNEL(DeviceType::kXPU) #endif +#endif + #undef REGISTER_PACK_KERNEL } // namespace diff --git a/oneflow/user/kernels/unpack_kernel.cpp b/oneflow/user/kernels/unpack_kernel.cpp index 482248c3add..828c97bf553 100644 --- a/oneflow/user/kernels/unpack_kernel.cpp +++ b/oneflow/user/kernels/unpack_kernel.cpp @@ -62,9 +62,17 @@ class UnpackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_UNPACK_KERNEL, DEVICE_TYPE_SEQ) -#if defined(WITH_MLU) + +#if defined(WITH_DEVICES) +#if defined(WITH_NPU) +REGISTER_UNPACK_KERNEL(DeviceType::kNPU) +#elif defined(WITH_MLU) REGISTER_UNPACK_KERNEL(DeviceType::kMLU) +#elif defined(WITH_XPU) +REGISTER_UNPACK_KERNEL(DeviceType::kXPU) +#endif #endif + #undef REGISTER_UNPACK_KERNEL } // namespace From 5abb6b10bd1432cbe8b5c15c35e7fc295030cf0f Mon Sep 17 00:00:00 2001 From: oneflow-ci-bot Date: Tue, 1 Jul 2025 12:59:05 +0000 Subject: [PATCH 4/5] auto format by CI --- cmake/oneflow.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index 19e30d66dda..8fd915e3893 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -371,7 +371,7 @@ message(STATUS "MLU support enabled.") if(DEVICES_ENABLED) # check WITH_DEVICES defined or not get_directory_property(EXISTING_DEFS COMPILE_DEFINITIONS) - + if(NOT "WITH_DEVICES" IN_LIST EXISTING_DEFS) add_definitions(-DWITH_DEVICES) message(STATUS "Added generic device support definition") From 972c6d77811ccf83b003b4802f5e92a256bed4c3 Mon Sep 17 00:00:00 2001 From: luyang Date: Wed, 2 Jul 2025 12:10:35 +0800 Subject: [PATCH 5/5] revert changes of 'add definition WITH_DEVICES' --- cmake/oneflow.cmake | 24 +------------------ oneflow/api/python/flags.cpp | 8 ------- .../core/auto_parallel/boxing_collector.cpp | 4 ++-- oneflow/core/functional/impl/nn_functor.cpp | 2 +- ...llective_boxing_sub_task_graph_builder.cpp | 12 +++------- ...erarchical_sub_task_graph_builder_impl.cpp | 2 +- oneflow/core/job/resource_desc.cpp | 10 +------- oneflow/core/job/runtime.cpp | 4 ++-- .../insert_nccl_logical_op_pass.cpp | 4 ++-- oneflow/core/job_rewriter/job_completer.cpp | 4 ++-- .../nccl_logical_chain_strict_order_pass.cpp | 4 ++-- .../nccl_logical_op_fusion_pass.cpp | 4 ++-- .../kernel/nccl_send_recv_boxing_kernel.cpp | 4 ++-- .../kernels/nccl_logical_2d_sbp_kernels.cpp | 4 ++-- .../kernels/nccl_logical_fusion_kernel.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_kernels.cpp | 4 ++-- .../kernels/nccl_logical_send_recv_kernel.cpp | 4 ++-- oneflow/user/kernels/pack_kernel.cpp | 9 +------ oneflow/user/kernels/unpack_kernel.cpp | 10 +------- 19 files changed, 31 insertions(+), 90 deletions(-) diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index 8fd915e3893..45ec5dcbf54 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -353,37 +353,15 @@ if(BUILD_CUDA) PROPERTIES COMPILE_FLAGS "-DCUDA_REAL_ARCHS=\"${CUDA_REAL_ARCHS}\"") endif() -# init flag of devices -set(DEVICES_ENABLED FALSE) - if(BUILD_NPU) add_definitions(-DWITH_NPU) - set(DEVICES_ENABLED TRUE) endif() -message(STATUS "NPU support enabled.") +message(STATUS "BUILD_NPU: ${BUILD_NPU}") if(BUILD_MLU) add_definitions(-DWITH_MLU) - set(DEVICES_ENABLED TRUE) endif() -message(STATUS "MLU support enabled.") - -if(DEVICES_ENABLED) - # check WITH_DEVICES defined or not - get_directory_property(EXISTING_DEFS COMPILE_DEFINITIONS) - - if(NOT "WITH_DEVICES" IN_LIST EXISTING_DEFS) - add_definitions(-DWITH_DEVICES) - message(STATUS "Added generic device support definition") - else() - message(STATUS "Generic device support already defined") - endif() -endif() - -# show all devices status message(STATUS "BUILD_MLU: ${BUILD_MLU}") -message(STATUS "BUILD_NPU: ${BUILD_NPU}") -message(STATUS "Generic device support: ${DEVICE_ENABLED}") if(BUILD_CUDA AND WITH_CUTLASS) if(CUDA_VERSION VERSION_GREATER_EQUAL "10.1") diff --git a/oneflow/api/python/flags.cpp b/oneflow/api/python/flags.cpp index 212fad21c21..eb3e8c838eb 100644 --- a/oneflow/api/python/flags.cpp +++ b/oneflow/api/python/flags.cpp @@ -29,14 +29,6 @@ ONEFLOW_API_PYBIND11_MODULE("flags", m) { #endif // WITH_CUDA }); - m.def("with_devices", []() { -#ifdef WITH_DEVICES - return true; -#else - return false; -#endif // WITH_DEVICES - }); - m.def("with_npu", []() { #ifdef WITH_NPU return true; diff --git a/oneflow/core/auto_parallel/boxing_collector.cpp b/oneflow/core/auto_parallel/boxing_collector.cpp index fdf1f6adb35..5fb8ba08518 100644 --- a/oneflow/core/auto_parallel/boxing_collector.cpp +++ b/oneflow/core/auto_parallel/boxing_collector.cpp @@ -581,7 +581,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const return Maybe::Ok(); } -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) // Use a general basic communication if no P in the consumer if (((Singleton::Get()->nccl_use_compute_stream() && producer_parallel_desc == consumer_parallel_desc) @@ -600,7 +600,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const // Otherwise, one-step transfer return Maybe::Ok(); } -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) if (JUST(ComputeLazyCopyCostBetweenNdSbp(sbp_producer, sbp_consumer, logical_blob_desc, producer_parallel_desc, consumer_parallel_desc, diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index d6833da44b4..f2d1a8a029b 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -2238,7 +2238,7 @@ class SparseSoftmaxCrossEntropyFunctor { const std::shared_ptr& label) const { if (!(logits->is_global() && label->is_global())) { return false; } // npu-implementation not support ms version yet -#if defined(WITH_DEVICES) +#if defined(WITH_NPU) || defined(WITH_MLU) return false; #endif diff --git a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp index ad2557a3ebf..ed19e1f2669 100644 --- a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp +++ b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp @@ -35,18 +35,12 @@ CollectiveBoxingSubTskGphBuilder::CollectiveBoxingSubTskGphBuilder() { if (collective_boxing_conf.nccl_enable_all_to_all()) { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kCUDA)); -#else - LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; -#endif - -#if defined(WITH_DEVICES) -#if defined(WITH_NPU) +#elif defined(WITH_NPU) builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kNPU)); #elif defined(WITH_MLU) builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kMLU)); -#elif defined(WITH_XPU) - builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kXPU)); -#endif +#else + LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; #endif } chain_builder_.reset(new ChainSubTskGphBuilder(builders)); diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index 20333a207fe..6362b2f2343 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -102,7 +102,7 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (in_parallel_desc.device_type() == out_parallel_desc.device_type() && in_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index e18ad03afaa..815459c993c 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -73,21 +73,13 @@ CollectiveBoxingConf ResourceDesc::collective_boxing_conf() const { bool ResourceDesc::nccl_use_compute_stream() const { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 return resource_.nccl_use_compute_stream(); -#else - return false; -#endif - -#if defined(WITH_DEVICES) -#if defined(WITH_NPU) +#elif defined(WITH_NPU) return resource_.nccl_use_compute_stream(); #elif defined(WITH_MLU) return resource_.nccl_use_compute_stream(); -#elif defined(WITH_XPU) - return resource_.nccl_use_compute_stream(); #else return false; #endif -#endif } void ResourceDesc::DumpCudnnConf(const JobConfigProto& job_conf) { diff --git a/oneflow/core/job/runtime.cpp b/oneflow/core/job/runtime.cpp index aaf677b7a76..1f23fedcbeb 100644 --- a/oneflow/core/job/runtime.cpp +++ b/oneflow/core/job/runtime.cpp @@ -70,7 +70,7 @@ Runtime::Runtime( Singleton::Get()->AddPlan(plan); collective_boxing_scheduler_plan_token_ = Singleton::Get()->AddPlan(plan); -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) const auto& vaild_ccl_comm_mgr_device_types = EagerCclCommMgrBuilder::Get().vaild_ccl_comm_mgr_device_types(); if (!vaild_ccl_comm_mgr_device_types.empty() && !Singleton::Get()) { @@ -78,7 +78,7 @@ Runtime::Runtime( EagerCclCommMgrBuilder::Get().NewCclCommMgr(vaild_ccl_comm_mgr_device_types.front())); } Singleton::Get()->CreateCommFromPlan(plan); -#endif // defined(WITH_CUDA) || WITH_DEVICES +#endif // defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) } std::vector source_tasks; source_tasks.reserve(plan.task().size()); diff --git a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp index 69fda0fb1f9..ed96998c2d1 100644 --- a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp +++ b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp @@ -16,7 +16,7 @@ limitations under the License. #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/common/util.h" #include "oneflow/core/job/nd_sbp_util.h" -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/framework/framework.h" #include "oneflow/core/framework/nd_sbp.h" #include "oneflow/core/framework/instructions_builder.h" @@ -883,4 +883,4 @@ REGISTER_JOB_PASS("InsertNcclLogicalOpPass", InsertNcclLogicalOpPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || defined(WITH_MLU) diff --git a/oneflow/core/job_rewriter/job_completer.cpp b/oneflow/core/job_rewriter/job_completer.cpp index fbf9ba27bc3..f6fe5a639af 100644 --- a/oneflow/core/job_rewriter/job_completer.cpp +++ b/oneflow/core/job_rewriter/job_completer.cpp @@ -153,7 +153,7 @@ Maybe JobCompleter::Complete(Job* job) { compile_tc->Count("[GraphCompile]" + job_name + " SystemOpFillJobNamePass", 1, true); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) if (Singleton::Get()->nccl_use_compute_stream()) { // NOTE(chengcheng): this pass need as last pass for insert correct op with nccl boxing. JUST(JobPass4Name("InsertNcclLogicalOpPass")(job, &job_pass_ctx)); @@ -169,7 +169,7 @@ Maybe JobCompleter::Complete(Job* job) { JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); } -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU JUST(JobPass4Name("LogicalChainPass")(job, &job_pass_ctx)); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); diff --git a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp index 788f6a230cc..d366602ebed 100644 --- a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -210,4 +210,4 @@ REGISTER_JOB_PASS("NcclLogicalChainStrictOrderPass", NcclLogicalChainStrictOrder } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp index 99b5bf51b79..0087b7446cd 100644 --- a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#if defined(WITH_CUDA) || defined(WITH_DEVICES) +#if defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -293,4 +293,4 @@ REGISTER_JOB_PASS("NcclLogicalOpFusionPass", NcclLogicalOpFusionPass); } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp index 7fc99a5a302..a4111f61ae9 100644 --- a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp +++ b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp @@ -22,7 +22,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -254,4 +254,4 @@ REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxin } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp index ae80efa66ab..a149319942c 100644 --- a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp @@ -25,7 +25,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -554,4 +554,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduc } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp index 10fed2a6649..ac22b2fa2f2 100644 --- a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "collective_communication/include/all_to_all.h" #include "collective_communication/include/reduce_scatter.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -712,4 +712,4 @@ REGISTER_USER_KERNEL("_nccl_logical_fusion") } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_kernels.cpp b/oneflow/user/kernels/nccl_logical_kernels.cpp index e4fc78f8dfd..b721f74d5e4 100644 --- a/oneflow/user/kernels/nccl_logical_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_kernels.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/broadcast.h" #include "oneflow/user/kernels/collective_communication/include/reduce.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -640,4 +640,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_s2s"); } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp index e8666e6f524..21a4e98a1fd 100644 --- a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp @@ -30,7 +30,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_DEVICES) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) || defined(WITH_MLU) namespace oneflow { @@ -294,4 +294,4 @@ REGISTER_USER_KERNEL("_nccl_logical_send_recv") } // namespace oneflow -#endif // WITH_CUDA || WITH_DEVICES +#endif // WITH_CUDA || WITH_NPU || WITH_MLU diff --git a/oneflow/user/kernels/pack_kernel.cpp b/oneflow/user/kernels/pack_kernel.cpp index 03ff31dcf62..c61923a27bb 100644 --- a/oneflow/user/kernels/pack_kernel.cpp +++ b/oneflow/user/kernels/pack_kernel.cpp @@ -69,16 +69,9 @@ class PackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_PACK_KERNEL, DEVICE_TYPE_SEQ) -#if defined(WITH_DEVICES) -#if defined(WITH_NPU) -REGISTER_PACK_KERNEL(DeviceType::kNPU) -#elif defined(WITH_MLU) +#if defined(WITH_MLU) REGISTER_PACK_KERNEL(DeviceType::kMLU) -#elif defined(WITH_XPU) -REGISTER_PACK_KERNEL(DeviceType::kXPU) #endif -#endif - #undef REGISTER_PACK_KERNEL } // namespace diff --git a/oneflow/user/kernels/unpack_kernel.cpp b/oneflow/user/kernels/unpack_kernel.cpp index 828c97bf553..482248c3add 100644 --- a/oneflow/user/kernels/unpack_kernel.cpp +++ b/oneflow/user/kernels/unpack_kernel.cpp @@ -62,17 +62,9 @@ class UnpackKernel final : public user_op::OpKernel { (user_op::HobDeviceType() == device)); OF_PP_FOR_EACH_TUPLE(REGISTER_UNPACK_KERNEL, DEVICE_TYPE_SEQ) - -#if defined(WITH_DEVICES) -#if defined(WITH_NPU) -REGISTER_UNPACK_KERNEL(DeviceType::kNPU) -#elif defined(WITH_MLU) +#if defined(WITH_MLU) REGISTER_UNPACK_KERNEL(DeviceType::kMLU) -#elif defined(WITH_XPU) -REGISTER_UNPACK_KERNEL(DeviceType::kXPU) -#endif #endif - #undef REGISTER_UNPACK_KERNEL } // namespace