proteus specializes only the outermost region

[jandrej]

See MFE

// A JIT variable that belongs to a *nested* JIT region is silently ignored:
// proteus specializes only the outermost region, and nothing warns about it.
// This holds for both interfaces -- registered lambdas and annotated functions.
//
// Run it and compare the traces:
//
//   PROTEUS_TRACE_OUTPUT="specialization;kernel-trace" ./build/nested-jit-vars
//
// All four variants compute the same, correct answers. The difference is what
// gets compiled: the hoisted variants specialize both v and w, so the three
// calls below -- same v, different w -- produce three specializations each. The
// nested variants specialize only v, so all three calls share one compilation
// and w stays an ordinary runtime value.

#include <cstdio>

#include <proteus/JitInterface.h>

template <typename F> void run(F &&Func) { proteus::register_lambda(Func)(); }

// Deliberately external linkage. A JIT'd body that touches an internal-linkage
// global (anything in an anonymous namespace, for instance) aborts ProteusPass
// at compile time with "Broken JIT module found" -- unrelated to nesting, and
// it hits annotated functions just as it hits lambdas.
int Result;

// --- registered lambdas ----------------------------------------------------

// Both runtime constants live in the capture list of the registered lambda.
void hoisted_lambda(int v, int w) {
  run([=, v = proteus::jit_variable(v), w = proteus::jit_variable(w)]()
          __attribute__((annotate("jit"))) { Result = v * 100 + w; });
}

// Same computation, but w is a jit_variable of a lambda registered *inside* the
// JIT'd body. That inner registration never reaches the runtime: the outer
// region is compiled as one unit, and w is compiled as a plain value.
void nested_lambda(int v, int w) {
  run([=, v = proteus::jit_variable(v)]() __attribute__((annotate("jit"))) {
    run([=, w = proteus::jit_variable(w)]()
            __attribute__((annotate("jit"))) { Result = v * 100 + w; });
  });
}

// --- annotated functions ---------------------------------------------------

// Both arguments are specialized by one annotation (positions are 1-based).
__attribute__((annotate("jit", 1, 2))) void hoisted_fn(int v, int w) {
  Result = v * 100 + w;
}

// The inner function carries its own annotation for w, and it is just as inert
// as the inner lambda above: called from inside a JIT'd region it is compiled
// as part of that region, never dispatched, never specialized.
__attribute__((annotate("jit", 1))) void inner_fn(int w) { Result += w; }

__attribute__((annotate("jit", 1))) void nested_fn(int v, int w) {
  Result = v * 100;
  inner_fn(w);
}

// ---------------------------------------------------------------------------

template <typename F> void run(const char *Label, F Variant) {
  printf("%s\n", Label);
  // Same v every time, so only w could trigger a new specialization.
  for (int w : {10, 20, 30}) {
    Variant(1, w);
    printf("  v = 1, w = %2d -> %d\n", w, Result);
  }
  printf("\n");
}

int main() {
  printf("Expect identical results from all four, but different traces:\n"
         "  hoisted: three specializations, w replaced each time\n"
         "  nested:  one specialization, no [LambdaSpec]/[ArgSpec] for w\n\n");

  run("lambda, v and w both in the outer capture list",
        [](int v, int w) { hoisted_lambda(v, w); });
  run("lambda, w in an inner registered lambda",
        [](int v, int w) { nested_lambda(v, w); });

  run("annotated function specializing both arguments",
        [](int v, int w) { hoisted_fn(v, w); });
  run("annotated function calling an annotated function that specializes w",
        [](int v, int w) { nested_fn(v, w); });

  return 0;
}