Skip to content

Commit 6eb273e

Browse files
scottmcmscottmcm
committed
[codegen] assume the tag, not the relative discriminant
1 parent adcb3d3 commit 6eb273e

File tree

3 files changed

+65
-47
lines changed

3 files changed

+65
-47
lines changed

‎compiler/rustc_codegen_ssa/src/mir/operand.rs

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,35 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
498498
bx.cx().const_uint(cast_to, niche_variants.start().as_u32() as u64);
499499
(is_niche, tagged_discr, 0)
500500
} else {
501+
// Thanks to parameter attributes and load metadata, LLVM already knows
502+
// the general valid range of the tag. It's possible, though, for there
503+
// to be an impossible value *in the middle*, which those ranges don't
504+
// communicate, so it's worth an `assume` to let the optimizer know.
505+
// Most importantly, this means when optimizing a variant test like
506+
// `SELECT(is_niche, complex, CONST) == CONST` it's ok to simplify that
507+
// to `!is_niche` because the `complex` part can't possibly match.
508+
//
509+
// This was previously asserted on `tagged_discr` below, where the
510+
// impossible value is more obvious, but that caused an intermediate
511+
// value to become multi-use and thus not optimize, so instead this
512+
// assumes on the original input which is always multi-use. See
513+
// <https://github.com/llvm/llvm-project/issues/134024#issuecomment-3131782555>
514+
//
515+
// FIXME: If we ever get range assume operand bundles in LLVM (so we
516+
// don't need the `icmp`s in the instruction stream any more), it
517+
// might be worth moving this back to being on the switch argument
518+
// where it's more obviously applicable.
519+
if niche_variants.contains(&untagged_variant)
520+
&& bx.cx().sess().opts.optimize != OptLevel::No
521+
{
522+
let impossible = niche_start
523+
.wrapping_add(u128::from(untagged_variant.as_u32()))
524+
.wrapping_sub(u128::from(niche_variants.start().as_u32()));
525+
let impossible = bx.cx().const_uint_big(tag_llty, impossible);
526+
let ne = bx.icmp(IntPredicate::IntNE, tag, impossible);
527+
bx.assume(ne);
528+
}
529+
501530
// With multiple niched variants we'll have to actually compute
502531
// the variant index from the stored tag.
503532
//
@@ -588,20 +617,6 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
588617
let untagged_variant_const =
589618
bx.cx().const_uint(cast_to, u64::from(untagged_variant.as_u32()));
590619

591-
// Thanks to parameter attributes and load metadata, LLVM already knows
592-
// the general valid range of the tag. It's possible, though, for there
593-
// to be an impossible value *in the middle*, which those ranges don't
594-
// communicate, so it's worth an `assume` to let the optimizer know.
595-
// Most importantly, this means when optimizing a variant test like
596-
// `SELECT(is_niche, complex, CONST) == CONST` it's ok to simplify that
597-
// to `!is_niche` because the `complex` part can't possibly match.
598-
if niche_variants.contains(&untagged_variant)
599-
&& bx.cx().sess().opts.optimize != OptLevel::No
600-
{
601-
let ne = bx.icmp(IntPredicate::IntNE, tagged_discr, untagged_variant_const);
602-
bx.assume(ne);
603-
}
604-
605620
let discr = bx.select(is_niche, tagged_discr, untagged_variant_const);
606621

607622
// In principle we could insert assumes on the possible range of `discr`, but

‎tests/codegen-llvm/enum/enum-discriminant-eq.rs

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -91,16 +91,16 @@ pub enum Mid<T> {
9191
pub fn mid_bool_eq_discr(a: Mid<bool>, b: Mid<bool>) -> bool {
9292
// CHECK-LABEL: @mid_bool_eq_discr(
9393

94+
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i8 %a, 3
95+
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
9496
// CHECK: %[[A_REL_DISCR:.+]] = add nsw i8 %a, -2
9597
// CHECK: %[[A_IS_NICHE:.+]] = icmp samesign ugt i8 %a, 1
96-
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i8 %[[A_REL_DISCR]], 1
97-
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
9898
// CHECK: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i8 %[[A_REL_DISCR]], i8 1
9999

100+
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i8 %b, 3
101+
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
100102
// CHECK: %[[B_REL_DISCR:.+]] = add nsw i8 %b, -2
101103
// CHECK: %[[B_IS_NICHE:.+]] = icmp samesign ugt i8 %b, 1
102-
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i8 %[[B_REL_DISCR]], 1
103-
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
104104
// CHECK: %[[B_DISCR:.+]] = select i1 %[[B_IS_NICHE]], i8 %[[B_REL_DISCR]], i8 1
105105

106106
// CHECK: ret i1 %[[R]]
@@ -111,16 +111,16 @@ pub fn mid_bool_eq_discr(a: Mid<bool>, b: Mid<bool>) -> bool {
111111
pub fn mid_ord_eq_discr(a: Mid<Ordering>, b: Mid<Ordering>) -> bool {
112112
// CHECK-LABEL: @mid_ord_eq_discr(
113113

114+
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i8 %a, 3
115+
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
114116
// CHECK: %[[A_REL_DISCR:.+]] = add nsw i8 %a, -2
115117
// CHECK: %[[A_IS_NICHE:.+]] = icmp sgt i8 %a, 1
116-
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i8 %[[A_REL_DISCR]], 1
117-
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
118118
// CHECK: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i8 %[[A_REL_DISCR]], i8 1
119119

120+
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i8 %b, 3
121+
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
120122
// CHECK: %[[B_REL_DISCR:.+]] = add nsw i8 %b, -2
121123
// CHECK: %[[B_IS_NICHE:.+]] = icmp sgt i8 %b, 1
122-
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i8 %[[B_REL_DISCR]], 1
123-
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
124124
// CHECK: %[[B_DISCR:.+]] = select i1 %[[B_IS_NICHE]], i8 %[[B_REL_DISCR]], i8 1
125125

126126
// CHECK: %[[R:.+]] = icmp eq i8 %[[A_DISCR]], %[[B_DISCR]]
@@ -140,16 +140,16 @@ pub fn mid_nz32_eq_discr(a: Mid<NonZero<u32>>, b: Mid<NonZero<u32>>) -> bool {
140140
pub fn mid_ac_eq_discr(a: Mid<AC>, b: Mid<AC>) -> bool {
141141
// CHECK-LABEL: @mid_ac_eq_discr(
142142

143-
// LLVM20: %[[A_REL_DISCR:.+]] = xor i8 %a, -128
144-
// CHECK: %[[A_IS_NICHE:.+]] = icmp slt i8 %a, 0
145143
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i8 %a, -127
146144
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
145+
// LLVM20: %[[A_REL_DISCR:.+]] = xor i8 %a, -128
146+
// CHECK: %[[A_IS_NICHE:.+]] = icmp slt i8 %a, 0
147147
// LLVM20: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i8 %[[A_REL_DISCR]], i8 1
148148

149-
// LLVM20: %[[B_REL_DISCR:.+]] = xor i8 %b, -128
150-
// CHECK: %[[B_IS_NICHE:.+]] = icmp slt i8 %b, 0
151149
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i8 %b, -127
152150
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
151+
// LLVM20: %[[B_REL_DISCR:.+]] = xor i8 %b, -128
152+
// CHECK: %[[B_IS_NICHE:.+]] = icmp slt i8 %b, 0
153153
// LLVM20: %[[B_DISCR:.+]] = select i1 %[[B_IS_NICHE]], i8 %[[B_REL_DISCR]], i8 1
154154

155155
// LLVM21: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i8 %a, i8 -127
@@ -166,21 +166,24 @@ pub fn mid_ac_eq_discr(a: Mid<AC>, b: Mid<AC>) -> bool {
166166
pub fn mid_giant_eq_discr(a: Mid<Giant>, b: Mid<Giant>) -> bool {
167167
// CHECK-LABEL: @mid_giant_eq_discr(
168168

169+
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i128 %a, 6
170+
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
169171
// CHECK: %[[A_TRUNC:.+]] = trunc nuw nsw i128 %a to i64
170-
// CHECK: %[[A_REL_DISCR:.+]] = add nsw i64 %[[A_TRUNC]], -5
172+
// LLVM20: %[[A_REL_DISCR:.+]] = add nsw i64 %[[A_TRUNC]], -5
171173
// CHECK: %[[A_IS_NICHE:.+]] = icmp samesign ugt i128 %a, 4
172-
// CHECK: %[[A_NOT_HOLE:.+]] = icmp ne i64 %[[A_REL_DISCR]], 1
173-
// CHECK: tail call void @llvm.assume(i1 %[[A_NOT_HOLE]])
174-
// CHECK: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i64 %[[A_REL_DISCR]], i64 1
174+
// LLVM20: %[[A_DISCR:.+]] = select i1 %[[A_IS_NICHE]], i64 %[[A_REL_DISCR]], i64 1
175+
// LLVM21: %[[A_MODIFIED_TAG:.+]] = select i1 %[[A_IS_NICHE]], i64 %[[A_TRUNC]], i64 6
175176

177+
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i128 %b, 6
178+
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
176179
// CHECK: %[[B_TRUNC:.+]] = trunc nuw nsw i128 %b to i64
177-
// CHECK: %[[B_REL_DISCR:.+]] = add nsw i64 %[[B_TRUNC]], -5
180+
// LLVM20: %[[B_REL_DISCR:.+]] = add nsw i64 %[[B_TRUNC]], -5
178181
// CHECK: %[[B_IS_NICHE:.+]] = icmp samesign ugt i128 %b, 4
179-
// CHECK: %[[B_NOT_HOLE:.+]] = icmp ne i64 %[[B_REL_DISCR]], 1
180-
// CHECK: tail call void @llvm.assume(i1 %[[B_NOT_HOLE]])
181-
// CHECK: %[[B_DISCR:.+]] = select i1 %[[B_IS_NICHE]], i64 %[[B_REL_DISCR]], i64 1
182+
// LLVM20: %[[B_DISCR:.+]] = select i1 %[[B_IS_NICHE]], i64 %[[B_REL_DISCR]], i64 1
183+
// LLVM21: %[[B_MODIFIED_TAG:.+]] = select i1 %[[B_IS_NICHE]], i64 %[[B_TRUNC]], i64 6
182184

183-
// CHECK: %[[R:.+]] = icmp eq i64 %[[A_DISCR]], %[[B_DISCR]]
185+
// LLVM20: %[[R:.+]] = icmp eq i64 %[[A_DISCR]], %[[B_DISCR]]
186+
// LLVM21: %[[R:.+]] = icmp eq i64 %[[A_MODIFIED_TAG]], %[[B_MODIFIED_TAG]]
184187
// CHECK: ret i1 %[[R]]
185188
discriminant_value(&a) == discriminant_value(&b)
186189
}

‎tests/codegen-llvm/enum/enum-match.rs

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -138,18 +138,18 @@ pub fn match3(e: Option<&u8>) -> i16 {
138138

139139
#[derive(PartialEq)]
140140
pub enum MiddleNiche {
141-
A,
142-
B,
143-
C(bool),
144-
D,
145-
E,
141+
A, // tag 2
142+
B, // tag 3
143+
C(bool), // untagged
144+
D, // tag 5
145+
E, // tag 6
146146
}
147147

148148
// CHECK-LABEL: define{{( dso_local)?}} noundef{{( range\(i8 -?[0-9]+, -?[0-9]+\))?}} i8 @match4(i8{{.+}}%0)
149149
// CHECK-NEXT: start:
150-
// CHECK-NEXT: %[[REL_VAR:.+]] = add{{( nsw)?}} i8 %0, -2
151-
// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %[[REL_VAR]], 2
150+
// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %0, 4
152151
// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
152+
// CHECK-NEXT: %[[REL_VAR:.+]] = add{{( nsw)?}} i8 %0, -2
153153
// CHECK-NEXT: %[[NOT_NICHE:.+]] = icmp{{( samesign)?}} ult i8 %0, 2
154154
// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[NOT_NICHE]], i8 2, i8 %[[REL_VAR]]
155155
// CHECK-NEXT: switch i8 %[[DISCR]]
@@ -443,19 +443,19 @@ pub enum HugeVariantIndex {
443443
V255(Never),
444444
V256(Never),
445445

446-
Possible257,
447-
Bool258(bool),
448-
Possible259,
446+
Possible257, // tag 2
447+
Bool258(bool), // untagged
448+
Possible259, // tag 4
449449
}
450450

451451
// CHECK-LABEL: define{{( dso_local)?}} noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match5(i8{{.+}}%0)
452452
// CHECK-NEXT: start:
453+
// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %0, 3
454+
// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
453455
// CHECK-NEXT: %[[REL_VAR:.+]] = add{{( nsw)?}} i8 %0, -2
454456
// CHECK-NEXT: %[[REL_VAR_WIDE:.+]] = zext i8 %[[REL_VAR]] to i64
455457
// CHECK-NEXT: %[[IS_NICHE:.+]] = icmp{{( samesign)?}} ugt i8 %0, 1
456458
// CHECK-NEXT: %[[NICHE_DISCR:.+]] = add nuw nsw i64 %[[REL_VAR_WIDE]], 257
457-
// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i64 %[[NICHE_DISCR]], 258
458-
// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
459459
// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[IS_NICHE]], i64 %[[NICHE_DISCR]], i64 258
460460
// CHECK-NEXT: switch i64 %[[DISCR]],
461461
// CHECK-NEXT: i64 257,

0 commit comments

Comments
 (0)