From c97cff9a7b0788990216bbb3b462094e20947e4f Mon Sep 17 00:00:00 2001 From: Rahul Joshi Date: Thu, 28 Jan 2021 10:04:59 -0800 Subject: [PATCH] [XLA:GPU] Migrate AllReduce thunk emission to MLIR. PiperOrigin-RevId: 354335704 --- include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td index 7b8dd94..ee39bfc 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td +++ b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td @@ -556,6 +556,10 @@ def LHLO_AllReduceOp : LHLO_Op<"all_reduce", [SameVariadicOperandSize]>, ); let regions = (region SizedRegion<1>:$computation); let verifier = [{ return Verify(*this); }]; + let extraClassDeclaration = [{ + // AllReduce is cross replica if channel_id is not set. + bool IsCrossReplica() { return !channel_id().hasValue(); } + }]; } def LHLO_CollectivePermuteOp: LHLO_Op<"collective_permute", [SameTypeOperands]>,