adaptation_of_mlu_device_based_on_refactor_xccl_primitive

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 "")

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 "")

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)

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;

oneflow/core/auto_parallel/boxing_collector.cpp

@@ -581,7 +581,7 @@ Maybe<void> BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const
     return Maybe<void>::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<ResourceDesc, ForSession>::Get()->nccl_use_compute_stream()
         && producer_parallel_desc == consumer_parallel_desc)
@@ -600,7 +600,7 @@ Maybe<void> BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const
     // Otherwise, one-step transfer
     return Maybe<void>::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,

oneflow/core/common/auto_registration_factory.h

@@ -68,7 +68,8 @@ struct AutoRegistrationFactory {
   bool has_creators() const { return creators_.get() != nullptr; }
 
   const HashMap<Key, Creator>& 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();
   }
 

oneflow/core/functional/impl/nn_functor.cpp

@@ -2238,7 +2238,7 @@ class SparseSoftmaxCrossEntropyFunctor {
                                const std::shared_ptr<one::Tensor>& 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
 

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

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);

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

oneflow/core/job/runtime.cpp

@@ -70,15 +70,15 @@ Runtime::Runtime(
     Singleton<RuntimeJobDescs>::Get()->AddPlan(plan);
     collective_boxing_scheduler_plan_token_ =
         Singleton<boxing::collective::Scheduler>::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<EagerCclCommMgr>::Get()) {
       Singleton<EagerCclCommMgr>::SetAllocated(
           EagerCclCommMgrBuilder::Get().NewCclCommMgr(vaild_ccl_comm_mgr_device_types.front()));
     }
     Singleton<EagerCclCommMgr>::Get()->CreateCommFromPlan(plan);
-#endif  // defined(WITH_CUDA) || defined(WITH_NPU)
+#endif  // defined(WITH_CUDA) || defined(WITH_NPU) || defined(WITH_MLU)
   }
   std::vector<const TaskProto*> source_tasks;
   source_tasks.reserve(plan.task().size());

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)

oneflow/core/job_rewriter/job_completer.cpp

@@ -153,7 +153,7 @@ Maybe<void> 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<ResourceDesc, ForSession>::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<void> 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));
 

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

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

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

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 {
 

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

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,12 +703,13 @@ size_t InferNcclLogicalFusionKernelTmpBufferSize(user_op::InferContext* ctx) {
 REGISTER_USER_KERNEL("_nccl_logical_fusion")
     .SetCreateFn<CclLogicalFusionKernel>()
     .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)
 }  // namespace
 
 }  // namespace oneflow
 
-#endif  // WITH_CUDA || WITH_NPU
+#endif  // WITH_CUDA || WITH_NPU || WITH_MLU

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

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<CclLogicalSendRecv>()
     .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

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

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