[Draft] Pass JITUserContext::user_context to custom handlers. Add UserContext param.

[mzient]

The draft intended as a live demonstration of ideas discussed in #6255. It is still not ready for review, let alone merging.

Problem recap: currently, it's impossible to pass per-launch user context and handlers to JIT compiled functions. This becomes problematic when the target is an accelerator (e.g. CUDA device), where it may be desirable o execute the code in user-provided context. It's especially important when running in multi-device systems.
It also decouples the user context passed to custom JIT handlers from JITUserContext.

This draft explores possibilities of passing user context to JIT-compiled code.

Usage:

struct MyMemPool {
  // ...
  void *allocate(size_t);
  void free(void *);

  static MyMemPool &get(int device_id);
};

void *my_malloc(void *pool, size_t size) {
  return ((MyMemPool*)pool)->allocate(size);
}

void my_free(void *pool, void *mem) {
  return ((MyMemPool*)pool)->free(mem);
}

int main() {
  Func f("my_func");
  Param<float> scale;
  // define the function
  f.set_custom_allocator(my_malloc, my_free);
  f.compile_jit();

  // now the function can run in parallel with 2 different allocators
  
  std::thread t1([&]() {
    UserContextParam uc;
    ParamMap pm1({
      { scale, 2 },
      { uc, &MyMemPool::get(0) }
    });
    Buffer<uint8_t> out_buf1( /* sizes here */);
    f.realize(out_buf1, pm1);
  });

  std::thread t2([&]() {
    ParamMap pm2({
      { scale, 2 },
      { uc, &MyMemPool::get(1) }
    });
    Buffer<uint8_t> out_buf2( /* sizes here */);
    f.realize(out_buf2, pm2);
  });

  t1.join();
  t2.join();
}

Signed-off-by: Michał Zientkiewicz [email protected]

[zvookin]

I need at least a description of how this is proposed to work. It likely should be in comments in the files extending the APIs. There is no test to show usage. Until this is addressed, I cannot do much with this PR.

I think the basic idea is ok, but the main reason you're having problems with the existing thing is that there wasn't enough thinking about use cases upfront. Even ParamMap is insufficiently documented that I keep having to justify its existence. (It's purpose is to provide an enduser accessible way to make JIT, realize, calls that take parameters in a threadsafe way.)

[mzient]

@zvookin Thanks for the comments.
This code is still very much in draft state - its intended purpose was to serve as an illustration for the discussion of the proposed feature rather than an actual pull request. When it's ready for review, I'll convert it to a regular PR. If this generates too much noise in your systems (I see that there were some CI jobs running, for example), please let me know - I'll close it and re-open at a later stage.

[zvookin]

Per "its intended purpose was to serve as an illustration for the discussion of the proposed feature rather than an actual pull request"...

Expecting folks to figure out the issues that need to be discussed from the implementation alone is a lot to ask. Some sentences explaining what you are trying to do, decisions you've made, and perhaps 10 or 20 lines of code showing how this will be used, even if it doesn't compile or work yet, would go a long way toward facilitating the discussion needed. I believe these would also assist you in your work as well.

[mzient]

I've added a usage example and linked the issue. We can also move the discussion here, if that's a preferred workflow in your project - or I can even close this draft until it reaches next readiness level.
I'm sorry if I made myself too much at home here, bringing with me the customs from my project, where drafts are rarely proactively looked at.

[steven-johnson]

Is this PR still active? Should it be closed?

[mzient]

Is this PR still active? Should it be closed?

I think it can be closed - similar functionality is provided by #6313, which has already been merged.