chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/python/mlir_wrapper/BUILD

@@ -0,0 +1,54 @@
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_python_pybind_extension(
+    name = "mlir_wrapper",
+    srcs = [
+        "attrs.cc",
+        "basic_classes.cc",
+        "builders.cc",
+        "mlir_wrapper.cc",
+        "mlir_wrapper.h",
+        "ops.cc",
+        "types.cc",
+    ],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "mlir_wrapper.pyi",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@llvm-project//llvm:FileCheckLib",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@pybind11",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "filecheck_wrapper",
+    srcs = ["filecheck_wrapper.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "filecheck_wrapper.pyi",
+    ],
+    starlark_only = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@llvm-project//llvm:FileCheckLib",
+        "@llvm-project//llvm:Support",
+        "@pybind11",
+    ],
+)

tensorflow/compiler/mlir/python/mlir_wrapper/attrs.cc

@@ -0,0 +1,26 @@
+/* Copyright 2020 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 <cstdint>
+
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_attrs(py::module& m) {
+  py::class_<mlir::IntegerAttr, mlir::Attribute>(m, "IntegerAttr")
+      .def("get",
+           py::overload_cast<mlir::Type, int64_t>(&mlir::IntegerAttr::get));
+}

tensorflow/compiler/mlir/python/mlir_wrapper/basic_classes.cc

@@ -0,0 +1,53 @@
+/* Copyright 2020 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 "llvm/FileCheck/FileCheck.h"
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_basic_classes(py::module& m) {
+  py::class_<mlir::MLIRContext>(m, "MLIRContext").def(py::init<>());
+
+  py::class_<mlir::Location>(m, "Location");
+
+  py::class_<mlir::UnknownLoc>(m, "UnknownLoc")
+      .def("get", [](mlir::MLIRContext* context) -> mlir::Location {
+        return mlir::UnknownLoc::get(context);
+      });
+
+  py::class_<mlir::Region>(m, "Region")
+      .def("back", &mlir::Region::back, py::return_value_policy::reference)
+      .def("front", &mlir::Region::front, py::return_value_policy::reference)
+      .def("add_block", [](mlir::Region& r) { r.push_back(new mlir::Block); })
+      .def("push_back", &mlir::Region::push_back)
+      .def("size", [](mlir::Region& r) { return r.getBlocks().size(); })
+      .def("front", &mlir::Region::front, py::return_value_policy::reference);
+  py::class_<mlir::Block::iterator>(m, "Block_Iterator");
+  py::class_<mlir::Block>(m, "Block")
+      .def("new", ([]() { return new mlir::Block; }),
+           py::return_value_policy::reference)
+      .def("end", &mlir::Block::end)
+      .def("addArgument", [](mlir::Block& block, mlir::Type type) {
+        return block.addArgument(type, block.getParent()->getLoc());
+      });
+
+  py::class_<mlir::Value>(m, "Value").def("getType", &mlir::Value::getType);
+  py::class_<mlir::OpResult, mlir::Value>(m, "OpResult");
+  py::class_<mlir::BlockArgument, mlir::Value>(m, "BlockArgument");
+}

tensorflow/compiler/mlir/python/mlir_wrapper/builders.cc

@@ -0,0 +1,55 @@
+/* Copyright 2020 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 "mlir/IR/Builders.h"  // from @llvm-project
+
+#include <vector>
+
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_builders(py::module& m) {
+  py::class_<mlir::Builder>(m, "Builder")
+      .def(py::init<mlir::MLIRContext*>())
+      .def("getFunctionType",
+           [](mlir::Builder& b, std::vector<mlir::Type> inputs,
+              std::vector<mlir::Type> outputs) {
+             return b.getFunctionType(llvm::ArrayRef<mlir::Type>(inputs),
+                                      llvm::ArrayRef<mlir::Type>(outputs));
+           });
+  py::class_<mlir::OpBuilder>(m, "OpBuilder")
+      .def(py::init<mlir::MLIRContext*>())
+      .def(py::init<mlir::Region&>())
+      .def(py::init<mlir::Operation*>())
+      .def(py::init<mlir::Block*, mlir::Block::iterator>())
+      .def("getUnknownLoc", &mlir::OpBuilder::getUnknownLoc)
+      .def("setInsertionPoint",
+           py::overload_cast<mlir::Block*, mlir::Block::iterator>(
+               &mlir::OpBuilder::setInsertionPoint))
+      .def("saveInsertionPoint", &mlir::OpBuilder::saveInsertionPoint)
+      .def("restoreInsertionPoint", &mlir::OpBuilder::restoreInsertionPoint)
+      .def(
+          "create",
+          [](mlir::OpBuilder& opb, mlir::OperationState& state) {
+            return opb.create(state);
+          },
+          py::return_value_policy::reference)
+      .def("getContext", &mlir::OpBuilder::getContext,
+           py::return_value_policy::reference);
+
+  py::class_<mlir::OpBuilder::InsertPoint>(m, "OpBuilder_InsertionPoint")
+      .def("getBlock", &mlir::OpBuilder::InsertPoint::getBlock);
+}

tensorflow/compiler/mlir/python/mlir_wrapper/filecheck_wrapper.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 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 <string>
+
+#include "llvm/FileCheck/FileCheck.h"
+#include "llvm/Support/SourceMgr.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+PYBIND11_MODULE(filecheck_wrapper, m) {
+  m.def("check", [](std::string input, std::string check) {
+    llvm::FileCheckRequest fcr;
+    llvm::FileCheck fc(fcr);
+    llvm::SourceMgr SM = llvm::SourceMgr();
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
+                          llvm::SMLoc());
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(check),
+                          llvm::SMLoc());
+    fc.readCheckFile(SM, llvm::StringRef(check));
+    return fc.checkInput(SM, llvm::StringRef(input));
+  });
+}

tensorflow/compiler/mlir/python/mlir_wrapper/filecheck_wrapper.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+def check(arg0: str, arg1: str) -> bool: ...

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.cc

@@ -0,0 +1,65 @@
+/* Copyright 2020 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/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+#include <string>
+
+#include "llvm/Support/SourceMgr.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+PYBIND11_MODULE(mlir_wrapper, m) {
+  m.def("preloadTensorFlowDialects", [](mlir::MLIRContext &context) {
+    mlir::DialectRegistry registry;
+    mlir::RegisterAllTensorFlowDialects(registry);
+    context.appendDialectRegistry(registry);
+    context.loadAllAvailableDialects();
+  });
+
+  m.def("verify", [](std::string input) {
+    llvm::SourceMgr SM = llvm::SourceMgr();
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
+                          llvm::SMLoc());
+    mlir::DialectRegistry registry;
+    mlir::RegisterAllTensorFlowDialects(registry);
+    mlir::MLIRContext ctx(registry);
+    ctx.loadAllAvailableDialects();
+    auto module = mlir::parseSourceFile<mlir::ModuleOp>(SM, &ctx);
+    if (!module) {
+      return false;
+    }
+    if (failed(mlir::verify(*module))) {
+      module->emitError("Invalid MLIR module: failed verification.");
+      return false;
+    }
+    return true;
+  });
+
+  init_basic_classes(m);
+  init_types(m);
+  init_builders(m);
+  init_ops(m);
+  init_attrs(m);
+}

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h

@@ -0,0 +1,30 @@
+/* Copyright 2020 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_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_
+#define TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+
+namespace py = pybind11;
+
+void init_basic_classes(py::module& m);
+void init_types(py::module& m);
+void init_builders(py::module& m);
+void init_ops(py::module& m);
+void init_attrs(py::module& m);
+
+#endif  // TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.pyi

@@ -0,0 +1,193 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+from typing import overload
+
+class Attribute:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Block:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def addArgument(self, *args, **kwargs): ...
+    def end(self) -> Block_Iterator: ...
+    def new(self) -> Block: ...
+
+class BlockArgument(Value):
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Block_Iterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Builder:
+    def __init__(self, arg0: MLIRContext) -> None: ...
+    def getFunctionType(self, arg0: list[Type], arg1: list[Type]) -> FunctionType: ...
+
+class FloatType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getBF16(self) -> FloatType: ...
+    def getF16(self) -> FloatType: ...
+    def getF32(self) -> FloatType: ...
+    def getF64(self) -> FloatType: ...
+
+class FuncOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: str, arg1: FunctionType) -> FuncOp: ...
+    def getArguments(self) -> list[BlockArgument]: ...
+    def getBody(self) -> Region: ...
+    def getName(self) -> str: ...
+    def getType(self) -> FunctionType: ...
+
+class FunctionType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getResults(self) -> list[Type]: ...
+
+class IntegerAttr(Attribute):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: int) -> IntegerAttr: ...
+
+class IntegerType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: int) -> IntegerType: ...
+
+class Location:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class MLIRContext:
+    def __init__(self) -> None: ...
+
+class ModuleOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self) -> ModuleOp: ...
+    def dump(self) -> None: ...
+    def getAsStr(self) -> str: ...
+    def push_back(self, arg0) -> None: ...
+
+class OpBuilder:
+    @overload
+    def __init__(self, arg0: MLIRContext) -> None: ...
+    @overload
+    def __init__(self, arg0: Region) -> None: ...
+    @overload
+    def __init__(self, arg0) -> None: ...
+    @overload
+    def __init__(self, arg0: Block, arg1: Block_Iterator) -> None: ...
+    def create(self, *args, **kwargs): ...
+    def getContext(self) -> MLIRContext: ...
+    def getUnknownLoc(self) -> Location: ...
+    def restoreInsertionPoint(self, arg0) -> None: ...
+    def saveInsertionPoint(self, *args, **kwargs): ...
+    def setInsertionPoint(self, arg0: Block, arg1: Block_Iterator) -> None: ...
+
+class OpBuilder_InsertionPoint:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getBlock(self) -> Block: ...
+
+class OpResult(Value):
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Operation:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def dump(self) -> None: ...
+    def getNumResults(self) -> int: ...
+    def getRegion(self, arg0: int) -> Region: ...
+    def getResult(self, arg0: int) -> OpResult: ...
+
+class OperationState:
+    def __init__(self, arg0: Location, arg1: str) -> None: ...
+    def addOperands(self, arg0: list[Value]) -> None: ...
+    def addRegion(self) -> Region: ...
+    def addTypes(self, arg0: list[Type]) -> None: ...
+
+class RankedTensorType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: Type) -> RankedTensorType: ...
+
+class Region:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def add_block(self) -> None: ...
+    def back(self, *args, **kwargs): ...
+    def front(self, *args, **kwargs): ...
+    def push_back(self, arg0) -> None: ...
+    def size(self) -> int: ...
+
+class ReturnOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: list[Value]) -> Operation: ...
+
+class Tf_AddV2Op:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_AnyOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value, arg3: bool) -> Operation: ...
+
+class Tf_ConstOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1) -> Operation: ...
+
+class Tf_EqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_GreaterEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_GreaterOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_LegacyCallOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: list[Type], arg2: list[Value], arg3: str) -> Operation: ...
+
+class Tf_LessEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_LessOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_NegOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value) -> Operation: ...
+
+class Tf_NotEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_SubOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Type:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class UnknownLoc:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self) -> Location: ...
+
+class UnrankedTensorType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self) -> UnrankedTensorType: ...
+
+class Value:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getType(self, *args, **kwargs): ...
+
+def preloadTensorFlowDialects(arg0) -> None: ...
+def verify(arg0: str) -> bool: ...

tensorflow/compiler/mlir/python/mlir_wrapper/ops.cc

@@ -0,0 +1,196 @@
+/* Copyright 2020 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 <memory>
+#include <string>
+#include <vector>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+void init_ops(py::module& m) {
+  py::class_<mlir::Operation, std::unique_ptr<mlir::Operation, py::nodelete>>(
+      m, "Operation")
+      .def("getRegion", &mlir::Operation::getRegion,
+           py::return_value_policy::reference)
+      .def("getResult", &mlir::Operation::getResult)
+      .def("dump", &mlir::Operation::dump)
+      .def("getNumResults", &mlir::Operation::getNumResults);
+
+  py::class_<mlir::OperationState>(m, "OperationState")
+      .def(py::init([](mlir::Location loc, std::string name) {
+        return mlir::OperationState(loc, llvm::StringRef(name));
+      }))
+      .def("addTypes",
+           [](mlir::OperationState& state, std::vector<mlir::Type> tys) {
+             state.addTypes(mlir::ArrayRef<mlir::Type>(tys));
+           })
+      .def("addOperands",
+           [](mlir::OperationState& os, std::vector<mlir::Value> ops) {
+             os.addOperands(mlir::ArrayRef<mlir::Value>(ops));
+           })
+      .def("addRegion", py::overload_cast<>(&mlir::OperationState::addRegion),
+           py::return_value_policy::reference);
+
+  py::class_<mlir::ModuleOp>(m, "ModuleOp")
+      .def("create",
+           [](mlir::Location loc) { return mlir::ModuleOp::create(loc); })
+      .def("push_back",
+           [](mlir::ModuleOp& m, mlir::func::FuncOp f) { m.push_back(f); })
+      .def("dump", &mlir::ModuleOp::dump)
+      .def("getAsStr", [](mlir::ModuleOp& m) {
+        std::string str;
+        llvm::raw_string_ostream os(str);
+        m.print(os);
+        return os.str();
+      });
+
+  py::class_<mlir::func::FuncOp>(m, "FuncOp")
+      .def("create",
+           [](mlir::Location location, std::string name,
+              mlir::FunctionType type) {
+             auto func = mlir::func::FuncOp::create(location, name, type);
+             func.addEntryBlock();
+             return func;
+           })
+      .def(
+          "getBody",
+          [](mlir::func::FuncOp& f) -> mlir::Region& { return f.getBody(); },
+          py::return_value_policy::reference)
+      .def("getArguments",
+           [](mlir::func::FuncOp& f) { return f.getArguments().vec(); })
+      .def("getName", [](mlir::func::FuncOp& f) { return f.getName().str(); })
+      .def("getType", &mlir::func::FuncOp::getFunctionType);
+
+  py::class_<mlir::func::ReturnOp>(m, "ReturnOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              std::vector<mlir::Value> values) -> mlir::Operation* {
+             return mlir::func::ReturnOp::create(
+                        opb, loc, mlir::ArrayRef<mlir::Value>(values))
+                 .getOperation();
+           });
+
+  // mlir::TF::AddOp
+  py::class_<mlir::TF::AddV2Op>(m, "Tf_AddV2Op")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::AddV2Op::create(opb, loc, x, y).getOperation();
+           });
+
+  py::class_<mlir::TF::AnyOp>(m, "Tf_AnyOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value input,
+              mlir::Value reduction_indices,
+              bool keep_dims = false) -> mlir::Operation* {
+             return mlir::TF::AnyOp::create(opb, loc, opb.getI1Type(), input,
+                                            reduction_indices, keep_dims)
+                 .getOperation();
+           });
+
+  // mlir::TF::ConstOp
+  py::class_<mlir::TF::ConstOp>(m, "Tf_ConstOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              mlir::Attribute value) -> mlir::Operation* {
+             return mlir::TF::ConstOp::create(opb, loc, value).getOperation();
+           });
+
+  // mlir::TF::EqualOp
+  py::class_<mlir::TF::EqualOp>(m, "Tf_EqualOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::EqualOp::create(opb, loc, x, y,
+                                              opb.getBoolAttr(true))
+                 .getOperation();
+           });
+
+  // mlir::TF::GreaterEqualOp
+  py::class_<mlir::TF::GreaterEqualOp>(m, "Tf_GreaterEqualOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::GreaterEqualOp::create(opb, loc, x, y)
+                 .getOperation();
+           });
+
+  // mlir::TF::GreaterOp
+  py::class_<mlir::TF::GreaterOp>(m, "Tf_GreaterOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::GreaterOp::create(opb, loc, x, y).getOperation();
+           });
+
+  // mlir::TF::LegacyCallOp
+  py::class_<mlir::TF::LegacyCallOp>(m, "Tf_LegacyCallOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              std::vector<mlir::Type> output, std::vector<mlir::Value> args,
+              std::string f) -> mlir::Operation* {
+             return mlir::TF::LegacyCallOp::create(
+                        opb, loc, mlir::ArrayRef<mlir::Type>(output),
+                        mlir::ArrayRef<mlir::Value>(args),
+                        /*args_attrs=*/nullptr,
+                        /*res_attrs=*/nullptr, mlir::StringRef(f))
+                 .getOperation();
+           });
+
+  // mlir::TF::LessEqualOp
+  py::class_<mlir::TF::LessEqualOp>(m, "Tf_LessEqualOp")
+      .def(
+          "create",
+          [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+             mlir::Value y) -> mlir::Operation* {
+            return mlir::TF::LessEqualOp::create(opb, loc, x, y).getOperation();
+          });
+
+  // mlir::TF::LessOp
+  py::class_<mlir::TF::LessOp>(m, "Tf_LessOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::LessOp::create(opb, loc, x, y).getOperation();
+           });
+
+  // mlir::TF::NegOp
+  py::class_<mlir::TF::NegOp>(m, "Tf_NegOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              mlir::Value x) -> mlir::Operation* {
+             return mlir::TF::NegOp::create(opb, loc, x).getOperation();
+           });
+
+  py::class_<mlir::TF::NotEqualOp>(m, "Tf_NotEqualOp")
+      .def("create", [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+                        mlir::Value y) {
+        return mlir::TF::NotEqualOp::create(
+                   opb, loc, x, y, mlir::BoolAttr::get(opb.getContext(), true))
+            .getOperation();
+      });
+
+  // mlir::TF::SubOp
+  py::class_<mlir::TF::SubOp>(m, "Tf_SubOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::SubOp::create(opb, loc, x, y).getOperation();
+           });
+}

tensorflow/compiler/mlir/python/mlir_wrapper/types.cc

@@ -0,0 +1,62 @@
+/* Copyright 2020 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 <cstdint>
+#include <vector>
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+void init_types(py::module& m) {
+  // Type
+  py::class_<mlir::Type> Type(m, "Type");
+
+  // Type Sub-classes
+  py::class_<mlir::FunctionType, mlir::Type>(m, "FunctionType")
+      .def("getResults",
+           [](mlir::FunctionType& ft) { return ft.getResults().vec(); });
+
+  py::class_<mlir::FloatType, mlir::Type>(m, "FloatType")
+      .def("getBF16",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::BFloat16Type::get(context);
+           })
+      .def("getF16",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::Float16Type::get(context);
+           })
+      .def("getF32",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::Float32Type::get(context);
+           })
+      .def("getF64", [](mlir::MLIRContext* context) -> mlir::FloatType {
+        return mlir::Float64Type::get(context);
+      });
+
+  py::class_<mlir::IntegerType, mlir::Type>(m, "IntegerType")
+      .def("get", [](mlir::MLIRContext* context, unsigned width) {
+        return mlir::IntegerType::get(context, width,
+                                      mlir::IntegerType::Signless);
+      });
+
+  py::class_<mlir::UnrankedTensorType, mlir::Type>(m, "UnrankedTensorType")
+      .def("get", &mlir::UnrankedTensorType::get);
+
+  py::class_<mlir::RankedTensorType, mlir::Type>(m, "RankedTensorType")
+      .def("get", [](std::vector<int64_t> shape, mlir::Type ty) {
+        return mlir::RankedTensorType::get(mlir::ArrayRef<int64_t>(shape), ty);
+      });
+}