27977445675

Committed 22 Jun 2026 07:11PM UTC coverage: 75.406% (-0.3%) from 75.677%

Build # 27977445675

Build Type

Pull #2304

github

Committed by

web-flow

Commit Message

Merge f7daf456f into 61d1fc0cd

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

Coverage Stats

869 of 1218 branches covered (71.35%)

Branch coverage included in aggregate %.

113 of 157 new or added lines in 4 files covered. (71.97%)

110 existing lines in 1 file now uncovered.

3356 of 4385 relevant lines covered (76.53%)

266.27 hits per line

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

83.46

/dpctl/program/utils/_utils.py

#                      Data Parallel Control (dpctl)
#
# Copyright 2020-2025 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:
            result_id = int(words[i + 1])
            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] = {
                "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] = {
                "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])
            if decoration == SpirvDecoration.SpecId:
                ids[target_id] = int(words[i + 3])
        elif opcode == SpirvOpCode.OpName:
            target_id = int(words[i + 1])
            name_bytes = words[i + 2 : i + word_count].tobytes()
            names[target_id] = name_bytes.split(b"\x00", 1)[0].decode("utf-8")

        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 2020-2025 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	result_id = int(words[i + 1])	1✔
121	types[result_id] = {"dtype": "?", "itemsize": 1}	1✔
122	elif opcode == SpirvOpCode.OpTypeInt:	1✔
123	result_id = int(words[i + 1])	1✔
124	width = int(words[i + 2])	1✔
125	signed = int(words[i + 3])	1✔
126	prefix = "i" if signed else "u"	1✔
127	types[result_id] = {	1✔
128	"dtype": f"{prefix}{width // 8}",
129	"itemsize": width // 8,
130	}
131	elif opcode == SpirvOpCode.OpTypeFloat:	1✔
132	result_id = int(words[i + 1])	1✔
133	width = int(words[i + 2])	1✔
134	types[result_id] = {	1✔
135	"dtype": f"f{width // 8}",
136	"itemsize": width // 8,
137	}
138	elif opcode == SpirvOpCode.OpSpecConstant:	1✔
139	type_id = int(words[i + 1])	1✔
140	result_id = int(words[i + 2])	1✔
141	constants[result_id] = type_id	1✔
142	literal_words = words[i + 3 : i + word_count]	1✔
143	defaults[result_id] = literal_words.tobytes()	1✔
144	elif opcode == SpirvOpCode.OpSpecConstantTrue:	1!
NEW 145	type_id = int(words[i + 1])	×
NEW 146	result_id = int(words[i + 2])	×
NEW 147	constants[result_id] = type_id	×
NEW 148	defaults[result_id] = True	×
149	elif opcode == SpirvOpCode.OpSpecConstantFalse:	1!
NEW 150	type_id = int(words[i + 1])	×
NEW 151	result_id = int(words[i + 2])	×
NEW 152	constants[result_id] = type_id	×
NEW 153	defaults[result_id] = False	×
154	elif opcode == SpirvOpCode.OpDecorate:	1✔
155	target_id = int(words[i + 1])	1✔
156	decoration = int(words[i + 2])	1✔
157	if decoration == SpirvDecoration.SpecId:	1✔
158	ids[target_id] = int(words[i + 3])	1✔
159	elif opcode == SpirvOpCode.OpName:	1✔
160	target_id = int(words[i + 1])	1✔
161	name_bytes = words[i + 2 : i + word_count].tobytes()	1✔
162	names[target_id] = name_bytes.split(b"\x00", 1)[0].decode("utf-8")	1✔
163
164	i += word_count	1✔
165
166	# a spec ID may appear multiple times in the same binary with different
167	# target IDs. We only need to keep one, so skip duplicates
168	unique_ids = set()	1✔
169	result = []	1✔
170	for target_id, spec_id in ids.items():	1✔
171	if spec_id in unique_ids:	1✔
172	continue	1✔
173	unique_ids.add(spec_id)	1✔
174	type_id = constants.get(target_id)	1✔
175	type_info = types.get(type_id, {"dtype": "unknown_type", "itemsize": 0})	1✔
176	name = names.get(target_id, f"unnamed_spec_const_{spec_id}")	1✔
177
178	dtype_str = type_info["dtype"]	1✔
179	raw_default = defaults.get(target_id)	1✔
180	default_value = None	1✔
181	if isinstance(raw_default, bool):	1!
NEW 182	default_value = raw_default	×
183	elif isinstance(raw_default, bytes) and dtype_str != "unknown_type":	1!
184	try:	1✔
185	default_value = np.frombuffer(raw_default, dtype=dtype_str)[	1✔
186	0
187	].item()
NEW 188	except (ValueError, TypeError):	×
NEW 189	default_value = None	×
190
191	result.append(	1✔
192	SpecializationConstantInfo(
193	spec_id=spec_id,
194	dtype=dtype_str,
195	name=name,
196	itemsize=type_info["itemsize"],
197	default_value=default_value,
198	)
199	)
200
201	return tuple(sorted(result, key=lambda x: x.spec_id))	1✔

IntelPython / dpctl / 27977445675

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