chore: import upstream snapshot with attribution

tensorflow/dtensor/mlir/ir/tf_dtensor.cc

@@ -0,0 +1,313 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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/dtensor/mlir/ir/tf_dtensor.h"
+
+#include <cassert>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h"
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+using namespace mlir;  // NOLINT
+
+namespace mlir {
+namespace TF {
+namespace {
+
+RankedTensorType GetRankedTensorType(mlir::Value val) {
+  mlir::Type type = val.getType();
+  if (auto type_with_subtype =
+          mlir::dyn_cast<mlir::TF::TensorFlowTypeWithSubtype>(
+              mlir::getElementTypeOrSelf(val))) {
+    if (type_with_subtype.GetSubtypes().size() == 1) {
+      type = type_with_subtype.GetSubtypes().front();
+    }
+  }
+  return mlir::dyn_cast_or_null<RankedTensorType>(type);
+}
+}  // namespace
+
+mlir::LogicalResult DTensorLayout::verify() {
+  DTensorLayout op = *this;
+  const auto& layout = op.getLayout();
+  if (layout.IsEmpty()) return mlir::success();
+  if (layout.IsSingleDevice()) return mlir::success();
+
+  auto input_value = op.getInput();
+
+  RankedTensorType type = GetRankedTensorType(input_value);
+
+  if (!type) return mlir::success();
+
+  const auto& num_shards = layout.num_shards();
+  if (num_shards.size() != type.getRank()) {
+    return op.emitOpError(llvm::formatv(
+        "requires matching rank for layout and input, but got {0} as suggested "
+        "rank from layout but {1} from shape.",
+        num_shards.size(), type.getRank()));
+  }
+
+  for (const auto& dim_and_index :
+       llvm::enumerate(llvm::zip(type.getShape(), num_shards))) {
+    const int dimension_index = dim_and_index.index();
+    const auto& dim_and_shards = dim_and_index.value();
+    const int dim = std::get<0>(dim_and_shards);
+    const int num_shard_for_dim = std::get<1>(dim_and_shards);
+    if (dim <= 0) continue;
+
+    if (dim % num_shard_for_dim != 0)
+      return op.emitOpError(llvm::formatv(
+          "requires dimension {0} to be divisible by sharding "
+          "specified in DTensorLayout, but got dimension size={1} is not "
+          "divisible by number of shards in layout for this dimension={2}.",
+          dimension_index, dim, num_shard_for_dim));
+  }
+
+  return mlir::success();
+}
+
+mlir::LogicalResult DTensorAllGatherOp::verify() {
+  DTensorAllGatherOp op = *this;
+  const tensorflow::dtensor::Layout input_layout = op.getInputLayout();
+  const tensorflow::dtensor::Layout output_layout = op.getOutputLayout();
+
+  if (input_layout.rank() != output_layout.rank())
+    return op.emitOpError()
+           << "received input and output layouts of unequal ranks "
+           << input_layout.rank() << " and " << output_layout.rank();
+
+  for (int32_t i = 0; i < input_layout.rank(); ++i) {
+    if (input_layout.sharding_spec(i) != output_layout.sharding_spec(i) &&
+        tensorflow::dtensor::Layout::IsShardedDimension(
+            output_layout.sharding_spec(i))) {
+      return op.emitOpError()
+             << "dimension " << i << " of output layout has sharding spec "
+             << output_layout.sharding_spec(i)
+             << " which is more sharded then the input layout spec "
+             << input_layout.sharding_spec(i);
+    }
+  }
+
+  RankedTensorType input_type =
+      mlir::dyn_cast<RankedTensorType>(op.getInput().getType());
+  if (!input_type) return mlir::success();
+
+  if (input_type.getRank() != input_layout.rank())
+    return op.emitOpError()
+           << "input layout rank " << input_layout.rank()
+           << " is not equal to input rank " << input_type.getRank();
+
+  RankedTensorType output_type =
+      mlir::dyn_cast<RankedTensorType>(op.getOutput().getType());
+  if (!output_type) return mlir::success();
+
+  if (output_type.getRank() != output_layout.rank())
+    return op.emitOpError()
+           << "output layout rank " << output_layout.rank()
+           << " is not equal to output rank " << output_type.getRank();
+
+  std::vector<int64_t> computed_output_shape =
+      output_layout.LocalShapeFromGlobalShape(
+          input_layout.GlobalShapeFromLocalShape(input_type.getShape()));
+
+  for (int32_t i = 0; i < computed_output_shape.size(); ++i) {
+    if (computed_output_shape[i] != output_type.getShape()[i]) {
+      return op.emitOpError()
+             << "computed output shape " << computed_output_shape[i]
+             << " at dimension " << i << " is not equal to actual output shape "
+             << output_type.getShape()[i];
+    }
+  }
+
+  return mlir::success();
+}
+
+mlir::LogicalResult DTensorAllScatterOp::verify() {
+  DTensorAllScatterOp op = *this;
+  const tensorflow::dtensor::Layout input_layout = op.getInputLayout();
+  const tensorflow::dtensor::Layout output_layout = op.getOutputLayout();
+
+  if (input_layout.rank() != output_layout.rank())
+    return op.emitOpError()
+           << "received input and output layouts of unequal ranks "
+           << input_layout.rank() << " and " << output_layout.rank();
+
+  for (int32_t i = 0; i < input_layout.rank(); ++i) {
+    if (input_layout.sharding_spec(i) != output_layout.sharding_spec(i) &&
+        tensorflow::dtensor::Layout::IsShardedDimension(
+            input_layout.sharding_spec(i))) {
+      return op.emitOpError()
+             << "dimension " << i << " of input layout has sharding spec "
+             << input_layout.sharding_spec(i)
+             << " which is more sharded then the output layout spec "
+             << output_layout.sharding_spec(i);
+    }
+  }
+
+  RankedTensorType input_type =
+      mlir::dyn_cast<RankedTensorType>(op.getInput().getType());
+  if (!input_type) return mlir::success();
+
+  if (input_type.getRank() != input_layout.rank())
+    return op.emitOpError()
+           << "input layout rank " << input_layout.rank()
+           << " is not equal to input rank " << input_type.getRank();
+
+  RankedTensorType output_type =
+      mlir::dyn_cast<RankedTensorType>(op.getOutput().getType());
+  if (!output_type) return mlir::success();
+
+  if (output_type.getRank() != output_layout.rank())
+    return op.emitOpError()
+           << "output layout rank " << output_layout.rank()
+           << " is not equal to output rank " << output_type.getRank();
+
+  std::vector<int64_t> computed_output_shape =
+      output_layout.LocalShapeFromGlobalShape(
+          input_layout.GlobalShapeFromLocalShape(input_type.getShape()));
+
+  for (int32_t i = 0; i < computed_output_shape.size(); ++i) {
+    if (computed_output_shape[i] != output_type.getShape()[i]) {
+      return op.emitOpError()
+             << "computed output shape " << computed_output_shape[i]
+             << " at dimension " << i << " is not equal to actual output shape "
+             << output_type.getShape()[i];
+    }
+  }
+
+  return mlir::success();
+}
+
+mlir::LogicalResult DTensorAllToAllOp::verify() {
+  DTensorAllToAllOp op = *this;
+  const tensorflow::dtensor::Layout input_layout = op.getInputLayout();
+  const tensorflow::dtensor::Layout output_layout = op.getOutputLayout();
+
+  if (input_layout.rank() != output_layout.rank())
+    return op.emitOpError()
+           << "received input and output layouts of unequal ranks "
+           << input_layout.rank() << " and " << output_layout.rank();
+
+  int32_t num_split_dims = 0;
+  int32_t num_concat_dims = 0;
+  std::string split_spec;
+  std::string concat_spec;
+  for (int32_t i = 0; i < input_layout.rank(); ++i) {
+    if (input_layout.sharding_spec(i) == output_layout.sharding_spec(i))
+      continue;
+    if (tensorflow::dtensor::Layout::IsUnshardedDimension(
+            input_layout.sharding_spec(i)) &&
+        tensorflow::dtensor::Layout::IsShardedDimension(
+            output_layout.sharding_spec(i))) {
+      num_split_dims++;
+      split_spec = output_layout.sharding_spec(i);
+    } else if (tensorflow::dtensor::Layout::IsShardedDimension(
+                   input_layout.sharding_spec(i)) &&
+               tensorflow::dtensor::Layout::IsUnshardedDimension(
+                   output_layout.sharding_spec(i))) {
+      num_concat_dims++;
+      concat_spec = input_layout.sharding_spec(i);
+    }
+  }
+  if (num_split_dims != 1 || num_concat_dims != 1 ||
+      split_spec != concat_spec) {
+    return op.emitOpError() << "must have one mesh dimension which is being "
+                               "unsharded in one axis and sharded in another";
+  }
+
+  RankedTensorType input_type =
+      mlir::dyn_cast<RankedTensorType>(op.getInput().getType());
+  if (!input_type) return mlir::success();
+
+  if (input_type.getRank() != input_layout.rank())
+    return op.emitOpError()
+           << "input layout rank " << input_layout.rank()
+           << " is not equal to input rank " << input_type.getRank();
+
+  RankedTensorType output_type =
+      mlir::dyn_cast<RankedTensorType>(op.getOutput().getType());
+  if (!output_type) return mlir::success();
+
+  if (output_type.getRank() != output_layout.rank())
+    return op.emitOpError()
+           << "output layout rank " << output_layout.rank()
+           << " is not equal to output rank " << output_type.getRank();
+
+  std::vector<int64_t> computed_output_shape =
+      output_layout.LocalShapeFromGlobalShape(
+          input_layout.GlobalShapeFromLocalShape(input_type.getShape()));
+
+  for (int32_t i = 0; i < computed_output_shape.size(); ++i) {
+    if (computed_output_shape[i] != output_type.getShape()[i]) {
+      return op.emitOpError()
+             << "computed output shape " << computed_output_shape[i]
+             << " at dimension " << i << " is not equal to actual output shape "
+             << output_type.getShape()[i];
+    }
+  }
+
+  return mlir::success();
+}
+
+LogicalResult DTensorLayout::inferReturnTypes(
+    MLIRContext* context, std::optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, mlir::PropertyRef, RegionRange regions,
+    SmallVectorImpl<Type>& inferredReturnTypes) {
+  assert(operands.size() == 1);
+  inferredReturnTypes.assign({operands[0].getType()});
+  return success();
+}
+
+void DTensorOpAdderHook(TensorFlowDialect& dialect) {
+  dialect.addOperations<
+#define GET_OP_LIST
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.cc.inc"
+      >();
+}
+
+int RegisterOnce() {
+  TF_DIALECT_REGISTER_ADDITIONAL_OPERATIONS(DTensorOpAdderHook)
+  return 0;
+}
+
+int RegisterDTensorTFOps() {
+  static int r = RegisterOnce();
+  return r;
+}
+
+}  // namespace TF
+}  // namespace mlir
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.cc.inc"

tensorflow/dtensor/mlir/ir/tf_dtensor.h

@@ -0,0 +1,54 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_H_
+#define TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_H_
+
+// Additional TF dialect operations for DTensor.
+#include "mlir/Dialect/Traits.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/DerivedAttributeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_verifiers.h"
+#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
+
+namespace mlir {
+namespace TF {
+
+// Add DTensor related operations into the TF dialect.
+int RegisterDTensorTFOps();
+
+}  // namespace TF
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h.inc"
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_H_

tensorflow/dtensor/mlir/ir/tf_dtensor.td

@@ -0,0 +1,315 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_OPS
+#define TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_OPS
+
+// Definitions for additional DTensor operations to add to the TF dialect.
+// To update
+// * Add //third_party/tensorflow/dtensor/cc:dtensor_ops to the internal dialect_generator_lib target.
+// * Run `patch_tf_dialect.sh OpA,OpB`
+// * Copy resulting output into this file.
+
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// DTensor attribute definitions
+//===----------------------------------------------------------------------===//
+
+class DTensor_DTensorAttr <string name, string description> :
+    Attr<CPred<"llvm::isa<mlir::dtensor::" # name # "Attr>($_self)">,
+         "DTensor " # description # " attribute">;
+
+def DTensor_LayoutAttr : DTensor_DTensorAttr<"Layout", "layout"> {
+  let returnType = "mlir::dtensor::LayoutAttr::Layout";
+  let convertFromStorage = "llvm::cast<mlir::dtensor::LayoutAttr>($_self).getValue()";
+}
+
+def DTensor_MeshAttr : DTensor_DTensorAttr<"Mesh", "mesh"> {
+  let returnType = "mlir::dtensor::MeshAttr::Mesh";
+  let convertFromStorage = "llvm::cast<mlir::dtensor::MeshAttr>($_self).getValue()";
+}
+
+//===----------------------------------------------------------------------===//
+// DTensor op definitions
+//===----------------------------------------------------------------------===//
+
+def Tf_DTensorSend : TF_Op<"DTensorSend", []> {
+  let summary = "Sends data to different mesh.";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    StrAttr:$key,
+    DTensor_MeshAttr:$target_mesh
+  );
+
+  let results = (outs);
+
+  TF_DerivedOperandTypeAttr Tinput = TF_DerivedOperandTypeAttr<0>;
+}
+
+def Tf_DTensorRecv : TF_Op<"DTensorRecv", []> {
+  let summary = "Receives data from different mesh.";
+
+  let arguments = (ins
+    StrAttr:$key,
+    TF_ShapeAttr:$shape,
+    DTensor_MeshAttr:$mesh
+ );
+ let results = (outs
+    TF_Tensor:$output
+ );
+
+ TF_DerivedResultTypeListAttr Toutputs = TF_DerivedResultTypeListAttr<0>;
+}
+
+def TF_DTensorLayout: TF_Op<"DTensorLayout", [DeclareOpInterfaceMethods<InferTypeOpInterface>, Pure, TF_AllTypesMatch<["input", "output"]>, TF_MustExecute, TF_NoConstantFold]> {
+  let summary = [{
+    Represents computational layout of an input tensor.
+    }];
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    DTensor_LayoutAttr:$layout,
+    TF_ShapeAttr:$global_shape
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+  let hasVerifier = 1;
+}
+
+def TF_RelayoutOp : TF_Op<"Relayout", [Pure, TF_AllTypesMatch<["input", "output"]>, TF_NoConstantFold]> {
+  let summary = "Change layout of input to target layout inside the same mesh cluster.";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+
+    StrAttr:$layout
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_RelayoutLikeOp : TF_Op<"RelayoutLike", [Pure, TF_AllTypesMatch<["input", "output"]>, TF_NoConstantFold]> {
+  let summary = "Change layout of the gradients to match the layout of the forward pass Relayout's input.";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+
+    TF_Tensor:$layout_input
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+  TF_DerivedOperandTypeAttr U = TF_DerivedOperandTypeAttr<1>;
+}
+
+def TF_CopyToMeshOp : TF_Op<"CopyToMesh", [Pure]> {
+  let summary = "";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    StrAttr:$mesh
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_CopyToMeshGradOp : TF_Op<"CopyToMeshGrad", [Pure]> {
+  let summary = "";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    TF_Tensor:$forward_input
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_DTensorAllReduceOp : TF_Op<"DTensorAllReduce", [Pure]> {
+  let summary = "";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    TF_Int32Tensor:$group_assignment,
+
+    TF_AnyStrAttrOf<["Min", "Max", "Mul", "Add", "Mean", "Any", "All"]>:$reduce_op,
+    StrAttr:$device_type
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_DTensorReduceScatterOp : TF_Op<"DTensorReduceScatter", [Pure]> {
+  let summary = "";
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    TF_Int32Tensor:$group_assignment,
+    TF_Int32Tensor:$scatter_dimension,
+
+    TF_AnyStrAttrOf<["Min", "Max", "Mul", "Add", "Mean", "Any", "All"]>:$reduce_op,
+    StrAttr:$device_type
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_DTensorShardedPrefixOp : TF_Op<"DTensorShardedPrefix", [Pure]> {
+  let summary = "Queries the generated shard prefix from DTensor SaveV2.";
+
+  let description = [{
+DTensor SPMD will generates multiple SaveV2 ops from a single SaveV2 op. Each
+generated SaveV2 op will have different prefix so that tensors can be saved
+into multiple different files. This Op would query `all` the prefix genearetd.
+Normally a MergeV2Checkpoints op would consume the output of this file so that
+the checkpoint can cover all files.
+  }];
+
+  let arguments = (ins
+    TF_StrTensor:$prefix,
+    TF_StrTensor:$tensor_names,
+    TF_StrTensor:$shape_and_slices,
+    TF_StrTensor:$mesh,
+    TF_StrTensor:$layouts,
+    Variadic<TF_Tensor>:$tensors
+  );
+
+  let results = (outs
+    TF_StrTensor:$generated_prefixes
+  );
+}
+
+def TF_DTensorRestoreV2Op : TF_Op<"DTensorRestoreV2", [Pure]> {
+  let summary = "DTensor RestoreV2 that takes extra shape and layouts.";
+
+  let description = [{
+DTensor RestoreV2 is needed so that we can restore tensors with layouts.
+Ideally this can be done in tf.function style restore but it is currently
+blocked by CopyToMesh. Note that this will be also used by name-based restore.
+  }];
+
+  let arguments = (ins
+    TF_StrTensor:$prefix,
+    TF_StrTensor:$tensor_names,
+    TF_StrTensor:$shape_and_slices,
+    TF_ShapeAttrArray:$input_shapes,
+    StrArrayAttr:$input_layouts
+  );
+
+  let results = (outs
+    Variadic<TF_Tensor>:$tensors
+  );
+
+  TF_DerivedResultTypeListAttr dtypes = TF_DerivedResultTypeListAttr<0>;
+}
+def TF_DTensorAllScatterOp : TF_Op<"DTensorAllScatter", [Pure]> {
+  let summary = "Slices the input to the given layout.";
+
+  let description = [{
+This op takes both an input and an output layout. The output layout must be more
+sharded than the input layout.
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    DTensor_LayoutAttr:$input_layout,
+    DTensor_LayoutAttr:$output_layout
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+    );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+
+  let hasVerifier = 1;
+}
+
+def TF_DTensorAllGatherOp : TF_Op<"DTensorAllGather", [Pure]> {
+  let summary = "Concatenates the input to match the given layout.";
+
+  let description = [{
+This op takes both an input and an output layout. The output layout must be less
+sharded than the input layout.
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    DTensor_LayoutAttr:$input_layout,
+    DTensor_LayoutAttr:$output_layout
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+    );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+
+  let hasVerifier = 1;
+}
+
+def TF_DTensorAllToAllOp : TF_Op<"DTensorAllToAll", [Pure]> {
+  let summary = "Mutually exchanges the input to match the given layout.";
+
+  let description = [{"
+This op takes both an input and an output layout. There can be one mesh
+dimension which is becoming unsharded in one axis while becoming sharded 
+in another axis."
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$input,
+    DTensor_LayoutAttr:$input_layout,
+    DTensor_LayoutAttr:$output_layout
+  );
+
+  let results = (outs
+    TF_Tensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+
+  let hasVerifier = 1;
+}
+
+#endif // TENSORFLOW_DTENSOR_MLIR_IR_TF_DTENSOR_OPS