/*   $Source: bitbucket.org:berkeleylab/gasnet.git/tests/testimm.c $
 * Description: GASNet Active Messages IMMEDIATE test
 * Copyright (c) 2018-2024, The Regents of the University of California
 * Terms of use are as specified in license.txt
 */

// TODOs:
//  + Additional schedule options:
//    - Randomized
//      May have static randomization at start of run, or dynamic per-operation
//    - Work-sharing option
//      Rather than statically partitioning the ops per-rank, just keep a total count
//    - "nbrhd-balanced" in which consecutive operations are to distinct nbrhds
//      This should reduce ingress bottlenecks, but may not "fit" if nbrhds have unequal sizes
//  + Synchronous (GEX_EVENT_NOW) operation(s)

#include <gasnetex.h>

static size_t param_SZ = 0;
#define TEST_SEGSZ_EXPR ((uintptr_t)param_SZ)

#include <test.h>

static gex_Client_t      myclient;
static gex_EP_t          myep;
static gex_TM_t          myteam;
static gex_Segment_t     mysegment;

gex_Rank_t myrank = 0;
gex_Rank_t numrank = 0;

gex_RankInfo_t *nbrhdinfo = NULL;
gex_Rank_t nbrhdsize; // size of the neighborhood, or 1 for PSHM-only

static int64_t param_Z = 0;
static long param_B = 0;
static long param_N = 0;

static int in_segment = 1;

static enum {
  TEST_POLL_NEXT,
  TEST_POLL_RETRY,
  TEST_POLL_LAZY,
  TEST_POLL_ALWAYS
} poll_mode;

static void *local_addr;
static long *remain;

static gasnett_atomic_t expect = gasnett_atomic_init(0);

#define hidx_expect_dec_handler 200
#define hidx_expect_set_handler 201

void expect_dec_handler(gex_Token_t token, void *buf, size_t nbytes) {
  // decrement the counter
  gasnett_atomic_decrement(&expect, 0);
}
void expect_set_handler(gex_Token_t token, void *buf, size_t nbytes) {
  // set the counter and reply with a decrement
  assert(nbytes == sizeof(gasnett_atomic_val_t));
  gasnett_atomic_set(&expect, *(gasnett_atomic_val_t*)buf, 0);
  gex_AM_ReplyMedium0(token, hidx_expect_dec_handler, NULL, 0, GEX_EVENT_NOW, 0);
}

void doMed(gex_Flags_t imm_flag);
void doLong(gex_Flags_t imm_flag);
void doPut(gex_Flags_t imm_flag);
void doGet(gex_Flags_t imm_flag);

int main(int argc, char **argv) {
  gex_AM_Entry_t htable[] = { 
    { hidx_expect_dec_handler, (gex_AM_Fn_t)expect_dec_handler, GEX_FLAG_AM_REQREP|GEX_FLAG_AM_MEDLONG, 0 },
    { hidx_expect_set_handler, (gex_AM_Fn_t)expect_set_handler, GEX_FLAG_AM_REQUEST|GEX_FLAG_AM_MEDIUM, 0 }
  };

  GASNET_Safe(gex_Client_Init(&myclient, &myep, &myteam, "testimm", &argc, &argv, 0));

  myrank = gex_TM_QueryRank(myteam);
  numrank = gex_TM_QuerySize(myteam);

  int enable_given = 0;
  int enable_med   = 0;
  int enable_long  = 0;
  int enable_put   = 0;
  int enable_get   = 0;

  int force_mixed  = 0;

  int help = 0;
  int arg = 1;
  while (argc > arg) {
    if (!strcmp(argv[arg], "-z")) {
      ++arg;
      if (argc > arg) { param_Z = atol(argv[arg]); arg++; }
      else help = 1;
      param_Z = MAX(param_Z, 0);
    } else if (!strcmp(argv[arg], "-b")) {
      ++arg;
      if (argc > arg) { param_B = atol(argv[arg]); arg++; }
      else help = 1;
    } else if (!strcmp(argv[arg], "-m")) {
      ++arg;
      enable_med = enable_given = 1;
    } else if (!strcmp(argv[arg], "-l")) {
      ++arg;
      enable_long = enable_given = 1;
    } else if (!strcmp(argv[arg], "-p")) {
      ++arg;
      enable_put = enable_given = 1;
    } else if (!strcmp(argv[arg], "-g")) {
      ++arg;
      enable_get = enable_given = 1;
    } else if (!strcmp(argv[arg], "-mixed")) {
      ++arg;
      force_mixed = 1;
    } else if (!strcmp(argv[arg], "-in")) {
      ++arg;
      in_segment = 1;
    } else if (!strcmp(argv[arg], "-out")) {
      ++arg;
      in_segment = 0;
    } else if (!strcmp(argv[arg], "-poll-next")) {
      ++arg;
      poll_mode = TEST_POLL_NEXT;
    } else if (!strcmp(argv[arg], "-poll-retry")) {
      ++arg;
      poll_mode = TEST_POLL_RETRY;
    } else if (!strcmp(argv[arg], "-poll-lazy")) {
      ++arg;
      poll_mode = TEST_POLL_LAZY;
    } else if (!strcmp(argv[arg], "-poll-always")) {
      ++arg;
      poll_mode = TEST_POLL_ALWAYS;
    } else if (argv[arg][0] == '-') {
      help = 1;
      ++arg;
    } else break;
  }

  if (!enable_given) {
    enable_med = enable_long = 1;
  }

  if (argc > arg) { param_N = atol(argv[arg]); ++arg; }
  if (!param_N) param_N = 10000;

  if (argc > arg) { param_SZ = atoi(argv[arg]); ++arg; }
  if (!param_SZ) { param_SZ = 1024*1024; }
  if (enable_med) {
    param_SZ = MIN(param_SZ, gex_AM_LUBRequestMedium());
  }
  if (enable_long) {
    param_SZ = MIN(param_SZ, gex_AM_LUBRequestLong());
  }

  if (!param_Z) param_Z = 500;

  if (!param_B) param_B = param_N; // what UPC would call indefinite layout

  GASNET_Safe(gex_Segment_Attach(&mysegment, myteam, TEST_SEGSZ_REQUEST));
  GASNET_Safe(gex_EP_RegisterHandlers(myep, htable, sizeof(htable)/sizeof(gex_AM_Entry_t)));

  test_init("testimm",1, "[options] (msgcnt) (msgsz)\n"
             "  Active rank 0 injects msgcnt operations of size msgsz to each passive peer.\n"
             "  Note that msgsz will be reduced if RequestMedium or RequestLong are\n"
             "  to be timed and msgsz would exceed the respective LUBRequest limit.\n"
             "  Options:\n"
             "  -in / -out\n"
             "        Selects whether the initiator's buffer is in the GASNet segment\n"
             "        or not (default is 'in').\n"
             "  -m / -l / -p / -g\n"
             "        Respectively enable timing of the RequestMedium, RequestLong, \n"
             "        PutNBI and GetNBI operations.\n"
             "        The default is to test all four operations, except omitting\n"
             "        Put and Get in runs with only shared-memory communication.\n"
             "  -mixed\n"
             "        By default this test excludes passive ranks that are reachable\n"
             "        via shared-memory communication, unless that is *all* ranks.\n"
             "        This option overrides this behavior, allowing the test to use\n"
             "        a mix of shared-memory and network communication.  This test\n"
             "        does not exercise \"best practices\" for this configuration.\n"
             "  -z <interval_us>\n"
             "        Sets the minimum interval (in us) between poll calls by passive\n"
             "        ranks (default is 500).  Actual inter-poll delays may be longer.\n"
             "  -b <blocksz>\n"
             "        Sets block size of nominal communication schedule, where\n"
             "        \"nominal\" means the schedule applied in the absence of\n"
             "        back-pressure indication due to GEX_FLAG_IMMEDIATE.\n"
             "           0 (default) sends all to each rank before advancing\n"
             "           1 sends only one per rank before advancing\n"
             "           N sends N per rank before advancing\n"
             "  -poll-* options select the polling policy (default is -poll-next):\n"
             "    -poll-next   poll upon back pressure, advancing to the next peer\n"
             "    -poll-retry  poll upon back pressure, retrying the same peer once\n"
             "    -poll-lazy   advance to the next peer upon back pressure,\n"
             "                 but poll only between loops over peers\n"
             "    -poll-always advance to the next peer upon back pressure,\n"
             "                 but poll before every IMMEDIATE operation\n"
           );
  if (help || argc > arg) test_usage();

  if (numrank < 2) {
    MSG0("WARNING: This test requires two or more ranks.  Test skipped.\n");
    gasnet_exit(0); /* exit 0 to prevent false negatives in test harness */
  }

  TEST_PRINT_CONDUITINFO();

  char *space = NULL;
  char nbrhd_warning[64] = "";
  gex_System_QueryNbrhdInfo(&nbrhdinfo, &nbrhdsize, NULL);
  int nbrhd_only = (nbrhdsize == numrank);
  if (!myrank) {
    if (nbrhdsize == 1) {
      // The passive ranks are all OUTSIDE our neighborhood
    } else if (nbrhdsize == numrank) {
      // The passive ranks are all INSIDE our neighborhood
      nbrhdinfo = NULL; // suppress filtering
      nbrhdsize = 1;    // and correct reported passive rank count
      assert(nbrhd_only);
    #if !GASNET_CONDUIT_SMP // would be "just noise" for smp-conduit
      strcpy(nbrhd_warning, "\n  WARNING: all ranks are reachable via shared-memory");
    #endif
    } else if (force_mixed) {
      // Mixed case: but commandline ask us not to omit ranks within our neighborhood
      snprintf(nbrhd_warning, sizeof(nbrhd_warning),
               "\n  WARNING: %d shared-memory rank%s allowed by -mixed option",
               (nbrhdsize-1), (nbrhdsize>2)?"s":"");
      nbrhdinfo = NULL; // suppress filtering
      nbrhdsize = 1;    // and correct reported passive rank count
    } else {
      // Mixed case: we will omit ranks within our neighborhood
      snprintf(nbrhd_warning, sizeof(nbrhd_warning),
               "\n  WARNING: %d shared-memory rank%s omitted",
               (nbrhdsize-1), (nbrhdsize>2)?"s":"");
    }

    if (in_segment) {
      local_addr = TEST_MYSEG();
    } else {
      space = test_malloc(param_SZ + PAGESZ);
      local_addr = alignup_ptr(space, PAGESZ);
    }
    remain = test_malloc(sizeof(long) * numrank);
  } // rank 0

  if (!enable_given) {
    enable_put = enable_get = !nbrhd_only;
  }

  MSG0("Running testimm with %d ranks\n"
       "  active rank issues %ld operations per rank, each of length %ld\n"
       "  active rank's local addresses %sside the segment\n"
       "  nominal schedule issues operations in blocks of %ld\n"
       "  polling policy is \"%s\"\n"
       "  %d passive rank%s sleep (at least) %"PRIi64"us between polls\n"
       "  operations timed:%s%s%s%s"
       "%s",
       numrank,
       param_N, (long)param_SZ,
       (in_segment ? "in" : "out"),
       param_B,
       (poll_mode==TEST_POLL_NEXT   ? "next" :
       (poll_mode==TEST_POLL_RETRY  ? "retry" :
       (poll_mode==TEST_POLL_ALWAYS ? "always" :
                                      "lazy"))),
       (numrank-nbrhdsize), (numrank-nbrhdsize>1)?"s":"", param_Z,
       enable_med?" RequestMedium":"", enable_long?" RequestLong":"",
       enable_put?" PutNBI":"", enable_get?" GetNBI":"",
       nbrhd_warning
      );

  if (enable_med) {
    doMed(0);
    doMed(GEX_FLAG_IMMEDIATE);
  }
  if (enable_long) {
    doLong(0);
    doLong(GEX_FLAG_IMMEDIATE);
  }
  if (enable_put) {
    doPut(0);
    doPut(GEX_FLAG_IMMEDIATE);
  }
  if (enable_get) {
    doGet(0);
    doGet(GEX_FLAG_IMMEDIATE);
  }

  if (!myrank) {
    test_free(space);
    test_free(remain);
  }

  BARRIER();
  MSG("done.");

  gasnet_exit(0);
  return 0;
}

void sleepy_barrier(uint64_t interval_ns) {
  gex_Event_t bar = gex_Coll_BarrierNB(myteam,0);
  do {
    gasnett_nsleep(interval_ns);
    gasnet_AMPoll();
  } while (gex_Event_Test(bar) != GASNET_OK);
}

void passive(void) {
  uint64_t interval_ns = 1000 * param_Z;
  sleepy_barrier(interval_ns); // start barrier (passive)
  do {
    gasnett_nsleep(interval_ns);
    gasnet_AMPoll();
  } while (gasnett_atomic_read(&expect, 0)); // await AM arrivals
  sleepy_barrier(interval_ns); // end barrier (passive)
}

#define ACTIVE(OPERATION, SYNC) do {                    \
  int done;                                             \
  do {                                                  \
    done = 1;                                           \
    for (gex_Rank_t r = 1; r < numrank; ++r) {          \
      long count, limit = MIN(param_B, remain[r]);      \
      for (count = 0; count < limit; ++count) {         \
        int did_retry = 0;                              \
        if ((poll_mode == TEST_POLL_ALWAYS) && imm_flag)\
          gasnet_AMPoll();                              \
        retry:                                          \
        if ( OPERATION ) { /* got IMM backpressure */   \
          switch (poll_mode) {                          \
            case TEST_POLL_RETRY:                       \
              if (!did_retry) {                         \
                gasnet_AMPoll();                        \
                did_retry = 1;                          \
                goto retry;                             \
              }                                         \
              break;                                    \
            case TEST_POLL_NEXT:                        \
              gasnet_AMPoll();                          \
              break;                                    \
            case TEST_POLL_LAZY: break;                 \
            case TEST_POLL_ALWAYS: break;               \
            default: gasnett_unreachable();             \
          }                                             \
          break; /* end injection to this peer */       \
        }                                               \
      }                                                 \
      remain[r] -= count;                               \
      done &= !remain[r];                               \
    } /* end loop over ranks */                         \
    if ((poll_mode == TEST_POLL_LAZY) && imm_flag)      \
      gasnet_AMPoll();                                  \
  } while (!done);                                      \
  gex_NBI_Wait(SYNC,0);                                 \
} while (0)

void init_remain(int isAMtest) {
  for (gex_Rank_t r = 1; r < numrank; ++r) remain[r] = param_N;
  if (nbrhdinfo) { // Optionally exclude PSHM-peers by zeroing their remain[]
    for (gex_Rank_t i = 0; i < nbrhdsize; ++i) remain[nbrhdinfo[i].gex_jobrank] = 0;
  }
  if (isAMtest) {
    gasnett_atomic_set(&expect, 0, 0);
    for (gex_Rank_t r = 1; r < numrank; ++r) {
      if (remain[r]) { // tell each passive rank what to expect
        gasnett_atomic_val_t val = (gasnett_atomic_val_t)remain[r];
        assert((long)val == remain[r]);
        gasnett_atomic_increment(&expect, 0);
        gex_AM_RequestMedium0(myteam, r, hidx_expect_set_handler, &val, sizeof(val), GEX_EVENT_NOW, 0);
      }
    }
    do { // await acknowledgment
      gasnet_AMPoll();
    } while (gasnett_atomic_read(&expect, 0));
  }
  gex_Event_Wait(gex_Coll_BarrierNB(myteam,0)); // start barrier (active)
}

void report(const char *name, gex_Flags_t imm_flag, double elapsed, double *prev) {
  char delta[32];
  if (*prev > 0.0) {
    snprintf(delta, sizeof(delta), ", %+.2f%% improvement", 100.*((*prev) - elapsed)/(*prev));
  } else {
    delta[0] = '\0';
  }
  double bw = ((double)param_N * param_SZ) / (elapsed * 1024 * 1024);
  printf("%-8s IMM=%c completed in %7.3f seconds (%7.3f MiB/s)%s\n",
         name, imm_flag?'Y':'N', elapsed, bw, delta);
  *prev = elapsed;
}

void doMed(gex_Flags_t imm_flag) {
  if (myrank) {
    passive();
  } else {
    init_remain(1);
    gasnett_tick_t start_ticks = gasnett_ticks_now();
    ACTIVE( gex_AM_RequestMedium0(myteam, r, hidx_expect_dec_handler, local_addr, param_SZ,
                                  GEX_EVENT_GROUP, imm_flag),
            GEX_EC_AM );
    gasnett_tick_t end_ticks = gasnett_ticks_now();
    double elapsed = 1e-9 * gasnett_ticks_to_ns(end_ticks - start_ticks);

    static double prev;
    report("MEDIUM:", imm_flag, elapsed, &prev);

    gex_Event_Wait(gex_Coll_BarrierNB(myteam,0)); // end barrier (active)
  }
}

void doLong(gex_Flags_t imm_flag) {
  if (myrank) {
    passive();
  } else {
    init_remain(1);
    gasnett_tick_t start_ticks = gasnett_ticks_now();
    ACTIVE( gex_AM_RequestLong0(myteam, r, hidx_expect_dec_handler, local_addr, param_SZ,
                                TEST_SEG(r), GEX_EVENT_GROUP, imm_flag),
            GEX_EC_AM );
    gasnett_tick_t end_ticks = gasnett_ticks_now();
    double elapsed = 1e-9 * gasnett_ticks_to_ns(end_ticks - start_ticks);

    static double prev;
    report("LONG:", imm_flag, elapsed, &prev);

    gex_Event_Wait(gex_Coll_BarrierNB(myteam,0)); // end barrier (active)
  }
}

void doPut(gex_Flags_t imm_flag) {
  if (myrank) {
    passive();
  } else {
    init_remain(0);
    gasnett_tick_t start_ticks = gasnett_ticks_now();
    ACTIVE( gex_RMA_PutNBI(myteam, r, TEST_SEG(r), local_addr, param_SZ,
                           GEX_EVENT_DEFER, imm_flag),
            GEX_EC_PUT );
    gasnett_tick_t end_ticks = gasnett_ticks_now();
    double elapsed = 1e-9 * gasnett_ticks_to_ns(end_ticks - start_ticks);

    static double prev;
    report("PUT:", imm_flag, elapsed, &prev);

    gex_Event_Wait(gex_Coll_BarrierNB(myteam,0)); // end barrier (active)
  }
}

void doGet(gex_Flags_t imm_flag) {
  if (myrank) {
    passive();
  } else {
    init_remain(0);
    gasnett_tick_t start_ticks = gasnett_ticks_now();
    ACTIVE( gex_RMA_GetNBI(myteam, local_addr, r, TEST_SEG(r), param_SZ, imm_flag),
            GEX_EC_GET );
    gasnett_tick_t end_ticks = gasnett_ticks_now();
    double elapsed = 1e-9 * gasnett_ticks_to_ns(end_ticks - start_ticks);

    static double prev;
    report("GET:", imm_flag, elapsed, &prev);

    gex_Event_Wait(gex_Coll_BarrierNB(myteam,0)); // end barrier (active)
  }
}