Add end2end examples for defining and running composite ops

PiperOrigin-RevId: 337520243
Change-Id: I32265c18b3736990765a9f6c416b00c03c93d335

Author: liufengdb

tensorflow/compiler/mlir/tfr/BUILD

@@ -1,5 +1,6 @@
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", "tf_py_test")
 load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")
+load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
 load(
     "//third_party/mlir:tblgen.bzl",
     "gentbl",
@@ -347,3 +348,17 @@ py_test(
         "//tensorflow/python:client_testlib",
     ],
 )
+
+py_library(
+    name = "test_utils",
+    srcs = ["python/test_utils.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow/python:client_testlib",
+    ],
+)
+
+gen_op_libraries(
+    name = "one_op",
+    src = "define_op_template.py",
+)

tensorflow/compiler/mlir/tfr/tfr.bzl renamed to tensorflow/compiler/mlir/tfr/build_defs.bzl

@@ -1,12 +1,12 @@
 """BUILD extension for TF composition project."""
 
-load("//tensorflow:tensorflow.bzl", "py_binary", "tf_gen_op_wrapper_py")
-load("//tensorflow:tensorflow.google.bzl", "pytype_library")
+load("//tensorflow:tensorflow.bzl", "py_binary", "tf_custom_op_library", "tf_gen_op_wrapper_py")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
 
 def gen_op_libraries(
         name,
         src,
-        deps,
+        deps = [],
         tags = [],
         test = False):
     """gen_op_libraries() generates all cc and py libraries for composite op source.
@@ -21,36 +21,33 @@ def gen_op_libraries(
     if not src.endswith(".py") or name == src[:-3]:
         fail("'src' %s conflicts with op Python wrapper. Rename it to be different from 'name'." % src)
 
-    gen_op_lib_exec = src[:-3]
+    gen_op_lib_exec = src[:-3]  # Strip off the .py
     py_binary(
         name = gen_op_lib_exec,
         srcs = [src],
         srcs_version = "PY2AND3",
         python_version = "PY3",
         deps = [
-            "//tensorflow/python:platform",
+            "//tensorflow/compiler/mlir/tfr:op_reg_gen",
+            "//tensorflow/compiler/mlir/tfr:tfr_gen",
+            "//tensorflow/compiler/mlir/tfr:composite",
         ] + deps,
     )
 
-    register_op = "register_" + name
+    registed_op = "registed_" + name
     native.genrule(
-        name = register_op,
+        name = registed_op,
         srcs = [],
         outs = [name + ".inc.cc"],
         cmd = "$(location %s) --output=$@ --gen_register_op=true" % gen_op_lib_exec,
         exec_tools = [":" + gen_op_lib_exec],
-        local = 1,
         tags = tags,
     )
 
     native.cc_library(
         name = name + "_cc",
         testonly = test,
-        srcs = [":" + register_op],
-        copts = [
-            "-Wno-unused-result",
-            "-Wno-unused-variable",
-        ],
+        srcs = [":" + registed_op],
         deps = [
             "//tensorflow/core:framework",
             "//tensorflow/core:lib",
@@ -59,62 +56,61 @@ def gen_op_libraries(
         alwayslink = 1,
     )
 
+    tf_custom_op_library(
+        name = name + ".so",
+        srcs = [":" + registed_op],
+    )
+
     tf_gen_op_wrapper_py(
-        name = name,
-        out = name + ".py",
+        name = "gen_" + name,
+        out = "gen_" + name + ".py",
         deps = [
             ":%s_cc" % name,
         ],
     )
 
-    pytype_library(
-        name = name + "_grads",
-        srcs = [
-            src,
-        ],
+    tf_custom_op_py_library(
+        name = name,
+        dso = [":%s.so" % name],
+        kernels = [":%s_cc" % name],
         srcs_version = "PY2AND3",
         deps = [
-            "//third_party/py/numpy",
-            "//third_party/py/tensorflow",
-        ] + deps,
-    )
-
-    pytype_library(
-        name = name + "_lib",
-        srcs = [
-            name + ".py",
+            ":gen_%s" % name,
         ],
-        srcs_version = "PY2AND3",
-        deps = [
-            ":%s" % name,
-            ":%s_cc" % name,
-            ":%s_grads" % name,
-            "//third_party/py/numpy",
-            "//third_party/py/tensorflow",
-        ] + deps,
     )
 
     # Link the register op and rebuild the binary
-    gen_tfr_lib_exec = gen_op_lib_exec + "_registered"
+    gen_tfr_lib_exec = gen_op_lib_exec + "_with_op_library"
     py_binary(
         name = gen_tfr_lib_exec,
         main = src,
         srcs = [src],
         srcs_version = "PY2AND3",
         python_version = "PY3",
         deps = [
-            "//tensorflow/python:platform",
-            ":%s" % name + "_cc",
+            "//tensorflow/compiler/mlir/tfr:op_reg_gen",
+            "//tensorflow/compiler/mlir/tfr:tfr_gen",
+            "//tensorflow/compiler/mlir/tfr:composite",
+            ":%s" % name,
         ] + deps,
     )
 
-    op_tfr = "composite_" + name
     native.genrule(
-        name = op_tfr,
+        name = name + "_mlir",
         srcs = [],
         outs = [name + ".mlir"],
         cmd = "$(location %s) --output=$@ --gen_register_op=false" % gen_tfr_lib_exec,
         exec_tools = [":" + gen_tfr_lib_exec],
-        local = 1,
         tags = tags,
     )
+
+    native.py_library(
+        name = name + "_py",
+        srcs = [src],
+        srcs_version = "PY2AND3",
+        deps = [
+            "//tensorflow/compiler/mlir/tfr:op_reg_gen",
+            "//tensorflow/compiler/mlir/tfr:tfr_gen",
+            "//tensorflow/compiler/mlir/tfr:composite",
+        ] + deps,
+    )

tensorflow/compiler/mlir/tfr/define_op_template.py

@@ -0,0 +1,64 @@
+# 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.
+"""A template to define composite ops."""
+
+# pylint: disable=g-direct-tensorflow-import
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+
+from tensorflow.compiler.mlir.tfr.python.composite import Composite
+from tensorflow.compiler.mlir.tfr.python.op_reg_gen import gen_register_op
+from tensorflow.compiler.mlir.tfr.python.tfr_gen import tfr_gen_from_module
+from tensorflow.python.platform import app
+from tensorflow.python.platform import flags
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'output', None,
+    'Path to write the genereated register op file and MLIR file.')
+
+flags.DEFINE_bool('gen_register_op', True,
+                  'Generate register op cc file or tfr mlir file.')
+
+flags.mark_flag_as_required('output')
+
+
+@Composite('TestRandom', derived_attrs=['T: numbertype'], outputs=['o: T'])
+def _composite_random_op():
+  pass
+
+
+def main(_):
+  if FLAGS.gen_register_op:
+    assert FLAGS.output.endswith('.cc')
+    generated_code = gen_register_op(sys.modules[__name__], '_composite_')
+  else:
+    assert FLAGS.output.endswith('.mlir')
+    generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_')
+
+  dirname = os.path.dirname(FLAGS.output)
+  if not os.path.exists(dirname):
+    os.makedirs(dirname)
+  with open(FLAGS.output, 'w') as f:
+    f.write(generated_code)
+
+
+if __name__ == '__main__':
+  app.run(main=main)

tensorflow/compiler/mlir/tfr/examples/mnist/BUILD

@@ -0,0 +1,60 @@
+load("//tensorflow:tensorflow.bzl", "py_binary")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
+load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
+
+package(
+    default_visibility = [
+        ":friends",
+    ],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+package_group(
+    name = "friends",
+    includes = ["//third_party/mlir:subpackages"],
+    packages = [
+        "//tensorflow/compiler/mlir/tfr/...",
+    ],
+)
+
+gen_op_libraries(
+    name = "mnist_ops",
+    src = "ops_defs.py",
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_py_test(
+    name = "mnist_ops_test",
+    size = "small",
+    srcs = ["mnist_ops_test.py"],
+    data = [":mnist_ops_mlir"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    tags = [
+        "no_pip",
+        "no_windows",  # TODO(b/170752141)
+        "nomac",  # TODO(b/170752141)
+    ],
+    deps = [
+        ":mnist_ops",
+        ":mnist_ops_py",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/compiler/mlir/tfr:test_utils",
+    ],
+)
+
+py_binary(
+    name = "mnist_train",
+    srcs = ["mnist_train.py"],
+    data = [":mnist_ops_mlir"],
+    python_version = "PY3",
+    deps = [
+        ":mnist_ops",
+        ":mnist_ops_py",
+        "//tensorflow:tensorflow_py",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+    ],
+)