Calls to JIT functions within JIT functions

[DavidPoliakoff]

Hey Hal. What I'm actually trying to do is make RAJA use JIT, and I think I have a path to it. However, I ran into a bug (some might refer to it as sanity), I don't think you can call JIT-marked functions from one another. A reproducer (no special compilation other than -fjit needed)

template<std::size_t N>
[[clang::jit]] void ping(std::size_t step);

template<std::size_t N>
[[clang::jit]] void pong(std::size_t step){
  std::cout << "Pong "<<N<<"\n";
}

template<std::size_t N>
[[clang::jit]] void ping(std::size_t step){
  std::cout << "Ping "<<N<<"\n";
  pong<N-step>(step);
}


int main(int argc, char* argv[]){
  std::size_t size = atoi(argv[1]);
  std::cout << "Size is "<<size<<"\n";
  ping<size>(size);

}

Note that this isn't really mutual recursion, it's just a call to ping, which calls pong. I think if we solve this we'll help allow constructs like

for(int rep = 0; rep<repeats;rep++){
    std::cout << "Loop iteration "<<rep<<"\n";
    RAJA::forall<RAJA::seq_exec>(RAJA::RangeSegment(0,size),[&](int i){
      out_matrix[i] = 5.0f;
      std::cout << "Made it in the outer\n";
      RAJA::forall<RAJA::seq_exec>(RAJA::RangeSegment(0,size),[&](int j){
      std::cout << "Made it in the middle\n";
        RAJA::forall<RAJA::seq_exec>(RAJA::RangeSegment(0,size),[&](int k){
          std::cout << "Made it in the inner\n";
          out_matrix[MAT2D(i,j,size)] += input_matrix1[MAT2D(i,k,size)] * input_matrix2[MAT2D(k,j,size)];
        });
      });
    });
  }

Which under the hood are calls to your JIT system, every call to RAJA::forall here is actually JITting its loop bounds (thanks to the bug you fixed yesterday)

[hfinkel]

Right now there are just two modes: During the ahead-of-time compile all JIT templates are delayed. At runtime, the system aggressively instantiates (the logic being to limit needless round trips into the compiler), and does this by more-or-less ignoring the JIT attributes. To make this work, there would need to be a fallback to generating more JIT calls (this requires some additional infrastructure that can be added, I think, but isn't there now).

[DavidPoliakoff]

@hfinkel Just to be clear, this one isn't needed on a short timescale. I file issues so we're aware of what's up. If we work with the -O2 (I'm testing this now) I should have RAJA autojitting its kernel (nested parallelism) loop bounds, this issue is for later