colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[XLA:GPU] Migrate AllReduce thunk emission to MLIR. 

PiperOrigin-RevId: 354335704
This commit is contained in:
Rahul Joshi 2021-01-28 10:04:59 -08:00 committed by TensorFlow MLIR Team
parent 30ce82790d
commit c97cff9a7b
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td
View File

@ -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]>,