commit c3d63c1d5831925cbd69818e07f6508abc48e0cb
author Ben Vanik <ben.vanik@gmail.com> Thu Oct 27 22:15:27 2022 +0000
committer GitHub <noreply@github.com> Thu Oct 27 15:15:27 2022 -0700
tree 47e0081d57203a1261643f9bdd31ecce3c35df31
parent 5601928597c3785d27162eba76cb0d7dd4bd47cc

Basic IPO pass for function cleanup. (#10911)

This removes unused args/results, inlines uniformly constant
args/results, and deduplicates the ones that remain. This is similar to
the block patterns we have and designed to be run in a fixed-point
iteration so that cse/canonicalization/block patterns can all help each
other.

In the future this should be extended to support globals by biasing
towards making loads/stores local with uses. Today we'll end up with a
lot of global loads -> call -> global stores and that's blocking the
naive global folding/fusion pass from kicking in. We'll probably want a
dialect/op interface like the inliner has to control this (and may be
able to just reuse the inliner interface).

Example call before + 726 globals:
```mlir
%4:270 = call @jit__train_step$main(%0, %1, %c3200000, %c0, %c3200000, %c3200000, %0, %2, %c6400000, %c0, %c6400000, %c6400000, %0, %3, %c200, %c0, %c200, %c200, %_train_state$0__timepoint, %_train_state$0, %_train_state$0__storage_size, %_train_state$0__offset, %_train_state$0__length, %_train_state$0__size, %_train_state$1__timepoint, %_train_state$1, %_train_state$1__storage_size, %_train_state$1__offset, %_train_state$1__length, %_train_state$1__size, %_train_state$2__timepoint, %_train_state$2, %_train_state$2__storage_size, %_train_state$2__offset, %_train_state$2__length, %_train_state$2__size, %_train_state$3__timepoint, %_train_state$3, %_train_state$3__storage_size, %_train_state$3__offset, %_train_state$3__length, %_train_state$3__size, %_train_state$4__timepoint, %_train_state$4, %_train_state$4__storage_size, %_train_state$4__offset, %_train_state$4__length, %_train_state$4__size, %_train_state$5__timepoint, %_train_state$5, %_train_state$5__storage_size, %_train_state$5__offset, %_train_state$5__length, %_train_state$5__size, %_train_state$6__timepoint, %_train_state$6, %_train_state$6__storage_size, %_train_state$6__offset, %_train_state$6__length, %_train_state$6__size, %_train_state$7__timepoint, %_train_state$7, %_train_state$7__storage_size, %_train_state$7__offset, %_train_state$7__length, %_train_state$7__size, %_train_state$8__timepoint, %_train_state$8, %_train_state$8__storage_size, %_train_state$8__offset, %_train_state$8__length, %_train_state$8__size, %_train_state$9__timepoint, %_train_state$9, %_train_state$9__storage_size, %_train_state$9__offset, %_train_state$9__length, %_train_state$9__size, %_train_state$10__timepoint, %_train_state$10, %_train_state$10__storage_size, %_train_state$10__offset, %_train_state$10__length, %_train_state$10__size, %_train_state$11__timepoint, %_train_state$11, %_train_state$11__storage_size, %_train_state$11__offset, %_train_state$11__length, %_train_state$11__size, %_train_state$12__timepoint, %_train_state$12, %_train_state$12__storage_size, %_train_state$12__offset, %_train_state$12__length, %_train_state$12__size, %_train_state$13__timepoint, %_train_state$13, %_train_state$13__storage_size, %_train_state$13__offset, %_train_state$13__length, %_train_state$13__size, %_train_state$14__timepoint, %_train_state$14, %_train_state$14__storage_size, %_train_state$14__offset, %_train_state$14__length, %_train_state$14__size, %_train_state$15__timepoint, %_train_state$15, %_train_state$15__storage_size, %_train_state$15__offset, %_train_state$15__length, %_train_state$15__size, %_train_state$16__timepoint, %_train_state$16, %_train_state$16__storage_size, %_train_state$16__offset, %_train_state$16__length, %_train_state$16__size, %_train_state$17__timepoint, %_train_state$17, %_train_state$17__storage_size, %_train_state$17__offset, %_train_state$17__length, %_train_state$17__size, %_train_state$18__timepoint, %_train_state$18, %_train_state$18__storage_size, %_train_state$18__offset, %_train_state$18__length, %_train_state$18__size, %_train_state$19__timepoint, %_train_state$19, %_train_state$19__storage_size, %_train_state$19__offset, %_train_state$19__length, %_train_state$19__size, %_train_state$20__timepoint, %_train_state$20, %_train_state$20__storage_size, %_train_state$20__offset, %_train_state$20__length, %_train_state$20__size, %_train_state$21__timepoint, %_train_state$21, %_train_state$21__storage_size, %_train_state$21__offset, %_train_state$21__length, %_train_state$21__size, %_train_state$22__timepoint, %_train_state$22, %_train_state$22__storage_size, %_train_state$22__offset, %_train_state$22__length, %_train_state$22__size, %_train_state$23__timepoint, %_train_state$23, %_train_state$23__storage_size, %_train_state$23__offset, %_train_state$23__length, %_train_state$23__size, %_train_state$24__timepoint, %_train_state$24, %_train_state$24__storage_size, %_train_state$24__offset, %_train_state$24__length, %_train_state$24__size, %_train_state$25__timepoint, %_train_state$25, %_train_state$25__storage_size, %_train_state$25__offset, %_train_state$25__length, %_train_state$25__size, %_train_state$26__timepoint, %_train_state$26, %_train_state$26__storage_size, %_train_state$26__offset, %_train_state$26__length, %_train_state$26__size, %_train_state$27__timepoint, %_train_state$27, %_train_state$27__storage_size, %_train_state$27__offset, %_train_state$27__length, %_train_state$27__size, %_train_state$28__timepoint, %_train_state$28, %_train_state$28__storage_size, %_train_state$28__offset, %_train_state$28__length, %_train_state$28__size, %_train_state$29__timepoint, %_train_state$29, %_train_state$29__storage_size, %_train_state$29__offset, %_train_state$29__length, %_train_state$29__size, %_train_state$30__timepoint, %_train_state$30, %_train_state$30__storage_size, %_train_state$30__offset, %_train_state$30__length, %_train_state$30__size, %_train_state$31__timepoint, %_train_state$31, %_train_state$31__storage_size, %_train_state$31__offset, %_train_state$31__length, %_train_state$31__size, %_train_state$32__timepoint, %_train_state$32, %_train_state$32__storage_size, %_train_state$32__offset, %_train_state$32__length, %_train_state$32__size, %_train_state$33__timepoint, %_train_state$33, %_train_state$33__storage_size, %_train_state$33__offset, %_train_state$33__length, %_train_state$33__size, %_train_state$34__timepoint, %_train_state$34, %_train_state$34__storage_size, %_train_state$34__offset, %_train_state$34__length, %_train_state$34__size, %_prng_key$0__timepoint, %_prng_key$0, %_prng_key$0__storage_size, %_prng_key$0__offset, %_prng_key$0__length, %_prng_key$0__size) : (!stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index) -> (!stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index)
  return %45, %2, %c10003776, %c10003712, %c4, %c4, %45, %2, %c10003776, %c6669248, %c512, %c512, %45, %2, %c10003776, %c6669760, %c3145728, %c3145728, %45, %2, %c10003776, %c9815488, %c1024, %c1024, %45, %2, %c10003776, %c9816512, %c131072, %c131072, %45, %2, %c10003776, %c9947584, %c256, %c256, %45, %2, %c10003776, %c9947840, %c10240, %c10240, %45, %2, %c10003776, %c9958080, %c512, %c512, %45, %2, %c10003776, %c9958592, %c32768, %c32768, %45, %2, %c10003776, %c9991360, %c48, %c48, %45, %2, %c10003776, %c9991424, %c12288, %c12288, %45, %2, %c10003776, %c64, %c4, %c4, %45, %2, %c10003776, %c128, %c4, %c4, %45, %2, %c10003776, %c192, %c4, %c4, %45, %2, %c10003776, %c256, %c512, %c512, %45, %2, %c10003776, %c768, %c3145728, %c3145728, %45, %2, %c10003776, %c3146496, %c1024, %c1024, %45, %2, %c10003776, %c3147520, %c131072, %c131072, %45, %2, %c10003776, %c3278592, %c256, %c256, %45, %2, %c10003776, %c3278848, %c10240, %c10240, %45, %2, %c10003776, %c3289088, %c512, %c512, %45, %2, %c10003776, %c3289600, %c32768, %c32768, %45, %2, %c10003776, %c3322368, %c48, %c48, %45, %2, %c10003776, %c3322432, %c12288, %c12288, %45, %2, %c10003776, %c3334720, %c512, %c512, %45, %2, %c10003776, %c3335232, %c3145728, %c3145728, %45, %2, %c10003776, %c6480960, %c1024, %c1024, %45, %2, %c10003776, %c6481984, %c131072, %c131072, %45, %2, %c10003776, %c6613056, %c256, %c256, %45, %2, %c10003776, %c6613312, %c10240, %c10240, %45, %2, %c10003776, %c6623552, %c512, %c512, %45, %2, %c10003776, %c6624064, %c32768, %c32768, %45, %2, %c10003776, %c6656832, %c48, %c48, %45, %2, %c10003776, %c6656896, %c12288, %c12288, %45, %2, %c10003776, %c6669184, %c4, %c4, %45, %2, %c10003776, %c0, %c8, %c8, %45, %1, %c320, %c64, %c4, %c4, %45, %1, %c320, %c128, %c4, %c4, %45, %1, %c320, %c0, %c4, %c4, %45, %1, %c320, %c64, %c4, %c4, %45, %1, %c320, %c0, %c4, %c4, %45, %1, %c320, %c192, %c4, %c4, %45, %1, %c320, %c0, %c4, %c4, %45, %1, %c320, %c0, %c4, %c4, %45, %1, %c320, %c256, %c4, %c4 : !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<variable>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index, !stream.timepoint, !stream.resource<external>, index, index, index, index
```

After IPO + 132 globals (offset args can be removed with future global
support in IPO):
```mlir
%3:3 = call @jit__train_step$main(%0, %1, %2, %_train_state$0__timepoint, %_train_state$0, %_train_state$0__storage_size, %_train_state$0__offset, %_train_state$1__offset, %_train_state$2, %_train_state$2__storage_size, %_train_state$2__offset, %_train_state$3__offset, %_train_state$4__offset, %_train_state$5__offset, %_train_state$6__offset, %_train_state$7__offset, %_train_state$8__offset, %_train_state$9__offset, %_train_state$10__offset, %_train_state$11__offset, %_train_state$12__offset, %_train_state$13__offset, %_train_state$14__offset, %_train_state$15__offset, %_train_state$16__offset, %_train_state$17__offset, %_train_state$18__offset, %_train_state$19__offset, %_train_state$20__offset, %_train_state$21__offset, %_train_state$22__offset, %_train_state$23__offset, %_train_state$24__offset, %_train_state$25__offset, %_train_state$26__offset, %_train_state$27__offset, %_train_state$28__offset, %_train_state$29__offset, %_train_state$30__offset, %_train_state$31__offset, %_train_state$32__offset, %_train_state$33__offset, %_train_state$34__offset, %_prng_key$0__offset) : (!stream.resource<external>, !stream.resource<external>, !stream.resource<external>, !stream.timepoint, !stream.resource<variable>, index, index, index, !stream.resource<variable>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index) -> (!stream.timepoint, !stream.resource<variable>, !stream.resource<external>)
  return %40, %2, %1 : !stream.timepoint, !stream.resource<variable>, !stream.resource<external>
```

Full IR (after/before):
https://gist.github.com/benvanik/1c8e73353e3d7a3f7fe2d8ac21765186

Fixes #10903 (by reducing timepoint count crossing calls).