Refactor specialiazation constants to use dataclass

also adds spec_id, itemsize, and default_value fields

Author: ndgrigorian

dpctl/program/utils/_utils.py

@@ -16,6 +16,7 @@
 
 """Implements various utilities for the dpctl.program module."""
 
+from dataclasses import dataclass
 from enum import IntEnum
 
 import numpy as np
@@ -31,6 +32,7 @@ class SpirvOpCode(IntEnum):
     OpSpecConstantTrue = 48
     OpSpecConstantFalse = 49
     OpSpecConstant = 50
+    OpFunction = 54
     OpDecorate = 71
 
 
@@ -40,10 +42,52 @@ class SpirvDecoration(IntEnum):
     SpecId = 1
 
 
+@dataclass(frozen=True)
+class SpecializationConstantInfo:
+    """Data class representing specialization constant information."""
+
+    spec_id: int
+    dtype: str
+    name: str
+    itemsize: int
+    default_value: int | float | bool | None
+
+
 def parse_spirv_specializations(
     spv_bytes: bytes | bytearray | memoryview,
-) -> dict[int, dict[str, str]]:
-    """Parses SPIR-V binary to extract specialization constants."""
+) -> tuple[SpecializationConstantInfo]:
+    """
+    Parses SPIR-V byte stream to extract information about specializations,
+    including the specialization IDs, types, and names.
+
+    Note that the dtype information may be imprecise, as the compiler may
+    choose to, for example, represent a bool as char, or may represent both
+    signed and unsigned integers as unsigned integer bit buckets of the same
+    length.
+
+    Args:
+        spv_bytes (Union[bytes, bytearray, memoryview]):
+            the SPIR-V byte stream.
+
+    Returns:
+        tuple[SpecializationConstantInfo]:
+            a tuple of parsed constants and their information represented by
+            `SpecializationConstantInfo` objects, sorted by their
+            specialization IDs. The length of the tuple is equal to the number
+            of specialization constants found. Each
+            `SpecializationConstantInfo` object contains the following
+            attributes:
+
+            - `spec_id` (int): The specialization ID.
+            - `dtype` (str): A NumPy style string representing the data type.
+            - `itemsize` (int): The size of the specialization constant in
+                bytes.
+            - `name` (str): The variable name. If not preserved in the binary,
+                a default name in the format `unnamed_spec_const_{spec_id}` is
+                used.
+            - `default_value` (int | float | bool | None): The default value of
+                the specialization constant. If not specified, `None` is used.
+    """
     words = np.frombuffer(spv_bytes, dtype=np.uint32)
 
     # verify magic number
@@ -54,6 +98,7 @@ def parse_spirv_specializations(
     ids = {}
     names = {}
     constants = {}
+    defaults = {}
 
     i = 5  # skip 5 word header
     while i < len(words):
@@ -64,27 +109,45 @@ def parse_spirv_specializations(
         if word_count == 0:
             raise ValueError(f"Invalid SPIR-V instruction at word index {i}")
 
-        if opcode == SpirvOpCode.OpTypeBool:
+        if opcode == SpirvOpCode.OpFunction:
+            # everything following is not relevant to specialization constant
+            # parsing, so we can stop parsing at this point
+            break
+        elif opcode == SpirvOpCode.OpTypeBool:
             result_id = int(words[i + 1])
-            types[result_id] = "?"
+            types[result_id] = {"dtype": "?", "itemsize": 1}
         elif opcode == SpirvOpCode.OpTypeInt:
             result_id = int(words[i + 1])
             width = int(words[i + 2])
             signed = int(words[i + 3])
             prefix = "i" if signed else "u"
-            types[result_id] = f"{prefix}{width // 8}"
+            types[result_id] = {
+                "dtype": f"{prefix}{width // 8}",
+                "itemsize": width // 8,
+            }
         elif opcode == SpirvOpCode.OpTypeFloat:
             result_id = int(words[i + 1])
             width = int(words[i + 2])
-            types[result_id] = f"f{width // 8}"
-        elif opcode in (
-            SpirvOpCode.OpSpecConstant,
-            SpirvOpCode.OpSpecConstantTrue,
-            SpirvOpCode.OpSpecConstantFalse,
-        ):
+            types[result_id] = {
+                "dtype": f"f{width // 8}",
+                "itemsize": width // 8,
+            }
+        elif opcode == SpirvOpCode.OpSpecConstant:
+            type_id = int(words[i + 1])
+            result_id = int(words[i + 2])
+            constants[result_id] = type_id
+            literal_words = words[i + 3 : i + word_count]
+            defaults[result_id] = literal_words.tobytes()
+        elif opcode == SpirvOpCode.OpSpecConstantTrue:
             type_id = int(words[i + 1])
             result_id = int(words[i + 2])
             constants[result_id] = type_id
+            defaults[result_id] = True
+        elif opcode == SpirvOpCode.OpSpecConstantFalse:
+            type_id = int(words[i + 1])
+            result_id = int(words[i + 2])
+            constants[result_id] = type_id
+            defaults[result_id] = False
         elif opcode == SpirvOpCode.OpDecorate:
             target_id = int(words[i + 1])
             decoration = int(words[i + 2])
@@ -97,15 +160,37 @@ def parse_spirv_specializations(
 
         i += word_count
 
-    result = {}
+    # a spec ID may appear multiple times in the same binary with different
+    # target IDs. We only need to keep one, so skip duplicates
+    unique_ids = set()
+    result = []
     for target_id, spec_id in ids.items():
+        if spec_id in unique_ids:
+            continue
+        unique_ids.add(spec_id)
         type_id = constants.get(target_id)
-        dtype_str = types.get(type_id, "unknown_type")
+        type_info = types.get(type_id, {"dtype": "unknown_type", "itemsize": 0})
         name = names.get(target_id, f"unnamed_spec_const_{spec_id}")
 
-        result[spec_id] = {
-            "name": name,
-            "dtype": dtype_str,
-        }
-
-    return result
+        dtype_str = type_info["dtype"]
+        raw_default = defaults.get(target_id)
+        default_value = None
+        if isinstance(raw_default, bytes):
+            try:
+                default_value = np.frombuffer(raw_default, dtype=dtype_str)[
+                    0
+                ].item()
+            except Exception:
+                default_value = None
+
+        result.append(
+            SpecializationConstantInfo(
+                spec_id=spec_id,
+                dtype=dtype_str,
+                name=name,
+                itemsize=type_info["itemsize"],
+                default_value=default_value,
+            )
+        )
+
+    return tuple(sorted(result, key=lambda x: x.spec_id))

dpctl/tests/input_files/specialization_constant_composite.spv

@@ -341,3 +341,49 @@ def test_create_kernel_bundle_with_composite_spec_const():
 
     # 1.0 * 10 + 2.5 = 12.5
     assert np.all(y == 12.5)
+
+
+def test_spirv_specializations_parser():
+    spirv_file = get_spirv_abspath("specialization_constant_kernel.spv")
+    with open(spirv_file, "rb") as spv:
+        spv_bytes = spv.read()
+    spec_consts = dpctl_prog.utils.parse_spirv_specializations(spv_bytes)
+    assert len(spec_consts) == 1
+    assert spec_consts[0].dtype == "u4"
+
+    spirv_file = get_spirv_abspath("specialization_constant_composite.spv")
+    with open(spirv_file, "rb") as spv:
+        spv_bytes = spv.read()
+
+    spec_consts = dpctl_prog.utils.parse_spirv_specializations(spv_bytes)
+    assert len(spec_consts) == 3
+    spec_const0, spec_const1, spec_const2 = spec_consts
+    assert spec_const0.dtype == "u4"
+    assert spec_const0.itemsize == 4
+    assert spec_const0.name == "unnamed_spec_const_0"
+    assert spec_const0.default_value == 1
+
+    assert spec_const1.dtype == "f4"
+    assert spec_const1.itemsize == 4
+    assert spec_const1.name == "unnamed_spec_const_1"
+    assert spec_const1.default_value == 0
+
+    # compiler translates bool to char
+    assert spec_const2.dtype == "u1"
+    assert spec_const2.itemsize == 1
+    assert spec_const2.name == "unnamed_spec_const_2"
+    assert spec_const2.default_value == 0
+
+
+def test_spirv_specializations_parser_no_spec_consts():
+    spirv_file = get_spirv_abspath("multi_kernel.spv")
+    with open(spirv_file, "rb") as spv:
+        spv_bytes = spv.read()
+    spec_consts = dpctl_prog.utils.parse_spirv_specializations(spv_bytes)
+    assert not spec_consts
+
+
+def test_spirv_specializations_parser_invalid_spirv():
+    invalid_spv = b"\x00\x01\x02\x03\x04\x05"
+    with pytest.raises(ValueError):
+        dpctl_prog.utils.parse_spirv_specializations(invalid_spv)