2407170feb
TFJS ops are registered as TF custom ops, and utilize export_graphdef.cc to build out the GraphDef object that could contain both TF and TFJS dialects. PiperOrigin-RevId: 311158257 Change-Id: I7313a5a01f12ef742a97fd5e9ff2bbffe8498b0c
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
# Copyright 2019 Google Inc. 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.
|
|
"""Lit runner site configuration."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import os
|
|
import platform
|
|
import lit.llvm
|
|
|
|
# Handle the test srcdir for platforms. On windows, things are weird with bazel.
|
|
if platform.system() == 'Windows':
|
|
srcdir = os.environ['TEST_SRCDIR']
|
|
real_test_srcdir = srcdir[:srcdir.find('tensorflow/compiler/mlir')]
|
|
external_srcdir = os.path.join(real_test_srcdir, 'external')
|
|
else:
|
|
real_test_srcdir = os.environ['TEST_SRCDIR']
|
|
external_srcdir = real_test_srcdir
|
|
|
|
# Lint for undefined variables is disabled as config is not defined inside this
|
|
# file, instead config is injected by lit.py. The structure is common for lit
|
|
# tests and intended to only persist temporarily (b/136126535).
|
|
# pylint: disable=undefined-variable
|
|
config.llvm_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'llvm')
|
|
config.mlir_obj_root = os.path.join(real_test_srcdir)
|
|
config.mlir_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'mlir')
|
|
# TODO(jpienaar): Replace with suffices in build rule.
|
|
config.suffixes = ['.td', '.mlir', '.pbtxt']
|
|
|
|
mlir_tf_tools_dirs = [
|
|
'tensorflow/compiler/mlir',
|
|
'tensorflow/compiler/mlir/lite',
|
|
'tensorflow/compiler/mlir/tensorflow',
|
|
'tensorflow/compiler/mlir/tfjs',
|
|
'tensorflow/compiler/mlir/xla',
|
|
'tensorflow/compiler/aot',
|
|
'tensorflow/compiler/xla/service/mlir_gpu',
|
|
]
|
|
config.mlir_tf_tools_dirs = [
|
|
os.path.join(real_test_srcdir, os.environ['TEST_WORKSPACE'], s)
|
|
for s in mlir_tf_tools_dirs
|
|
]
|
|
test_dir = os.environ['TEST_TARGET']
|
|
test_dir = test_dir.strip('/').rsplit(':', 1)[0]
|
|
config.mlir_test_dir = os.path.join(real_test_srcdir,
|
|
os.environ['TEST_WORKSPACE'], test_dir)
|
|
|
|
if platform.system() == 'Windows':
|
|
# Configure this to work with msys2, TF's preferred windows bash.
|
|
config.lit_tools_dir = '/usr/bin'
|
|
|
|
lit.llvm.initialize(lit_config, config)
|
|
|
|
# Let the main config do the real work.
|
|
lit_config.load_config(
|
|
config,
|
|
os.path.join(
|
|
os.path.join(real_test_srcdir, os.environ['TEST_WORKSPACE'],
|
|
'tensorflow/compiler/mlir/runlit.cfg.py')))
|
|
# pylint: enable=undefined-variable
|