Koren/v3 msm

.gitignore

@@ -17,4 +17,4 @@
 **/Cargo.lock
 **/icicle/build/
 **/wrappers/rust/icicle-cuda-runtime/src/bindings.rs
-**/build/*
+**/build/*

icicle/backend/cpu/include/tasks_manager.h

@@ -1,7 +1,6 @@
 #pragma once
 #include <atomic>
 #include <thread>
-#include <stdexcept>
 #include <cassert>
 
 #define LOG_TASKS_PER_THREAD 3
@@ -13,7 +12,8 @@
 /**
  * @class TaskBase
  * @brief abstract base for a task supported by `TasksManager`.
- * Important
+ * Important - the user does not manually create these tasks - they are part of the manager and are accessed by the
+ * various get_task functions.
  */
 class TaskBase
 {
@@ -91,12 +91,19 @@ class TasksManager
    */
   Task* get_idle_or_completed_task();
 
+  /**
+   * @brief Get idle task to be dispatched without handling previous results (As it holds no previous result). This is
+   * not a blocking function,
+   * @return Task* - pointer to an idle task. nullptr if no task is available (All are either running or completed).
+   */
+  Task* get_idle_task();
+
   /**
    * @brief Get task that holds previous result to be handled by the user. This function blocks the code until a
    * completed task is found or all tasks are idle with no result.
    * @return Task* - pointer to a completed task. nullptr if no task is available (all are idle without results).
    * NOTE: The task's status should be updated if new tasks are to be assigned / other completed tasks are requested.
-   * Use dispatch
+   * Use dispatch if setting a new task, or set_idle if to just mark the task result as handled.
    */
   Task* get_completed_task();
 
@@ -140,9 +147,16 @@ class TasksManager
      */
     Task* get_idle_or_completed_task();
 
+    /**
+     * @brief Get idle task to be dispatched without handling previous results (As it holds no previous result). This
+     * isn't a blocking function - it checks all worker's tasks and returns.
+     * @return Task* - pointer to an idle task. nullptr if no task is available.
+     */
+    Task* get_idle_task();
+
     /**
      * @brief Get task that holds previous result to be handled by the user. This isn't a blocking function - it checks
-     * all worker's tasks and if no completed one is found a nullptr is returned.
+     * all worker's tasks and returns.
      * @param is_idle - boolean flag indicating if all worker's tasks are idle.
      * @return Task* - pointer to a completed task. nullptr if no task is available.
      * NOTE: if using completed_task to assign additional tasks, the existing result must be handled before hand.
@@ -213,6 +227,17 @@ Task* TasksManager<Task>::Worker::get_idle_or_completed_task()
   return nullptr;
 }
 
+template <class Task>
+Task* TasksManager<Task>::Worker::get_idle_task()
+{
+  for (int i = 0; i < m_tasks.size(); i++) {
+    m_next_task_idx = (1 + m_next_task_idx) & TASK_IDX_MASK;
+
+    if (m_tasks[m_next_task_idx].is_idle()) { return &m_tasks[m_next_task_idx]; }
+  }
+  return nullptr;
+}
+
 template <class Task>
 Task* TasksManager<Task>::Worker::get_completed_task(bool& is_idle)
 {
@@ -255,6 +280,20 @@ Task* TasksManager<Task>::get_idle_or_completed_task()
   }
 }
 
+template <class Task>
+Task* TasksManager<Task>::get_idle_task()
+{
+  Task* idle_task = nullptr;
+  for (int i = 0; i < m_workers.size(); i++) {
+    m_next_worker_idx = (m_next_worker_idx < m_workers.size() - 1) ? m_next_worker_idx + 1 : 0;
+
+    idle_task = m_workers[m_next_worker_idx].get_idle_task();
+    if (idle_task != nullptr) { return idle_task; }
+  }
+  // No completed tasks were found in the loop - return null.
+  return nullptr;
+}
+
 template <class Task>
 Task* TasksManager<Task>::get_completed_task()
 {

icicle/backend/cpu/src/curve/cpu_msm.cpp

@@ -1,43 +1,13 @@
 
+#include "cpu_msm.hpp"
 #include "icicle/backend/msm_backend.h"
-#include "icicle/errors.h"
-#include "icicle/runtime.h"
 
-#include "icicle/curves/projective.h"
-#include "icicle/curves/curve_config.h"
+// This file's purpose is just to register the hpp functions to the Icicle API.
 
-using namespace curve_config;
-using namespace icicle;
-
-template <typename S, typename A, typename P>
-eIcicleError
-cpu_msm(const Device& device, const S* scalars, const A* bases, int msm_size, const MSMConfig& config, P* results)
-{
-  for (auto batch_idx = 0; batch_idx < config.batch_size; ++batch_idx) {
-    P res = P::zero();
-    const S* batch_scalars = scalars + msm_size * batch_idx;
-    const A* batch_bases = config.are_bases_shared ? bases : bases + msm_size * batch_idx;
-    for (auto i = 0; i < msm_size; ++i) {
-      res = res + P::from_affine(batch_bases[i]) * batch_scalars[i];
-    }
-    results[batch_idx] = res;
-  }
-  return eIcicleError::SUCCESS;
-}
-
-template <typename A>
-eIcicleError cpu_msm_precompute_bases(
-  const Device& device, const A* input_bases, int nof_bases, const MSMConfig& config, A* output_bases)
-{
-  ICICLE_ASSERT(!config.are_points_on_device && !config.are_scalars_on_device);
-  memcpy(output_bases, input_bases, sizeof(A) * nof_bases);
-  return eIcicleError::SUCCESS;
-}
-
-REGISTER_MSM_BACKEND("CPU", (cpu_msm<scalar_t, affine_t, projective_t>));
-REGISTER_MSM_PRE_COMPUTE_BASES_BACKEND("CPU", cpu_msm_precompute_bases<affine_t>);
+REGISTER_MSM_PRE_COMPUTE_BASES_BACKEND("CPU", (cpu_msm_precompute_bases<affine_t, projective_t>));
+REGISTER_MSM_BACKEND("CPU", (cpu_msm<affine_t, projective_t>));
 
 #ifdef G2
-REGISTER_MSM_G2_BACKEND("CPU", (cpu_msm<scalar_t, g2_affine_t, g2_projective_t>));
-REGISTER_MSM_G2_PRE_COMPUTE_BASES_BACKEND("CPU", cpu_msm_precompute_bases<g2_affine_t>);
-#endif // G2
+REGISTER_MSM_G2_PRE_COMPUTE_BASES_BACKEND("CPU", (cpu_msm_precompute_bases<g2_affine_t, g2_projective_t>));
+REGISTER_MSM_G2_BACKEND("CPU", (cpu_msm<g2_affine_t, g2_projective_t>));
+#endif