From 1be1123c70bc0c64d627b3ff35811bc6726ec672 Mon Sep 17 00:00:00 2001
From: Rahul Joshi <jurahul@google.com>
Date: Fri, 29 Jan 2021 13:30:59 -0800
Subject: [PATCH] [XLA:GPU] Add support for PartitionId

PiperOrigin-RevId: 354599221
---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td | 8 ++++++++
 include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td     | 4 ++++
 2 files changed, 12 insertions(+)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
index 57fdfb6..a948c32 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
@@ -637,6 +637,14 @@ class BASE_HLO_ReplicaIdOp {
   }];
 }
 
+class BASE_HLO_PartitionIdOp {
+  string summary = "PartitionId operator";
+
+  string description = [{
+    Returns the unique ID (int32 scalar) of the partition.
+  }];
+}
+
 
 class BASE_HLO_AllReduceOp {
   string summary = "AllReduce operator";
diff --git a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td
index ee39bfc..9706473 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td
@@ -608,6 +608,10 @@ def LHLO_ReplicaIdOp : LHLO_Op<"replica_id", []>, BASE_HLO_ReplicaIdOp {
   let arguments = (ins Arg<MemRefOf<[UI32]>, "", [MemWrite]>);
 }
 
+def LHLO_PartitionIdOp : LHLO_Op<"partition_id", []>, BASE_HLO_PartitionIdOp {
+  let arguments = (ins Arg<MemRefOf<[UI32]>, "", [MemWrite]>);
+}
+
 def LHLO_TriangularSolveOp: LHLO_Op<"triangular_solve", [SameOperandsElementType]>,
                             BASE_HLO_TriangularSolveOp {
   let arguments = (ins