bscpkgs/garlic/exp/hpcg/common.nix

{
  stdenv
, lib
, stdexp
, bsc
, stages
, callPackage
}:

with lib;

rec {

  checkInput = {nextStage, conf, ...}: stages.exec {
    inherit nextStage;
    pre = optionalString (! (conf.enableGen or false)) (
    let
      gen = callPackage ./gen.nix { };
      inputTre = gen.getInputTre conf;
      exp = inputTre.experiment;
      unit = elemAt exp.units 0;
      expName = baseNameOf (toString exp);
      unitName = baseNameOf (toString unit);
      inputPath = "$GARLIC_OUT/${expName}/${unitName}/1";
    in
      ''
        # Force the generation of the input resultTree as a dependency:
        # ${toString inputTre.result}

        # Ensure the input dataset is still available
        export HPCG_INPUT_PATH="${toString inputPath}"

        if [ ! -e "$HPCG_INPUT_PATH" ]; then
          >&2 echo "Missing input dataset: $HPCG_INPUT_PATH"
          exit 1
        fi
      ''
    );
  };

  getSizePerTask = cpusPerTask: sizePerCpu:
    mapAttrs (name: val: val * cpusPerTask) sizePerCpu;

  exec = {nextStage, conf, ...}: let
    actionArg = if (conf.enableGen or false)
    then "--store=."
    else "--load=\"$HPCG_INPUT_PATH\"";

  in stages.exec {
    inherit nextStage;
    argv = [
      "--nx=${toString conf.sizePerTask.x}"
      "--ny=${toString conf.sizePerTask.y}"
      "--nz=${toString conf.sizePerTask.z}"
      "--npx=${toString conf.nprocs.x}"
      "--npy=${toString conf.nprocs.y}"
      "--npz=${toString conf.nprocs.z}"
      "--nblocks=${toString conf.nblocks}"
      "--ncomms=${toString conf.ncomms}"
      # The input symlink is generated by the input stage, which is generated by
      # the genInput function.
      actionArg
    ] ++ optional (conf.disableAspectRatio or false) "--no-ar=1";
  };

  program = {nextStage, conf, ...}: bsc.apps.hpcg.override {
    inherit (conf) gitBranch;
  };

  pipeline = stdexp.stdPipeline ++ [
    checkInput
    exec program ];
}