Removed duplicate per-backend tf and iree input mlir files (#2521)
The `iree_input__<backend>.mlir` files are all identical, and the `tf_input__<backend>.mlir` files are all identical up to constants in their function names. For example:
```mlir
< func @__inference_simple_matmul_70(%arg0: tensor<128x3072xf32> {tf._user_specified_name = "a", tf_saved_model.index_path = [0]}, %arg1: tensor<3072x256xf32> {tf._user_specified_name = "b", tf_saved_model.index_path = [1]}) -> (tensor<128x256xf32> {tf_saved_model.index_path = []}) attributes {tf._input_shapes = [#tf.shape<128x3072>, #tf.shape<3072x256>], tf_saved_model.exported_names = ["simple_matmul"]} {
---
> func @__inference_simple_matmul_670(%arg0: tensor<128x3072xf32> {tf._user_specified_name = "a", tf_saved_model.index_path = [0]}, %arg1: tensor<3072x256xf32> {tf._user_specified_name = "b", tf_saved_model.index_path = [1]}) -> (tensor<128x256xf32> {tf_saved_model.index_path = []}) attributes {tf._input_shapes = [#tf.shape<128x3072>, #tf.shape<3072x256>], tf_saved_model.exported_names = ["simple_matmul"]} {
```diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
index 8f5ec8d..46a3785 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
@@ -39,6 +39,16 @@
np.random.seed(seed)
+def backends_to_str(target_backends):
+ """Creates a flattened and normalized string representing target_backends."""
+ normalized_backends = []
+ for backend in target_backends:
+ # Remove unusual characters and ensure names don't end or start in "_".
+ backend = re.sub("[^0-9a-zA-Z_]+", "_", backend)
+ normalized_backends.append(backend.strip("_"))
+ return "__".join(normalized_backends)
+
+
def compile_tf_module(tf_module,
target_backends=(),
exported_names=(),
@@ -52,9 +62,9 @@
saved_model:
A TF SavedModel directory containing the files used translate the
tf.Module into an IREE module.
- tf_input__backends.mlir:
+ tf_input.mlir:
MLIR for the module in TF's input dialect.
- iree_input__backends.mlir:
+ iree_input.mlir:
The MLIR above translated to IREE via compiler.TF_IMPORT_PASS_PIPELINE.
compiled__backends.vmfb:
A VM FlatBuffer compiled to the target backends from the IREE MLIR above.
@@ -77,14 +87,6 @@
# We break up the compilation here so we can save intermediary artifacts.
compiler_context = compiler.Context()
- if artifacts_dir is not None:
- normalized_backends = []
- for backend in target_backends:
- # Remove unusual characters and ensure names don't end or start in "_".
- backend = re.sub("[^0-9a-zA-Z_]+", "_", backend)
- normalized_backends.append(backend.strip("_"))
- backends_string = "__".join(normalized_backends)
-
# Convert the tf_module into raw TF input MLIR.
compiler_module = compiler.tf_load_saved_model(
sm_path,
@@ -93,8 +95,7 @@
pass_pipeline=())
if artifacts_dir is not None:
- tf_mlir_path = os.path.join(artifacts_dir,
- f"tf_input__{backends_string}.mlir")
+ tf_mlir_path = os.path.join(artifacts_dir, "tf_input.mlir")
logging.info("Saving raw TF input MLIR to: %s", tf_mlir_path)
with open(tf_mlir_path, "w") as f:
f.write(compiler_module.to_asm())
@@ -103,16 +104,15 @@
compiler_module.run_pass_pipeline(compiler.TF_IMPORT_PASS_PIPELINE)
if artifacts_dir is not None:
- iree_mlir_path = os.path.join(artifacts_dir,
- f"iree_input__{backends_string}.mlir")
+ iree_mlir_path = os.path.join(artifacts_dir, "iree_input.mlir")
logging.info("Saving IREE input MLIR to: %s", iree_mlir_path)
with open(iree_mlir_path, "w") as f:
f.write(compiler_module.to_asm())
compiled_module = compiler_module.compile(target_backends=target_backends)
if artifacts_dir is not None:
- compiled_path = os.path.join(artifacts_dir,
- f"compiled__{backends_string}.vmfb")
+ compiled_name = f"compiled__{backends_to_str(target_backends)}.vmfb"
+ compiled_path = os.path.join(artifacts_dir, compiled_name)
logging.info("Saving compiled IREE module to: %s", compiled_path)
with open(compiled_path, "wb") as f:
f.write(compiled_module)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
index 25645df..fe25f97 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
@@ -17,6 +17,7 @@
import os
import tempfile
+from absl import logging
from absl.testing import parameterized
from pyiree.tf.support import tf_utils
import tensorflow as tf
@@ -52,7 +53,7 @@
},
{
'testcase_name': 'multiple_backends',
- 'target_backends': ['vmla', 'llvm'],
+ 'target_backends': ['vmla', 'llvm-ir'],
},
])
def test_artifact_saving(self, target_backends):
@@ -64,13 +65,13 @@
artifacts_dir=artifacts_dir)
artifacts_to_check = [
- 'saved_model',
- f'tf_input__{"__".join(target_backends)}.mlir',
- f'iree_input__{"__".join(target_backends)}.mlir',
- f'compiled__{"__".join(target_backends)}.vmfb',
+ 'saved_model', 'tf_input.mlir', 'iree_input.mlir',
+ f'compiled__{tf_utils.backends_to_str(target_backends)}.vmfb',
]
for artifact in artifacts_to_check:
- self.assertTrue(os.path.exists(os.path.join(artifacts_dir, artifact)))
+ artifact_path = os.path.join(artifacts_dir, artifact)
+ logging.info('Checking path: %s', artifact_path)
+ self.assertTrue(os.path.exists(artifact_path))
@parameterized.named_parameters([
{
