Support bitwise allreduce operations in the communicator (#8623)
This commit is contained in:
@@ -23,6 +23,17 @@ struct FHelper {
|
||||
}
|
||||
};
|
||||
|
||||
template<typename DType>
|
||||
struct FHelper<op::BitAND, DType> {
|
||||
static void
|
||||
Allreduce(DType *,
|
||||
size_t ,
|
||||
void (*)(void *arg),
|
||||
void *) {
|
||||
utils::Error("DataType does not support bitwise AND operation");
|
||||
}
|
||||
};
|
||||
|
||||
template<typename DType>
|
||||
struct FHelper<op::BitOR, DType> {
|
||||
static void
|
||||
@@ -30,7 +41,18 @@ struct FHelper<op::BitOR, DType> {
|
||||
size_t ,
|
||||
void (*)(void *arg),
|
||||
void *) {
|
||||
utils::Error("DataType does not support bitwise or operation");
|
||||
utils::Error("DataType does not support bitwise OR operation");
|
||||
}
|
||||
};
|
||||
|
||||
template<typename DType>
|
||||
struct FHelper<op::BitXOR, DType> {
|
||||
static void
|
||||
Allreduce(DType *,
|
||||
size_t ,
|
||||
void (*)(void *arg),
|
||||
void *) {
|
||||
utils::Error("DataType does not support bitwise XOR operation");
|
||||
}
|
||||
};
|
||||
|
||||
@@ -111,12 +133,24 @@ void Allreduce(void *sendrecvbuf,
|
||||
count, enum_dtype,
|
||||
prepare_fun, prepare_arg);
|
||||
return;
|
||||
case kBitwiseAND:
|
||||
Allreduce<op::BitAND>
|
||||
(sendrecvbuf,
|
||||
count, enum_dtype,
|
||||
prepare_fun, prepare_arg);
|
||||
return;
|
||||
case kBitwiseOR:
|
||||
Allreduce<op::BitOR>
|
||||
(sendrecvbuf,
|
||||
count, enum_dtype,
|
||||
prepare_fun, prepare_arg);
|
||||
return;
|
||||
case kBitwiseXOR:
|
||||
Allreduce<op::BitXOR>
|
||||
(sendrecvbuf,
|
||||
count, enum_dtype,
|
||||
prepare_fun, prepare_arg);
|
||||
return;
|
||||
default: utils::Error("unknown enum_op");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user