Skip to content

Commit 62799d0

Browse files
khyperiakhyperia
authored
Make BuilderSpirv::has_capability and has_extension O(1) (#688)
1 parent ca53cbe commit 62799d0

File tree

5 files changed

+99
-49
lines changed

5 files changed

+99
-49
lines changed

‎crates/rustc_codegen_spirv/src/builder/ext_inst.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,12 @@ impl ExtInst {
5555
}
5656
if !bx
5757
.builder
58-
.has_extension("SPV_INTEL_shader_integer_functions2")
58+
.has_extension(bx.sym.spv_intel_shader_integer_functions2)
5959
{
60-
bx.zombie(to_zombie, "extension IntegerFunctions2INTEL is required");
60+
bx.zombie(
61+
to_zombie,
62+
"extension SPV_INTEL_shader_integer_functions2 is required",
63+
);
6164
}
6265
}
6366
}

‎crates/rustc_codegen_spirv/src/builder_spirv.rs

Lines changed: 75 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
use crate::builder;
22
use crate::codegen_cx::CodegenCx;
33
use crate::spirv_type::SpirvType;
4+
use crate::symbols::Symbols;
45
use crate::target::SpirvTarget;
56
use crate::target_feature::TargetFeature;
67
use rspirv::dr::{Block, Builder, Module, Operand};
78
use rspirv::spirv::{AddressingModel, Capability, MemoryModel, Op, StorageClass, Word};
89
use rspirv::{binary::Assemble, binary::Disassemble};
9-
use rustc_data_structures::fx::FxHashMap;
10+
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
1011
use rustc_middle::bug;
12+
use rustc_span::symbol::Symbol;
1113
use rustc_span::{Span, DUMMY_SP};
1214
use std::cell::{RefCell, RefMut};
1315
use std::rc::Rc;
@@ -319,52 +321,111 @@ pub struct BuilderSpirv {
319321
const_to_id: RefCell<FxHashMap<WithType<SpirvConst>, WithConstLegality<Word>>>,
320322
id_to_const: RefCell<FxHashMap<Word, WithConstLegality<SpirvConst>>>,
321323
string_cache: RefCell<FxHashMap<String, Word>>,
324+
325+
enabled_capabilities: FxHashSet<Capability>,
326+
enabled_extensions: FxHashSet<Symbol>,
322327
}
323328

324329
impl BuilderSpirv {
325-
pub fn new(target: &SpirvTarget, features: &[TargetFeature]) -> Self {
330+
pub fn new(
331+
sym: &Symbols,
332+
target: &SpirvTarget,
333+
features: &[TargetFeature],
334+
bindless: bool,
335+
) -> Self {
326336
let version = target.spirv_version();
327337
let memory_model = target.memory_model();
328338

329339
let mut builder = Builder::new();
330340
builder.set_version(version.0, version.1);
331341

342+
let mut enabled_capabilities = FxHashSet::default();
343+
let mut enabled_extensions = FxHashSet::default();
344+
345+
fn add_cap(
346+
builder: &mut Builder,
347+
enabled_capabilities: &mut FxHashSet<Capability>,
348+
cap: Capability,
349+
) {
350+
// This should be the only callsite of Builder::capability (aside from tests), to make
351+
// sure the hashset stays in sync.
352+
builder.capability(cap);
353+
enabled_capabilities.insert(cap);
354+
}
355+
fn add_ext(builder: &mut Builder, enabled_extensions: &mut FxHashSet<Symbol>, ext: Symbol) {
356+
// This should be the only callsite of Builder::extension (aside from tests), to make
357+
// sure the hashset stays in sync.
358+
builder.extension(&*ext.as_str());
359+
enabled_extensions.insert(ext);
360+
}
361+
332362
for feature in features {
333-
match feature {
334-
TargetFeature::Capability(cap) => builder.capability(*cap),
335-
TargetFeature::Extension(ext) => builder.extension(&*ext.as_str()),
363+
match *feature {
364+
TargetFeature::Capability(cap) => {
365+
add_cap(&mut builder, &mut enabled_capabilities, cap);
366+
}
367+
TargetFeature::Extension(ext) => {
368+
add_ext(&mut builder, &mut enabled_extensions, ext);
369+
}
336370
}
337371
}
338372

339373
if target.is_kernel() {
340-
builder.capability(Capability::Kernel);
374+
add_cap(&mut builder, &mut enabled_capabilities, Capability::Kernel);
341375
} else {
342-
builder.capability(Capability::Shader);
376+
add_cap(&mut builder, &mut enabled_capabilities, Capability::Shader);
343377
if memory_model == MemoryModel::Vulkan {
344378
if version < (1, 5) {
345-
builder.extension("SPV_KHR_vulkan_memory_model");
379+
add_ext(
380+
&mut builder,
381+
&mut enabled_extensions,
382+
sym.spv_khr_vulkan_memory_model,
383+
);
346384
}
347-
builder.capability(Capability::VulkanMemoryModel);
385+
add_cap(
386+
&mut builder,
387+
&mut enabled_capabilities,
388+
Capability::VulkanMemoryModel,
389+
);
348390
}
349391
}
350392

351393
// The linker will always be ran on this module
352-
builder.capability(Capability::Linkage);
394+
add_cap(&mut builder, &mut enabled_capabilities, Capability::Linkage);
353395

354396
let addressing_model = if target.is_kernel() {
355-
builder.capability(Capability::Addresses);
397+
add_cap(
398+
&mut builder,
399+
&mut enabled_capabilities,
400+
Capability::Addresses,
401+
);
356402
AddressingModel::Physical32
357403
} else {
358404
AddressingModel::Logical
359405
};
360406

361407
builder.memory_model(addressing_model, memory_model);
362408

409+
if bindless {
410+
add_ext(
411+
&mut builder,
412+
&mut enabled_extensions,
413+
sym.spv_ext_descriptor_indexing,
414+
);
415+
add_cap(
416+
&mut builder,
417+
&mut enabled_capabilities,
418+
Capability::RuntimeDescriptorArray,
419+
);
420+
}
421+
363422
Self {
364423
builder: RefCell::new(builder),
365424
const_to_id: Default::default(),
366425
id_to_const: Default::default(),
367426
string_cache: Default::default(),
427+
enabled_capabilities,
428+
enabled_extensions,
368429
}
369430
}
370431

@@ -410,27 +471,11 @@ impl BuilderSpirv {
410471
}
411472

412473
pub fn has_capability(&self, capability: Capability) -> bool {
413-
self.builder
414-
.borrow()
415-
.module_ref()
416-
.capabilities
417-
.iter()
418-
.any(|inst| {
419-
inst.class.opcode == Op::Capability
420-
&& inst.operands[0].unwrap_capability() == capability
421-
})
474+
self.enabled_capabilities.contains(&capability)
422475
}
423476

424-
pub fn has_extension(&self, extension: &str) -> bool {
425-
self.builder
426-
.borrow()
427-
.module_ref()
428-
.extensions
429-
.iter()
430-
.any(|inst| {
431-
inst.class.opcode == Op::Extension
432-
&& inst.operands[0].unwrap_literal_string() == extension
433-
})
477+
pub fn has_extension(&self, extension: Symbol) -> bool {
478+
self.enabled_extensions.contains(&extension)
434479
}
435480

436481
pub fn select_function_by_id(&self, id: Word) -> BuilderCursor {

‎crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

Lines changed: 9 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -312,23 +312,17 @@ impl<'tcx> CodegenCx<'tcx> {
312312
bx.call(entry_func, &call_args, None);
313313
bx.ret_void();
314314

315-
if self.bindless() {
316-
self.emit_global().extension("SPV_EXT_descriptor_indexing");
317-
self.emit_global()
318-
.capability(Capability::RuntimeDescriptorArray);
319-
320-
if self.target.spirv_version() > (1, 3) {
321-
let sets = self.bindless_descriptor_sets.borrow().unwrap();
315+
if self.bindless() && self.target.spirv_version() > (1, 3) {
316+
let sets = self.bindless_descriptor_sets.borrow().unwrap();
322317

323-
op_entry_point_interface_operands.push(sets.buffers);
318+
op_entry_point_interface_operands.push(sets.buffers);
324319

325-
//op_entry_point_interface_operands
326-
// .push(sets.sampled_image_1d);
327-
// op_entry_point_interface_operands
328-
// .push(sets.sampled_image_2d);
329-
//op_entry_point_interface_operands
330-
//.push(sets.sampled_image_3d);
331-
}
320+
//op_entry_point_interface_operands
321+
// .push(sets.sampled_image_1d);
322+
// op_entry_point_interface_operands
323+
// .push(sets.sampled_image_2d);
324+
//op_entry_point_interface_operands
325+
//.push(sets.sampled_image_3d);
332326
}
333327

334328
let stub_fn_id = stub_fn.def_cx(self);

‎crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ impl<'tcx> CodegenCx<'tcx> {
132132
let result = Self {
133133
tcx,
134134
codegen_unit,
135-
builder: BuilderSpirv::new(&target, &features),
135+
builder: BuilderSpirv::new(&sym, &target, &features, bindless),
136136
instances: Default::default(),
137137
function_parameter_values: Default::default(),
138138
type_cache: Default::default(),

‎crates/rustc_codegen_spirv/src/symbols.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ pub struct Symbols {
2121
pub libm: Symbol,
2222
pub num_traits: Symbol,
2323
pub entry_point_name: Symbol,
24+
pub spv_intel_shader_integer_functions2: Symbol,
25+
pub spv_khr_vulkan_memory_model: Symbol,
26+
pub spv_ext_descriptor_indexing: Symbol,
2427
descriptor_set: Symbol,
2528
binding: Symbol,
2629
input_attachment_index: Symbol,
@@ -366,11 +369,16 @@ impl Symbols {
366369
Self {
367370
fmt_decimal: Symbol::intern("fmt_decimal"),
368371

369-
entry_point_name: Symbol::intern("entry_point_name"),
370372
spirv: Symbol::intern("spirv"),
371373
spirv_std: Symbol::intern("spirv_std"),
372374
libm: Symbol::intern("libm"),
373375
num_traits: Symbol::intern("num_traits"),
376+
entry_point_name: Symbol::intern("entry_point_name"),
377+
spv_intel_shader_integer_functions2: Symbol::intern(
378+
"SPV_INTEL_shader_integer_functions2",
379+
),
380+
spv_khr_vulkan_memory_model: Symbol::intern("SPV_KHR_vulkan_memory_model"),
381+
spv_ext_descriptor_indexing: Symbol::intern("SPV_EXT_descriptor_indexing"),
374382
descriptor_set: Symbol::intern("descriptor_set"),
375383
binding: Symbol::intern("binding"),
376384
input_attachment_index: Symbol::intern("input_attachment_index"),

0 commit comments

Comments
 (0)