limitations under the License.
==============================================================================*/

include "mlir/Pass/PassBase.td"

def EmbeddingPipeliningPass : Pass<"tf-embedding-pipelining", "mlir::ModuleOp"> {
  let summary = "Rewrite graph for embedding pipelining";
  let constructor = "TFDevice::CreateEmbeddingPipeliningPass()";
    let description = [{

  pm.addNestedPass<FuncOp>(mlir::TF::CreateDecomposeReduceDatasetPass());
  // Only one of EmbeddingSequencing and EmbeddingPipelining will actually
  // run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
  // runs, embedding attributes are stripped and the sequencing pass will have
  // no effect. If the pipelining pass doesn't run, embedding attributes are

    if (mlir::isa<TF::ResourceType>(tensor_type.getElementType())) {
      return true;
    }
  }
  return false;
}

struct EmbeddingPipeliningPass
    : public ::impl::EmbeddingPipeliningPassBase<EmbeddingPipeliningPass> {
  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;
};