diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2f339b5fd..c476f7eb8 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -177,6 +177,9 @@ add_library(triton SHARED ${PYTHON_SRC})
 
 # find_package(PythonLibs REQUIRED)
 
+set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+set(TRITON_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
+
 target_link_libraries(triton
   ${PYTHON_LIBRARIES}
   TritonIR
diff --git a/python/setup.py b/python/setup.py
index e8dd815de..05f008d2c 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -79,6 +79,8 @@ class CMakeBuild(build_ext):
 
     def build_extension(self, ext):
         llvm_include_dir, llvm_library_dir = get_llvm()
+        # lit is used by the test suite
+        lit_dir = shutil.which('lit')
         self.debug = True
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
         # create build directories
@@ -99,7 +101,8 @@ class CMakeBuild(build_ext):
             # '-DPYTHON_EXECUTABLE=' + sys.executable,
             # '-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON',
             "-DTRITON_LLVM_BUILD_DIR=" + llvm_build_dir,
-            "-DPYTHON_INCLUDE_DIRS=" + ";".join(python_include_dirs)
+            "-DPYTHON_INCLUDE_DIRS=" + ";".join(python_include_dirs),
+            "-DLLVM_EXTERNAL_LIT=" + lit_dir
         ]
         # configuration
         cfg = "Debug" if self.debug else "Release"
@@ -119,6 +122,11 @@ class CMakeBuild(build_ext):
         subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=self.build_temp, env=env)
         subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp)
 
+        # run tests. Note: this depends on llvm-lit
+        # -DLLVM_EXTERNAL_LIT=<path-to-lit.py>
+        # Note: get_llvm_lit_path(...) in llvm/cmake/modules/AddLLVM.cmake
+        subprocess.call(["cmake", "--build", ".", "--target", "check-triton"], cwd=self.build_temp, env=env)
+
 
 setup(
     name="triton",
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 80f217775..0619e0030 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -10,7 +10,7 @@ configure_lit_site_cfg(
 )
 
 set(TRITON_TEST_DEPENDS
-  FileCheck count not
+  # FileCheck count not
   triton-opt
 )
 
diff --git a/test/Triton/vecadd.mlir b/test/Triton/vecadd.mlir
index 8e4e23b69..d99c3ab1f 100644
--- a/test/Triton/vecadd.mlir
+++ b/test/Triton/vecadd.mlir
@@ -1,4 +1,4 @@
-// RUN: triton-opt %s -tritongpu-verifier -verify-diagnostics
+// RUN: triton-opt %s -verify-diagnostics
 
 module {
   func @add_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32__(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32, %arg5: i32) {
@@ -11,9 +11,9 @@ module {
     %5 = tt.broadcast %arg3 : (i32) -> tensor<256xi32>
     %6 = arith.cmpi slt, %4, %5 : tensor<256xi32>
     %7 = tt.broadcast %arg0 : (!tt.ptr<f32>) -> tensor<256x!tt.ptr<f32>>
-    %8 = tt.getelementptr %7, %4, : tensor<256x!tt.ptr<f32>>
+    %8 = tt.getelementptr %7, %4 : tensor<256x!tt.ptr<f32>>
     %9 = tt.broadcast %arg1 : (!tt.ptr<f32>) -> tensor<256x!tt.ptr<f32>>
-    %10 = tt.getelementptr %9, %4, : tensor<256x!tt.ptr<f32>>
+    %10 = tt.getelementptr %9, %4 : tensor<256x!tt.ptr<f32>>
     %cst = arith.constant 0.000000e+00 : f32
     %11 = tt.broadcast %cst : (f32) -> tensor<256xf32>
     %c0_i32 = arith.constant 0 : i32
@@ -31,13 +31,13 @@ module {
       %22 = arith.addf %19, %21 : tensor<256xf32>
       %23 = arith.addf %arg7, %22 : tensor<256xf32>
       %24 = tt.broadcast %arg5 : (i32) -> tensor<256xi32>
-      %25 = tt.getelementptr %arg8, %24, : tensor<256x!tt.ptr<f32>>
+      %25 = tt.getelementptr %arg8, %24 : tensor<256x!tt.ptr<f32>>
       %26 = tt.broadcast %arg5 : (i32) -> tensor<256xi32>
-      %27 = tt.getelementptr %arg9, %26, : tensor<256x!tt.ptr<f32>>
+      %27 = tt.getelementptr %arg9, %26 : tensor<256x!tt.ptr<f32>>
       scf.yield %23, %25, %27 : tensor<256xf32>, tensor<256x!tt.ptr<f32>>, tensor<256x!tt.ptr<f32>>
     }
     %16 = tt.broadcast %arg2 : (!tt.ptr<f32>) -> tensor<256x!tt.ptr<f32>>
-    %17 = tt.getelementptr %16, %4, : tensor<256x!tt.ptr<f32>>
+    %17 = tt.getelementptr %16, %4 : tensor<256x!tt.ptr<f32>>
     tt.store %17, %15#0, %6, : tensor<256xf32>
     return
   }
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in
index 69a2a4a6e..1118ed36b 100644
--- a/test/lit.site.cfg.py.in
+++ b/test/lit.site.cfg.py.in
@@ -17,7 +17,7 @@ config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
 
 
 import lit.llvm
-lit.llvminitialize(lit_config, config)
+lit.llvm.initialize(lit_config, config)
 
 # Let the main config do the real work
 lit_config.load_config(config, "@TRITON_SOURCE_DIR@/test/lit.cfg.py")