Revert "Remove size-based partitioning heuristic (#256)" (#376)

* Revert "Remove size-based partitioning heuristic (#256)"

This reverts commit cfe828e.

* Clang-tidy fixes

Author: manopapad

legate/settings.py

@@ -81,6 +81,48 @@ class LegateRuntimeSettings(Settings):
         """,
     )
 
+    min_gpu_chunk: EnvOnlySetting[int] = EnvOnlySetting(
+        "min_gpu_chunk",
+        "LEGATE_MIN_GPU_CHUNK",
+        default=1048576,  # 1 << 20
+        test_default=2,
+        convert=convert_int,
+        help="""
+        If using GPUs, any task operating on arrays smaller than this will
+        not be parallelized across more than one GPU.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
+    min_cpu_chunk: EnvOnlySetting[int] = EnvOnlySetting(
+        "min_cpu_chunk",
+        "LEGATE_MIN_CPU_CHUNK",
+        default=16384,  # 1 << 14
+        test_default=2,
+        convert=convert_int,
+        help="""
+        If using CPUs, any task operating on arrays smaller than this will
+        not be parallelized across more than one core.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
+    min_omp_chunk: EnvOnlySetting[int] = EnvOnlySetting(
+        "min_omp_chunk",
+        "LEGATE_MIN_OMP_CHUNK",
+        default=131072,  # 1 << 17
+        test_default=2,
+        convert=convert_int,
+        help="""
+        If using OpenMP, any task operating on arrays smaller than this will
+        not be parallelized across more than one OpenMP group.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
     window_size: EnvOnlySetting[int] = EnvOnlySetting(
         "window_size",
         "LEGATE_WINDOW_SIZE",
@@ -94,6 +136,43 @@ class LegateRuntimeSettings(Settings):
         """,
     )
 
+    max_pending_exceptions: EnvOnlySetting[int] = EnvOnlySetting(
+        "max_pending_exceptions",
+        "LEGATE_MAX_PENDING_EXCEPTIONS",
+        default=64,
+        test_default=1,
+        convert=convert_int,
+        help="""
+        How many possibly-exception-throwing tasks to emit before blocking.
+        Legate by default does not wait for operations to complete, but instead
+        "runs ahead" and continues launching work, which will complete
+        asynchronously. If an operation throws an exception, then by the time
+        an exception is reported execution may have progressed beyond the
+        launch of the faulting operation. If you need to check for an exception
+        at the exact point where it might get thrown (e.g. to catch it and
+        recover gracefully), set this to 1. Note that this will introduce more
+        blocking in the control logic of your program, likely reducing overall
+        utilization.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
+    precise_exception_trace: EnvOnlySetting[bool] = EnvOnlySetting(
+        "precise_exception_trace",
+        "LEGATE_PRECISE_EXCEPTION_TRACE",
+        default=False,
+        test_default=False,
+        convert=convert_bool,
+        help="""
+        Whether to capture the stacktrace at the point when a potentially
+        faulting operation is launched, so a more accurate error location can
+        be reported in case an exception is thrown.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
     field_reuse_frac: EnvOnlySetting[int] = EnvOnlySetting(
         "field_reuse_frac",
         "LEGATE_FIELD_REUSE_FRAC",
@@ -125,6 +204,23 @@ class LegateRuntimeSettings(Settings):
         """,
     )
 
+    max_lru_length: EnvOnlySetting[int] = EnvOnlySetting(
+        "max_lru_length",
+        "LEGATE_MAX_LRU_LENGTH",
+        default=5,
+        test_default=1,
+        convert=convert_int,
+        help="""
+        Once the last Store of a given shape is garbage collected, the
+        resources associated with it are placed on an LRU queue, rather than
+        getting freed immediately, in case the program creates a Store of the
+        same shape in the near future. This setting controls the length of that
+        LRU queue.
+
+        This is a read-only environment variable setting used by the runtime.
+        """,
+    )
+
     disable_mpi: EnvOnlySetting[bool] = EnvOnlySetting(
         "disable_mpi",
         "LEGATE_DISABLE_MPI",

src/core/legate_c.h

@@ -66,7 +66,10 @@ typedef enum legate_core_tunable_t {
   LEGATE_CORE_TUNABLE_TOTAL_GPUS,
   LEGATE_CORE_TUNABLE_NUM_NODES,
   LEGATE_CORE_TUNABLE_HAS_SOCKET_MEM,
+  LEGATE_CORE_TUNABLE_MIN_SHARD_VOLUME,
+  LEGATE_CORE_TUNABLE_WINDOW_SIZE,
   LEGATE_CORE_TUNABLE_FIELD_REUSE_SIZE,
+  LEGATE_CORE_TUNABLE_MAX_LRU_LENGTH,
 } legate_core_tunable_t;
 
 typedef enum legate_core_variant_t {

src/core/mapping/detail/core_mapper.cc

@@ -58,8 +58,19 @@ class CoreMapper final : public Mapper {
 
  private:
   const LocalMachine machine{};
+  // TODO(wonchanl): Some of these should be moved to legate::detail::Config
+  const int64_t min_gpu_chunk{
+    extract_env("LEGATE_MIN_GPU_CHUNK", MIN_GPU_CHUNK_DEFAULT, MIN_GPU_CHUNK_TEST)};
+  const int64_t min_cpu_chunk{
+    extract_env("LEGATE_MIN_CPU_CHUNK", MIN_CPU_CHUNK_DEFAULT, MIN_CPU_CHUNK_TEST)};
+  const int64_t min_omp_chunk{
+    extract_env("LEGATE_MIN_OMP_CHUNK", MIN_OMP_CHUNK_DEFAULT, MIN_OMP_CHUNK_TEST)};
+  const uint32_t window_size{
+    extract_env("LEGATE_WINDOW_SIZE", WINDOW_SIZE_DEFAULT, WINDOW_SIZE_TEST)};
   const uint32_t field_reuse_frac{
     extract_env("LEGATE_FIELD_REUSE_FRAC", FIELD_REUSE_FRAC_DEFAULT, FIELD_REUSE_FRAC_TEST)};
+  const uint32_t max_lru_length{
+    extract_env("LEGATE_MAX_LRU_LENGTH", MAX_LRU_LENGTH_DEFAULT, MAX_LRU_LENGTH_TEST)};
 };
 
 void CoreMapper::set_machine(const legate::mapping::MachineQueryInterface* /*m*/) {}
@@ -91,9 +102,25 @@ Scalar CoreMapper::tunable_value(TunableID tunable_id)
     case LEGATE_CORE_TUNABLE_NUM_NODES: {
       return Scalar{static_cast<int32_t>(machine.total_nodes)};
     }
+    case LEGATE_CORE_TUNABLE_MIN_SHARD_VOLUME: {
+      // TODO(wonchanl): make these profile guided
+      if (machine.has_gpus()) {
+        // Make sure we can get at least 1M elements on each GPU
+        return Scalar{min_gpu_chunk};
+      }
+      if (machine.has_omps()) {
+        // Make sure we get at least 128K elements on each OpenMP
+        return Scalar{min_omp_chunk};
+      }
+      // Make sure we can get at least 8KB elements on each CPU
+      return Scalar{min_cpu_chunk};
+    }
     case LEGATE_CORE_TUNABLE_HAS_SOCKET_MEM: {
       return Scalar{machine.has_socket_memory()};
     }
+    case LEGATE_CORE_TUNABLE_WINDOW_SIZE: {
+      return Scalar{window_size};
+    }
     case LEGATE_CORE_TUNABLE_FIELD_REUSE_SIZE: {
       // Multiply this by the total number of nodes and then scale by the frac
       const uint64_t global_mem_size =
@@ -102,6 +129,9 @@ Scalar CoreMapper::tunable_value(TunableID tunable_id)
                                                           : machine.system_memory().capacity());
       return Scalar{global_mem_size / field_reuse_frac};
     }
+    case LEGATE_CORE_TUNABLE_MAX_LRU_LENGTH: {
+      return Scalar{max_lru_length};
+    }
     default: break;
   }
   // Illegal tunable variable

src/core/runtime/detail/partition_manager.cc

@@ -23,6 +23,18 @@
 
 namespace legate::detail {
 
+PartitionManager::PartitionManager(Runtime* runtime)
+{
+  auto mapper_id = runtime->core_library()->get_mapper_id();
+
+  min_shard_volume_ =
+    runtime->get_tunable<int64_t>(mapper_id, LEGATE_CORE_TUNABLE_MIN_SHARD_VOLUME);
+
+  if (LegateDefined(LEGATE_USE_DEBUG)) {
+    assert(min_shard_volume_ > 0);
+  }
+}
+
 const std::vector<uint32_t>& PartitionManager::get_factors(const mapping::detail::Machine& machine)
 {
   auto curr_num_pieces = machine.count();
@@ -64,7 +76,7 @@ tuple<uint64_t> PartitionManager::compute_launch_shape(const mapping::detail::Ma
   // Prune out any dimensions that are 1
   std::vector<size_t> temp_shape{};
   std::vector<uint32_t> temp_dims{};
-  size_t volume = 1;
+  int64_t volume = 1;
 
   temp_dims.reserve(shape.size());
   temp_shape.reserve(shape.size());
@@ -76,23 +88,28 @@ tuple<uint64_t> PartitionManager::compute_launch_shape(const mapping::detail::Ma
     }
     temp_shape.push_back(extent);
     temp_dims.push_back(dim);
-    volume *= extent;
+    volume *= static_cast<int64_t>(extent);
   }
 
-  if (temp_shape.empty()) {
+  // Figure out how many shards we can make with this array
+  int64_t max_pieces = (volume + min_shard_volume_ - 1) / min_shard_volume_;
+  assert(volume == 0 || max_pieces > 0);
+  // If we can only make one piece return that now
+  if (max_pieces <= 1) {
     return {};
   }
 
   // TODO(wonchanl): We need a better heuristic
-  auto max_pieces = curr_num_pieces;
+  max_pieces = curr_num_pieces;
 
   // First compute the N-th root of the number of pieces
   const auto ndim = temp_shape.size();
+  assert(ndim > 0);
   std::vector<size_t> temp_result{};
 
   if (1 == ndim) {
     // Easy one dimensional case
-    temp_result.push_back(std::min<size_t>(temp_shape.front(), max_pieces));
+    temp_result.push_back(std::min<size_t>(temp_shape.front(), static_cast<size_t>(max_pieces)));
   } else if (2 == ndim) {
     if (volume < max_pieces) {
       // TBD: Once the max_pieces heuristic is fixed, this should never happen
@@ -112,11 +129,11 @@ tuple<uint64_t> PartitionManager::compute_launch_shape(const mapping::detail::Ma
       // try rounding n both up and down
       constexpr auto EPSILON = 1e-12;
 
-      auto n1 = std::max<size_t>(1, static_cast<size_t>(std::floor(n + EPSILON)));
+      auto n1 = std::max<int64_t>(1, static_cast<int64_t>(std::floor(n + EPSILON)));
       while (max_pieces % n1 != 0) {
         --n1;
       }
-      auto n2 = std::max<size_t>(1, static_cast<size_t>(std::floor(n - EPSILON)));
+      auto n2 = std::max<int64_t>(1, static_cast<int64_t>(std::floor(n - EPSILON)));
       while (max_pieces % n2 != 0) {
         ++n2;
       }

src/core/runtime/detail/partition_manager.h

@@ -30,8 +30,12 @@ struct Machine;
 
 namespace legate::detail {
 
+class Runtime;
+
 class PartitionManager {
  public:
+  explicit PartitionManager(Runtime* runtime);
+
   [[nodiscard]] const std::vector<uint32_t>& get_factors(const mapping::detail::Machine& machine);
 
   [[nodiscard]] tuple<uint64_t> compute_launch_shape(const mapping::detail::Machine& machine,
@@ -67,6 +71,7 @@ class PartitionManager {
                                   Legion::FieldID field_id);
 
  private:
+  int64_t min_shard_volume_{};
   std::unordered_map<uint32_t, std::vector<uint32_t>> all_factors_{};
 
   using TilingCacheKey = std::pair<Legion::IndexSpace, Tiling>;

src/core/runtime/detail/runtime.cc

@@ -89,7 +89,6 @@ constexpr const char* const TOPLEVEL_NAME     = "Legate Core Toplevel Task";
 
 Runtime::Runtime()
   : legion_runtime_{Legion::Runtime::get_runtime()},
-    window_size_{extract_env("LEGATE_WINDOW_SIZE", WINDOW_SIZE_DEFAULT, WINDOW_SIZE_TEST)},
     field_reuse_freq_{
       extract_env("LEGATE_FIELD_REUSE_FREQ", FIELD_REUSE_FREQ_DEFAULT, FIELD_REUSE_FREQ_TEST)},
     force_consensus_match_{!!extract_env("LEGATE_CONSENSUS", CONSENSUS_DEFAULT, CONSENSUS_TEST)}
@@ -198,7 +197,7 @@ void Runtime::initialize(Legion::Context legion_context)
   core_library_   = find_library(CORE_LIBRARY_NAME, false /*can_fail*/);
   // TODO(jfaibussowit): Use smart pointers for these
   communicator_manager_ = new CommunicatorManager{};
-  partition_manager_    = new PartitionManager{};
+  partition_manager_    = new PartitionManager{this};
   machine_manager_      = new MachineManager{};
   provenance_manager_   = new ProvenanceManager{};
   Config::has_socket_mem =

src/core/runtime/detail/runtime.h

@@ -359,7 +359,7 @@ class Runtime {
     registered_shardings_{};
 
   std::vector<InternalSharedPtr<Operation>> operations_;
-  size_t window_size_{};
+  size_t window_size_{1};
   uint64_t next_unique_id_{};
 
   using RegionFieldID = std::pair<Legion::LogicalRegion, Legion::FieldID>;

src/core/task/detail/return.h

@@ -50,7 +50,7 @@ struct ReturnedException {
   ReturnedException() = default;
   ReturnedException(int32_t index, std::string_view error_message);
 
-  static inline constexpr auto MAX_MESSAGE_SIZE = 2048;
+  static inline constexpr auto MAX_MESSAGE_SIZE = 256;
 
   [[nodiscard]] bool raised() const;
 

src/env_defaults.h

@@ -17,15 +17,43 @@
 
 #pragma once
 
+#include "legate_defines.h"
+
+// 1 << 20 (need actual number for python to parse)
+#define MIN_GPU_CHUNK_DEFAULT 1048576
+#define MIN_GPU_CHUNK_TEST 2
+
+// 1 << 14 (need actual number for python to parse)
+#define MIN_CPU_CHUNK_DEFAULT 16384
+#define MIN_CPU_CHUNK_TEST 2
+
+// 1 << 17 (need actual number for python to parse)
+#define MIN_OMP_CHUNK_DEFAULT 131072
+#define MIN_OMP_CHUNK_TEST 2
+
 #define WINDOW_SIZE_DEFAULT 1
 #define WINDOW_SIZE_TEST 1
 
+#if LegateDefined(LEGATE_USE_DEBUG)
+// In debug mode, the default is always block on tasks that can throw exceptions
+#define MAX_PENDING_EXCEPTIONS_DEFAULT 1
+#else
+#define MAX_PENDING_EXCEPTIONS_DEFAULT 64
+#endif
+#define MAX_PENDING_EXCEPTIONS_TEST 1
+
+#define PRECISE_EXCEPTION_TRACE_DEFAULT 0
+#define PRECISE_EXCEPTION_TRACE_TEST 0
+
 #define FIELD_REUSE_FRAC_DEFAULT 256
 #define FIELD_REUSE_FRAC_TEST 1
 
 #define FIELD_REUSE_FREQ_DEFAULT 32
 #define FIELD_REUSE_FREQ_TEST 8
 
+#define MAX_LRU_LENGTH_DEFAULT 5
+#define MAX_LRU_LENGTH_TEST 1
+
 #define DISABLE_MPI_DEFAULT 0
 #define DISABLE_MPI_TEST 0
 

tests/unit/legate/core/util/task_util.py

@@ -200,9 +200,7 @@ def deep_clone() -> TestFunction[_P, _T]:
             fn_copy.self = fn_copy
 
             def mark_called() -> None:
-                # FIXME: This assertion doesn't hold if more than one Python
-                # task gets launched.
-                # assert not fn_copy.self.called, f"{fn}, {fn_copy.self}"
+                assert not fn_copy.self.called, f"{fn}, {fn_copy.self}"
                 fn_copy.self.called = True
 
             fn_copy.self.mark_called = mark_called