Skip to content

Commit 8969965

Browse files
robamlerteoxoy
robamler
authored and
teoxoy
committed
[naga] Vectorize [un]pack4x{I, U}8[Clamp] on msl
Implements more direct conversions between 32-bit integers and 4x8-bit integer vectors using bit casting to/from `packed_[u]char4` when on MSL 2.1+ (older versions of MSL don't seem to support these bit casts). - `unpack4x{I, U}8(x)` becomes `[u]int4(as_type<packed_[u]char4>(x))`; - `pack4x{I, U}8(x)` becomes `as_type<uint>(packed_[u]char4(x))`; and - `pack4x{I, U}8Clamp(x)` becomes `as_type<uint>(packed_uchar4(metal::clamp(x, 0, 255)))`. These bit casts match the WGSL spec for these functions because Metal runs on little-endian machines.
1 parent b32eb4a commit 8969965

File tree

4 files changed

+303
-47
lines changed

4 files changed

+303
-47
lines changed

‎naga/src/back/msl/writer.rs

Lines changed: 93 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1524,6 +1524,58 @@ impl<W: Write> Writer<W> {
15241524
Ok(())
15251525
}
15261526

1527+
/// Emit code for the WGSL functions `pack4x{I, U}8[Clamp]`.
1528+
fn put_pack4x8(
1529+
&mut self,
1530+
arg: Handle<crate::Expression>,
1531+
context: &ExpressionContext<'_>,
1532+
was_signed: bool,
1533+
clamp_bounds: Option<(&str, &str)>,
1534+
) -> Result<(), Error> {
1535+
let write_arg = |this: &mut Self| -> BackendResult {
1536+
if let Some((min, max)) = clamp_bounds {
1537+
// Clamping with scalar bounds works (component-wise) even for packed_[u]char4.
1538+
write!(this.out, "{NAMESPACE}::clamp(")?;
1539+
this.put_expression(arg, context, true)?;
1540+
write!(this.out, ", {min}, {max})")?;
1541+
} else {
1542+
this.put_expression(arg, context, true)?;
1543+
}
1544+
Ok(())
1545+
};
1546+
1547+
if context.lang_version >= (2, 1) {
1548+
let packed_type = if was_signed {
1549+
"packed_char4"
1550+
} else {
1551+
"packed_uchar4"
1552+
};
1553+
// Metal uses little endian byte order, which matches what WGSL expects here.
1554+
write!(self.out, "as_type<uint>({packed_type}(")?;
1555+
write_arg(self)?;
1556+
write!(self.out, "))")?;
1557+
} else {
1558+
// MSL < 2.1 doesn't support `as_type` casting between packed chars and scalars.
1559+
if was_signed {
1560+
write!(self.out, "uint(")?;
1561+
}
1562+
write!(self.out, "(")?;
1563+
write_arg(self)?;
1564+
write!(self.out, "[0] & 0xFF) | ((")?;
1565+
write_arg(self)?;
1566+
write!(self.out, "[1] & 0xFF) << 8) | ((")?;
1567+
write_arg(self)?;
1568+
write!(self.out, "[2] & 0xFF) << 16) | ((")?;
1569+
write_arg(self)?;
1570+
write!(self.out, "[3] & 0xFF) << 24)")?;
1571+
if was_signed {
1572+
write!(self.out, ")")?;
1573+
}
1574+
}
1575+
1576+
Ok(())
1577+
}
1578+
15271579
/// Emit code for the isign expression.
15281580
///
15291581
fn put_isign(
@@ -2490,53 +2542,41 @@ impl<W: Write> Writer<W> {
24902542
write!(self.out, "{fun_name}")?;
24912543
self.put_call_parameters(iter::once(arg), context)?;
24922544
}
2493-
fun @ (Mf::Pack4xI8 | Mf::Pack4xU8 | Mf::Pack4xI8Clamp | Mf::Pack4xU8Clamp) => {
2494-
let was_signed = matches!(fun, Mf::Pack4xI8 | Mf::Pack4xI8Clamp);
2495-
let clamp_bounds = match fun {
2496-
Mf::Pack4xI8Clamp => Some(("-128", "127")),
2497-
Mf::Pack4xU8Clamp => Some(("0", "255")),
2498-
_ => None,
2499-
};
2500-
if was_signed {
2501-
write!(self.out, "uint(")?;
2502-
}
2503-
let write_arg = |this: &mut Self| -> BackendResult {
2504-
if let Some((min, max)) = clamp_bounds {
2505-
write!(this.out, "{NAMESPACE}::clamp(")?;
2506-
this.put_expression(arg, context, true)?;
2507-
write!(this.out, ", {min}, {max})")?;
2508-
} else {
2509-
this.put_expression(arg, context, true)?;
2510-
}
2511-
Ok(())
2512-
};
2513-
write!(self.out, "(")?;
2514-
write_arg(self)?;
2515-
write!(self.out, "[0] & 0xFF) | ((")?;
2516-
write_arg(self)?;
2517-
write!(self.out, "[1] & 0xFF) << 8) | ((")?;
2518-
write_arg(self)?;
2519-
write!(self.out, "[2] & 0xFF) << 16) | ((")?;
2520-
write_arg(self)?;
2521-
write!(self.out, "[3] & 0xFF) << 24)")?;
2522-
if was_signed {
2523-
write!(self.out, ")")?;
2524-
}
2545+
Mf::Pack4xI8 => self.put_pack4x8(arg, context, true, None)?,
2546+
Mf::Pack4xU8 => self.put_pack4x8(arg, context, false, None)?,
2547+
Mf::Pack4xI8Clamp => {
2548+
self.put_pack4x8(arg, context, true, Some(("-128", "127")))?
2549+
}
2550+
Mf::Pack4xU8Clamp => {
2551+
self.put_pack4x8(arg, context, false, Some(("0", "255")))?
25252552
}
25262553
fun @ (Mf::Unpack4xI8 | Mf::Unpack4xU8) => {
2527-
write!(self.out, "(")?;
2528-
if matches!(fun, Mf::Unpack4xU8) {
2529-
write!(self.out, "u")?;
2554+
let sign_prefix = if matches!(fun, Mf::Unpack4xU8) {
2555+
"u"
2556+
} else {
2557+
""
2558+
};
2559+
2560+
if context.lang_version >= (2, 1) {
2561+
// Metal uses little endian byte order, which matches what WGSL expects here.
2562+
write!(
2563+
self.out,
2564+
"{sign_prefix}int4(as_type<packed_{sign_prefix}char4>("
2565+
)?;
2566+
self.put_expression(arg, context, true)?;
2567+
write!(self.out, "))")?;
2568+
} else {
2569+
// MSL < 2.1 doesn't support `as_type` casting between packed chars and scalars.
2570+
write!(self.out, "({sign_prefix}int4(")?;
2571+
self.put_expression(arg, context, true)?;
2572+
write!(self.out, ", ")?;
2573+
self.put_expression(arg, context, true)?;
2574+
write!(self.out, " >> 8, ")?;
2575+
self.put_expression(arg, context, true)?;
2576+
write!(self.out, " >> 16, ")?;
2577+
self.put_expression(arg, context, true)?;
2578+
write!(self.out, " >> 24) << 24 >> 24)")?;
25302579
}
2531-
write!(self.out, "int4(")?;
2532-
self.put_expression(arg, context, true)?;
2533-
write!(self.out, ", ")?;
2534-
self.put_expression(arg, context, true)?;
2535-
write!(self.out, " >> 8, ")?;
2536-
self.put_expression(arg, context, true)?;
2537-
write!(self.out, " >> 16, ")?;
2538-
self.put_expression(arg, context, true)?;
2539-
write!(self.out, " >> 24) << 24 >> 24)")?;
25402580
}
25412581
Mf::QuantizeToF16 => {
25422582
match *context.resolve_type(arg) {
@@ -3279,14 +3319,20 @@ impl<W: Write> Writer<W> {
32793319
self.need_bake_expressions.insert(arg);
32803320
self.need_bake_expressions.insert(arg1.unwrap());
32813321
}
3282-
crate::MathFunction::FirstLeadingBit
3283-
| crate::MathFunction::Pack4xI8
3322+
crate::MathFunction::FirstLeadingBit => {
3323+
self.need_bake_expressions.insert(arg);
3324+
}
3325+
crate::MathFunction::Pack4xI8
32843326
| crate::MathFunction::Pack4xU8
32853327
| crate::MathFunction::Pack4xI8Clamp
32863328
| crate::MathFunction::Pack4xU8Clamp
32873329
| crate::MathFunction::Unpack4xI8
32883330
| crate::MathFunction::Unpack4xU8 => {
3289-
self.need_bake_expressions.insert(arg);
3331+
// On MSL < 2.1, we emit a polyfill for these functions that uses the
3332+
// argument multiple times. This is no longer necessary on MSL >= 2.1.
3333+
if context.lang_version < (2, 1) {
3334+
self.need_bake_expressions.insert(arg);
3335+
}
32903336
}
32913337
crate::MathFunction::ExtractBits => {
32923338
// Only argument 1 is re-used.

‎naga/tests/in/wgsl/bits-optimized-msl.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
targets = "METAL"
2+
3+
[msl]
4+
lang_version = [2, 1]

‎naga/tests/in/wgsl/bits-optimized-msl.wgsl

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Keep in sync with `bits_downlevel` and `bits_downlevel_webgl`
2+
3+
@compute @workgroup_size(1)
4+
fn main() {
5+
var i = 0;
6+
var i2 = vec2<i32>(0);
7+
var i3 = vec3<i32>(0);
8+
var i4 = vec4<i32>(0);
9+
var u = 0u;
10+
var u2 = vec2<u32>(0u);
11+
var u3 = vec3<u32>(0u);
12+
var u4 = vec4<u32>(0u);
13+
var f2 = vec2<f32>(0.0);
14+
var f4 = vec4<f32>(0.0);
15+
u = pack4x8snorm(f4);
16+
u = pack4x8unorm(f4);
17+
u = pack2x16snorm(f2);
18+
u = pack2x16unorm(f2);
19+
u = pack2x16float(f2);
20+
u = pack4xI8(i4);
21+
u = pack4xU8(u4);
22+
u = pack4xI8Clamp(i4);
23+
u = pack4xU8Clamp(u4);
24+
f4 = unpack4x8snorm(u);
25+
f4 = unpack4x8unorm(u);
26+
f2 = unpack2x16snorm(u);
27+
f2 = unpack2x16unorm(u);
28+
f2 = unpack2x16float(u);
29+
i4 = unpack4xI8(u);
30+
u4 = unpack4xU8(u);
31+
i = insertBits(i, i, 5u, 10u);
32+
i2 = insertBits(i2, i2, 5u, 10u);
33+
i3 = insertBits(i3, i3, 5u, 10u);
34+
i4 = insertBits(i4, i4, 5u, 10u);
35+
u = insertBits(u, u, 5u, 10u);
36+
u2 = insertBits(u2, u2, 5u, 10u);
37+
u3 = insertBits(u3, u3, 5u, 10u);
38+
u4 = insertBits(u4, u4, 5u, 10u);
39+
i = extractBits(i, 5u, 10u);
40+
i2 = extractBits(i2, 5u, 10u);
41+
i3 = extractBits(i3, 5u, 10u);
42+
i4 = extractBits(i4, 5u, 10u);
43+
u = extractBits(u, 5u, 10u);
44+
u2 = extractBits(u2, 5u, 10u);
45+
u3 = extractBits(u3, 5u, 10u);
46+
u4 = extractBits(u4, 5u, 10u);
47+
i = firstTrailingBit(i);
48+
u2 = firstTrailingBit(u2);
49+
i3 = firstLeadingBit(i3);
50+
u3 = firstLeadingBit(u3);
51+
i = firstLeadingBit(i);
52+
u = firstLeadingBit(u);
53+
i = countOneBits(i);
54+
i2 = countOneBits(i2);
55+
i3 = countOneBits(i3);
56+
i4 = countOneBits(i4);
57+
u = countOneBits(u);
58+
u2 = countOneBits(u2);
59+
u3 = countOneBits(u3);
60+
u4 = countOneBits(u4);
61+
i = reverseBits(i);
62+
i2 = reverseBits(i2);
63+
i3 = reverseBits(i3);
64+
i4 = reverseBits(i4);
65+
u = reverseBits(u);
66+
u2 = reverseBits(u2);
67+
u3 = reverseBits(u3);
68+
u4 = reverseBits(u4);
69+
}

‎naga/tests/out/msl/wgsl-bits-optimized-msl.msl

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
// language: metal2.1
2+
#include <metal_stdlib>
3+
#include <simd/simd.h>
4+
5+
using metal::uint;
6+
7+
8+
kernel void main_(
9+
) {
10+
int i = 0;
11+
metal::int2 i2_ = metal::int2(0);
12+
metal::int3 i3_ = metal::int3(0);
13+
metal::int4 i4_ = metal::int4(0);
14+
uint u = 0u;
15+
metal::uint2 u2_ = metal::uint2(0u);
16+
metal::uint3 u3_ = metal::uint3(0u);
17+
metal::uint4 u4_ = metal::uint4(0u);
18+
metal::float2 f2_ = metal::float2(0.0);
19+
metal::float4 f4_ = metal::float4(0.0);
20+
metal::float4 _e28 = f4_;
21+
u = metal::pack_float_to_snorm4x8(_e28);
22+
metal::float4 _e30 = f4_;
23+
u = metal::pack_float_to_unorm4x8(_e30);
24+
metal::float2 _e32 = f2_;
25+
u = metal::pack_float_to_snorm2x16(_e32);
26+
metal::float2 _e34 = f2_;
27+
u = metal::pack_float_to_unorm2x16(_e34);
28+
metal::float2 _e36 = f2_;
29+
u = as_type<uint>(half2(_e36));
30+
metal::int4 _e38 = i4_;
31+
u = as_type<uint>(packed_char4(_e38));
32+
metal::uint4 _e40 = u4_;
33+
u = as_type<uint>(packed_uchar4(_e40));
34+
metal::int4 _e42 = i4_;
35+
u = as_type<uint>(packed_char4(metal::clamp(_e42, -128, 127)));
36+
metal::uint4 _e44 = u4_;
37+
u = as_type<uint>(packed_uchar4(metal::clamp(_e44, 0, 255)));
38+
uint _e46 = u;
39+
f4_ = metal::unpack_snorm4x8_to_float(_e46);
40+
uint _e48 = u;
41+
f4_ = metal::unpack_unorm4x8_to_float(_e48);
42+
uint _e50 = u;
43+
f2_ = metal::unpack_snorm2x16_to_float(_e50);
44+
uint _e52 = u;
45+
f2_ = metal::unpack_unorm2x16_to_float(_e52);
46+
uint _e54 = u;
47+
f2_ = float2(as_type<half2>(_e54));
48+
uint _e56 = u;
49+
i4_ = int4(as_type<packed_char4>(_e56));
50+
uint _e58 = u;
51+
u4_ = uint4(as_type<packed_uchar4>(_e58));
52+
int _e60 = i;
53+
int _e61 = i;
54+
i = metal::insert_bits(_e60, _e61, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
55+
metal::int2 _e65 = i2_;
56+
metal::int2 _e66 = i2_;
57+
i2_ = metal::insert_bits(_e65, _e66, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
58+
metal::int3 _e70 = i3_;
59+
metal::int3 _e71 = i3_;
60+
i3_ = metal::insert_bits(_e70, _e71, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
61+
metal::int4 _e75 = i4_;
62+
metal::int4 _e76 = i4_;
63+
i4_ = metal::insert_bits(_e75, _e76, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
64+
uint _e80 = u;
65+
uint _e81 = u;
66+
u = metal::insert_bits(_e80, _e81, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
67+
metal::uint2 _e85 = u2_;
68+
metal::uint2 _e86 = u2_;
69+
u2_ = metal::insert_bits(_e85, _e86, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
70+
metal::uint3 _e90 = u3_;
71+
metal::uint3 _e91 = u3_;
72+
u3_ = metal::insert_bits(_e90, _e91, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
73+
metal::uint4 _e95 = u4_;
74+
metal::uint4 _e96 = u4_;
75+
u4_ = metal::insert_bits(_e95, _e96, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
76+
int _e100 = i;
77+
i = metal::extract_bits(_e100, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
78+
metal::int2 _e104 = i2_;
79+
i2_ = metal::extract_bits(_e104, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
80+
metal::int3 _e108 = i3_;
81+
i3_ = metal::extract_bits(_e108, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
82+
metal::int4 _e112 = i4_;
83+
i4_ = metal::extract_bits(_e112, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
84+
uint _e116 = u;
85+
u = metal::extract_bits(_e116, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
86+
metal::uint2 _e120 = u2_;
87+
u2_ = metal::extract_bits(_e120, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
88+
metal::uint3 _e124 = u3_;
89+
u3_ = metal::extract_bits(_e124, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
90+
metal::uint4 _e128 = u4_;
91+
u4_ = metal::extract_bits(_e128, metal::min(5u, 32u), metal::min(10u, 32u - metal::min(5u, 32u)));
92+
int _e132 = i;
93+
i = (((metal::ctz(_e132) + 1) % 33) - 1);
94+
metal::uint2 _e134 = u2_;
95+
u2_ = (((metal::ctz(_e134) + 1) % 33) - 1);
96+
metal::int3 _e136 = i3_;
97+
i3_ = metal::select(31 - metal::clz(metal::select(_e136, ~_e136, _e136 < 0)), int3(-1), _e136 == 0 || _e136 == -1);
98+
metal::uint3 _e138 = u3_;
99+
u3_ = metal::select(31 - metal::clz(_e138), uint3(-1), _e138 == 0 || _e138 == -1);
100+
int _e140 = i;
101+
i = metal::select(31 - metal::clz(metal::select(_e140, ~_e140, _e140 < 0)), int(-1), _e140 == 0 || _e140 == -1);
102+
uint _e142 = u;
103+
u = metal::select(31 - metal::clz(_e142), uint(-1), _e142 == 0 || _e142 == -1);
104+
int _e144 = i;
105+
i = metal::popcount(_e144);
106+
metal::int2 _e146 = i2_;
107+
i2_ = metal::popcount(_e146);
108+
metal::int3 _e148 = i3_;
109+
i3_ = metal::popcount(_e148);
110+
metal::int4 _e150 = i4_;
111+
i4_ = metal::popcount(_e150);
112+
uint _e152 = u;
113+
u = metal::popcount(_e152);
114+
metal::uint2 _e154 = u2_;
115+
u2_ = metal::popcount(_e154);
116+
metal::uint3 _e156 = u3_;
117+
u3_ = metal::popcount(_e156);
118+
metal::uint4 _e158 = u4_;
119+
u4_ = metal::popcount(_e158);
120+
int _e160 = i;
121+
i = metal::reverse_bits(_e160);
122+
metal::int2 _e162 = i2_;
123+
i2_ = metal::reverse_bits(_e162);
124+
metal::int3 _e164 = i3_;
125+
i3_ = metal::reverse_bits(_e164);
126+
metal::int4 _e166 = i4_;
127+
i4_ = metal::reverse_bits(_e166);
128+
uint _e168 = u;
129+
u = metal::reverse_bits(_e168);
130+
metal::uint2 _e170 = u2_;
131+
u2_ = metal::reverse_bits(_e170);
132+
metal::uint3 _e172 = u3_;
133+
u3_ = metal::reverse_bits(_e172);
134+
metal::uint4 _e174 = u4_;
135+
u4_ = metal::reverse_bits(_e174);
136+
return;
137+
}

0 commit comments

Comments
 (0)