Catalog: parallelize all-sky tile prefetch

src/catalog.c

@@ -6,6 +6,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <omp.h>
+
 #define PI 3.14159265358979323846
 
 enum {
@@ -108,26 +110,40 @@ static size_t tile_index(int ra_index, int dec_index)
     return (size_t)dec_index * CATALOG_ALL_SKY_RA_TILES + ra_index;
 }
 
+static int load_tile_file(const StarCatalog *catalog, int ra_index,
+                          int dec_index, Star **stars, size_t *count)
+{
+    char path[PATH_MAX];
+    int written;
+    if (catalog == NULL || stars == NULL || count == NULL)
+        return -1;
+    *stars = NULL;
+    *count = 0;
+    written = snprintf(path, sizeof path, "%s/tile_ra%03d_dec%03d.csv",
+                       catalog->all_sky_root, ra_index, dec_index);
+    if (written < 0 || (size_t)written >= sizeof path)
+        return -1;
+    StarCatalog temporary = {0};
+    if (catalog_load_csv(&temporary, path))
+        return -1;
+    *stars = temporary.stars;
+    *count = temporary.count;
+    return 0;
+}
+
 static int load_tile(StarCatalog *catalog, int ra_index, int dec_index)
 {
     CatalogTile *tile = &catalog->tiles[tile_index(ra_index, dec_index)];
-    char path[PATH_MAX];
-    int written;
+    Star *stars;
+    size_t count;
     if (tile->state != 0)
         return tile->state == 1 ? 0 : -1;
-    written = snprintf(path, sizeof path, "%s/tile_ra%03d_dec%03d.csv",
-                       catalog->all_sky_root, ra_index, dec_index);
-    if (written < 0 || (size_t)written >= sizeof path) {
+    if (load_tile_file(catalog, ra_index, dec_index, &stars, &count)) {
         tile->state = -1;
         return -1;
     }
-    StarCatalog temporary = {0};
-    if (catalog_load_csv(&temporary, path)) {
-        tile->state = -1;
-        return -1;
-    }
-    tile->stars = temporary.stars;
-    tile->count = temporary.count;
+    tile->stars = stars;
+    tile->count = count;
     tile->state = 1;
     catalog->count += tile->count;
     return 0;
@@ -255,15 +271,15 @@ static int tile_is_fully_contained(const double direction[3][3], int ra_index,
     return 1;
 }
 
-int catalog_visit_source_triangle(StarCatalog *catalog,
-                                  const double direction[3][3],
-                                  int load_missing, CatalogTileVisitor visitor,
-                                  void *context)
+typedef int (*CatalogTileIndexVisitor)(int ra_index, int dec_index,
+                                       void *context);
+
+static int visit_source_triangle_tile_indices(
+    const double direction[3][3], CatalogTileIndexVisitor visitor,
+    void *context)
 {
-    if (catalog == NULL || direction == NULL || visitor == NULL)
+    if (direction == NULL || visitor == NULL)
         return -1;
-    if (catalog->kind == STAR_CATALOG_MEMORY)
-        return visitor(catalog->stars, catalog->count, 0, context);
     double longitude[3], latitude[3], unwrapped[3];
     for (int i = 0; i < 3; ++i) {
         lon_lat_from_direction(direction[i], &longitude[i], &latitude[i]);
@@ -311,17 +327,149 @@ int catalog_visit_source_triangle(StarCatalog *catalog,
         for (int raw_ra = all_ra ? 0 : ra_first;
              raw_ra <= (all_ra ? 359 : ra_last); ++raw_ra) {
             const int ra = (raw_ra % 360 + 360) % 360;
-            CatalogTile *tile = &catalog->tiles[tile_index(ra, dec)];
-            if (load_missing && load_tile(catalog, ra, dec))
-                continue; /* Downloader has not finished this tile yet. */
-            if (tile->state != 1 ||
-                visitor(tile->stars, tile->count,
-                        tile_is_fully_contained(direction, ra, dec), context))
+            if (visitor(ra, dec, context))
                 return -1;
         }
     return 0;
 }
 
+typedef struct {
+    StarCatalog *catalog;
+    const double (*direction)[3];
+    int load_missing;
+    CatalogTileVisitor visitor;
+    void *context;
+} CatalogVisitContext;
+
+static int visit_catalog_tile(int ra_index, int dec_index, void *opaque)
+{
+    CatalogVisitContext *context = opaque;
+    CatalogTile *tile =
+        &context->catalog->tiles[tile_index(ra_index, dec_index)];
+    if (context->load_missing && load_tile(context->catalog, ra_index, dec_index))
+        return 0; /* Downloader has not finished this tile yet. */
+    if (tile->state != 1)
+        return 0;
+    return context->visitor(
+        tile->stars, tile->count,
+        tile_is_fully_contained(context->direction, ra_index, dec_index),
+        context->context);
+}
+
+typedef struct {
+    unsigned char *requested;
+} CatalogMarkContext;
+
+static int mark_catalog_tile(int ra_index, int dec_index, void *opaque)
+{
+    CatalogMarkContext *context = opaque;
+    context->requested[tile_index(ra_index, dec_index)] = 1;
+    return 0;
+}
+
+int catalog_mark_source_triangle_tiles(
+    const double direction[3][3],
+    unsigned char requested[CATALOG_ALL_SKY_TILE_COUNT])
+{
+    CatalogMarkContext context = {.requested = requested};
+    if (requested == NULL)
+        return -1;
+    return visit_source_triangle_tile_indices(direction, mark_catalog_tile,
+                                              &context);
+}
+
+typedef struct {
+    size_t tile_id;
+    Star *stars;
+    size_t count;
+    int loaded;
+} CatalogPendingTile;
+
+int catalog_prefetch_marked_tiles(
+    StarCatalog *catalog,
+    const unsigned char requested[CATALOG_ALL_SKY_TILE_COUNT],
+    int worker_count, CatalogPrefetchStats *stats)
+{
+    CatalogPendingTile *pending = NULL;
+    size_t pending_count = 0;
+    double load_start;
+    int result = -1;
+    if (stats != NULL)
+        *stats = (CatalogPrefetchStats){0};
+    if (catalog == NULL || requested == NULL || worker_count <= 0)
+        return -1;
+    if (catalog->kind != STAR_CATALOG_ALL_SKY)
+        return 0;
+    for (size_t tile_id = 0; tile_id < CATALOG_ALL_SKY_TILE_COUNT; ++tile_id) {
+        if (!requested[tile_id])
+            continue;
+        if (stats != NULL)
+            ++stats->requested_tiles;
+        if (catalog->tiles[tile_id].state == 0)
+            ++pending_count;
+    }
+    if (pending_count == 0)
+        return 0;
+    pending = calloc(pending_count, sizeof *pending);
+    if (pending == NULL)
+        return -1;
+    size_t pending_index = 0;
+    for (size_t tile_id = 0; tile_id < CATALOG_ALL_SKY_TILE_COUNT; ++tile_id)
+        if (requested[tile_id] && catalog->tiles[tile_id].state == 0)
+            pending[pending_index++].tile_id = tile_id;
+    if (worker_count > (int)pending_count)
+        worker_count = (int)pending_count;
+    load_start = omp_get_wtime();
+#pragma omp parallel for num_threads(worker_count) schedule(static)
+    for (size_t i = 0; i < pending_count; ++i) {
+        const int ra_index =
+            (int)(pending[i].tile_id % CATALOG_ALL_SKY_RA_TILES);
+        const int dec_index =
+            (int)(pending[i].tile_id / CATALOG_ALL_SKY_RA_TILES);
+        pending[i].loaded = !load_tile_file(
+            catalog, ra_index, dec_index, &pending[i].stars, &pending[i].count);
+    }
+    if (stats != NULL)
+        stats->load_seconds = omp_get_wtime() - load_start;
+    /* Only this serial commit mutates the shared catalog cache. */
+    for (size_t i = 0; i < pending_count; ++i) {
+        CatalogTile *tile = &catalog->tiles[pending[i].tile_id];
+        if (pending[i].loaded) {
+            tile->stars = pending[i].stars;
+            tile->count = pending[i].count;
+            tile->state = 1;
+            catalog->count += tile->count;
+            if (stats != NULL) {
+                ++stats->newly_loaded_tiles;
+                stats->newly_loaded_stars += tile->count;
+            }
+        } else {
+            tile->state = -1;
+            if (stats != NULL)
+                ++stats->unavailable_tiles;
+        }
+    }
+    result = 0;
+    free(pending);
+    return result;
+}
+
+int catalog_visit_source_triangle(StarCatalog *catalog,
+                                  const double direction[3][3],
+                                  int load_missing, CatalogTileVisitor visitor,
+                                  void *context)
+{
+    if (catalog == NULL || direction == NULL || visitor == NULL)
+        return -1;
+    if (catalog->kind == STAR_CATALOG_MEMORY)
+        return visitor(catalog->stars, catalog->count, 0, context);
+    CatalogVisitContext visit_context = {
+        .catalog = catalog, .direction = direction, .load_missing = load_missing,
+        .visitor = visitor, .context = context};
+    return visit_source_triangle_tile_indices(direction, visit_catalog_tile,
+                                              &visit_context);
+}
+
 void catalog_destroy(StarCatalog *catalog)
 {
     if (catalog->tiles != NULL)