25360310698

Committed 05 May 2026 05:55AM UTC coverage: 75.3%. First build

Build # 25360310698

Build Type

Pull #2304

github

Committed by

web-flow

Commit Message

Merge 17fff5a96 into 2b2056d35

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

Coverage Stats

858 of 1198 branches covered (71.62%)

Branch coverage included in aggregate %.

102 of 136 new or added lines in 4 files covered. (75.0%)

3288 of 4308 relevant lines covered (76.32%)

266.01 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):
    """SPIR-V operation codes used in the dpctl.program module."""

    OpName = 5
    OpTypeBool = 20
    OpTypeInt = 21
    OpTypeFloat = 22
    OpSpecConstantTrue = 48
    OpSpecConstantFalse = 49
    OpSpecConstant = 50
    OpFunction = 54
    OpDecorate = 71


class SpirvDecoration(IntEnum):
    """SPIR-V decoration codes used in the dpctl.program module."""

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

IntelPython / dpctl / 25360310698

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