Use llvm filecheck instead of filecheck python lib. (#7070)

The filecheck python dependency isn't available in python3.9. This
change removes it in lieu of the llvm FileCheck binary, which is now
packaged into the wheel.

Author: BabakkGraphcore

CMakeLists.txt

@@ -296,6 +296,20 @@ if(TRITON_BUILD_PYTHON_MODULE)
     target_link_libraries(triton PRIVATE z)
   endif()
   target_link_options(triton PRIVATE ${LLVM_LDFLAGS})
+
+  if (NOT DEFINED LLVM_SYSPATH)
+      message(FATAL_ERROR "LLVM_SYSPATH must be set.")
+  endif()
+
+  if (NOT DEFINED TRITON_WHEEL_DIR)
+      message(FATAL_ERROR "TRITON_WHEEL_DIR must be set.")
+  endif()
+
+  configure_file(
+    "${LLVM_SYSPATH}/bin/FileCheck"
+    "${TRITON_WHEEL_DIR}/FileCheck"
+    COPYONLY)
+
 endif()
 
 if (UNIX AND NOT APPLE)

python/test-requirements.txt

@@ -6,5 +6,4 @@ pytest-forked
 pytest-xdist
 scipy>=1.7.1
 llnl-hatchet
-filecheck
 expecttest

python/test/unit/test_filecheck.py

@@ -32,5 +32,5 @@ def test_kernel():
         # CHECK: %c42_i32
         anchor(scalar)
 
-    with pytest.raises(ValueError, match="Couldn't match \"%c42_i32\""):
+    with pytest.raises(ValueError, match="expected string not found in input\n # CHECK: %c42_i32"):
         run_filecheck_test(test_kernel)

python/triton/_filecheck.py

@@ -1,12 +1,7 @@
-import sys
 import os
-import io
 import inspect
-
-from filecheck.options import Options
-from filecheck.finput import FInput
-from filecheck.parser import Parser, pattern_for_opts
-from filecheck.matcher import Matcher
+import subprocess
+import tempfile
 
 import triton
 from triton.compiler import ASTSource, make_backend
@@ -22,8 +17,8 @@
 stub_target = GPUTarget("cuda", 100, 32)
 stub_backend = make_backend(stub_target)
 
-llvm_bin_dir = os.path.join(os.path.dirname(sys.executable), "bin")
-filecheck_path = os.path.join(llvm_bin_dir, "FileCheck")
+triton_dir = os.path.dirname(__file__)
+filecheck_path = os.path.join(triton_dir, "FileCheck")
 
 
 class MatchError(ValueError):
@@ -37,14 +32,21 @@ def __str__(self):
 
 
 def run_filecheck(name, module_str, check_template):
-    options = Options(match_filename=name)
-    fin = FInput(name, module_str)
-    ops = io.StringIO(check_template)
-    parser = Parser(options, ops, *pattern_for_opts(options))
-    matcher = Matcher(options, fin, parser)
-    matcher.stderr = io.StringIO()
-    if matcher.run() != 0:
-        raise MatchError(matcher.stderr.getvalue(), module_str)
+    with tempfile.TemporaryDirectory() as tempdir:
+        temp_module = os.path.join(tempdir, "module")
+        with open(temp_module, "w") as temp:
+            temp.write(module_str)
+
+        temp_expected = os.path.join(tempdir, "expected")
+        with open(temp_expected, "w") as temp:
+            temp.write(check_template)
+
+        try:
+            subprocess.check_output([filecheck_path, temp_expected, "--input-file", temp_module],
+                                    stderr=subprocess.STDOUT)
+        except subprocess.CalledProcessError as error:
+            decoded = error.output.decode('unicode_escape')
+            raise ValueError(decoded)
 
 
 def run_parser(kernel_fn):

setup.py

@@ -311,6 +311,8 @@ def get_thirdparty_packages(packages: list):
             thirdparty_cmake_args.append(f"-D{p.include_flag}={package_dir}/include")
         if p.lib_flag:
             thirdparty_cmake_args.append(f"-D{p.lib_flag}={package_dir}/lib")
+        if p.syspath_var_name:
+            thirdparty_cmake_args.append(f"-D{p.syspath_var_name}={package_dir}")
         if p.sym_name is not None:
             sym_link_path = os.path.join(package_root_dir, p.sym_name)
             update_symlink(sym_link_path, package_dir)
@@ -435,6 +437,8 @@ def build_extension(self, ext):
         thirdparty_cmake_args = get_thirdparty_packages([get_llvm_package_info()])
         thirdparty_cmake_args += self.get_pybind11_cmake_args()
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
+        wheeldir = os.path.dirname(extdir)
+
         # create build directories
         if not os.path.exists(self.build_temp):
             os.makedirs(self.build_temp)
@@ -448,7 +452,8 @@ def build_extension(self, ext):
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir, "-DTRITON_BUILD_PYTHON_MODULE=ON",
             "-DPython3_EXECUTABLE:FILEPATH=" + sys.executable, "-DPython3_INCLUDE_DIR=" + python_include_dir,
             "-DTRITON_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends if not b.is_external]),
-            "-DTRITON_PLUGIN_DIRS=" + ';'.join([b.src_dir for b in backends if b.is_external])
+            "-DTRITON_PLUGIN_DIRS=" + ';'.join([b.src_dir for b in backends if b.is_external]),
+            "-DTRITON_WHEEL_DIR=" + wheeldir
         ]
         if lit_dir is not None:
             cmake_args.append("-DLLVM_EXTERNAL_LIT=" + lit_dir)