spirv.py 1.08 KB
Newer Older
1 2 3 4
"""Utility for Interacting with SPIRV Tools"""
import subprocess
import os
from . import util
5
from .._ffi.base import py_str
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39

def optimize(spv_bin):
    """Optimize SPIRV using spirv-opt via CLI

    Note that the spirv-opt is still experimental.

    Parameters
    ----------
    spv_bin : bytearray
        The spirv file

    Return
    ------
    cobj_bin : bytearray
        The HSA Code Object
    """

    tmp_dir = util.tempdir()
    tmp_in = tmp_dir.relpath("input.spv")
    tmp_out = tmp_dir.relpath("output.spv")
    with open(tmp_in, "wb") as out_file:
        out_file.write(bytes(spv_bin))

    sdk = os.environ.get("VULKAN_SDK", None)
    cmd = os.path.join(sdk, "bin/spirv-opt") if sdk else "spirv-opt"
    args = [cmd, "-O", tmp_in, "-o", tmp_out]
    proc = subprocess.Popen(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    (out, _) = proc.communicate()

    if proc.returncode != 0:
        msg = "Opitmizationerror using spirv-opt:\n"
40
        msg += py_str(out)
41 42 43
        raise RuntimeError(msg)

    return bytearray(open(tmp_out, "rb").read())