[isl-spatial-analysis] Add unicast hops calculation

NVlabs · Mar 27, 2024 · 48a7ebf · 48a7ebf
1 parent 3a1179e
commit 48a7ebf
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 18 deletions.
diff --git a/include/loop-analysis/spatial-analysis.hpp b/include/loop-analysis/spatial-analysis.hpp
@@ -26,6 +26,7 @@ struct TransferInfo
   {
     double accesses;
     double hops;
+    double unicast_hops;
   };
   std::map<std::pair<uint64_t, uint64_t>, AccessStats> compat_access_stats;
 

diff --git a/src/SConscript b/src/SConscript
@@ -303,6 +303,7 @@ applications/looptree-model/main.cpp
 
 unit_test_sources = Split("""
 unit-tests/test-main.cpp
+unit-tests/test-multicast.cpp
 unit-tests/test-simple-link-transfer.cpp
 """)
 

diff --git a/src/loop-analysis/isl-analysis/isl-to-legacy-adaptor.cpp b/src/loop-analysis/isl-analysis/isl-to-legacy-adaptor.cpp
@@ -190,7 +190,8 @@ CompoundDataMovementNest GenerateCompoundDataMovementNest(
         {
           tile.access_stats.stats[key] = AccessStats{
             .accesses = access_stats.accesses / tile.replication_factor,
-            .hops = access_stats.hops
+            .hops = access_stats.hops,
+            .unicast_hops = access_stats.unicast_hops
           };
         }
 

diff --git a/src/loop-analysis/spatial-analysis.cpp b/src/loop-analysis/spatial-analysis.cpp
@@ -171,6 +171,10 @@ SimpleMulticastModel::Apply(const Fill& fill, const Occupancy& occ) const
   auto p_data_to_max_x = isl_map_lexmax(
     isl_map_project_out(wrapped_fill.copy(), isl_dim_out, 1, 1)
   );
+  auto p_data_to_x = isl_map_project_out(wrapped_fill.copy(),
+                                         isl_dim_out,
+                                         1,
+                                         1);
   auto p_x_hops_cost = isl_pw_qpolynomial_from_qpolynomial(
     isl_qpolynomial_add(
       isl_qpolynomial_var_on_domain(
@@ -183,12 +187,14 @@ SimpleMulticastModel::Apply(const Fill& fill, const Occupancy& occ) const
       )
     )
   );
-  auto p_x_to_data = isl_map_reverse(p_data_to_max_x);
-  p_data_per_dst = isl_map_card(isl_map_copy(p_x_to_data));
-  auto p_x_hops_per_dst = isl_pw_qpolynomial_mul(p_x_hops_cost,
-                                                 p_data_per_dst);
+  auto p_max_x_to_data = isl_map_reverse(p_data_to_max_x);
+  p_data_per_dst = isl_map_card(isl_map_copy(p_max_x_to_data));
+  auto p_x_hops_per_dst = isl_pw_qpolynomial_mul(
+    isl_pw_qpolynomial_copy(p_x_hops_cost),
+    p_data_per_dst
+  );
   double total_x_hops = 0;
-  auto p_x_domain = isl_map_domain(p_x_to_data);
+  auto p_x_domain = isl_map_domain(p_max_x_to_data);
   if (isl_set_is_singleton(isl_set_copy(p_x_domain)))
   {
     auto p_sample = isl_set_sample_point(p_x_domain);
@@ -208,6 +214,20 @@ SimpleMulticastModel::Apply(const Fill& fill, const Occupancy& occ) const
     ));
   }
 
+  auto p_x_to_data = isl_map_reverse(p_data_to_x);
+  p_data_per_dst = isl_map_card(isl_map_copy(p_x_to_data));
+  auto p_x_unicast_hops_per_dst = isl_pw_qpolynomial_mul(
+    p_x_hops_cost,
+    p_data_per_dst
+  );
+  double total_x_unicast_hops = 0;
+  total_x_unicast_hops = isl::val_to_double(isl_qpolynomial_get_constant_val(
+    isl_pw_qpolynomial_as_qpolynomial(isl_set_apply_pw_qpolynomial(
+      isl_map_domain(p_x_to_data),
+      p_x_unicast_hops_per_dst
+    ))
+  ));
+
   auto p_total_accesses = isl_pw_qpolynomial_sum(isl_map_card(isl_set_unwrap(
     wrapped_fill.domain().release()
   )));
@@ -218,9 +238,11 @@ SimpleMulticastModel::Apply(const Fill& fill, const Occupancy& occ) const
   );
 
   auto total_hops = total_x_hops + total_y_hops;
+  auto total_unicast_hops = total_x_unicast_hops + total_y_hops;
   auto& stats = transfer_info.compat_access_stats[std::make_pair(1, 1)];
   stats.accesses = accesses;
   stats.hops = total_hops / accesses;
+  stats.unicast_hops = total_unicast_hops / accesses;
 
   transfer_info.total_child_accesses = isl::val_to_double(
     isl::get_val_from_singular(

diff --git a/src/unit-tests/test-multicast.cpp b/src/unit-tests/test-multicast.cpp
@@ -8,10 +8,11 @@ BOOST_AUTO_TEST_CASE(TestSimpleMulticastModel)
   using namespace analysis;
 
   /* Graphically:
-   *  ------------ t
-   *  ---- x     ---- x 
+   *  ------------> t
+   *  ----> x     ----> x 
    * | [0] [1]   | [1] [2]
    * | [1] [2]   | [2] [3]
+   * v           v
    * y           y        
    */
   auto fill = Fill(
@@ -29,12 +30,9 @@ BOOST_AUTO_TEST_CASE(TestSimpleMulticastModel)
   auto info = multicast_model.Apply(fill, occ);
 
   auto& singles = info.compat_access_stats.at(std::make_pair(1, 1));
-  BOOST_CHECK(singles.accesses == 2);
-  BOOST_CHECK(singles.hops == 4);
-
-  auto& doubles = info.compat_access_stats.at(std::make_pair(2, 1));
-  BOOST_CHECK(doubles.accesses == 4);
-  BOOST_CHECK(doubles.hops == 3.5);
+  BOOST_CHECK_CLOSE(singles.accesses, 6, 0.1);
+  BOOST_CHECK_CLOSE(singles.hops, 22.0/6, 0.1);
+  BOOST_CHECK_CLOSE(singles.unicast_hops, 4, 0.1);
 }
 
 BOOST_AUTO_TEST_CASE(TestSimpleMulticastModel_FlattenedCoords)
@@ -64,8 +62,7 @@ BOOST_AUTO_TEST_CASE(TestSimpleMulticastModel_FlattenedCoords)
   auto info = multicast_model.Apply(fill, occ);
 
   auto& singles = info.compat_access_stats.at(std::make_pair(1, 1));
-  BOOST_CHECK(singles.accesses == 4);
-
-  auto& doubles = info.compat_access_stats.at(std::make_pair(2, 1));
-  BOOST_CHECK(doubles.accesses == 4);
+  BOOST_CHECK_CLOSE(singles.accesses, 6, 0.1);
+  BOOST_CHECK_CLOSE(singles.hops, 29.0/6, 0.1);
+  BOOST_CHECK_CLOSE(singles.unicast_hops, 32.0/6, 0.1);
 }