Generalize module format in AOT tool (#4818)

This PR modifies the aot-tool to potentially support several ze-module
formats (aside from spirv and native binary).

The `compiler.cpp` now has a string `format_name` template parameter
instead of a boolean `is_spv` that contains a format name ("native",
"spirv", ...). The string is then matched to an actual format in
`get_module_format()`.

Signed-off-by: dchigarev <[email protected]>

Author: dchigarev

python/triton/tools/compile.py

@@ -205,12 +205,16 @@ def constexpr(s):
         "_placeholder": "",
     }
     if is_xpu():
+        if args.generate_native_code:
+            format_name = "native"
+        else:
+            format_name = "spirv"
         params |= {
             "arg_types": ", ".join(ty_to_cpp(arg) for arg in arg_types_not_1),
             "grf_mode": args.grf_mode,
             "build_flags": ccinfo.metadata.build_flags,
             "threads_per_warp": args.threads_per_warp,
-            "is_spv": "false" if args.generate_native_code else "true",
+            "format_name": format_name,
         }
     output_files = []
     backend_name = target.backend

third_party/intel/tools/intel/compile.cpp

@@ -25,10 +25,10 @@ static inline void gpuAssert(ze_result_t code, const char *file, int line) {{
 }}
 
 // ze globals
-#define SPV_NAME {kernel_name}_spv
+#define BIN_NAME {kernel_name}_bin
 ze_module_handle_t {kernel_name}_mod = NULL;
 ze_kernel_handle_t {kernel_name}_func = NULL;
-unsigned char SPV_NAME[{bin_size}] = {{ {bin_data} }};
+unsigned char BIN_NAME[{bin_size}] = {{ {bin_data} }};
 // sycl globals
 const sycl::device sycl_device;
 const auto ctx =
@@ -42,15 +42,25 @@ void unload_{kernel_name}(void) {{
     // Not implemeted
 }}
 
+static ze_module_format_t get_module_format(const std::string& format_name) {{
+  if (format_name == "spirv") {{
+    return ZE_MODULE_FORMAT_IL_SPIRV;
+  }} else if (format_name == "native") {{
+    return ZE_MODULE_FORMAT_NATIVE;
+  }} else {{
+    throw std::runtime_error("Unsupported module format");
+  }}
+}}
+
 void load_{kernel_name}() {{
-    uint8_t *binary_ptr = (uint8_t *)&SPV_NAME;
+    uint8_t *binary_ptr = (uint8_t *)&BIN_NAME;
     size_t binary_size = {bin_size};
 
-    const bool is_spv = {is_spv};
+    const std::string format_name = "{format_name}";
 
     ze_module_desc_t module_description {{}};
     module_description.stype = ZE_STRUCTURE_TYPE_MODULE_DESC;
-    module_description.format = is_spv ? ZE_MODULE_FORMAT_IL_SPIRV : ZE_MODULE_FORMAT_NATIVE;
+    module_description.format = get_module_format(format_name);
     module_description.inputSize = static_cast<uint32_t>(binary_size);
     module_description.pInputModule = binary_ptr;
     module_description.pBuildFlags = "{build_flags}";