Add _Raw.replicaId() op. (#1018)

Author: pschuh

Sources/CX10/xla_tensor_wrapper.cc

@@ -572,6 +572,9 @@ OpaqueXLATensor* XLATensor_repeat(OpaqueXLATensor* input,
   return new XLATensor(
       XLATensor::repeat(*input, XlaHelpers::I64List(repeats.slice())));
 }
+OpaqueXLATensor* XLATensor_replica_id(const struct CDevice device) {
+  return new XLATensor(XLATensor::xla_replica_id(ConvertDevice(device)));
+}
 OpaqueXLATensor* XLATensor_resize_value(OpaqueXLATensor* a, Int64ArrayRef arr) {
   return new XLATensor(
       XLATensor::resize_value(*a, XlaHelpers::I64List(arr.slice())));

Sources/CX10/xla_tensor_wrapper.h

@@ -341,6 +341,7 @@ XLA_API OpaqueXLATensor* XLATensor_relu(OpaqueXLATensor* a);
 XLA_API OpaqueXLATensor* XLATensor_rem(OpaqueXLATensor* a, OpaqueXLATensor* b);
 XLA_API OpaqueXLATensor* XLATensor_repeat(OpaqueXLATensor* input,
                                           Int64ArrayRef repeats);
+XLA_API OpaqueXLATensor* XLATensor_replica_id(const struct CDevice device);
 XLA_API OpaqueXLATensor*
 XLATensor_resize_value(OpaqueXLATensor* a, Int64ArrayRef arr);
 XLA_API OpaqueXLATensor* XLATensor_round_to_even(OpaqueXLATensor* a);

Sources/x10/swift_bindings/XLATensor.swift

@@ -743,6 +743,10 @@ extension XLATensor {
     return XLATensor(_handle: XLATensor_relu(a.handle))
   }
 
+  static func replica_id(_ device: Device) -> XLATensor {
+    return XLATensor(_handle: XLATensor_replica_id(device.cdevice));
+  }
+
   static func resize_value(_ value: XLATensor, _ dims: [Int64]) -> XLATensor {
     defer { _fixLifetime(value) }
     return dims.withArrayRef { dims in

Sources/x10/swift_bindings/apis/RawOpsManual.swift

@@ -2929,6 +2929,10 @@ public enum _RawXLA {
     return Tensor(_xla: XLATensor.threshold_backward(gradients.xlaTensor, features.xlaTensor, 0))
   }
 
+  public static func replicaId(_ device: Device) -> Tensor<Int32> {
+    return Tensor(_xla: XLATensor.replica_id(device))
+  }
+
   /// Reshapes a tensor.
   ///
   /// Given `tensor`, this operation returns a tensor that has the same values

Sources/x10/xla_tensor/aten_compat.h

@@ -757,6 +757,7 @@
   _(aten, xla_max_pool_grad)                                \
   _(aten, xla_pad)                                          \
   _(aten, xla_rem)                                          \
+  _(aten, xla_replica_id)                                   \
   _(aten, xla_slice)                                        \
   _(aten, xla_truncated_normal)                             \
   _(aten, xla_is_finite)                                    \

Sources/x10/xla_tensor/ops/replica_id.cpp

@@ -0,0 +1,46 @@
+// Copyright 2020 TensorFlow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tensorflow/compiler/tf2xla/xla_tensor/ops/replica_id.h"
+
+#include "tensorflow/compiler/tf2xla/xla_tensor/lowering_context.h"
+
+namespace swift_xla {
+namespace ir {
+namespace ops {
+
+ReplicaId::ReplicaId() : Node(
+    ir::OpKind(at::aten::xla_replica_id),
+    {},
+    xla::ShapeUtil::MakeShape(xla::S32, {}),
+    1, 0x18923728) {}
+
+NodePtr ReplicaId::Clone(OpList operands) const {
+  return MakeNode<ReplicaId>();
+}
+
+XlaOpVector ReplicaId::Lower(LoweringContext* loctx) const {
+  return ReturnOp(xla::ConvertElementType(
+      xla::ReplicaId(loctx->builder()), xla::S32), loctx);
+}
+
+std::string ReplicaId::ToString() const {
+  std::stringstream ss;
+  ss << Node::ToString();
+  return ss.str();
+}
+
+}  // namespace ops
+}  // namespace ir
+}  // namespace swift_xla

Sources/x10/xla_tensor/ops/replica_id.h

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2020 TensorFlow Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorflow/compiler/tf2xla/xla_tensor/ir.h"
+
+namespace swift_xla {
+namespace ir {
+namespace ops {
+
+class ReplicaId : public Node {
+ public:
+  ReplicaId();
+
+  std::string ToString() const override;
+
+  NodePtr Clone(OpList operands) const override;
+
+  XlaOpVector Lower(LoweringContext* loctx) const override;
+};
+
+}  // namespace ops
+}  // namespace ir
+}  // namespace swift_xla

Sources/x10/xla_tensor/tensor.h

@@ -1264,6 +1264,8 @@ class XLATensor {
 
   static XLATensor xla_truncated_normal(const XLATensor& input);
 
+  static XLATensor xla_replica_id(const Device& device);
+
  private:
   struct SyncTensorsConfig {
     // Whether we want to force XLA data on the target tensors (hence trimming

Sources/x10/xla_tensor/tensor_methods.cpp

@@ -95,6 +95,7 @@
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/reflection_pad2d.h"
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/reflection_pad2d_backward.h"
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/repeat.h"
+#include "tensorflow/compiler/tf2xla/xla_tensor/ops/replica_id.h"
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/replication_pad.h"
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/replication_pad_backward.h"
 #include "tensorflow/compiler/tf2xla/xla_tensor/ops/resize.h"
@@ -3037,4 +3038,8 @@ XLATensor XLATensor::xla_truncated_normal(const XLATensor& input) {
   return input.CreateFrom(ir::ops::XlaTruncatedNormal(input.GetIrValue()));
 }
 
+XLATensor XLATensor::xla_replica_id(const Device& device) {
+  return XLATensor::Create(ir::MakeNode<ir::ops::ReplicaId>(), device);
+}
+
 }  // namespace swift_xla