Set up llvm lit test with bazel

# Copyright 2019 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.
"""Lit runner configuration."""

import os
import platform
import sys
import lit.formats
from lit.llvm import llvm_config
from lit.llvm.subst import ToolSubst

# Lint for undefined variables is disabled as config is not defined inside this
# file, instead config is injected by way of evaluating runlit.cfg.py from
# runlit.site.cfg.py which in turn is evaluated by lit.py. The structure is
# common for lit tests and intended to only persist temporarily (b/136126535).
# pylint: disable=undefined-variable
# Configuration file for the 'lit' test runner.

# name: The name of this test suite.
config.name = 'MLIR ' + os.path.basename(config.mlir_test_dir)

config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)

# suffixes: A list of file extensions to treat as test files.
config.suffixes = ['.cc', '.hlo', '.hlotxt', '.json', '.mlir', '.pbtxt', '.py']

# test_source_root: The root path where tests are located.
config.test_source_root = config.mlir_test_dir

# test_exec_root: The root path where tests should be run.
config.test_exec_root = os.environ['RUNFILES_DIR']

if platform.system() == 'Windows':
    tool_patterns = [
        ToolSubst('FileCheck.exe', unresolved='fatal'),
        #  Handle these specially as they are strings searched for during testing.
        ToolSubst('count.exe', unresolved='fatal'),
        ToolSubst('not.exe', unresolved='fatal')
    ]

    llvm_config.config.substitutions.append(
        ('%python', '"%s"' % (sys.executable)))

    llvm_config.add_tool_substitutions(tool_patterns,
                                       [llvm_config.config.llvm_tools_dir])
else:
    llvm_config.use_default_substitutions()

llvm_config.config.substitutions.append(
    ('%tfrt_bindir', 'tensorflow/compiler/aot'))

# Tweak the PATH to include the tools dir.
llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)

tool_dirs = config.mlir_custom_tools_dirs + [
    config.mlir_tools_dir, config.llvm_tools_dir
]

# need modify
tool_names = [
    ## add your own tools
    "mlir-opt",
    "mlir-translate"
]
tools = [ToolSubst(s, unresolved='ignore') for s in tool_names]
llvm_config.add_tool_substitutions(tools, tool_dirs)
# pylint: enable=undefined-variable

import os
import platform
import lit.llvm

## ================ need modify ================
mlir_custom_tools_dirs = [
    'src
]
# the folder contains all test files
tests_root_dir = 'src/'
# runlit.cfg.py file path
runlit_cfg_py_path = "src/runlit.cfg.py"
## ================ need modify end ================

# 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(tests_root_dir)]
  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']

config.mlir_custom_tools_dirs = [
    os.path.join(real_test_srcdir, os.environ['TEST_WORKSPACE'], s)
    for s in mlir_custom_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'],
                     runlit_cfg_py_path)))
# pylint: enable=undefined-variable

# Test definitions for Lit, the LLVM test runner.
#
# This is reusing the LLVM Lit test runner in the interim until the new build
# rules are upstreamed.
# TODO(b/136126535): remove this custom rule.
"""Lit runner globbing test
"""

load("@bazel_skylib//lib:paths.bzl", "paths")

# Default values used by the test runner.
_default_test_file_exts = ["mlir", ".pbtxt", ".td"]
_default_driver = "@llvm-project//mlir:run_lit.sh"
_default_size = "small"
_default_tags = []

# These are patterns which we should never match, for tests, subdirectories, or
# test input data files.
_ALWAYS_EXCLUDE = [
 "**/LICENSE.txt",
 "**/README.txt",
 "**/lit.local.cfg",
 # Exclude input files that have spaces in their names, since bazel
 # cannot cope with such "targets" in the srcs list.
 "**/* *",
 "**/* */**",
]

def _run_lit_test(name, data, lit_path, size, tags, driver, features, exec_properties):
 """Runs lit on all tests it can find in `data` under tensorflow/compiler/mlir.

 Note that, due to Bazel's hermetic builds, lit only sees the tests that
 are included in the `data` parameter, regardless of what other tests might
 exist in the directory searched.

 Args:
   name: str, the name of the test, including extension.
   data: [str], the data input to the test.
   lit_path: [str], the path to put lit files, start from root folder
   size: str, the size of the test.
   tags: [str], tags to attach to the test.
   driver: str, label of the driver shell script.
           Note: use of a custom driver is not currently supported
           and specifying a default driver will abort the tests.
   features: [str], list of extra features to enable.
 """

 # Disable tests on windows for now, to enable testing rest of all xla and mlir.
 native.py_test(
     name = name,
     srcs = ["@llvm-project//llvm:lit"],
     tags = tags + ["no_pip", "no_windows"],
     args = [
         lit_path + "/" + paths.basename(data[-1]) + " --config-prefix=runlit -v",
     ] + features,
     data = data + [
         "//{}:litfiles".format(lit_path),
         "@llvm-project//llvm:FileCheck",
         "@llvm-project//llvm:count",
         "@llvm-project//llvm:not",
     ],
     #        deps = ["@pypi_lit//:pkg"],
     size = size,
     main = "lit.py",
     exec_properties = exec_properties,
 )

def glob_lit_tests(
     name = None,
     lit_path = None,
     exclude = [],
     test_file_exts = _default_test_file_exts,
     default_size = _default_size,
     size_override = {},
     data = [],
     per_test_extra_data = {},
     default_tags = _default_tags,
     tags_override = {},
     driver = _default_driver,
     features = [],
     exec_properties = {}):
 """Creates all plausible Lit tests (and their inputs) under this directory.

 Args:
   name: str, name of the test_suite rule to generate for running all tests.
   exclude: [str], paths to exclude (for tests and inputs).
   test_file_exts: [str], extensions for files that are tests.
   default_size: str, the test size for targets not in "size_override".
   size_override: {str: str}, sizes to use for specific tests.
   data: [str], additional input data to the test.
   per_test_extra_data: {str: [str]}, extra data to attach to a given file.
   default_tags: [str], additional tags to attach to the test.
   tags_override: {str: str}, tags to add to specific tests.
   driver: str, label of the driver shell script.
           Note: use of a custom driver is not currently supported
           and specifying a default driver will abort the tests.
   features: [str], list of extra features to enable.
   exec_properties: a dictionary of properties to pass on.
 """

 # Ignore some patterns by default for tests and input data.
 exclude = _ALWAYS_EXCLUDE + exclude

 tests = native.glob(
     ["*." + ext for ext in test_file_exts],
     exclude = exclude,
 )

 # Run tests individually such that errors can be attributed to a specific
 # failure.
 all_tests = []
 for curr_test in tests:
     all_tests.append(curr_test + ".test")

     # Instantiate this test with updated parameters.
     _run_lit_test(
         name = curr_test + ".test",
         data = data + [curr_test] + per_test_extra_data.get(curr_test, []),
         lit_path = lit_path,
         size = size_override.get(curr_test, default_size),
         tags = default_tags + tags_override.get(curr_test, []),
         driver = driver,
         features = features,
         exec_properties = exec_properties,
     )

 # TODO: remove this check after making it a required param.
 if name:
     native.test_suite(
         name = name,
         tests = all_tests,
         tags = ["manual"],
     )