Implement Integer Rings in RNS Representation

icicle/backend/cpu/src/field/cpu_ntt.cpp

@@ -3,24 +3,24 @@
 using namespace field_config;
 using namespace icicle;
 
-eIcicleError
-cpu_ntt_init_domain(const Device& device, const scalar_t& primitive_root, const NTTInitDomainConfig& config)
+template <typename S = scalar_t>
+eIcicleError cpu_ntt_init_domain(const Device& device, const S& primitive_root, const NTTInitDomainConfig& config)
 {
-  auto err = ntt_cpu::CpuNttDomain<scalar_t>::cpu_ntt_init_domain(device, primitive_root, config);
+  auto err = ntt_cpu::CpuNttDomain<S>::cpu_ntt_init_domain(device, primitive_root, config);
   return err;
 }
 
 template <typename S = scalar_t>
 eIcicleError cpu_ntt_release_domain(const Device& device, const S& dummy)
 {
-  auto err = ntt_cpu::CpuNttDomain<scalar_t>::cpu_ntt_release_domain(device);
+  auto err = ntt_cpu::CpuNttDomain<S>::cpu_ntt_release_domain(device);
   return err;
 }
 
 template <typename S = scalar_t>
 eIcicleError cpu_get_root_of_unity_from_domain(const Device& device, uint64_t logn, S* rou)
 {
-  auto err = ntt_cpu::CpuNttDomain<scalar_t>::get_root_of_unity_from_domain(device, logn, rou);
+  auto err = ntt_cpu::CpuNttDomain<S>::get_root_of_unity_from_domain(device, logn, rou);
   return err;
 }
 
@@ -32,11 +32,18 @@ cpu_ntt(const Device& device, const E* input, uint64_t size, NTTDir dir, const N
   return err;
 }
 
-REGISTER_NTT_INIT_DOMAIN_BACKEND("CPU", (cpu_ntt_init_domain));
+REGISTER_NTT_INIT_DOMAIN_BACKEND("CPU", (cpu_ntt_init_domain<scalar_t>));
 REGISTER_NTT_RELEASE_DOMAIN_BACKEND("CPU", cpu_ntt_release_domain<scalar_t>);
 REGISTER_NTT_GET_ROU_FROM_DOMAIN_BACKEND("CPU", cpu_get_root_of_unity_from_domain<scalar_t>);
 REGISTER_NTT_BACKEND("CPU", (cpu_ntt<scalar_t, scalar_t>));
 
 #ifdef EXT_FIELD
 REGISTER_NTT_EXT_FIELD_BACKEND("CPU", (cpu_ntt<scalar_t, extension_t>));
 #endif // EXT_FIELD
+
+#ifdef RING
+REGISTER_NTT_INIT_DOMAIN_RING_RNS_BACKEND("CPU", (cpu_ntt_init_domain<scalar_rns_t>));
+REGISTER_NTT_RELEASE_DOMAIN_RING_RNS_BACKEND("CPU", cpu_ntt_release_domain<scalar_rns_t>);
+REGISTER_NTT_GET_ROU_FROM_DOMAIN_RING_RNS_BACKEND("CPU", cpu_get_root_of_unity_from_domain<scalar_rns_t>);
+REGISTER_NTT_RING_RNS_BACKEND("CPU", (cpu_ntt<scalar_rns_t, scalar_rns_t>));
+#endif // RING

icicle/backend/cpu/src/field/cpu_vec_ops.cpp

@@ -733,9 +733,6 @@ eIcicleError cpu_matrix_transpose(
 }
 
 REGISTER_MATRIX_TRANSPOSE_BACKEND("CPU", cpu_matrix_transpose<scalar_t>);
-#ifdef EXT_FIELD
-REGISTER_MATRIX_TRANSPOSE_EXT_FIELD_BACKEND("CPU", cpu_matrix_transpose<extension_t>);
-#endif // EXT_FIELD
 
 /*********************************** BIT REVERSE ***********************************/
 template <typename T>
@@ -765,9 +762,6 @@ cpu_bit_reverse(const Device& device, const T* vec_in, uint64_t size, const VecO
 }
 
 REGISTER_BIT_REVERSE_BACKEND("CPU", cpu_bit_reverse<scalar_t>);
-#ifdef EXT_FIELD
-REGISTER_BIT_REVERSE_EXT_FIELD_BACKEND("CPU", cpu_bit_reverse<extension_t>);
-#endif // EXT_FIELD
 
 /*********************************** SLICE ***********************************/
 
@@ -802,9 +796,6 @@ eIcicleError cpu_slice(
 }
 
 REGISTER_SLICE_BACKEND("CPU", cpu_slice<scalar_t>);
-#ifdef EXT_FIELD
-REGISTER_SLICE_EXT_FIELD_BACKEND("CPU", cpu_slice<extension_t>);
-#endif // EXT_FIELD
 
 /*********************************** Highest non-zero idx ***********************************/
 template <typename T>
@@ -993,6 +984,9 @@ eIcicleError cpu_poly_divide(
 REGISTER_POLYNOMIAL_DIVISION("CPU", cpu_poly_divide<scalar_t>);
 
 #ifdef EXT_FIELD
+REGISTER_MATRIX_TRANSPOSE_EXT_FIELD_BACKEND("CPU", cpu_matrix_transpose<extension_t>);
+REGISTER_BIT_REVERSE_EXT_FIELD_BACKEND("CPU", cpu_bit_reverse<extension_t>);
+REGISTER_SLICE_EXT_FIELD_BACKEND("CPU", cpu_slice<extension_t>);
 REGISTER_VECTOR_ADD_EXT_FIELD_BACKEND("CPU", cpu_vector_add<extension_t>);
 REGISTER_VECTOR_ACCUMULATE_EXT_FIELD_BACKEND("CPU", cpu_vector_accumulate<extension_t>);
 REGISTER_VECTOR_SUB_EXT_FIELD_BACKEND("CPU", cpu_vector_sub<extension_t>);
@@ -1006,4 +1000,54 @@ REGISTER_SCALAR_MUL_VEC_EXT_FIELD_BACKEND("CPU", cpu_scalar_mul<extension_t>);
 REGISTER_SCALAR_ADD_VEC_EXT_FIELD_BACKEND("CPU", cpu_scalar_add<extension_t>);
 REGISTER_SCALAR_SUB_VEC_EXT_FIELD_BACKEND("CPU", cpu_scalar_sub<extension_t>);
 REGISTER_EXECUTE_PROGRAM_EXT_FIELD_BACKEND("CPU", cpu_execute_program<extension_t>);
-#endif // EXT_FIELD
+#endif // EXT_FIELD
+
+#ifdef RING
+// Register APIs for rns type
+REGISTER_MATRIX_TRANSPOSE_RING_RNS_BACKEND("CPU", cpu_matrix_transpose<scalar_rns_t>);
+REGISTER_BIT_REVERSE_RING_RNS_BACKEND("CPU", cpu_bit_reverse<scalar_rns_t>);
+REGISTER_SLICE_RING_RNS_BACKEND("CPU", cpu_slice<scalar_rns_t>);
+REGISTER_VECTOR_ADD_RING_RNS_BACKEND("CPU", cpu_vector_add<scalar_rns_t>);
+REGISTER_VECTOR_ACCUMULATE_RING_RNS_BACKEND("CPU", cpu_vector_accumulate<scalar_rns_t>);
+REGISTER_VECTOR_SUB_RING_RNS_BACKEND("CPU", cpu_vector_sub<scalar_rns_t>);
+REGISTER_VECTOR_MUL_RING_RNS_BACKEND("CPU", (cpu_vector_mul<scalar_rns_t, scalar_rns_t>));
+REGISTER_VECTOR_DIV_RING_RNS_BACKEND("CPU", cpu_vector_div<scalar_rns_t>);
+REGISTER_CONVERT_MONTGOMERY_RING_RNS_BACKEND("CPU", cpu_convert_montgomery<scalar_rns_t>);
+REGISTER_VECTOR_SUM_RING_RNS_BACKEND("CPU", cpu_vector_sum<scalar_rns_t>);
+REGISTER_VECTOR_PRODUCT_RING_RNS_BACKEND("CPU", cpu_vector_product<scalar_rns_t>);
+REGISTER_SCALAR_MUL_VEC_RING_RNS_BACKEND("CPU", cpu_scalar_mul<scalar_rns_t>);
+REGISTER_SCALAR_ADD_VEC_RING_RNS_BACKEND("CPU", cpu_scalar_add<scalar_rns_t>);
+REGISTER_SCALAR_SUB_VEC_RING_RNS_BACKEND("CPU", cpu_scalar_sub<scalar_rns_t>);
+
+// RNS conversion
+template <typename SrcType, typename DstType, bool into_rns>
+eIcicleError
+cpu_convert_rns(const Device& device, const SrcType* input, uint64_t size, const VecOpsConfig& config, DstType* output)
+{
+  tf::Taskflow taskflow;
+  tf::Executor executor;
+  const uint64_t total_nof_operations = size * config.batch_size;
+
+  const int nof_workers = get_nof_workers(config);
+  const uint64_t worker_task_size = (total_nof_operations + nof_workers - 1) / nof_workers; // round up
+
+  for (uint64_t start_idx = 0; start_idx < total_nof_operations; start_idx += worker_task_size) {
+    taskflow.emplace([=]() {
+      const uint64_t end_idx = std::min(start_idx + worker_task_size, total_nof_operations);
+      for (uint64_t idx = start_idx; idx < end_idx; ++idx) {
+        if constexpr (into_rns) {
+          DstType::convert_direct_to_rns(&input[idx].limbs_storage, &output[idx].limbs_storage);
+        } else {
+          SrcType::convert_rns_to_direct(&input[idx].limbs_storage, &output[idx].limbs_storage);
+        }
+      }
+    });
+  }
+
+  executor.run(taskflow).wait();
+  taskflow.clear();
+  return eIcicleError::SUCCESS;
+}
+REGISTER_CONVERT_TO_RNS_BACKEND("CPU", (cpu_convert_rns<scalar_t, scalar_rns_t, true /*into rns*/>));
+REGISTER_CONVERT_FROM_RNS_BACKEND("CPU", (cpu_convert_rns<scalar_rns_t, scalar_t, false /*from rns*/>));
+#endif // RING

icicle/cmake/target_editor.cmake

@@ -23,6 +23,7 @@ function(handle_ring TARGET)
   target_sources(${TARGET} PRIVATE
     src/fields/ffi_extern.cpp
     src/vec_ops.cpp
+    src/rings/rns_vec_ops.cpp
     src/matrix_ops.cpp
   )
 endfunction()

icicle/include/icicle/backend/ntt_backend.h

@@ -89,4 +89,68 @@ namespace icicle {
       return true;                                                                                                     \
     }();                                                                                                               \
   }
+
+#ifdef RING
+  /*************************** NTT ***************************/
+  using NttRingRnsImpl = std::function<eIcicleError(
+    const Device& device,
+    const scalar_rns_t* input,
+    int size,
+    NTTDir dir,
+    const NTTConfig<scalar_rns_t>& config,
+    scalar_rns_t* output)>;
+
+  void register_ring_rns_ntt(const std::string& deviceType, NttRingRnsImpl impl);
+
+  #define REGISTER_NTT_RING_RNS_BACKEND(DEVICE_TYPE, FUNC)                                                             \
+    namespace {                                                                                                        \
+      static bool UNIQUE(_reg_ntt_ring_rns) = []() -> bool {                                                           \
+        register_ring_rns_ntt(DEVICE_TYPE, FUNC);                                                                      \
+        return true;                                                                                                   \
+      }();                                                                                                             \
+    }
+
+  /*************************** INIT DOMAIN ***************************/
+  using NttInitDomainRingRnsImpl = std::function<eIcicleError(
+    const Device& device, const scalar_rns_t& primitive_root, const NTTInitDomainConfig& config)>;
+
+  void register_ring_rns_ntt_init_domain(const std::string& deviceType, NttInitDomainRingRnsImpl);
+
+  #define REGISTER_NTT_INIT_DOMAIN_RING_RNS_BACKEND(DEVICE_TYPE, FUNC)                                                 \
+    namespace {                                                                                                        \
+      static bool UNIQUE(_reg_ntt_init_domain_ring_rns) = []() -> bool {                                               \
+        register_ring_rns_ntt_init_domain(DEVICE_TYPE, FUNC);                                                          \
+        return true;                                                                                                   \
+      }();                                                                                                             \
+    }
+
+  /*************************** RELEASE DOMAIN ***************************/
+  // Note: 'phantom' is a workaround for the function required per field but need to differentiate by type when
+  // calling.
+  using NttReleaseDomainRingRnsImpl = std::function<eIcicleError(const Device& device, const scalar_rns_t& phantom)>;
+
+  void register_ring_rns_ntt_release_domain(const std::string& deviceType, NttReleaseDomainRingRnsImpl);
+
+  #define REGISTER_NTT_RELEASE_DOMAIN_RING_RNS_BACKEND(DEVICE_TYPE, FUNC)                                              \
+    namespace {                                                                                                        \
+      static bool UNIQUE(_reg_ntt_release_domain_ring_rns) = []() -> bool {                                            \
+        register_ring_rns_ntt_release_domain(DEVICE_TYPE, FUNC);                                                       \
+        return true;                                                                                                   \
+      }();                                                                                                             \
+    }
+
+  /*************************** GET ROU FROM DOMAIN ***************************/
+  using NttGetRouFromDomainRingRnsImpl =
+    std::function<eIcicleError(const Device& device, uint64_t logn, scalar_rns_t* rou)>;
+
+  void register_ring_rns_ntt_get_rou_from_domain(const std::string& deviceType, NttGetRouFromDomainRingRnsImpl);
+
+  #define REGISTER_NTT_GET_ROU_FROM_DOMAIN_RING_RNS_BACKEND(DEVICE_TYPE, FUNC)                                         \
+    namespace {                                                                                                        \
+      static bool UNIQUE(_reg_ntt_get_rou_from_domain_ring_rns) = []() -> bool {                                       \
+        register_ring_rns_ntt_get_rou_from_domain(DEVICE_TYPE, FUNC);                                                  \
+        return true;                                                                                                   \
+      }();                                                                                                             \
+    }
+#endif // RING
 } // namespace icicle