From 7515cc6da4df033009046c83a8a07d1cf8427c9f Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa@m.titech.ac.jp>
Date: Sat, 24 Jun 2023 02:27:16 +0900
Subject: [PATCH] Fix tests

---
 test/polygeist-opt/pgo.mlir                              | 4 ++--
 tools/cgeist/Test/Verification/{ => CUDA}/cudasymbols.cu | 2 +-
 tools/cgeist/Test/lit.cfg                                | 7 ++++++-
 tools/cgeist/Test/lit.site.cfg.in                        | 2 ++
 4 files changed, 11 insertions(+), 4 deletions(-)
 rename tools/cgeist/Test/Verification/{ => CUDA}/cudasymbols.cu (98%)

diff --git a/test/polygeist-opt/pgo.mlir b/test/polygeist-opt/pgo.mlir
index 82b8dfd789a3..698ee2960829 100644
--- a/test/polygeist-opt/pgo.mlir
+++ b/test/polygeist-opt/pgo.mlir
@@ -23,8 +23,8 @@ module {
 // CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(2 : i32) : i32
 // CHECK:           %[[VAL_1:.*]] = llvm.mlir.constant(0 : i32) : i32
 // CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(1 : i32) : i32
-// CHECK:           %[[VAL_3:.*]] = llvm.mlir.addressof @kernelId.0 : !llvm.ptr<array<70 x i8>>
-// CHECK:           %[[VAL_4:.*]] = llvm.getelementptr %[[VAL_3]][0, 0] : (!llvm.ptr<array<70 x i8>>) -> !llvm.ptr<i8>
+// CHECK:           %[[VAL_3:.*]] = llvm.mlir.addressof @kernelId.0 : !llvm.ptr<array<[[LEN:.*]] x i8>>
+// CHECK:           %[[VAL_4:.*]] = llvm.getelementptr %[[VAL_3]][0, 0] : (!llvm.ptr<array<[[LEN]] x i8>>) -> !llvm.ptr<i8>
 // CHECK:           %[[VAL_5:.*]] = llvm.call @mgpurtPGOGetAlternative(%[[VAL_4]], %[[VAL_0]]) : (!llvm.ptr<i8>, i32) -> i32
 // CHECK:           %[[VAL_6:.*]] = llvm.icmp "eq" %[[VAL_5]], %[[VAL_1]] : i32
 // CHECK:           llvm.cond_br %[[VAL_6]], ^bb1, ^bb2
diff --git a/tools/cgeist/Test/Verification/cudasymbols.cu b/tools/cgeist/Test/Verification/CUDA/cudasymbols.cu
similarity index 98%
rename from tools/cgeist/Test/Verification/cudasymbols.cu
rename to tools/cgeist/Test/Verification/CUDA/cudasymbols.cu
index 4c86e83859fa..94b5b9512112 100644
--- a/tools/cgeist/Test/Verification/cudasymbols.cu
+++ b/tools/cgeist/Test/Verification/CUDA/cudasymbols.cu
@@ -2,7 +2,7 @@
 // RUN: cgeist %s --cuda-lower --cuda-gpu-arch=sm_60 -nocudalib -nocudainc %resourcedir --function=* -S -emit-llvm-dialect -output-intermediate-gpu=1 -emit-cuda -c | FileCheck %s
 // TODO only do this test if we have a cuda build
 
-#include "Inputs/cuda.h"
+#include "../Inputs/cuda.h"
 #include "__clang_cuda_builtin_vars.h"
 
 __device__ float dev_array[2];
diff --git a/tools/cgeist/Test/lit.cfg b/tools/cgeist/Test/lit.cfg
index eed85ef89fa6..e32339fce937 100644
--- a/tools/cgeist/Test/lit.cfg
+++ b/tools/cgeist/Test/lit.cfg
@@ -17,11 +17,16 @@ config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # suffixes: A list of file extensions to treat as test files. This is overriden
 # by individual lit.local.cfg files in the test subdirectories.
+
 config.suffixes = ['.c', '.cpp', '.cu']
 
 # excludes: A list of directories or files to exclude from the testsuite even
 # if they match the suffixes pattern.
-#config.excludes = ['Inputs']
+config.excludes = []
+if config.polygeist_enable_cuda == "0":
+    config.excludes += ['CUDA']
+if config.polygeist_enable_rocm == "0":
+    config.excludes += ['ROCm']
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
diff --git a/tools/cgeist/Test/lit.site.cfg.in b/tools/cgeist/Test/lit.site.cfg.in
index e5af5086fd1e..d252d7977f74 100644
--- a/tools/cgeist/Test/lit.site.cfg.in
+++ b/tools/cgeist/Test/lit.site.cfg.in
@@ -8,6 +8,8 @@ config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
 config.mlir_clang_obj_root = "@MLIR_CLANG_BINARY_DIR@"
 config.target_triple = "@TARGET_TRIPLE@"
 config.llvm_obj_root = path(r"@LLVM_BINARY_DIR@")
+config.polygeist_enable_cuda = "@POLYGEIST_ENABLE_CUDA@"
+config.polygeist_enable_rocm = "@POLYGEIST_ENABLE_ROCM@"
 
 # Support substitution of the tools and build_mode with user parameters.
 # This is used when we can't determine the tool dir at configuration time.