Remove fewer Storage calls in copy_prop

[ohadravid]

Modify the copy_prop MIR optimization pass to remove fewer Storage{Live,Dead} calls, allowing for better optimizations by LLVM - see #141649.

Details

This is my attempt to fix the mentioned issue (this is the first part, I also implemented a similar solution for GVN in this branch).

The idea is to use the MaybeStorageDead a new MaybeUninitializedLocals analysis and remove only the storage calls of heads that are maybe-uninit when the associated local is accessed (or, conversely, keep the storage of heads that are for-sure initialized in every relevant access).

When combined with the GVN change, the final example in the issue (#141649 (comment)) is optimized as expected by LLVM. I also measured the effect on a few functions in rav1d (where I originally saw the issue) and observed reduced stack usage in several of them.

This is my first attempt at working with MIR optimizations, so it's possible this isn't the right approach — but all tests pass, and the resulting diffs appear correct.

r? tmiasko

since he commented on the issue and pointed to these passes.

[rustbot]

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

[matthiaskrgr]

@bors try @rust-timer queue

[bors]

⌛ Trying commit d24d035 with merge ef7d206...

[bors]

☀️ Try build successful - checks-actions
Build commit: ef7d206 (ef7d20666974f0dac45b03e051f2e283f9d9f090)

[rust-timer]

Finished benchmarking commit (ef7d206): comparison URL.

Overall result: ❌ regressions - please read the text below

Benchmarking this pull request means it may be perf-sensitive – we'll automatically label it not fit for rolling up. You can override this, but we strongly advise not to, due to possible changes in compiler perf.

Next Steps: If you can justify the regressions found in this try perf run, please do so in sufficient writing along with @rustbot label: +perf-regression-triaged. If not, please fix the regressions and do another perf run. If its results are neutral or positive, the label will be automatically removed.

@bors rollup=never
@rustbot label: -S-waiting-on-perf +perf-regression

Instruction count

Our most reliable metric. Used to determine the overall result above. However, even this metric can be noisy.

	mean	range	count
Regressions ❌ (primary)	0.3%	[0.2%, 0.4%]	8
Regressions ❌ (secondary)	0.3%	[0.2%, 0.4%]	7
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	0.3%	[0.2%, 0.4%]	8

Max RSS (memory usage)

Results (primary 0.7%, secondary 3.4%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	3.5%	[1.8%, 5.0%]	5
Regressions ❌ (secondary)	3.4%	[3.4%, 3.4%]	1
Improvements ✅ (primary)	-3.9%	[-6.5%, -2.0%]	3
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	0.7%	[-6.5%, 5.0%]	8

Cycles

Results (primary -0.6%, secondary -0.1%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	3.8%	[3.8%, 3.8%]	1
Improvements ✅ (primary)	-0.6%	[-0.6%, -0.6%]	1
Improvements ✅ (secondary)	-4.1%	[-4.1%, -4.1%]	1
All ❌✅ (primary)	-0.6%	[-0.6%, -0.6%]	1

Binary size

Results (primary 0.0%, secondary 0.0%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	0.2%	[0.0%, 0.8%]	10
Regressions ❌ (secondary)	0.1%	[0.0%, 0.1%]	5
Improvements ✅ (primary)	-0.2%	[-0.8%, -0.0%]	8
Improvements ✅ (secondary)	-0.2%	[-0.2%, -0.2%]	1
All ❌✅ (primary)	0.0%	[-0.8%, 0.8%]	18

Bootstrap: 757.399s -> 756.065s (-0.18%)
Artifact size: 372.20 MiB -> 372.12 MiB (-0.02%)

[ohadravid]

@matthiaskrgr - I updated the impl to stop re-checking once a head is found to be maybe-dead, which should be a bit better

[matthiaskrgr]

@bors try @rust-timer queue

[bors]

⌛ Trying commit 905e968 with merge c0a2949...

[cjgillot]

Should this check happen in Replacer::visit_local, and move the replacement of storage statements to a dedicated cleanup visitor?

[bors]

☀️ Try build successful - checks-actions
Build commit: c0a2949 (c0a294957df10fc3880e1677c72c0cf122485509)

[ohadravid]

Should this check happen in Replacer::visit_local

I'm not sure how to make this work: using ResultsCursor requires a &body, but it's not possible to have that while running a MutVisitor since it requires a &mut body.

Is there a different way to do this?

[rust-timer]

Finished benchmarking commit (c0a2949): comparison URL.

Overall result: ❌✅ regressions and improvements - please read the text below

Benchmarking this pull request means it may be perf-sensitive – we'll automatically label it not fit for rolling up. You can override this, but we strongly advise not to, due to possible changes in compiler perf.

Next Steps: If you can justify the regressions found in this try perf run, please do so in sufficient writing along with @rustbot label: +perf-regression-triaged. If not, please fix the regressions and do another perf run. If its results are neutral or positive, the label will be automatically removed.

@bors rollup=never
@rustbot label: -S-waiting-on-perf +perf-regression

Instruction count

Our most reliable metric. Used to determine the overall result above. However, even this metric can be noisy.

	mean	range	count
Regressions ❌ (primary)	0.3%	[0.2%, 0.4%]	9
Regressions ❌ (secondary)	0.3%	[0.2%, 0.4%]	7
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-0.2%	[-0.2%, -0.2%]	1
All ❌✅ (primary)	0.3%	[0.2%, 0.4%]	9

Max RSS (memory usage)

Results (primary -0.1%, secondary -1.3%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	4.2%	[3.4%, 5.8%]	4
Regressions ❌ (secondary)	3.1%	[3.1%, 3.1%]	1
Improvements ✅ (primary)	-4.4%	[-6.6%, -1.8%]	4
Improvements ✅ (secondary)	-5.8%	[-5.8%, -5.8%]	1
All ❌✅ (primary)	-0.1%	[-6.6%, 5.8%]	8

Cycles

Results (secondary -1.0%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	2.3%	[2.3%, 2.3%]	1
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-2.6%	[-2.6%, -2.5%]	2
All ❌✅ (primary)	-	-	0

Binary size

Results (primary -0.0%, secondary 0.0%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	0.2%	[0.0%, 0.8%]	10
Regressions ❌ (secondary)	0.1%	[0.0%, 0.1%]	5
Improvements ✅ (primary)	-0.2%	[-0.8%, -0.0%]	8
Improvements ✅ (secondary)	-0.2%	[-0.2%, -0.2%]	1
All ❌✅ (primary)	-0.0%	[-0.8%, 0.8%]	18

Bootstrap: 756.494s -> 757.685s (0.16%)
Artifact size: 372.15 MiB -> 372.11 MiB (-0.01%)

[tmiasko]

SsaLocals::borrowed_locals()[local] describes whether local is borrowed, not whether any local in its copy class is borrowed. An example that doesn't work as it supposed to:

#![feature(custom_mir, core_intrinsics, freeze)]
extern crate core;
use core::intrinsics::mir::*;
use core::marker::Freeze;

#[custom_mir(dialect = "runtime")]
pub fn f<T: Copy + Freeze>(_1: (T, T)) -> T {
    mir! {
        let _2: T;
        let _3: T;
        let _4: &T;
        {
            StorageLive(_2);
            _2 = _1.0;
            _3 = _2;
            _4 = &_3;
            StorageDead(_2);
            RET = *_4;
            Return()
        }
    }
}

$ rustc +stage1 b.rs --crate-type=lib -Zmir-opt-level=0 -Zmir-enable-passes=+CopyProp -Zunpretty=mir -Copt-level=1
fn f(_1: (T, T)) -> T {
    let mut _0: T;
    let mut _2: T;
    let mut _3: T;
    let mut _4: &T;

    bb0: {
        StorageLive(_2);
        _2 = copy (_1.0: T);
        _4 = &_2;
        StorageDead(_2);
        _0 = copy (*_4);
        return;
    }
}

I think it should be fine to allow the head itself to be borrowed (it is only all other locals from the copy class that cannot be borrowed).

[ohadravid]

Oh I see, I read the compute_copy_classes comment but assumed it referred to ssa.borrowed_locals(), not the one it computes internally.

[ohadravid]

Add this test and pushed the correct fix.
I do get reordered output for some tests (like tests/mir-opt/pre-codegen/derived_ord.rs) and I'm not sure why, but the fixed impl now produces the correct output for this test.

[rust-log-analyzer]

The job x86_64-gnu-llvm-19 failed! Check out the build log: (web) (plain)

Click to see the possible cause of the failure (guessed by this bot)

#18 exporting to docker image format
#18 sending tarball 26.4s done
#18 DONE 42.0s
##[endgroup]
Setting extra environment values for docker:  --env ENABLE_GCC_CODEGEN=1 --env GCC_EXEC_PREFIX=/usr/lib/gcc/
[CI_JOB_NAME=x86_64-gnu-llvm-19]
[CI_JOB_NAME=x86_64-gnu-llvm-19]
debug: `DISABLE_CI_RUSTC_IF_INCOMPATIBLE` configured.
---
sccache: Listening on address 127.0.0.1:4226
##[group]Configure the build
configure: processing command line
configure: 
configure: build.configure-args := ['--build=x86_64-unknown-linux-gnu', '--llvm-root=/usr/lib/llvm-19', '--enable-llvm-link-shared', '--set', 'rust.randomize-layout=true', '--set', 'rust.thin-lto-import-instr-limit=10', '--set', 'build.print-step-timings', '--enable-verbose-tests', '--set', 'build.metrics', '--enable-verbose-configure', '--enable-sccache', '--disable-manage-submodules', '--enable-locked-deps', '--enable-cargo-native-static', '--set', 'rust.codegen-units-std=1', '--set', 'dist.compression-profile=balanced', '--dist-compression-formats=xz', '--set', 'rust.lld=false', '--disable-dist-src', '--release-channel=nightly', '--enable-debug-assertions', '--enable-overflow-checks', '--enable-llvm-assertions', '--set', 'rust.verify-llvm-ir', '--set', 'rust.codegen-backends=llvm,cranelift,gcc', '--set', 'llvm.static-libstdcpp', '--set', 'gcc.download-ci-gcc=true', '--enable-new-symbol-mangling']
configure: build.build          := x86_64-unknown-linux-gnu
configure: target.x86_64-unknown-linux-gnu.llvm-config := /usr/lib/llvm-19/bin/llvm-config
configure: llvm.link-shared     := True
configure: rust.randomize-layout := True
configure: rust.thin-lto-import-instr-limit := 10
---
48           _4 = copy _1;
49           StorageLive(_10);
+           StorageLive(_11);
+           StorageLive(_12);
50           _10 = discriminant(_4);
51           switchInt(move _10) -> [0: bb7, 1: bb6, otherwise: bb1];
52       }

90       }
91   
92       bb5: {
+           StorageDead(_12);
+           StorageDead(_11);
93           StorageDead(_10);
94           StorageDead(_4);
95           _5 = discriminant(_3);

98       }
99   
100       bb6: {
-           StorageLive(_12);
102           _12 = move ((_4 as Err).0: i32);
103           StorageLive(_13);
104           _13 = Result::<Infallible, i32>::Err(copy _12);

105           _3 = ControlFlow::<Result<Infallible, i32>, i32>::Break(move _13);
106           StorageDead(_13);
-           StorageDead(_12);
108 -         goto -> bb5;
109 +         goto -> bb8;
110       }

111   
112       bb7: {
-           StorageLive(_11);
114           _11 = move ((_4 as Ok).0: i32);
115           _3 = ControlFlow::<Result<Infallible, i32>, i32>::Continue(copy _11);
-           StorageDead(_11);
117           goto -> bb5;
118 +     }
119 + 

120 +     bb8: {
+ +         StorageDead(_12);
+ +         StorageDead(_11);
121 +         StorageDead(_10);
122 +         StorageDead(_4);
123 +         _5 = discriminant(_3);


thread '[mir-opt] tests/mir-opt/jump_threading.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/jump_threading.identity.JumpThreading.panic-unwind.diff
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

---- [mir-opt] tests/mir-opt/pre-codegen/derived_ord.rs stdout ----
37     }
38 
---
79     }
80 
81     bb3: {
-         StorageLive(_12);
83         _12 = move ((_11 as Some).0: std::cmp::Ordering);
84         StorageLive(_13);
85         _13 = discriminant(_12);

86         _0 = Le(move _13, const 0_i8);
87         StorageDead(_13);
-         StorageDead(_12);
89         StorageDead(_11);
+         StorageDead(_12);
90         return;
91     }
92 }


thread '[mir-opt] tests/mir-opt/pre-codegen/derived_ord.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/pre-codegen/derived_ord.demo_le.PreCodegen.after.mir

---- [mir-opt] tests/mir-opt/pre-codegen/loops.rs stdout ----
54     bb2: {
55         StorageLive(_13);
56         _5 = &mut _4;
+         StorageLive(_8);
57         StorageLive(_7);
58         StorageLive(_6);
59         _6 = &mut (_4.0: impl Iterator<Item = T>);

62 
63     bb3: {
64         StorageDead(_6);
-         StorageLive(_8);
66         _8 = &mut (_4.1: impl Fn(T) -> U);
67         StorageLive(_9);
+         StorageLive(_10);
68         _9 = discriminant(_7);
69         switchInt(move _9) -> [0: bb4, 1: bb6, otherwise: bb9];
70     }

71 
72     bb4: {
+         StorageDead(_10);
73         StorageDead(_9);
-         StorageDead(_8);
75         StorageDead(_7);
+         StorageDead(_8);
76         StorageDead(_13);
77         drop(_4) -> [return: bb5, unwind continue];
78     }

83     }
84 
85     bb6: {
-         StorageLive(_10);
87         _10 = move ((_7 as Some).0: T);
88         StorageLive(_12);
89         StorageLive(_11);

97         StorageDead(_12);
98         StorageDead(_10);
99         StorageDead(_9);
-         StorageDead(_8);
101         StorageDead(_7);
+         StorageDead(_8);
102         _14 = move ((_13 as Some).0: U);
103         _15 = opaque::<U>(move _14) -> [return: bb8, unwind: bb10];
104     }


thread '[mir-opt] tests/mir-opt/pre-codegen/loops.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/pre-codegen/loops.mapped.PreCodegen.after.mir

---- [mir-opt] tests/mir-opt/pre-codegen/simple_option_map.rs stdout ----
29         StorageLive(_6);
30         StorageLive(_4);
31         StorageLive(_2);
+         StorageLive(_3);
32         _2 = discriminant(_1);
33         switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb4];
34     }

35 
36     bb1: {
+         StorageDead(_3);
---
42     }
43 
44     bb2: {
-         StorageLive(_3);
46         _3 = copy ((_1 as Some).0: i32);
47         _4 = ControlFlow::<Option<Infallible>, i32>::Continue(copy _3);
48         StorageDead(_3);


thread '[mir-opt] tests/mir-opt/pre-codegen/simple_option_map.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/pre-codegen/simple_option_map.map_via_question_mark.PreCodegen.after.mir

---- [mir-opt] tests/mir-opt/pre-codegen/slice_index.rs stdout ----
36         StorageLive(_11);
37         StorageLive(_5);
38         _5 = &raw mut (*_1);
+         StorageLive(_8);
39         StorageLive(_6);
40         _6 = PtrMetadata(copy _1);
41         _7 = <std::ops::Range<usize> as SliceIndex<[T]>>::get_unchecked_mut::precondition_check(copy _3, copy _4, move _6) -> [return: bb1, unwind unreachable];

43 
44     bb1: {
45         StorageDead(_6);
-         StorageLive(_8);
47         _8 = SubUnchecked(copy _4, copy _3);
48         StorageLive(_9);
-         _9 = copy _5 as *mut u32 (PtrToPtr);
50         StorageLive(_10);
+         _9 = copy _5 as *mut u32 (PtrToPtr);
51         _10 = Offset(copy _9, copy _3);
52         _11 = *mut [u32] from (copy _10, copy _8);
53         StorageDead(_10);


thread '[mir-opt] tests/mir-opt/pre-codegen/slice_index.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/pre-codegen/slice_index.slice_get_unchecked_mut_range.PreCodegen.after.panic-unwind.mir

---- [mir-opt] tests/mir-opt/separate_const_switch.rs stdout ----
39       bb0: {
40           StorageLive(_2);
41           StorageLive(_6);
+           StorageLive(_7);
+           StorageLive(_8);
42           _6 = discriminant(_1);
43           switchInt(move _6) -> [0: bb6, 1: bb5, otherwise: bb1];
44       }

63       }
64   
65       bb4: {
+           StorageDead(_8);
+           StorageDead(_7);
66           StorageDead(_6);
67           _3 = discriminant(_2);
68 -         switchInt(move _3) -> [0: bb2, 1: bb3, otherwise: bb1];

70       }
71   
72       bb5: {
-           StorageLive(_8);
74           _8 = copy ((_1 as Err).0: i32);
75           StorageLive(_9);
76           _9 = Result::<Infallible, i32>::Err(copy _8);

77           _2 = ControlFlow::<Result<Infallible, i32>, i32>::Break(move _9);
78           StorageDead(_9);
-           StorageDead(_8);
80 -         goto -> bb4;
81 +         goto -> bb7;
82       }

83   
84       bb6: {
-           StorageLive(_7);
86           _7 = copy ((_1 as Ok).0: i32);
87           _2 = ControlFlow::<Result<Infallible, i32>, i32>::Continue(copy _7);
-           StorageDead(_7);
89           goto -> bb4;
90 +     }
91 + 

92 +     bb7: {
+ +         StorageDead(_8);
+ +         StorageDead(_7);
93 +         StorageDead(_6);
94 +         _3 = discriminant(_2);
95 +         goto -> bb3;


thread '[mir-opt] tests/mir-opt/separate_const_switch.rs' panicked at src/tools/compiletest/src/runtest/mir_opt.rs:84:21:
Actual MIR output differs from expected MIR output /checkout/tests/mir-opt/separate_const_switch.identity.JumpThreading.diff


failures:
    [mir-opt] tests/mir-opt/jump_threading.rs
    [mir-opt] tests/mir-opt/pre-codegen/derived_ord.rs