Skip to content

Instantly share code, notes, and snippets.

@johari
Last active January 31, 2025 18:25
Show Gist options
  • Select an option

  • Save johari/6bd40261dcfabc0f9906e06f696be8a4 to your computer and use it in GitHub Desktop.

Select an option

Save johari/6bd40261dcfabc0f9906e06f696be8a4 to your computer and use it in GitHub Desktop.
Nima's Legion notes

Legion is task graph. So is llbuild2fx.

(Task-based system) The task model heavily resemble's llbuild2fx's FXKey, computeValue and fi.request mechanism.

(Memmory model) llbuild2fx has a closer relationship with CAS for storage/memory. Legion's emphasis is on logical region (as an abstraction over memory) and Rect (to express the shape/dimensions of the computation).

(Decoupling the task from underbelly of compute) Legion has the following API

  • Runtime::preregister_task_variant<int, sum_task>(registrar, "sum", AUTO_GENERATE_ID); and
  • registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));

that resemble FXAction and FXExecutor in llbuild2fx.

Beyond one node?

  • llbuild2fx can be coupled with AWS-lambda style compute, but does not maintain a pool of workers itself. You have to wire it up to a compute engine via providing your own FXExecutor implementation.
  • Legion relies on GASNet (see tutorial). I've never worked with GASNet but I'm guessing it is some sort of gRPC (+ a unified memory address space over network) that supercomputers use. I found this talk on YouTube that describes GASNet.

Composable units of computation

Watch Alex Aiken hinting at compositional advantages of task-based systems.

To get a good sense of Legion's API and with a simple but non-trivial example, start with this link.

Papers

#include <cstdio>
#include <cassert>
#include <cstdlib>
#include "legion.h"
using namespace Legion;
enum TaskIDs {
TOP_LEVEL_TASK_ID,
FIBONACCI_TASK_ID,
SUM_TASK_ID,
};
void top_level_task(const Task *task,
const std::vector<PhysicalRegion> &regions,
Context ctx, Runtime *runtime) {
int num_fibonacci = 7; // Default value
const InputArgs &command_args = Runtime::get_input_args();
for (int i = 1; i < command_args.argc; i++) {
// Skip any legion runtime configuration parameters
if (command_args.argv[i][0] == '-') {
i++;
continue;
}
num_fibonacci = atoi(command_args.argv[i]);
assert(num_fibonacci >= 0);
break;
}
printf("Computing the first %d Fibonacci numbers...\n", num_fibonacci);
std::vector<Future> fib_results;
for (int i = 0; i < num_fibonacci; i++) {
TaskLauncher launcher(FIBONACCI_TASK_ID, TaskArgument(&i,sizeof(i)));
fib_results.push_back(runtime->execute_task(ctx, launcher));
}
for (int i = 0; i < num_fibonacci; i++) {
int result = fib_results[i].get_result<int>();
printf("Fibonacci(%d) = %d\n", i, result);
}
fib_results.clear();
}
int fibonacci_task(const Task *task,
const std::vector<PhysicalRegion> &regions,
Context ctx, Runtime *runtime) {
assert(task->arglen == sizeof(int));
int fib_num = *(const int*)task->args;
if (fib_num == 0)
return 0;
if (fib_num == 1)
return 1;
// Launch fib-1
const int fib1 = fib_num-1;
TaskLauncher t1(FIBONACCI_TASK_ID, TaskArgument(&fib1,sizeof(fib1)));
Future f1 = runtime->execute_task(ctx, t1);
// Launch fib-2
const int fib2 = fib_num-2;
TaskLauncher t2(FIBONACCI_TASK_ID, TaskArgument(&fib2,sizeof(fib2)));
Future f2 = runtime->execute_task(ctx, t2);
TaskLauncher sum(SUM_TASK_ID, TaskArgument(NULL, 0));
sum.add_future(f1);
sum.add_future(f2);
Future result = runtime->execute_task(ctx, sum);
return result.get_result<int>();
}
int sum_task(const Task *task,
const std::vector<PhysicalRegion> &regions,
Context ctx, Runtime *runtime) {
assert(task->futures.size() == 2);
Future f1 = task->futures[0];
int r1 = f1.get_result<int>();
Future f2 = task->futures[1];
int r2 = f2.get_result<int>();
return (r1 + r2);
}
int main(int argc, char **argv) {
Runtime::set_top_level_task_id(TOP_LEVEL_TASK_ID);
{
TaskVariantRegistrar registrar(TOP_LEVEL_TASK_ID, "top_level");
registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
Runtime::preregister_task_variant<top_level_task>(registrar, "top_level");
}
{
TaskVariantRegistrar registrar(FIBONACCI_TASK_ID, "fibonacci");
registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
Runtime::preregister_task_variant<int, fibonacci_task>(registrar, "fibonacci");
}
{
TaskVariantRegistrar registrar(SUM_TASK_ID, "sum");
registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
registrar.set_leaf(true);
Runtime::preregister_task_variant<int, sum_task>(registrar, "sum", AUTO_GENERATE_ID);
}
return Runtime::start(argc, argv);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment