Switch HSA/ROCR dispatch to hsaco/kernel-object loading

amd_triton_npu/backend/driver.py

@@ -1288,30 +1288,32 @@ def _get_output_format(runtime=None):
     If not set, defaults to "elf" on npu2 and "xclbin" on npu1.
     ELF format is only supported on npu2 (AIE2P) devices.
 
-    Under the HSA runtime the artifact must be PDI + insts, since the ROCR AIE
-    dispatch path consumes a raw ``aie.pdi`` and its ``insts.bin`` sidecar. In
-    that case "pdi" is forced and an explicit ELF/xclbin request is rejected.
+    Under the HSA runtime, "pdi" is rejected: a raw ``aie.pdi`` has no
+    ``AIE_PARTITION`` section, so the hsaco packer (which needs the PDI's
+    partition/column-count metadata, not just its bytes) has nothing to read
+    it from. Unlike XRT, which can dispatch a raw PDI + insts.bin directly,
+    HSA always packs through ``aie.compiler.hsaco.pack``, which only
+    understands the xclbin and full-ELF shapes.
     """
     if runtime is None:
         runtime = npu_config.runtime
     npu_version = detect_npu_version(runtime)
     configured_format = npu_config.output_format
-    if runtime == "hsa":
-        if configured_format is not None and configured_format != "pdi":
-            raise RuntimeError(
-                f"ROCR requires the 'pdi' output format, but "
-                f"output_format={configured_format!r} was requested. Unset "
-                "AMD_TRITON_NPU_OUTPUT_FORMAT (or npu_config.output_format), or "
-                "set it to 'pdi'."
-            )
-        return "pdi"
     if configured_format is not None:
         if configured_format == "elf" and npu_version == "npu1":
             raise RuntimeError(
                 "ELF output format is not supported on npu1 (AIE2) devices. "
                 "Unset or change AMD_TRITON_NPU_OUTPUT_FORMAT, or set "
                 "npu_config.output_format to 'xclbin', 'pdi', or None."
             )
+        if configured_format == "pdi" and runtime == "hsa":
+            raise RuntimeError(
+                "output_format='pdi' is not supported under the HSA runtime: "
+                "a raw PDI carries no partition/column-count metadata, which "
+                "the hsaco packer needs. Unset AMD_TRITON_NPU_OUTPUT_FORMAT "
+                "(or npu_config.output_format) to auto-detect xclbin/elf, or "
+                "set it to one of those explicitly."
+            )
         return configured_format
     # Auto-detect: ELF for npu2, xclbin for npu1.
     return "elf" if npu_version == "npu2" else "xclbin"
@@ -2450,6 +2452,64 @@ def _put_aircc_artifacts(cache, artifacts, output_format):
     return {"bin_path": bin_path, "insts_path": insts_path}
 
 
+# Inverse of NPU_AGENT_NAMES (driver.py's existing {"aie2": "npu1", "aie2p":
+# "npu2"} map): the hsaco section an npu_version packs into.
+_NPU_VERSION_TO_ARCH = {v: k for k, v in NPU_AGENT_NAMES.items()}
+
+
+def _pack_hsaco(output_format, artifacts, arch, kernel_name):
+    """Pack aircc's artifacts (see ``_aircc_compile``) into hsaco bytes.
+
+    ``kernel_name`` is the string the packed kernel is looked up by at
+    dispatch time: the instance-qualified "<kernel>:<instance>" form
+    ``_extract_elf_kernel_name`` already produces for a full ELF (the exact
+    string XRT's own ELF launcher already looks its kernel up by), or the
+    plain Triton kernel name for PDI+insts -- packing and the later
+    ``hsa_executable_get_symbol_by_name`` only have to agree with each
+    other, not with any real ELF symbol.
+
+    Returns:
+        bytes: the packed hsaco.
+
+    Raises:
+        ValueError: via aie.compiler.hsaco.pack, if the xclbin/full-ELF
+            artifacts are malformed (e.g. a missing partition column count, or
+            ``kernel_name`` not found in a full ELF's kernel list).
+        RuntimeError: via aie.compiler.hsaco.pack, if xclbinutil/llvm-objcopy
+            cannot be located or fail.
+    """
+    try:
+        from aie.compiler.hsaco import pack
+    except ImportError as e:
+        raise RuntimeError(
+            "could not import aie.compiler.hsaco: the installed mlir-aie "
+            "does not have hsaco packing support yet. This needs an "
+            "mlir-aie build from (or after) branch "
+            "ypapadop-amd/hsa-hsaco-loading -- utils/mlir-air-hash.txt is not "
+            "yet pinned to one (see this change's PR description)."
+        ) from e
+
+    if output_format == "elf":
+        (kernel,) = pack.kernels_from_full_elf(
+            artifacts["elf_path"], names=[kernel_name]
+        )
+    else:
+        pdi, num_cols = pack.partition_from_xclbin(artifacts["bin_path"])
+        insts = Path(artifacts["insts_path"]).read_bytes()
+        kernel = {
+            "name": kernel_name,
+            "pdi": pdi,
+            "num_cols": num_cols,
+            "insts": insts,
+        }
+    section = pack.build_section(arch, [kernel])
+    with tempfile.TemporaryDirectory() as scratch:
+        path = os.path.join(scratch, "kernel.hsaco")
+        pack.ensure_hsaco(path)
+        pack.inject(path, arch, section)
+        return Path(path).read_bytes()
+
+
 def compile_module(
     launcher_src,
     kernel_placeholder_name,
@@ -2463,8 +2523,7 @@ def compile_module(
     ``link_profile`` selects the runtime the generated launcher links against:
 
     * ``"xrt"``  -- XRT (xclbin/ELF); links xrt_coreutil + uuid.
-    * ``"hsa"``  -- HSA AIE dispatch via ROCR; links libhsa-runtime64. Requires
-      ``output_format == "pdi"``.
+    * ``"hsa"``  -- HSA AIE dispatch via ROCR; links libhsa-runtime64.
     """
     if link_profile not in ("xrt", "hsa"):
         raise ValueError(f"link_profile must be 'xrt' or 'hsa'; got {link_profile!r}")
@@ -2622,6 +2681,22 @@ def launch(
             else:
                 cache_bin_path = cached_artifacts["bin_path"]
                 cache_insts_path = cached_artifacts["insts_path"]
+            if link_profile == "hsa":
+                cache_hsaco_path = cache.get_file("kernel.hsaco")
+                if cache_hsaco_path is None:
+                    # Force a real rebuild: cache_path gates that branch below,
+                    # not cached_artifacts, so both have to be cleared or a
+                    # cache dir with a compiled .so but no kernel.hsaco (e.g.
+                    # left over from before this hsaco path existed) would
+                    # still take the cache-hit branch and reach the final
+                    # dispatch with no cache_hsaco_path/hsaco_kernel_name set.
+                    cached_artifacts = None
+                    cache_path = None
+                elif output_format == "elf":
+                    with open(cache_elf_kernel_path) as f:
+                        hsaco_kernel_name = f.read().strip()
+                else:
+                    hsaco_kernel_name = kernel_name
 
         if cache_path is None:
             with tempfile.TemporaryDirectory() as tmpdir:
@@ -2719,6 +2794,19 @@ def launch(
                 else:
                     cache_bin_path = cached_artifacts["bin_path"]
                     cache_insts_path = cached_artifacts["insts_path"]
+                if link_profile == "hsa":
+                    arch = _NPU_VERSION_TO_ARCH[npu_version]
+                    hsaco_kernel_name = (
+                        artifacts["elf_kernel_name"]
+                        if output_format == "elf"
+                        else kernel_name
+                    )
+                    hsaco_bytes = _pack_hsaco(
+                        output_format, artifacts, arch, hsaco_kernel_name
+                    )
+                    cache_hsaco_path = cache.put(
+                        hsaco_bytes, "kernel.hsaco", binary=True
+                    )
                 with open(so_path, "rb") as f:
                     cache_path = cache.put(f.read(), filename, binary=True)
 
@@ -2756,7 +2844,12 @@ def launch(
             raise RuntimeError(f"Cannot find {name} module in {cache_path}")
         mod = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(mod)
-        if output_format == "elf":
+        if link_profile == "hsa":
+            hsaco_path_str = cache_hsaco_path
+            if IS_WINDOWS and hsaco_path_str.startswith("\\\\?\\"):
+                hsaco_path_str = hsaco_path_str[4:]
+            mod.set_paths(hsaco_path_str, hsaco_kernel_name)
+        elif output_format == "elf":
             # Read the cached kernel name
             with open(cache_elf_kernel_path) as f:
                 elf_kernel_name = f.read().strip()
@@ -2809,16 +2902,20 @@ class NPULauncher(object):
     def __init__(self, src, metadata, runtime="xrt"):
         """Build the host launcher for ``src`` targeting ``runtime``.
 
-        ``runtime`` is "xrt" (xclbin/ELF via XRT) or "hsa" (PDI + insts via
-        HSA/ROCR); it is bound by ``NPUDriver`` when it constructs the launcher.
+        ``runtime`` is "xrt" (xclbin/ELF via XRT) or "hsa" (hsaco-packed
+        PDI+insts or full ELF via HSA/ROCR); it is bound by ``NPUDriver`` when
+        it constructs the launcher.
         """
         constants, signature = extract_signature_and_constants(src)
         if runtime == "hsa":
-            # HSA consumes PDI + insts and links ROCR. Import lazily so the
-            # (large) HSA codegen module is only loaded when the HSA path is used.
+            # HSA consumes a packed hsaco (xclbin/PDI+insts or full ELF) and
+            # links ROCR. Import lazily so the (large) HSA codegen module is
+            # only loaded when the HSA path is used.
             from .hsa_launcher import _generate_hsa_launcher
 
-            self.output_format = "pdi"
+            # Detect output format the same way the XRT branch does: ELF for
+            # npu2, xclbin for npu1. Both are packed into an hsaco before dispatch.
+            self.output_format = _get_output_format(runtime=runtime)
             launcher_src = _generate_hsa_launcher(
                 constants,
                 signature,
@@ -2875,11 +2972,14 @@ def get_npu_cache_dir(compiled_kernel):
 
     The NPU backend stores hardware-specific artifacts in a separate cache
     directory from Triton's main compiler cache. Depending on the selected
-    output format, the directory contains either:
+    output format and link profile, the directory contains:
 
-    * xclbin output: ``aie.xclbin``, ``insts.bin``, and
+    * xclbin output (xrt): ``aie.xclbin``, ``insts.bin``, and
+      ``__npu_dispatch.so``
+    * elf output (xrt): ``aie.elf``, ``elf_kernel_name.txt``, and
       ``__npu_dispatch.so``
-    * elf output: ``aie.elf``, ``elf_kernel_name.txt``, and
+    * hsa (either output format, packed into one hsaco): the above
+      format-specific artifacts, plus ``kernel.hsaco`` and
       ``__npu_dispatch.so``
 
     This function returns the path to that directory.
@@ -2965,7 +3065,7 @@ def __init__(self, runtime=None):
         ``runtime`` selects how kernels are dispatched:
 
         * ``"xrt"`` -- XRT (xclbin on npu1, ELF on npu2).
-        * ``"hsa"`` -- HSA/ROCR AIE dispatch (PDI + insts).
+        * ``"hsa"`` -- HSA/ROCR AIE dispatch (hsaco-packed PDI+insts or full ELF).
 
         Defaults to ``npu_config.runtime`` (the ``AMD_TRITON_NPU_RUNTIME`` env
         var, itself defaulting to ``"xrt"``), so ``NPUDriver()`` honors the

amd_triton_npu/backend/hsa_launcher.py

@@ -9,8 +9,9 @@
 the shared runtime library ``libtriton_npu_hsa.so`` (see
 ``include/HsaRuntime/HsaRuntime.{h,cpp}``) via a small C ABI:
 
-* ``triton_npu_hsa_prepare(pdi, insts)`` -> opaque program handle (once, from
-  ``set_paths``).
+* ``triton_npu_hsa_prepare(hsaco, kernel_name)`` -> opaque program handle
+  (once, from ``set_paths``), loaded from an hsaco and resolved by kernel
+  name.
 * ``triton_npu_hsa_dispatch(program, n, ptrs, sizes)`` (per launch).
 
 Because every launcher links the same shared library, the ``HsaRuntime``
@@ -34,7 +35,7 @@
 def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) -> str:
     """Generate the thin C++ CPython launcher that dispatches via HSA/ROCR.
 
-    The generated module exposes ``set_paths(pdi_path, insts_path)`` (which calls
+    The generated module exposes ``set_paths(hsaco_path, kernel_name)`` (which calls
     ``triton_npu_hsa_prepare`` once and stashes the handle) and ``launch(...)``
     (which marshals the tensor pointers/sizes and calls
     ``triton_npu_hsa_dispatch`` with the GIL released). All HSA state lives in
@@ -44,8 +45,9 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) ->
     (``codegen.written_pointer_args``), or None when that is not known.
 
     ``_kernel_name`` is accepted for signature parity with the XRT launcher
-    generators but is unused: the HSA path selects work by PDI/insts address,
-    not by a kernel symbol name.
+    generators but is unused here: the real kernel name is threaded through at
+    runtime via ``set_paths``'s second argument (see driver.py's
+    ``compile_module``), not baked into the generated source.
     """
     args_format = "".join(format_of(extracted_type(ty)) for ty in signature.values())
     fmt = "iiiOOOO" + args_format
@@ -114,19 +116,20 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) ->
 // Number of tensor kernel arguments for this specialized launcher.
 static constexpr std::uint32_t NUM_KERNARGS = {num_ptr_args};
 
-// Handle to this module's prepared (pdi, insts) program in the shared runtime.
+// Handle to this module's prepared hsaco kernel in the shared runtime.
 static triton_npu_hsa_program_t g_program = nullptr;
 
-// Python-callable set_paths(pdi_path, insts_path): prepare (load + cache) the
-// program in the shared runtime and stash its handle. Called once per module.
+// Python-callable set_paths(hsaco_path, kernel_name): pack-load the hsaco
+// in the shared runtime and stash its program handle. Called once per
+// module.
 static PyObject* py_set_paths(PyObject* self, PyObject* args) {{
-  const char* pdi;
-  const char* insts;
-  if (!PyArg_ParseTuple(args, "ss", &pdi, &insts)) {{
+  const char* hsaco_path;
+  const char* kernel_name;
+  if (!PyArg_ParseTuple(args, "ss", &hsaco_path, &kernel_name)) {{
     return NULL;
   }}
   char err[512];
-  g_program = triton_npu_hsa_prepare(pdi, insts, err, sizeof(err));
+  g_program = triton_npu_hsa_prepare(hsaco_path, kernel_name, err, sizeof(err));
   if (g_program == nullptr) {{
     PyErr_SetString(PyExc_RuntimeError, err);
     return NULL;
@@ -198,7 +201,7 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) ->
 // Methods exported by this per-signature dispatch extension module.
 static PyMethodDef ModuleMethods[] = {{
   {{"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}},
-  {{"set_paths", py_set_paths, METH_VARARGS, "Set paths to aie.pdi and insts.bin"}},
+  {{"set_paths", py_set_paths, METH_VARARGS, "Set path to kernel.hsaco and the kernel name"}},
   {{NULL, NULL, 0, NULL}}  // sentinel
 }};