mac compilation fix and api to query device availability

ingonyama-zk · Jun 16, 2024 · 919efcf · 919efcf
1 parent 0b47758
commit 919efcf
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 7 deletions.
diff --git a/icicle_v3/include/icicle/runtime.h b/icicle_v3/include/icicle/runtime.h
@@ -158,3 +158,13 @@ extern "C" eIcicleError icicle_device_synchronize();
  * @return eIcicleError Status of the properties query.
  */
 extern "C" eIcicleError icicle_get_device_properties(DeviceProperties& properties);
+
+/**
+ * @brief Checks if the specified device is available.
+ *
+ * @param dev The device to check for availability.
+ * @return eIcicleError Status of the device availability check.
+ *         - `SUCCESS` if the device is available.
+ *         - `INVALID_DEVICE` if the device is not available.
+ */
+extern "C" eIcicleError icicle_is_device_avialable(const Device& dev);
diff --git a/icicle_v3/src/curves/montgomery_conversion.cpp b/icicle_v3/src/curves/montgomery_conversion.cpp
@@ -1,3 +1,4 @@
+#include "icicle/vec_ops.h"
 #include "icicle/curves/montgomery_conversion.h"
 #include "icicle/dispatcher.h"
 #include "icicle/curves/curve_config.h"
@@ -10,14 +11,14 @@ namespace icicle {
   ICICLE_DISPATCHER_INST(AffineConvertMont, affine_convert_montgomery, AffineConvertMontImpl);
 
   extern "C" eIcicleError CONCAT_EXPAND(CURVE, affine_convert_montgomery)(
-    const affine_t* input, size_t n, bool is_into, const VecOpsConfig& config, affine_t* output)
+    const affine_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, affine_t* output)
   {
     return AffineConvertMont::execute(input, n, is_into, config, output);
   }
 
   template <>
   eIcicleError
-  convert_montgomery(const affine_t* input, size_t n, bool is_into, const VecOpsConfig& config, affine_t* output)
+  convert_montgomery(const affine_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, affine_t* output)
   {
     return CONCAT_EXPAND(CURVE, affine_convert_montgomery)(input, n, is_into, config, output);
   }
@@ -44,14 +45,14 @@ namespace icicle {
   ICICLE_DISPATCHER_INST(ProjectiveConvertMont, projective_convert_montgomery, ProjectiveConvertMontImpl);
 
   extern "C" eIcicleError CONCAT_EXPAND(CURVE, projective_convert_montgomery)(
-    const projective_t* input, size_t n, bool is_into, const VecOpsConfig& config, projective_t* output)
+    const projective_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, projective_t* output)
   {
     return ProjectiveConvertMont::execute(input, n, is_into, config, output);
   }
 
   template <>
   eIcicleError convert_montgomery(
-    const projective_t* input, size_t n, bool is_into, const VecOpsConfig& config, projective_t* output)
+    const projective_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, projective_t* output)
   {
     return CONCAT_EXPAND(CURVE, projective_convert_montgomery)(input, n, is_into, config, output);
   }
@@ -60,14 +61,14 @@ namespace icicle {
   ICICLE_DISPATCHER_INST(ProjectiveG2ConvertMont, projective_g2_convert_montgomery, ProjectiveG2ConvertMontImpl);
 
   extern "C" eIcicleError CONCAT_EXPAND(CURVE, projective_g2_convert_montgomery)(
-    const g2_projective_t* input, size_t n, bool is_into, const VecOpsConfig& config, g2_projective_t* output)
+    const g2_projective_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, g2_projective_t* output)
   {
     return ProjectiveG2ConvertMont::execute(input, n, is_into, config, output);
   }
 
   template <>
   eIcicleError convert_montgomery(
-    const g2_projective_t* input, size_t n, bool is_into, const VecOpsConfig& config, g2_projective_t* output)
+    const g2_projective_t* input, uint64_t n, bool is_into, const VecOpsConfig& config, g2_projective_t* output)
   {
     return CONCAT_EXPAND(CURVE, projective_g2_convert_montgomery)(input, n, is_into, config, output);
   }

diff --git a/icicle_v3/src/runtime.cpp b/icicle_v3/src/runtime.cpp
@@ -76,6 +76,11 @@ extern "C" eIcicleError icicle_get_device_properties(DeviceProperties& propertie
   return DeviceAPI::get_thread_local_deviceAPI()->get_device_properties(properties);
 }
 
+extern "C" eIcicleError icicle_is_device_avialable(const Device& dev)
+{
+  return is_device_registered(dev.type) ? eIcicleError::SUCCESS : eIcicleError::INVALID_DEVICE;
+}
+
 extern "C" eIcicleError icicle_create_stream(icicleStreamHandle* stream)
 {
   return DeviceAPI::get_thread_local_deviceAPI()->create_stream(stream);

diff --git a/icicle_v3/tests/test_device_api.cpp b/icicle_v3/tests/test_device_api.cpp
@@ -85,7 +85,10 @@ TEST_F(DeviceApiTest, ApiError)
 
 TEST_F(DeviceApiTest, AvailableMemory)
 {
-  icicle::Device dev = {"CUDA", 0}; // TODO Yuval: implement for CPU too
+  icicle::Device dev = {"CUDA", 0};
+  const bool is_cuda_registered = eIcicleError::SUCCESS == icicle_is_device_avialable(dev);
+  if (!is_cuda_registered) { return; } // TODO Yuval: implement for CPU too
+
   icicle_set_device(dev);
   size_t total, free;
   ASSERT_EQ(eIcicleError::SUCCESS, icicle_get_available_memory(total, free));