Cache isResharding during segmentation.

[wujingyue]

Compile time spent time on isResharding sticks out after #3482. #4121 improves this a bit but it still takes 20% of the compile time and I believe there are low-hanging fruits.

profile001.pdf

One opportunity is to cache isResharding during segmentation. To quantify the potential saving, I made the following change

diff --git a/csrc/fusion_segmenter.cpp b/csrc/fusion_segmenter.cpp
index ea403117..55ed1ea6 100644
--- a/csrc/fusion_segmenter.cpp
+++ b/csrc/fusion_segmenter.cpp
@@ -3941,6 +3941,8 @@ SegmentCandidateFinder::SegmentCandidateFinder(
     bool multi_device)
     : options_(options), runtime_inputs_(inputs) {
   FUSER_PERF_SCOPE("SegmentCandidateFinder::SegmentCandidateFinder");
+  std::cerr << "enter SegmentCandidateFinder::SegmentCandidateFinder"
+            << std::endl;
   NVF_ERROR(
       !options_.only_segment_resharding_exprs ||
           (!options_.run_translate_welford &&
@@ -3956,6 +3958,8 @@ SegmentCandidateFinder::SegmentCandidateFinder(

   privatizeUpcast();
   findSegments();
+  std::cerr << "exit SegmentCandidateFinder::SegmentCandidateFinder"
+            << std::endl;
 }

 // Add runtimeInfo() accessor with validation
diff --git a/csrc/multidevice/utils.cpp b/csrc/multidevice/utils.cpp
index c74db90f..53eb42d7 100644
--- a/csrc/multidevice/utils.cpp
+++ b/csrc/multidevice/utils.cpp
@@ -536,6 +536,15 @@ bool isResharding(const Expr* expr) {
     return false;
   }

+  static int64_t count = 0;
+  static auto& unique_exprs = *(new std::unordered_set<const Expr*>());
+  unique_exprs.insert(expr);
+  if (count % 100 == 0) {
+    std::cerr << "isResharding has been called for " << count << " times with "
+              << unique_exprs.size() << " unique Exprs." << std::endl;
+  }
+  count++;
+
   // We don't use getTvsWithDifferentSharding because it creates a computeAtMap,
   // which is too costly
   for (auto* input : ir_utils::filterByType<TensorView>(expr->inputs())) {

to collect how many times isResharding is called during segmentation and for how many unique Exprs.

$ mpirun -np 1 --output-filename /tmp/test_multidevice bin/test_multidevice --gtest_filter=DistributedTransformerTest.Backward/__half

[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from DistributedTransformerTest
[ RUN      ] DistributedTransformerTest.Backward/__half
isResharding has been called for 0 times with 1 unique Exprs.
isResharding has been called for 100 times with 101 unique Exprs.
isResharding has been called for 200 times with 157 unique Exprs.
isResharding has been called for 300 times with 157 unique Exprs.
isResharding has been called for 400 times with 157 unique Exprs.
isResharding has been called for 500 times with 157 unique Exprs.
isResharding has been called for 600 times with 157 unique Exprs.
isResharding has been called for 700 times with 157 unique Exprs.
isResharding has been called for 800 times with 157 unique Exprs.
enter SegmentCandidateFinder::SegmentCandidateFinder
isResharding has been called for 900 times with 158 unique Exprs.
isResharding has been called for 1000 times with 163 unique Exprs.
isResharding has been called for 1100 times with 163 unique Exprs.
isResharding has been called for 1200 times with 163 unique Exprs.
isResharding has been called for 1300 times with 163 unique Exprs.
isResharding has been called for 1400 times with 163 unique Exprs.
isResharding has been called for 1500 times with 163 unique Exprs.
isResharding has been called for 1600 times with 165 unique Exprs.
isResharding has been called for 1700 times with 167 unique Exprs.
isResharding has been called for 1800 times with 167 unique Exprs.
isResharding has been called for 1900 times with 167 unique Exprs.
isResharding has been called for 2000 times with 168 unique Exprs.
isResharding has been called for 2100 times with 168 unique Exprs.
isResharding has been called for 2200 times with 169 unique Exprs.
isResharding has been called for 2300 times with 170 unique Exprs.
isResharding has been called for 2400 times with 170 unique Exprs.
isResharding has been called for 2500 times with 170 unique Exprs.
isResharding has been called for 2600 times with 170 unique Exprs.
isResharding has been called for 2700 times with 170 unique Exprs.
isResharding has been called for 2800 times with 170 unique Exprs.
isResharding has been called for 2900 times with 170 unique Exprs.
isResharding has been called for 3000 times with 170 unique Exprs.
isResharding has been called for 3100 times with 170 unique Exprs.
isResharding has been called for 3200 times with 170 unique Exprs.
isResharding has been called for 3300 times with 170 unique Exprs.
isResharding has been called for 3400 times with 170 unique Exprs.
isResharding has been called for 3500 times with 170 unique Exprs.
isResharding has been called for 3600 times with 170 unique Exprs.
isResharding has been called for 3700 times with 170 unique Exprs.
isResharding has been called for 3800 times with 170 unique Exprs.
isResharding has been called for 3900 times with 170 unique Exprs.
isResharding has been called for 4000 times with 170 unique Exprs.
isResharding has been called for 4100 times with 170 unique Exprs.
isResharding has been called for 4200 times with 170 unique Exprs.
isResharding has been called for 4300 times with 170 unique Exprs.
isResharding has been called for 4400 times with 170 unique Exprs.
isResharding has been called for 4500 times with 170 unique Exprs.
isResharding has been called for 4600 times with 170 unique Exprs.
isResharding has been called for 4700 times with 170 unique Exprs.
isResharding has been called for 4800 times with 170 unique Exprs.
isResharding has been called for 4900 times with 170 unique Exprs.
isResharding has been called for 5000 times with 170 unique Exprs.
isResharding has been called for 5100 times with 170 unique Exprs.
isResharding has been called for 5200 times with 170 unique Exprs.
isResharding has been called for 5300 times with 170 unique Exprs.
isResharding has been called for 5400 times with 170 unique Exprs.
isResharding has been called for 5500 times with 170 unique Exprs.
isResharding has been called for 5600 times with 170 unique Exprs.
isResharding has been called for 5700 times with 170 unique Exprs.
isResharding has been called for 5800 times with 170 unique Exprs.
isResharding has been called for 5900 times with 170 unique Exprs.
isResharding has been called for 6000 times with 170 unique Exprs.
isResharding has been called for 6100 times with 170 unique Exprs.
isResharding has been called for 6200 times with 170 unique Exprs.
isResharding has been called for 6300 times with 170 unique Exprs.
isResharding has been called for 6400 times with 170 unique Exprs.
isResharding has been called for 6500 times with 170 unique Exprs.
isResharding has been called for 6600 times with 170 unique Exprs.
isResharding has been called for 6700 times with 170 unique Exprs.
isResharding has been called for 6800 times with 170 unique Exprs.
isResharding has been called for 6900 times with 170 unique Exprs.
isResharding has been called for 7000 times with 170 unique Exprs.
isResharding has been called for 7100 times with 170 unique Exprs.
isResharding has been called for 7200 times with 170 unique Exprs.
isResharding has been called for 7300 times with 170 unique Exprs.
isResharding has been called for 7400 times with 170 unique Exprs.
isResharding has been called for 7500 times with 170 unique Exprs.
isResharding has been called for 7600 times with 170 unique Exprs.
isResharding has been called for 7700 times with 170 unique Exprs.
isResharding has been called for 7800 times with 170 unique Exprs.
isResharding has been called for 7900 times with 170 unique Exprs.
isResharding has been called for 8000 times with 170 unique Exprs.
isResharding has been called for 8100 times with 170 unique Exprs.
isResharding has been called for 8200 times with 170 unique Exprs.
isResharding has been called for 8300 times with 170 unique Exprs.
isResharding has been called for 8400 times with 170 unique Exprs.
isResharding has been called for 8500 times with 170 unique Exprs.
isResharding has been called for 8600 times with 170 unique Exprs.
isResharding has been called for 8700 times with 170 unique Exprs.
isResharding has been called for 8800 times with 170 unique Exprs.
isResharding has been called for 8900 times with 170 unique Exprs.
isResharding has been called for 9000 times with 170 unique Exprs.
isResharding has been called for 9100 times with 170 unique Exprs.
isResharding has been called for 9200 times with 170 unique Exprs.
isResharding has been called for 9300 times with 170 unique Exprs.
isResharding has been called for 9400 times with 170 unique Exprs.
isResharding has been called for 9500 times with 170 unique Exprs.
isResharding has been called for 9600 times with 170 unique Exprs.
isResharding has been called for 9700 times with 170 unique Exprs.
isResharding has been called for 9800 times with 170 unique Exprs.
isResharding has been called for 9900 times with 170 unique Exprs.
isResharding has been called for 10000 times with 170 unique Exprs.
isResharding has been called for 10100 times with 170 unique Exprs.
isResharding has been called for 10200 times with 170 unique Exprs.
isResharding has been called for 10300 times with 170 unique Exprs.
isResharding has been called for 10400 times with 170 unique Exprs.
isResharding has been called for 10500 times with 170 unique Exprs.
isResharding has been called for 10600 times with 170 unique Exprs.
isResharding has been called for 10700 times with 170 unique Exprs.
isResharding has been called for 10800 times with 170 unique Exprs.
isResharding has been called for 10900 times with 170 unique Exprs.
isResharding has been called for 11000 times with 170 unique Exprs.
isResharding has been called for 11100 times with 170 unique Exprs.
isResharding has been called for 11200 times with 170 unique Exprs.
isResharding has been called for 11300 times with 170 unique Exprs.
isResharding has been called for 11400 times with 170 unique Exprs.
isResharding has been called for 11500 times with 170 unique Exprs.
isResharding has been called for 11600 times with 170 unique Exprs.
isResharding has been called for 11700 times with 170 unique Exprs.
isResharding has been called for 11800 times with 170 unique Exprs.
isResharding has been called for 11900 times with 170 unique Exprs.
isResharding has been called for 12000 times with 170 unique Exprs.
isResharding has been called for 12100 times with 170 unique Exprs.
isResharding has been called for 12200 times with 170 unique Exprs.
isResharding has been called for 12300 times with 170 unique Exprs.
isResharding has been called for 12400 times with 170 unique Exprs.
isResharding has been called for 12500 times with 170 unique Exprs.
isResharding has been called for 12600 times with 170 unique Exprs.
isResharding has been called for 12700 times with 170 unique Exprs.
isResharding has been called for 12800 times with 170 unique Exprs.
isResharding has been called for 12900 times with 170 unique Exprs.
isResharding has been called for 13000 times with 170 unique Exprs.
isResharding has been called for 13100 times with 170 unique Exprs.
isResharding has been called for 13200 times with 170 unique Exprs.
isResharding has been called for 13300 times with 170 unique Exprs.
isResharding has been called for 13400 times with 170 unique Exprs.
isResharding has been called for 13500 times with 170 unique Exprs.
isResharding has been called for 13600 times with 170 unique Exprs.
isResharding has been called for 13700 times with 170 unique Exprs.
isResharding has been called for 13800 times with 170 unique Exprs.
isResharding has been called for 13900 times with 170 unique Exprs.
isResharding has been called for 14000 times with 170 unique Exprs.
isResharding has been called for 14100 times with 170 unique Exprs.
isResharding has been called for 14200 times with 170 unique Exprs.
isResharding has been called for 14300 times with 170 unique Exprs.
isResharding has been called for 14400 times with 170 unique Exprs.
isResharding has been called for 14500 times with 170 unique Exprs.
isResharding has been called for 14600 times with 170 unique Exprs.
isResharding has been called for 14700 times with 170 unique Exprs.
isResharding has been called for 14800 times with 170 unique Exprs.
isResharding has been called for 14900 times with 170 unique Exprs.
isResharding has been called for 15000 times with 170 unique Exprs.
isResharding has been called for 15100 times with 170 unique Exprs.
isResharding has been called for 15200 times with 170 unique Exprs.
isResharding has been called for 15300 times with 170 unique Exprs.
isResharding has been called for 15400 times with 170 unique Exprs.
isResharding has been called for 15500 times with 170 unique Exprs.
isResharding has been called for 15600 times with 170 unique Exprs.
isResharding has been called for 15700 times with 170 unique Exprs.
isResharding has been called for 15800 times with 170 unique Exprs.
isResharding has been called for 15900 times with 170 unique Exprs.
isResharding has been called for 16000 times with 170 unique Exprs.
isResharding has been called for 16100 times with 170 unique Exprs.
isResharding has been called for 16200 times with 170 unique Exprs.
isResharding has been called for 16300 times with 170 unique Exprs.
isResharding has been called for 16400 times with 170 unique Exprs.
isResharding has been called for 16500 times with 170 unique Exprs.
isResharding has been called for 16600 times with 170 unique Exprs.
isResharding has been called for 16700 times with 170 unique Exprs.
isResharding has been called for 16800 times with 170 unique Exprs.
isResharding has been called for 16900 times with 170 unique Exprs.
isResharding has been called for 17000 times with 170 unique Exprs.
isResharding has been called for 17100 times with 170 unique Exprs.
isResharding has been called for 17200 times with 170 unique Exprs.
isResharding has been called for 17300 times with 170 unique Exprs.
isResharding has been called for 17400 times with 170 unique Exprs.
isResharding has been called for 17500 times with 170 unique Exprs.
isResharding has been called for 17600 times with 170 unique Exprs.
isResharding has been called for 17700 times with 171 unique Exprs.
isResharding has been called for 17800 times with 171 unique Exprs.
isResharding has been called for 17900 times with 171 unique Exprs.
isResharding has been called for 18000 times with 171 unique Exprs.
isResharding has been called for 18100 times with 171 unique Exprs.
isResharding has been called for 18200 times with 171 unique Exprs.
isResharding has been called for 18300 times with 171 unique Exprs.
isResharding has been called for 18400 times with 171 unique Exprs.
isResharding has been called for 18500 times with 171 unique Exprs.
exit SegmentCandidateFinder::SegmentCandidateFinder
isResharding has been called for 18600 times with 234 unique Exprs.
isResharding has been called for 18700 times with 305 unique Exprs.
isResharding has been called for 18800 times with 369 unique Exprs.
isResharding has been called for 18900 times with 466 unique Exprs.

For this test case, isResharding was called >17600 times during segmentation with only ~170 unique Expr*. This means a potential cache hit rate of 90%