Support bitwise allreduce operations in the communicator (#8623)

This commit is contained in:
Rong Ou
2022-12-24 14:40:05 -08:00
committed by GitHub
parent c7e82b5914
commit 77b069c25d
10 changed files with 207 additions and 26 deletions

View File

@@ -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");
}
}