Skip to content

Commit 07c74d2

Browse files
39aliFirestar99
39ali
authored and
Firestar99
committed
add support for CooperativeMatrix
1 parent df22665 commit 07c74d2

11 files changed

Lines changed: 377 additions & 47 deletions

File tree

‎crates/rustc_codegen_spirv/src/abi.rs‎

Lines changed: 68 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
44
use crate::attr::{AggregatedSpirvAttributes, IntrinsicType};
55
use crate::codegen_cx::CodegenCx;
6+
use crate::maybe_pqp_cg_ssa::traits::ConstCodegenMethods as _;
67
use crate::spirv_type::SpirvType;
78
use itertools::Itertools;
89
use rspirv::spirv::{Dim, ImageFormat, StorageClass, Word};
@@ -885,6 +886,48 @@ fn trans_intrinsic_type<'tcx>(
885886
args: GenericArgsRef<'tcx>,
886887
intrinsic_type_attr: IntrinsicType,
887888
) -> Result<Word, ErrorGuaranteed> {
889+
trait FromScalarInt: Sized {
890+
fn from_scalar_int(n: ScalarInt) -> Option<Self>;
891+
}
892+
893+
impl FromScalarInt for u32 {
894+
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
895+
Some(n.try_to_bits(Size::from_bits(32)).ok()?.try_into().unwrap())
896+
}
897+
}
898+
899+
impl FromScalarInt for Dim {
900+
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
901+
Dim::from_u32(u32::from_scalar_int(n)?)
902+
}
903+
}
904+
905+
impl FromScalarInt for ImageFormat {
906+
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
907+
ImageFormat::from_u32(u32::from_scalar_int(n)?)
908+
}
909+
}
910+
911+
fn const_int_value<'tcx, P: FromScalarInt>(
912+
cx: &CodegenCx<'tcx>,
913+
const_: Const<'tcx>,
914+
) -> Result<P, ErrorGuaranteed> {
915+
let ty::Value {
916+
ty: const_ty,
917+
valtree: const_val,
918+
} = const_.to_value();
919+
assert!(const_ty.is_integral());
920+
const_val
921+
.try_to_scalar()
922+
.and_then(|scalar| scalar.try_to_scalar_int().ok())
923+
.and_then(P::from_scalar_int)
924+
.ok_or_else(|| {
925+
cx.tcx
926+
.dcx()
927+
.err(format!("invalid value for const generic: {const_}"))
928+
})
929+
}
930+
888931
match intrinsic_type_attr {
889932
IntrinsicType::GenericImageType => {
890933
// see SpirvType::sizeof
@@ -948,48 +991,6 @@ fn trans_intrinsic_type<'tcx>(
948991
// let image_format: spirv::ImageFormat =
949992
// type_from_variant_discriminant(cx, args.const_at(6));
950993

951-
trait FromScalarInt: Sized {
952-
fn from_scalar_int(n: ScalarInt) -> Option<Self>;
953-
}
954-
955-
impl FromScalarInt for u32 {
956-
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
957-
Some(n.try_to_bits(Size::from_bits(32)).ok()?.try_into().unwrap())
958-
}
959-
}
960-
961-
impl FromScalarInt for Dim {
962-
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
963-
Dim::from_u32(u32::from_scalar_int(n)?)
964-
}
965-
}
966-
967-
impl FromScalarInt for ImageFormat {
968-
fn from_scalar_int(n: ScalarInt) -> Option<Self> {
969-
ImageFormat::from_u32(u32::from_scalar_int(n)?)
970-
}
971-
}
972-
973-
fn const_int_value<'tcx, P: FromScalarInt>(
974-
cx: &CodegenCx<'tcx>,
975-
const_: Const<'tcx>,
976-
) -> Result<P, ErrorGuaranteed> {
977-
let ty::Value {
978-
ty: const_ty,
979-
valtree: const_val,
980-
} = const_.to_value();
981-
assert!(const_ty.is_integral());
982-
const_val
983-
.try_to_scalar()
984-
.and_then(|scalar| scalar.try_to_scalar_int().ok())
985-
.and_then(P::from_scalar_int)
986-
.ok_or_else(|| {
987-
cx.tcx
988-
.dcx()
989-
.err(format!("invalid value for Image const generic: {const_}"))
990-
})
991-
}
992-
993994
let dim = const_int_value(cx, args.const_at(1))?;
994995
let depth = const_int_value(cx, args.const_at(2))?;
995996
let arrayed = const_int_value(cx, args.const_at(3))?;
@@ -1019,6 +1020,31 @@ fn trans_intrinsic_type<'tcx>(
10191020
Ok(SpirvType::AccelerationStructureKhr.def(span, cx))
10201021
}
10211022
IntrinsicType::RayQueryKhr => Ok(SpirvType::RayQueryKhr.def(span, cx)),
1023+
IntrinsicType::CooperativeMatrixKhr => {
1024+
if ty.size != Size::from_bytes(4) {
1025+
return Err(cx.tcx.dcx().err("cooperative_matrix type must have size 4"));
1026+
}
1027+
1028+
// Generic arg 0: component type T
1029+
let component_type = cx.layout_of(args.type_at(0)).spirv_type(span, cx);
1030+
// Const generic 1: USE (MatrixA=0, MatrixB=1, MatrixAccumulator=2)
1031+
let usage = cx.const_u32(const_int_value(cx, args.const_at(1))?).def_cx(cx);
1032+
// Const generic 2: ROWS
1033+
let rows = cx.const_u32(const_int_value(cx, args.const_at(2))?).def_cx(cx);
1034+
// Const generic 3: COLS
1035+
let columns = cx.const_u32(const_int_value(cx, args.const_at(3))?).def_cx(cx);
1036+
// Scope: Subgroup = 3
1037+
let scope = cx.const_u32(3).def_cx(cx);
1038+
1039+
Ok(SpirvType::CooperativeMatrixKhr {
1040+
component_type,
1041+
usage,
1042+
rows,
1043+
columns,
1044+
scope,
1045+
}
1046+
.def(span, cx))
1047+
}
10221048
IntrinsicType::SampledImage => {
10231049
// see SpirvType::sizeof
10241050
if ty.size != Size::from_bytes(4) {

‎crates/rustc_codegen_spirv/src/attr.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ pub enum IntrinsicType {
6969
TypedBuffer,
7070
Matrix,
7171
Vector,
72+
CooperativeMatrixKhr,
7273
}
7374

7475
#[derive(Copy, Clone, Debug, PartialEq, Eq)]

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
454454
self.fatal("cannot memset acceleration structure")
455455
}
456456
SpirvType::RayQueryKhr => self.fatal("cannot memset ray query"),
457+
SpirvType::CooperativeMatrixKhr { .. } => {
458+
self.fatal("cannot memset cooperative matrix")
459+
}
457460
}
458461
}
459462

@@ -511,6 +514,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
511514
self.fatal("cannot memset acceleration structure")
512515
}
513516
SpirvType::RayQueryKhr => self.fatal("cannot memset ray query"),
517+
SpirvType::CooperativeMatrixKhr { .. } => {
518+
self.fatal("cannot memset cooperative matrix")
519+
}
514520
}
515521
}
516522

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -623,7 +623,8 @@ impl<'tcx> CodegenCx<'tcx> {
623623
| SpirvType::SampledImage { .. }
624624
| SpirvType::InterfaceBlock { .. }
625625
| SpirvType::AccelerationStructureKhr
626-
| SpirvType::RayQueryKhr => {
626+
| SpirvType::RayQueryKhr
627+
| SpirvType::CooperativeMatrixKhr { .. } => {
627628
let result = self.undef(ty);
628629
self.zombie_no_span(
629630
result.def_cx(self),

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ impl BaseTypeCodegenMethods for CodegenCx<'_> {
209209
| SpirvType::SampledImage { .. }
210210
| SpirvType::AccelerationStructureKhr
211211
| SpirvType::RayQueryKhr
212+
| SpirvType::CooperativeMatrixKhr { .. }
212213
=> TypeKind::Token,
213214
}
214215
}

‎crates/rustc_codegen_spirv/src/spirv_type.rs‎

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,13 @@ pub enum SpirvType<'tcx> {
9090

9191
AccelerationStructureKhr,
9292
RayQueryKhr,
93+
CooperativeMatrixKhr {
94+
component_type: Word,
95+
usage: Word,
96+
rows: Word,
97+
columns: Word,
98+
scope: Word,
99+
},
93100
}
94101

95102
impl SpirvType<'_> {
@@ -193,6 +200,20 @@ impl SpirvType<'_> {
193200
cx.emit_global().type_acceleration_structure_khr_id(id)
194201
}
195202
Self::RayQueryKhr => cx.emit_global().type_ray_query_khr_id(id),
203+
Self::CooperativeMatrixKhr {
204+
component_type,
205+
scope,
206+
rows,
207+
columns,
208+
usage: use_,
209+
} => cx.emit_global().type_cooperative_matrix_khr_id(
210+
id,
211+
component_type,
212+
scope,
213+
rows,
214+
columns,
215+
use_,
216+
),
196217
Self::SampledImage { image_type } => {
197218
cx.emit_global().type_sampled_image_id(id, image_type)
198219
}
@@ -339,6 +360,7 @@ impl SpirvType<'_> {
339360
Self::Image { .. }
340361
| Self::AccelerationStructureKhr
341362
| Self::RayQueryKhr
363+
| Self::CooperativeMatrixKhr { .. }
342364
| Self::Sampler
343365
| Self::SampledImage { .. }
344366
| Self::InterfaceBlock { .. } => Size::from_bytes(4),
@@ -361,6 +383,7 @@ impl SpirvType<'_> {
361383
Self::Image { .. }
362384
| Self::AccelerationStructureKhr
363385
| Self::RayQueryKhr
386+
| Self::CooperativeMatrixKhr { .. }
364387
| Self::Sampler
365388
| Self::SampledImage { .. }
366389
| Self::InterfaceBlock { .. } => Align::from_bytes(4).unwrap(),
@@ -389,7 +412,10 @@ impl SpirvType<'_> {
389412
Self::InterfaceBlock { .. } | Self::RayQueryKhr | Self::SampledImage { .. } => None,
390413

391414
// Descriptor types
392-
Self::Image { .. } | Self::AccelerationStructureKhr | Self::Sampler => None,
415+
Self::Image { .. }
416+
| Self::AccelerationStructureKhr
417+
| Self::CooperativeMatrixKhr { .. }
418+
| Self::Sampler => None,
393419

394420
// Primitive types
395421
ty => ty.sizeof(cx),
@@ -455,6 +481,19 @@ impl SpirvType<'_> {
455481
SpirvType::InterfaceBlock { inner_type } => SpirvType::InterfaceBlock { inner_type },
456482
SpirvType::AccelerationStructureKhr => SpirvType::AccelerationStructureKhr,
457483
SpirvType::RayQueryKhr => SpirvType::RayQueryKhr,
484+
SpirvType::CooperativeMatrixKhr {
485+
component_type,
486+
scope,
487+
rows,
488+
columns,
489+
usage: use_,
490+
} => SpirvType::CooperativeMatrixKhr {
491+
component_type,
492+
scope,
493+
rows,
494+
columns,
495+
usage: use_,
496+
},
458497

459498
// Only these variants have any slices to arena-allocate.
460499
SpirvType::Adt {
@@ -644,6 +683,20 @@ impl fmt::Debug for SpirvTypePrinter<'_, '_> {
644683
.finish(),
645684
SpirvType::AccelerationStructureKhr => f.debug_struct("AccelerationStructure").finish(),
646685
SpirvType::RayQueryKhr => f.debug_struct("RayQuery").finish(),
686+
SpirvType::CooperativeMatrixKhr {
687+
component_type,
688+
scope,
689+
rows,
690+
columns,
691+
usage: use_,
692+
} => f
693+
.debug_struct("CooperativeMatrix")
694+
.field("component_type", &self.cx.debug_type(component_type))
695+
.field("scope", &self.cx.debug_type(scope))
696+
.field("rows", &self.cx.debug_type(rows))
697+
.field("columns", &self.cx.debug_type(columns))
698+
.field("use_", &self.cx.debug_type(use_))
699+
.finish(),
647700
};
648701
{
649702
let mut debug_stack = DEBUG_STACK.lock().unwrap();
@@ -797,6 +850,7 @@ impl SpirvTypePrinter<'_, '_> {
797850
}
798851
SpirvType::AccelerationStructureKhr => f.write_str("AccelerationStructureKhr"),
799852
SpirvType::RayQueryKhr => f.write_str("RayQuery"),
853+
SpirvType::CooperativeMatrixKhr { .. } => f.write_str("CooperativeMatrixKhr"),
800854
}
801855
}
802856
}

‎crates/rustc_codegen_spirv/src/spirv_type_constraints.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1086,10 +1086,10 @@ pub fn instruction_signatures(op: Op) -> Option<&'static [InstSig<'static>]> {
10861086
}
10871087
// SPV_KHR_cooperative_matrix
10881088
Op::TypeCooperativeMatrixKHR
1089-
| Op::CooperativeMatrixLoadKHR
1089+
| Op::CooperativeMatrixLengthKHR
10901090
| Op::CooperativeMatrixStoreKHR
1091-
| Op::CooperativeMatrixMulAddKHR
1092-
| Op::CooperativeMatrixLengthKHR => reserved!(SPV_KHR_cooperative_matrix),
1091+
| Op::CooperativeMatrixLoadKHR => {}
1092+
Op::CooperativeMatrixMulAddKHR => sig! { (_, _, T) -> T },
10931093
// SPV_QCOM_image_processing
10941094
Op::ImageSampleWeightedQCOM
10951095
| Op::ImageBoxFilterQCOM

‎crates/rustc_codegen_spirv/src/symbols.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -356,6 +356,10 @@ impl Symbols {
356356
"ray_query",
357357
SpirvAttribute::IntrinsicType(IntrinsicType::RayQueryKhr),
358358
),
359+
(
360+
"cooperative_matrix",
361+
SpirvAttribute::IntrinsicType(IntrinsicType::CooperativeMatrixKhr),
362+
),
359363
("block", SpirvAttribute::Block),
360364
("flat", SpirvAttribute::Flat),
361365
("invariant", SpirvAttribute::Invariant),

0 commit comments

Comments
 (0)