Merge pull request #1642 from RossBrunton/ross/icpx

[CMake] Parse kernel names using integration header rather than IR

Author: RossBrunton

scripts/generate_kernel_header.py

@@ -62,44 +62,32 @@ def generate_header(output_file, kernel_name_dict):
         fout.write(rendered)
 
 
-def get_mangled_names(dpcxx_path, source_file, output_header):
+def get_mangled_names(source_file, output_header):
     """Return a list of all the entry point names from a given sycl source file.
 
     Filters out wrapper and offset handler entry points.
     """
     output_dir = os.path.dirname(output_header)
-    il_file = os.path.join(output_dir, os.path.basename(source_file) + ".ll")
-    generate_il_command = f"""\
-        {dpcxx_path} -S -fsycl -fsycl-device-code-split=off \
-        -fsycl-device-only -o {il_file} {source_file}"""
-    subprocess.run(generate_il_command, shell=True)
-    kernel_line_regex = re.compile("define.*spir_kernel")
-    definition_lines = []
-    with open(il_file) as f:
+    name = os.path.splitext(os.path.basename(source_file))[0]
+    ih_file = os.path.join(output_dir, name, name + ".ih")
+    definitions = []
+    writing = False
+    with open(ih_file) as f:
         lines = f.readlines()
         for line in lines:
-            if kernel_line_regex.search(line) is not None:
-                definition_lines.append(line)
+            if "}" in line and writing:
+                break
+            # __pf_kernel_wrapper seems to be an internal function used by dpcpp
+            if writing and "19__pf_kernel_wrapper" not in line:
+                definitions.append(line.replace(",", "").strip()[1:-1])
+            if "const char* const kernel_names[] = {" in line:
+                    writing = True
 
-    entry_point_names = []
-    kernel_name_regex = re.compile(r"@(.*?)\(")
-    for line in definition_lines:
-        if kernel_name_regex.search(line) is None:
-            continue
-        match = kernel_name_regex.search(line)
-        assert isinstance(match, re.Match)
-        kernel_name = match.group(1)
-        if "kernel_wrapper" not in kernel_name and "with_offset" not in kernel_name:
-            entry_point_names.append(kernel_name)
-
-    os.remove(il_file)
-    return entry_point_names
+    return definitions
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--dpcxx_path",
-                        help="Full path to dpc++ compiler executable.")
     parser.add_argument(
         "-o",
         "--output",
@@ -112,7 +100,7 @@ def main():
     for source_file in args.source_files:
         program_name = os.path.splitext(os.path.basename(source_file))[0]
         mangled_names[program_name] = get_mangled_names(
-            args.dpcxx_path, source_file, args.output)
+            source_file, args.output)
     generate_header(args.output, mangled_names)
 
 

test/conformance/device_code/CMakeLists.txt

@@ -10,6 +10,12 @@ else()
     set(AMD_ARCH "${UR_CONFORMANCE_AMD_ARCH}")
 endif()
 
+if (WIN32)
+  set(NULDEV NUL)
+else()
+  set(NULDEV /dev/null)
+endif()
+
 cmake_path(GET UR_DPCXX EXTENSION EXE)
 cmake_path(REPLACE_FILENAME UR_DPCXX "clang-offload-extract${EXE}" OUTPUT_VARIABLE DEFAULT_EXTRACTOR_NAME)
 set(UR_DEVICE_CODE_EXTRACTOR "${DEFAULT_EXTRACTOR_NAME}" CACHE PATH "Path to clang-offload-extract")
@@ -100,6 +106,17 @@ macro(add_device_binary SOURCE_FILE)
         add_custom_target(generate_${KERNEL_NAME}_${TRIPLE} DEPENDS ${BIN_PATH})
         add_dependencies(generate_device_binaries generate_${KERNEL_NAME}_${TRIPLE})
     endforeach()
+
+    set(IH_PATH "${DEVICE_BINARY_DIR}/${KERNEL_NAME}.ih")
+    add_custom_command(OUTPUT "${IH_PATH}"
+        COMMAND ${UR_DPCXX} -fsycl -fsycl-device-code-split=off
+        -fsycl-device-only -c -Xclang -fsycl-int-header="${IH_PATH}"
+        ${DPCXX_BUILD_FLAGS_LIST} ${SOURCE_FILE} -o ${NULDEV}
+
+        WORKING_DIRECTORY "${DEVICE_BINARY_DIR}"
+        DEPENDS ${SOURCE_FILE}
+    )
+    list(APPEND DEVICE_IHS ${IH_PATH})
     list(APPEND DEVICE_CODE_SOURCES ${SOURCE_FILE})
 endmacro()
 
@@ -127,8 +144,8 @@ set(KERNEL_HEADER ${UR_CONFORMANCE_DEVICE_BINARIES_DIR}/kernel_entry_points.h)
 add_custom_command(OUTPUT ${KERNEL_HEADER}
     WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/scripts
     COMMAND ${Python3_EXECUTABLE} generate_kernel_header.py
-    --dpcxx_path ${UR_DPCXX} -o ${KERNEL_HEADER} ${DEVICE_CODE_SOURCES}
+    -o ${KERNEL_HEADER} ${DEVICE_CODE_SOURCES}
     DEPENDS ${PROJECT_SOURCE_DIR}/scripts/generate_kernel_header.py
-    ${DEVICE_CODE_SOURCES})
+    ${DEVICE_CODE_SOURCES} ${DEVICE_IHS})
 add_custom_target(kernel_names_header DEPENDS ${KERNEL_HEADER})
 add_dependencies(generate_device_binaries kernel_names_header)