chore: import upstream snapshot with attribution

tensorflow/dtensor/mlir/sparse_expansions/dynamic_enqueue_sparse_expander.cc

@@ -0,0 +1,116 @@
+/* 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/sparse_expansions/dynamic_enqueue_sparse_expander.h"
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/mlir/sparse_expander_common.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+// Indices tensor should be transformed from a shape [?, 2] tensor to a
+// [?, 3] tensor padded with 0's because
+// EnqueueTPUEmbeddingArbitraryTensorBatchOp expects a [?, 3] indices tensor.
+StatusOr<mlir::Value> ExpandIndices(mlir::OpBuilder& builder,
+                                    mlir::Value indices) {
+  int64_t num_dim =
+      mlir::dyn_cast<mlir::RankedTensorType>(indices.getType()).getDimSize(1);
+  if (num_dim != 2)
+    return absl::UnimplementedError(
+        "Sparse tensors with dense rank not equal to 2 is not yet supported in "
+        "DTensor.");
+  mlir::Location loc = indices.getLoc();
+  auto indices_padded_type = mlir::RankedTensorType::get(
+      {mlir::ShapedType::kDynamic, 3},
+      mlir::dyn_cast<mlir::RankedTensorType>(indices.getType())
+          .getElementType());
+  // Little trick to make a rank-2 tensor of [[0,0], [0,1]] using rank 1
+  // constants.
+  mlir::Value indices_padding = mlir::TF::ReshapeOp::create(
+      builder, loc,
+      mlir::TF::collection_ops_util::GetR1Const({0, 0, 0, 1}, builder, loc),
+      mlir::TF::collection_ops_util::GetR1Const({2, 2}, builder, loc));
+  mlir::Value indices_padded =
+      mlir::TF::PadOp::create(builder, loc, indices_padded_type,
+                              /*input=*/indices,
+                              /*paddings=*/indices_padding);
+  return indices_padded;
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> DynamicEnqueueSparseExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::TF::DynamicEnqueueTPUEmbeddingArbitraryTensorBatchOp dense_enqueue_op =
+      mlir::cast<mlir::TF::DynamicEnqueueTPUEmbeddingArbitraryTensorBatchOp>(
+          op);
+
+  mlir::OpBuilder builder(dense_enqueue_op);
+  mlir::Location location = dense_enqueue_op->getLoc();
+
+  mlir::OperandRange feature = dense_enqueue_op.getEmbeddingIndices();
+  llvm::SmallVector<mlir::Value, 4> indices;
+  llvm::SmallVector<mlir::Value, 4> values;
+
+  for (mlir::Value sparse_feature_value : feature) {
+    if (!IsSparseValue(sparse_feature_value)) {
+      return absl::InternalError(
+          "Expected feature input to DynamicEnqueueOp to be a sparse input, "
+          "but was not. This should not happen.");
+    }
+    // Indices tensor may need to be expanded to a different shape
+    // for Enqueue op to work properly.
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value expanded_indices,
+        ExpandIndices(
+            builder, GetIndicesFromSparseTensor(sparse_feature_value).value()));
+    indices.push_back(expanded_indices);
+    values.push_back(GetValuesFromSparseTensor(sparse_feature_value).value());
+  }
+  // Insert a new op with new sparse operands, and delete the old one.
+  // This op does not have a return value so we do not need to replace any
+  // consumers.
+  mlir::Operation* sparse_enqueue_op =
+      mlir::TF::DynamicEnqueueTPUEmbeddingArbitraryTensorBatchOp::create(
+          builder, location,
+          /*sample_indices_or_row_splits=*/indices,
+          /*embedding_indices=*/values,
+          /*aggregation_weights=*/dense_enqueue_op.getAggregationWeights(),
+          /*mode_override=*/
+          dense_enqueue_op.getModeOverride(),
+          /*device_ordinal=*/dense_enqueue_op.getDeviceOrdinal(),
+          /*combiners=*/dense_enqueue_op.getCombiners());
+  dense_enqueue_op.erase();
+  return sparse_enqueue_op;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/sparse_expansions/dynamic_enqueue_sparse_expander.h

@@ -0,0 +1,34 @@
+/* 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_SPARSE_EXPANSIONS_DYNAMIC_ENQUEUE_SPARSE_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANSIONS_DYNAMIC_ENQUEUE_SPARSE_EXPANDER_H_
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/mlir/sparse_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class DynamicEnqueueSparseExpander : public SparseExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANSIONS_DYNAMIC_ENQUEUE_SPARSE_EXPANDER_H_

tensorflow/dtensor/mlir/sparse_expansions/matmul_sparse_expander.cc

@@ -0,0 +1,59 @@
+/* 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/sparse_expansions/matmul_sparse_expander.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/mlir/sparse_expander_common.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> MatMulSparseExpander::ExpandOp(mlir::Operation* op) {
+  mlir::TF::MatMulOp mm = mlir::cast<mlir::TF::MatMulOp>(op);
+  // If any of the transpose attributes are true, then return original op.
+  if (mm.getTransposeA() || mm.getTransposeB()) return op;
+
+  // Expand to SparseTensorDenseMatMul Op only if the left operand
+  // is a SparseTensor.
+  if (IsSparseValue(op->getOperand(0)) && !IsSparseValue(op->getOperand(1))) {
+    mlir::OpBuilder builder(op);
+    // Since operand 0 is a SparseValue, we don't need to check that
+    // the indices, values, and dense_shapes exist.
+    mlir::TF::SparseTensorDenseMatMulOp new_op =
+        mlir::TF::SparseTensorDenseMatMulOp::create(
+            builder, op->getLoc(), op->getResultTypes(),
+            mlir::ValueRange{
+                GetIndicesFromSparseTensor(op->getOperand(0)).value(),
+                GetValuesFromSparseTensor(op->getOperand(0)).value(),
+                GetDenseShapesFromSparseTensor(op->getOperand(0)).value(),
+                op->getOperand(1)});
+
+    op->getResult(0).replaceAllUsesWith(new_op.getResult());
+    op->erase();
+    return new_op.getOperation();
+  }
+
+  // Any other case, return the original op.
+  return op;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/sparse_expansions/matmul_sparse_expander.h

@@ -0,0 +1,34 @@
+/* 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_SPARSE_EXPANSIONS_MATMUL_SPARSE_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANSIONS_MATMUL_SPARSE_EXPANDER_H_
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/mlir/sparse_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class MatMulSparseExpander : public SparseExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANSIONS_MATMUL_SPARSE_EXPANDER_H_