Add check message for unsupport allreduce op (#204)

* Add check message for unsupport allreduce op

Author: Chao1Han

src/dispatch_stub.cpp

@@ -554,6 +554,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::allreduce(std::v
                                                                        ProcessGroupCCL& pg_ccl) {
   checkSameType(tensors[0], tensors);
   c10::DeviceType dev_type = tensors[0].device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->allreduce_(tensors, opts, pg_ccl);
 }
 
@@ -562,6 +563,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::allreduce_coales
                                                                        ProcessGroupCCL& pg_ccl) {
   checkSameType(tensors[0], tensors);
   c10::DeviceType dev_type = tensors[0].device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->allreduce_coalesced_(tensors, opts, pg_ccl);
 }
 
@@ -570,6 +572,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::reduce(std::vect
                                                              ProcessGroupCCL& pg_ccl) {
   checkSameType(tensors[0], tensors);
   c10::DeviceType dev_type = tensors[0].device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->reduce_(tensors, opts, pg_ccl);
 }
 
@@ -638,6 +641,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::reduce_scatter(s
                                                                 const ReduceScatterOptions& opts,
                                                                 ProcessGroupCCL& pg_ccl) {
   c10::DeviceType dev_type = outputTensors[0].device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->reduce_scatter_(outputTensors, inputTensors, opts, pg_ccl);
 }
 
@@ -646,6 +650,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::_reduce_scatter_
                                                                 const ReduceScatterOptions& opts,
                                                                 ProcessGroupCCL& pg_ccl) {
   c10::DeviceType dev_type = inputTensor.device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->_reduce_scatter_base_(outputTensor, inputTensor, opts, pg_ccl);
 }
 
@@ -657,6 +662,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> DispatchStub::reduce_scatter_t
   checkSameType(inputTensors[0], inputTensors);
   checkSameType(outputTensors[0], outputTensors);
   c10::DeviceType dev_type = inputTensors[0].device().type();
+  check_supported_reduce_op(dev_type, opts.reduceOp);
   return get_ccl_stub(dev_type)->reduce_scatter_tensor_coalesced_(outputTensors, inputTensors, opts, pg_ccl);
 }
 

src/dispatch_stub.h

@@ -290,4 +290,48 @@ class DispatchStub {
   }
 };
 
+} // namespace oneccl_bindings_for_pytorch
+
+namespace {
+
+std::string reduce_op_to_string(c10d::ReduceOp op) {
+  switch (op) {
+    case c10d::ReduceOp::SUM:
+      return "SUM";
+    case c10d::ReduceOp::PRODUCT:
+      return "PRODUCT";
+    case c10d::ReduceOp::MIN:
+      return "MIN";
+    case c10d::ReduceOp::MAX:
+      return "MAX";
+    case c10d::ReduceOp::BAND:
+      return "BAND";
+    case c10d::ReduceOp::BOR:
+      return "BOR";
+    case c10d::ReduceOp::BXOR:
+      return "BXOR";
+    case c10d::ReduceOp::AVG:
+      return "AVG";
+    default:
+      return "UNKNOWN";
+  }
 }
+
+void check_supported_reduce_op(c10::DeviceType dev_type, c10d::ReduceOp op) {
+  if (dev_type == c10::DeviceType::XPU) {
+    switch (op) {
+      case c10d::ReduceOp::BAND:
+      case c10d::ReduceOp::BOR:
+      case c10d::ReduceOp::BXOR:
+      case c10d::ReduceOp::AVG:
+      case c10d::ReduceOp::PREMUL_SUM:
+      case c10d::ReduceOp::UNUSED:
+        TORCH_CHECK(false, ("Cannot use ReduceOp." + reduce_op_to_string(op) + " with XPU"));
+      default:
+        // No action needed for supported operations
+        break;
+    }
+  }
+}
+
+} // namespace