28947730040

Committed 08 Jul 2026 01:48PM UTC coverage: 74.758% (-0.9%) from 75.677%

Build # 28947730040

Build Type

Pull #2304

github

Committed by

web-flow

Commit Message

Merge 9c9a7c4c3 into baf89c997

Pull Request Pull Request #2304: Add support for specialization constants

Coverage Stats

880 of 1248 branches covered (70.51%)

Branch coverage included in aggregate %.

133 of 222 new or added lines in 4 files covered. (59.91%)

1 existing line in 1 file now uncovered.

3370 of 4437 relevant lines covered (75.95%)

262.97 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

73.26

/dpctl/program/utils/_utils.py

#                      Data Parallel Control (dpctl)
#
# Copyright 2026 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Implements various utilities for the dpctl.program module."""

from dataclasses import dataclass
from enum import IntEnum

import numpy as np


# these constants come from the SPIR-V spec:
# https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html
class SpirvOpCode(IntEnum):
    OpName = 5
    OpTypeBool = 20
    OpTypeInt = 21
    OpTypeFloat = 22
    OpSpecConstantTrue = 48
    OpSpecConstantFalse = 49
    OpSpecConstant = 50
    OpFunction = 54
    OpDecorate = 71


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,
) -> tuple[SpecializationConstantInfo]:
    """
    Parses SPIR-V byte stream to extract information about specializations,
    including the specialization IDs, types, names, and default values.

    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 (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
    if len(words) < 5 or words[0] != 0x07230203:
        raise ValueError("Invalid SPIR-V binary")

    types = {}
    ids = {}
    names = {}
    constants = {}
    defaults = {}

    i = 5  # skip 5 word header
    while i < len(words):
        word = words[i]
        opcode = word & 0xFFFF
        word_count = word >> 16

        if word_count == 0:
            raise ValueError(f"Invalid SPIR-V instruction at word index {i}")
        if i + word_count > len(words):
            raise ValueError(
                f"Invalid SPIR-V instruction at offset {i} (extends beyond "
                "buffer)"
            )

        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:
            if word_count < 2:
                raise ValueError(f"Truncated OpTypeBool at word index {i}")
            result_id = int(words[i + 1])
            types[result_id] = {"dtype": "?", "itemsize": 1}
        elif opcode == SpirvOpCode.OpTypeInt:
            if word_count < 4:
                raise ValueError(f"Truncated OpTypeInt at word index {i}")
            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] = {
                "dtype": f"{prefix}{width // 8}",
                "itemsize": width // 8,
            }
        elif opcode == SpirvOpCode.OpTypeFloat:
            if word_count < 3:
                raise ValueError(f"Truncated OpTypeFloat at word index {i}")
            result_id = int(words[i + 1])
            width = int(words[i + 2])
            types[result_id] = {
                "dtype": f"f{width // 8}",
                "itemsize": width // 8,
            }
        elif opcode == SpirvOpCode.OpSpecConstant:
            if word_count < 4:
                raise ValueError(f"Truncated OpSpecConstant at word index {i}")
            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:
            if word_count < 3:
                raise ValueError(
                    f"Truncated OpSpecConstantTrue at word index {i}"
                )
            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:
            if word_count < 3:
                raise ValueError(
                    f"Truncated OpSpecConstantFalse at word index {i}"
                )
            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:
            if word_count < 3:
                raise ValueError(f"Truncated OpDecorate at word index {i}")
            target_id = int(words[i + 1])
            decoration = int(words[i + 2])
            if decoration == SpirvDecoration.SpecId:
                if word_count < 4:
                    raise ValueError(
                        f"Truncated OpDecorate SpecId at word index {i}"
                    )
                ids[target_id] = int(words[i + 3])
        elif opcode == SpirvOpCode.OpName:
            if word_count < 2:
                raise ValueError(f"Truncated OpName at word index {i}")
            target_id = int(words[i + 1])
            name_bytes = words[i + 2 : i + word_count].tobytes()
            try:
                names[target_id] = name_bytes.split(b"\x00", 1)[0].decode(
                    "utf-8"
                )
            except UnicodeDecodeError:
                raise ValueError(f"Invalid UTF-8 in OpName at word index {i}")

        i += word_count

    # 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)
        type_info = types.get(type_id, {"dtype": "unknown_type", "itemsize": 0})
        name = names.get(target_id, f"unnamed_spec_const_{spec_id}")

        dtype_str = type_info["dtype"]
        raw_default = defaults.get(target_id)
        default_value = None
        if isinstance(raw_default, bool):
            default_value = raw_default
        elif isinstance(raw_default, bytes) and dtype_str != "unknown_type":
            try:
                default_value = np.frombuffer(raw_default, dtype=dtype_str)[
                    0
                ].item()
            except (ValueError, TypeError):
                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))

1	# Data Parallel Control (dpctl)
2	#
3	# Copyright 2026 Intel Corporation
4	#
5	# Licensed under the Apache License, Version 2.0 (the "License");
6	# you may not use this file except in compliance with the License.
7	# You may obtain a copy of the License at
8	#
9	# http://www.apache.org/licenses/LICENSE-2.0
10	#
11	# Unless required by applicable law or agreed to in writing, software
12	# distributed under the License is distributed on an "AS IS" BASIS,
13	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14	# See the License for the specific language governing permissions and
15	# limitations under the License.
16
17	"""Implements various utilities for the dpctl.program module."""
18
19	from dataclasses import dataclass	1✔
20	from enum import IntEnum	1✔
21
22	import numpy as np	1✔
23
24
25	# these constants come from the SPIR-V spec:
26	# https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html
27	class SpirvOpCode(IntEnum):	1✔
28	OpName = 5	1✔
29	OpTypeBool = 20	1✔
30	OpTypeInt = 21	1✔
31	OpTypeFloat = 22	1✔
32	OpSpecConstantTrue = 48	1✔
33	OpSpecConstantFalse = 49	1✔
34	OpSpecConstant = 50	1✔
35	OpFunction = 54	1✔
36	OpDecorate = 71	1✔
37
38
39	class SpirvDecoration(IntEnum):	1✔
40	SpecId = 1	1✔
41
42
43	@dataclass(frozen=True)	1✔
44	class SpecializationConstantInfo:	1✔
45	"""Data class representing specialization constant information."""
46
47	spec_id: int
48	dtype: str
49	name: str
50	itemsize: int
51	default_value: int \| float \| bool \| None
52
53
54	def parse_spirv_specializations(	1✔
55	spv_bytes: bytes \| bytearray \| memoryview,
56	) -> tuple[SpecializationConstantInfo]:
57	"""
58	Parses SPIR-V byte stream to extract information about specializations,
59	including the specialization IDs, types, names, and default values.
60
61	Note that the dtype information may be imprecise, as the compiler may
62	choose to, for example, represent a bool as char, or may represent both
63	signed and unsigned integers as unsigned integer bit buckets of the same
64	length.
65
66	Args:
67	spv_bytes (bytes \| bytearray \| memoryview):
68	the SPIR-V byte stream.
69
70	Returns:
71	tuple[SpecializationConstantInfo]:
72	a tuple of parsed constants and their information represented by
73	`SpecializationConstantInfo` objects, sorted by their
74	specialization IDs. The length of the tuple is equal to the number
75	of specialization constants found. Each
76	`SpecializationConstantInfo` object contains the following
77	attributes:
78
79	- `spec_id` (int): The specialization ID.
80	- `dtype` (str): A NumPy style string representing the data type.
81	- `itemsize` (int): The size of the specialization constant in
82	bytes.
83	- `name` (str): The variable name. If not preserved in the binary,
84	a default name in the format `unnamed_spec_const_{spec_id}` is
85	used.
86	- `default_value` (int \| float \| bool \| None): The default value of
87	the specialization constant. If not specified, `None` is used.
88	"""
89	words = np.frombuffer(spv_bytes, dtype=np.uint32)	1✔
90
91	# verify magic number
92	if len(words) < 5 or words[0] != 0x07230203:	1!
NEW 93	raise ValueError("Invalid SPIR-V binary")	×
94
95	types = {}	1✔
96	ids = {}	1✔
97	names = {}	1✔
98	constants = {}	1✔
99	defaults = {}	1✔
100
101	i = 5 # skip 5 word header	1✔
102	while i < len(words):	1!
103	word = words[i]	1✔
104	opcode = word & 0xFFFF	1✔
105	word_count = word >> 16	1✔
106
107	if word_count == 0:	1!
NEW 108	raise ValueError(f"Invalid SPIR-V instruction at word index {i}")	×
109	if i + word_count > len(words):	1!
NEW 110	raise ValueError(	×
111	f"Invalid SPIR-V instruction at offset {i} (extends beyond "
112	"buffer)"
113	)
114
115	if opcode == SpirvOpCode.OpFunction:	1✔
116	# everything following is not relevant to specialization constant
117	# parsing, so we can stop parsing at this point
118	break	1✔
119	elif opcode == SpirvOpCode.OpTypeBool:	1✔
120	if word_count < 2:	1!
NEW 121	raise ValueError(f"Truncated OpTypeBool at word index {i}")	×
122	result_id = int(words[i + 1])	1✔
123	types[result_id] = {"dtype": "?", "itemsize": 1}	1✔
124	elif opcode == SpirvOpCode.OpTypeInt:	1✔
125	if word_count < 4:	1!
NEW 126	raise ValueError(f"Truncated OpTypeInt at word index {i}")	×
127	result_id = int(words[i + 1])	1✔
128	width = int(words[i + 2])	1✔
129	signed = int(words[i + 3])	1✔
130	prefix = "i" if signed else "u"	1✔
131	types[result_id] = {	1✔
132	"dtype": f"{prefix}{width // 8}",
133	"itemsize": width // 8,
134	}
135	elif opcode == SpirvOpCode.OpTypeFloat:	1✔
136	if word_count < 3:	1!
NEW 137	raise ValueError(f"Truncated OpTypeFloat at word index {i}")	×
138	result_id = int(words[i + 1])	1✔
139	width = int(words[i + 2])	1✔
140	types[result_id] = {	1✔
141	"dtype": f"f{width // 8}",
142	"itemsize": width // 8,
143	}
144	elif opcode == SpirvOpCode.OpSpecConstant:	1✔
145	if word_count < 4:	1!
NEW 146	raise ValueError(f"Truncated OpSpecConstant at word index {i}")	×
147	type_id = int(words[i + 1])	1✔
148	result_id = int(words[i + 2])	1✔
149	constants[result_id] = type_id	1✔
150	literal_words = words[i + 3 : i + word_count]	1✔
151	defaults[result_id] = literal_words.tobytes()	1✔
152	elif opcode == SpirvOpCode.OpSpecConstantTrue:	1!
NEW 153	if word_count < 3:	×
NEW 154	raise ValueError(	×
155	f"Truncated OpSpecConstantTrue at word index {i}"
156	)
NEW 157	type_id = int(words[i + 1])	×
NEW 158	result_id = int(words[i + 2])	×
NEW 159	constants[result_id] = type_id	×
NEW 160	defaults[result_id] = True	×
161	elif opcode == SpirvOpCode.OpSpecConstantFalse:	1!
NEW 162	if word_count < 3:	×
NEW 163	raise ValueError(	×
164	f"Truncated OpSpecConstantFalse at word index {i}"
165	)
NEW 166	type_id = int(words[i + 1])	×
NEW 167	result_id = int(words[i + 2])	×
NEW 168	constants[result_id] = type_id	×
NEW 169	defaults[result_id] = False	×
170	elif opcode == SpirvOpCode.OpDecorate:	1✔
171	if word_count < 3:	1!
NEW 172	raise ValueError(f"Truncated OpDecorate at word index {i}")	×
173	target_id = int(words[i + 1])	1✔
174	decoration = int(words[i + 2])	1✔
175	if decoration == SpirvDecoration.SpecId:	1✔
176	if word_count < 4:	1!
NEW 177	raise ValueError(	×
178	f"Truncated OpDecorate SpecId at word index {i}"
179	)
180	ids[target_id] = int(words[i + 3])	1✔
181	elif opcode == SpirvOpCode.OpName:	1✔
182	if word_count < 2:	1!
NEW 183	raise ValueError(f"Truncated OpName at word index {i}")	×
184	target_id = int(words[i + 1])	1✔
185	name_bytes = words[i + 2 : i + word_count].tobytes()	1✔
186	try:	1✔
187	names[target_id] = name_bytes.split(b"\x00", 1)[0].decode(	1✔
188	"utf-8"
189	)
NEW 190	except UnicodeDecodeError:	×
NEW 191	raise ValueError(f"Invalid UTF-8 in OpName at word index {i}")	×
192
193	i += word_count	1✔
194
195	# a spec ID may appear multiple times in the same binary with different
196	# target IDs. We only need to keep one, so skip duplicates
197	unique_ids = set()	1✔
198	result = []	1✔
199	for target_id, spec_id in ids.items():	1✔
200	if spec_id in unique_ids:	1✔
201	continue	1✔
202	unique_ids.add(spec_id)	1✔
203	type_id = constants.get(target_id)	1✔
204	type_info = types.get(type_id, {"dtype": "unknown_type", "itemsize": 0})	1✔
205	name = names.get(target_id, f"unnamed_spec_const_{spec_id}")	1✔
206
207	dtype_str = type_info["dtype"]	1✔
208	raw_default = defaults.get(target_id)	1✔
209	default_value = None	1✔
210	if isinstance(raw_default, bool):	1!
NEW 211	default_value = raw_default	×
212	elif isinstance(raw_default, bytes) and dtype_str != "unknown_type":	1!
213	try:	1✔
214	default_value = np.frombuffer(raw_default, dtype=dtype_str)[	1✔
215	0
216	].item()
NEW 217	except (ValueError, TypeError):	×
NEW 218	default_value = None	×
219
220	result.append(	1✔
221	SpecializationConstantInfo(
222	spec_id=spec_id,
223	dtype=dtype_str,
224	name=name,
225	itemsize=type_info["itemsize"],
226	default_value=default_value,
227	)
228	)
229
230	return tuple(sorted(result, key=lambda x: x.spec_id))	1✔

IntelPython / dpctl / 28947730040

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous