提交 8d719b06 编辑于 作者: Stella Laurenzo's avatar Stella Laurenzo 提交者: tensorflow-copybara
浏览文件

Remove pyiree.compiler and switch everything to pyiree.compiler2.

* Also remove the pyiree.tf.compiler and remove build support for pyiree.xla.compiler (will be transitioned after).
* Also removes bazel build support for python packages.

PiperOrigin-RevId: 346446168
上级 67a7f0e6
......@@ -31,7 +31,6 @@ py_test(
python_version = "PY3",
srcs_version = "PY3",
tags = [
"manual",
"nokokoro",
"nopresubmit",
],
......@@ -64,7 +63,6 @@ py_test(
python_version = "PY3",
srcs_version = "PY3",
tags = [
"manual",
"nokokoro",
"nopresubmit",
],
......@@ -101,7 +99,6 @@ py_test(
python_version = "PY3",
srcs_version = "PY3",
tags = [
"manual",
"nokokoro",
"nopresubmit",
],
......@@ -144,7 +141,6 @@ py_test(
python_version = "PY3",
srcs_version = "PY3",
tags = [
"manual",
"nokokoro",
"nopresubmit",
],
......@@ -181,7 +177,6 @@ py_test(
python_version = "PY3",
srcs_version = "PY3",
tags = [
"manual",
"nokokoro",
"nopresubmit",
],
......
......@@ -15,9 +15,9 @@
import tempfile
from pyiree.compiler2 import tf as iree_compiler_tf
import tensorflow as tf
from iree.integrations.tensorflow.bindings.python.pyiree.tf import compiler as iree_compiler
from tensorflow_federated.proto.v0 import computation_pb2 as pb
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.common_libs import serialization_utils
......@@ -124,6 +124,10 @@ def import_tensorflow_computation(comp, name='fn'):
legacy_init_op=initializer,
strip_default_attrs=True)
builder.save()
iree_module = iree_compiler.tf_signature_def_saved_model_to_compiler_module(
model_dir, saved_model_tags=set(['unused']), exported_names=[name])
iree_module = iree_compiler_tf.compile_saved_model(
model_dir,
import_type='SIGNATURE_DEF',
import_only=True,
saved_model_tags=set(['unused']),
exported_names=[name])
return computation_module.ComputationModule(iree_module, name, type_spec)
......@@ -223,7 +223,8 @@ class CompilerTest(tf.test.TestCase):
comp_proto = computation_impl.ComputationImpl.get_proto(comp)
module = compiler.import_tensorflow_computation(comp_proto)
self.assertIsInstance(module, computation_module.ComputationModule)
mlir = module.compiler_module.to_asm(large_element_limit=100)
assert isinstance(module.compiler_module, bytes)
mlir = module.compiler_module.decode('utf-8')
return module, mlir
def _assert_mlir_contains_pattern(self, actual_mlir, expected_list):
......
......@@ -13,7 +13,6 @@
# limitations under the License.
"""The medium for exchanging executable logic between compiler and runtime."""
from iree.integrations.tensorflow.bindings.python.pyiree.tf import compiler as iree_compiler
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.api import typed_object
......@@ -42,8 +41,6 @@ class ComputationModule(typed_object.TypedObject):
always as an instance of `computation_types.FunctionType`, possibly with
a `None` parameter in case it takes no arguments.
"""
py_typecheck.check_type(compiler_module,
iree_compiler.binding.CompilerModule)
py_typecheck.check_type(function_name, str)
py_typecheck.check_type(type_signature, computation_types.FunctionType)
self._compiler_module = compiler_module
......
......@@ -12,9 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from pyiree.compiler2 import tf as iree_compiler_tf
import tensorflow as tf
from iree.integrations.tensorflow.bindings.python.pyiree.tf import compiler as iree_compiler
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.backends.iree import computation_module
......@@ -29,8 +29,11 @@ class ComputationModuleTest(tf.test.TestCase):
model_dir = '/tmp/foo'
save_options = tf.saved_model.SaveOptions(save_debug_info=True)
tf.saved_model.save(tf_module, model_dir, options=save_options)
iree_compiler_module = iree_compiler.tf_saved_model_to_compiler_module(
model_dir, exported_names=['foo'])
iree_compiler_module = iree_compiler_tf.compile_saved_model(
model_dir,
import_only=True,
exported_names=['foo'],
target_backends=iree_compiler_tf.DEFAULT_TESTING_BACKENDS)
my_computation_module = computation_module.ComputationModule(
iree_compiler_module, 'foo',
computation_types.FunctionType(tf.float32, tf.float32))
......
......@@ -15,7 +15,9 @@
import threading
from iree.bindings.python.pyiree import rt as iree_runtime
from pyiree import compiler2 as iree_compiler
from pyiree import rt as iree_runtime
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.core.api import typed_object
from tensorflow_federated.python.core.backends.iree import backend_info
......@@ -67,8 +69,8 @@ class ComputationCallable(typed_object.TypedObject):
"""
py_typecheck.check_type(module, computation_module.ComputationModule)
py_typecheck.check_type(backend, backend_info.BackendInfo)
flatbuffer_blob = module.compiler_module.compile(
target_backends=[backend.target_name])
flatbuffer_blob = iree_compiler.compile_str(
module.compiler_module, target_backends=[backend.target_name])
# TODO(b/153499219): Find a way to name the modules somehow differently
# for debugging. Right now, module names come from the implicit "module {}"
# that wraps anything parsed from ASM that lacks an explicit module
......
......@@ -15,10 +15,12 @@
import tempfile
import numpy as np
from pyiree import rt as iree_runtime
from pyiree.compiler2 import tf as iree_compiler_tf
import tensorflow as tf
from iree.bindings.python.pyiree import rt as iree_runtime
from iree.integrations.tensorflow.bindings.python.pyiree.tf import compiler as iree_compiler
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.backends.iree import backend_info
from tensorflow_federated.python.core.backends.iree import computation_module
......@@ -44,8 +46,8 @@ class RuntimeTest(tf.test.TestCase):
with tempfile.TemporaryDirectory() as model_dir:
save_options = tf.saved_model.SaveOptions(save_debug_info=True)
tf.saved_model.save(tf_module, model_dir, options=save_options)
iree_compiler_module = iree_compiler.tf_saved_model_to_compiler_module(
model_dir, exported_names=['foo'])
iree_compiler_module = iree_compiler_tf.compile_saved_model(
model_dir, import_only=True)
my_computation_module = computation_module.ComputationModule(
iree_compiler_module, 'foo',
computation_types.FunctionType(tf.float32, tf.float32))
......
Supports Markdown
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册