25470260757

Committed 07 May 2026 01:12AM UTC coverage: 75.295%. First build

Build # 25470260757

Build Type

Pull #2304

github

Committed by

web-flow

Commit Message

Merge 97871d969 into 1a0a91064

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

Coverage Stats

858 of 1198 branches covered (71.62%)

Branch coverage included in aggregate %.

101 of 135 new or added lines in 4 files covered. (74.81%)

3287 of 4307 relevant lines covered (76.32%)

264.67 hits per line

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

85.6

/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


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 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, 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))

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	class SpirvOpCode(IntEnum):	1✔
26	OpName = 5	1✔
27	OpTypeBool = 20	1✔
28	OpTypeInt = 21	1✔
29	OpTypeFloat = 22	1✔
30	OpSpecConstantTrue = 48	1✔
31	OpSpecConstantFalse = 49	1✔
32	OpSpecConstant = 50	1✔
33	OpFunction = 54	1✔
34	OpDecorate = 71	1✔
35
36
37	class SpirvDecoration(IntEnum):	1✔
38	SpecId = 1	1✔
39
40
41	@dataclass(frozen=True)	1✔
42	class SpecializationConstantInfo:	1✔
43	"""Data class representing specialization constant information."""
44
45	spec_id: int
46	dtype: str
47	name: str
48	itemsize: int
49	default_value: int \| float \| bool \| None
50
51
52	def parse_spirv_specializations(	1✔
53	spv_bytes: bytes \| bytearray \| memoryview,
54	) -> tuple[SpecializationConstantInfo]:
55	"""
56	Parses SPIR-V byte stream to extract information about specializations,
57	including the specialization IDs, types, names, and default values.
58
59	Note that the dtype information may be imprecise, as the compiler may
60	choose to, for example, represent a bool as char, or may represent both
61	signed and unsigned integers as unsigned integer bit buckets of the same
62	length.
63
64	Args:
65	spv_bytes (bytes \| bytearray \| memoryview):
66	the SPIR-V byte stream.
67
68	Returns:
69	tuple[SpecializationConstantInfo]:
70	a tuple of parsed constants and their information represented by
71	`SpecializationConstantInfo` objects, sorted by their
72	specialization IDs. The length of the tuple is equal to the number
73	of specialization constants found. Each
74	`SpecializationConstantInfo` object contains the following
75	attributes:
76
77	- `spec_id` (int): The specialization ID.
78	- `dtype` (str): A NumPy style string representing the data type.
79	- `itemsize` (int): The size of the specialization constant in
80	bytes.
81	- `name` (str): The variable name. If not preserved in the binary,
82	a default name in the format `unnamed_spec_const_{spec_id}` is
83	used.
84	- `default_value` (int \| float \| bool \| None): The default value of
85	the specialization constant. If not specified, `None` is used.
86	"""
87	words = np.frombuffer(spv_bytes, dtype=np.uint32)	1✔
88
89	# verify magic number
90	if len(words) < 5 or words[0] != 0x07230203:	1!
NEW 91	raise ValueError("Invalid SPIR-V binary")	×
92
93	types = {}	1✔
94	ids = {}	1✔
95	names = {}	1✔
96	constants = {}	1✔
97	defaults = {}	1✔
98
99	i = 5 # skip 5 word header	1✔
100	while i < len(words):	1!
101	word = words[i]	1✔
102	opcode = word & 0xFFFF	1✔
103	word_count = word >> 16	1✔
104
105	if word_count == 0:	1!
NEW 106	raise ValueError(f"Invalid SPIR-V instruction at word index {i}")	×
107
108	if opcode == SpirvOpCode.OpFunction:	1✔
109	# everything following is not relevant to specialization constant
110	# parsing, so we can stop parsing at this point
111	break	1✔
112	elif opcode == SpirvOpCode.OpTypeBool:	1✔
113	result_id = int(words[i + 1])	1✔
114	types[result_id] = {"dtype": "?", "itemsize": 1}	1✔
115	elif opcode == SpirvOpCode.OpTypeInt:	1✔
116	result_id = int(words[i + 1])	1✔
117	width = int(words[i + 2])	1✔
118	signed = int(words[i + 3])	1✔
119	prefix = "i" if signed else "u"	1✔
120	types[result_id] = {	1✔
121	"dtype": f"{prefix}{width // 8}",
122	"itemsize": width // 8,
123	}
124	elif opcode == SpirvOpCode.OpTypeFloat:	1✔
125	result_id = int(words[i + 1])	1✔
126	width = int(words[i + 2])	1✔
127	types[result_id] = {	1✔
128	"dtype": f"f{width // 8}",
129	"itemsize": width // 8,
130	}
131	elif opcode == SpirvOpCode.OpSpecConstant:	1✔
132	type_id = int(words[i + 1])	1✔
133	result_id = int(words[i + 2])	1✔
134	constants[result_id] = type_id	1✔
135	literal_words = words[i + 3 : i + word_count]	1✔
136	defaults[result_id] = literal_words.tobytes()	1✔
137	elif opcode == SpirvOpCode.OpSpecConstantTrue:	1!
NEW 138	type_id = int(words[i + 1])	×
NEW 139	result_id = int(words[i + 2])	×
NEW 140	constants[result_id] = type_id	×
NEW 141	defaults[result_id] = True	×
142	elif opcode == SpirvOpCode.OpSpecConstantFalse:	1!
NEW 143	type_id = int(words[i + 1])	×
NEW 144	result_id = int(words[i + 2])	×
NEW 145	constants[result_id] = type_id	×
NEW 146	defaults[result_id] = False	×
147	elif opcode == SpirvOpCode.OpDecorate:	1✔
148	target_id = int(words[i + 1])	1✔
149	decoration = int(words[i + 2])	1✔
150	if decoration == SpirvDecoration.SpecId:	1✔
151	ids[target_id] = int(words[i + 3])	1✔
152	elif opcode == SpirvOpCode.OpName:	1✔
153	target_id = int(words[i + 1])	1✔
154	name_bytes = words[i + 2 : i + word_count].tobytes()	1✔
155	names[target_id] = name_bytes.split(b"\x00", 1)[0].decode("utf-8")	1✔
156
157	i += word_count	1✔
158
159	# a spec ID may appear multiple times in the same binary with different
160	# target IDs. We only need to keep one, so skip duplicates
161	unique_ids = set()	1✔
162	result = []	1✔
163	for target_id, spec_id in ids.items():	1✔
164	if spec_id in unique_ids:	1✔
165	continue	1✔
166	unique_ids.add(spec_id)	1✔
167	type_id = constants.get(target_id)	1✔
168	type_info = types.get(type_id, {"dtype": "unknown_type", "itemsize": 0})	1✔
169	name = names.get(target_id, f"unnamed_spec_const_{spec_id}")	1✔
170
171	dtype_str = type_info["dtype"]	1✔
172	raw_default = defaults.get(target_id)	1✔
173	default_value = None	1✔
174	if isinstance(raw_default, bytes):	1!
175	try:	1✔
176	default_value = np.frombuffer(raw_default, dtype=dtype_str)[	1✔
177	0
178	].item()
NEW 179	except Exception:	×
NEW 180	default_value = None	×
181
182	result.append(	1✔
183	SpecializationConstantInfo(
184	spec_id=spec_id,
185	dtype=dtype_str,
186	name=name,
187	itemsize=type_info["itemsize"],
188	default_value=default_value,
189	)
190	)
191
192	return tuple(sorted(result, key=lambda x: x.spec_id))	1✔

IntelPython / dpctl / 25470260757

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