Macro expansion during JIT compilation?

[NamjaeChoi]

Hello,

I'm wondering if JIT compilation can properly expand macros. I guess it is somehow treating macros since CEED_QFUNCTION is also a macro, but I'm not sure if you have a special rule to specifically handle it or it is being handled by the compiler. The reason for this question is because we are defining QFunctions using macros like:

#define CEED_KERNEL(name, parent)                                                                  \
  class name final : public parent<name>                                                           \
  {                                                                                                \
  public:                                                                                          \
    static CEEDQFunctionPointer _qf;                                                               \
    auto getQFunction() { return _qf; }                                                            \
    std::string getQFunctionPath() { return std::string(__FILE__) + ":" + #name + "_qf"; }
#define CEED_KERNEL_END(name)                                                                      \
  }                                                                                                \
  ;                                                                                                \
                                                                                                   \
  CEED_QFUNCTION(name##_qf)                                                                        \
  (void * ctx, const CeedInt Q, const CeedScalar * const * in, CeedScalar * const * out)           \
  {                                                                                                \
    auto kernel = static_cast<name *>(ctx);                                                        \
    kernel->computeResidualInternal(Q, in, out);                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  CEEDQFunctionPointer name::_qf = name##_qf;

which is used in the user code like:

CEED_KERNEL(CEEDDiffusionKernel, CEEDDiffusion)
public:
static InputParameters validParams();

CEEDDiffusionKernel(const InputParameters & parameters);
CEED_KERNEL_END(CEEDDiffusionKernel)

But this leads to a JIT error that the QFunction was not defined:

/data/choin/moose/libCEED/backends/cuda/ceed-cuda-compile.cpp:125 in CeedCompile_Cuda(): NVRTC_ERROR_COMPILATION
default_program(63222): error: identifier "CEEDDiffusionKernel_qf" is undefined
      CEEDDiffusionKernel_qf(ctx, 1, inputs, outputs);

[jeremylt]

I'll try to make a running example of this to reproduce and understand what's happening, but it will take me time to do so.

[jeremylt]

If you aren't familiar, a minimal reproducible example greatly helps library maintainers in seeing the issue you are seeing.

[NamjaeChoi]

Let me try. This might be easy to make an example.

[jeremylt]

For CUDA JiT, the CEED_QFUNCTION macro is defined as

#define CEED_QFUNCTION(name) inline __device__ int name

I would guess you are getting a mixup with macro composition?

[jeremylt]

You can set the environment variable CEED_DEBUG=1 to see the code being JiTed, though that will potentially go away in the future as a result of #1695

[NamjaeChoi]

Hmm that just give me the code before expanding macro.

[jeremylt]

Ah, that's right. I'm still suspicious of chaining string manipulation macros. That's very touchy in my experience

[NamjaeChoi]

I just tried chaining macros in ex1-volume but seems to work fine. Let look into this more.

[NamjaeChoi]

Ah I think I know the reason. All our CEED source files are protected with a preprocessor, which is defined by our static compiler. But I guess JIT compiler doesn't know that? Is there a way to provide a preprocessor to the JIT compiler?

[NamjaeChoi]

A potential way to realize complicated QFunctions could be leveraging JIT link. I created a small repository demonstrating the concept, which is based on the nvJitLink example (https://docs.nvidia.com/cuda/nvjitlink/index.html#example-device-lto-link-time-optimization).

https://github.com/NamjaeChoi/jit_test

In the example, a kernel saxpy is JIT-compiled, and it calls a device function compute which is statically compiled. offline.cu that contains the body of compute also contains some host codes, and it is compiled as a standard object file offline.o. It is then JIT-linked with the JIT-compiled kernel saxpy as well as statically linked to the executable. In this way, users can use QFunctions as simple entry points that merely call the statically-compiled user device functions which can have all the complicated things like our object-oriented design.

This will require libCEED to be able to receive the list of object files and do JIT-link. Also, a capability to receive QFunction source strings instead of source files in libCEED would be favorable, so that we don't have to create dummy source files that contain the QFunctions merely calling the user device functions.

This is just a quick thought that I wanted to share and not a request of any urgent implementation.

[jeremylt]

Hmm, how would that work without breaking the ability to select between CPU and GPU at runtime?

[NamjaeChoi]

Given the user provided proper function qualifiers (__host__ __device__) to his precompiled function, I wonder if there is a way for the CPU JIT compiler to link with the precompiled function.

[jeremylt]

CPU backends do not have JiT

[NamjaeChoi]

Then I guess it should work without any special treatment, because the object file is always statically linked. JIT link is an additional phase for GPU.