search

inducer / arraycontext

Choose your favorite numpy-workalike!
6 stars 11 forks source link

Add `compile_trace_callback` to pytato array contexts #173

Closed inducer closed 2 years ago

inducer commented 2 years ago

I'm trying to see how we can make the compilation process a bit less black-box-y. This came to mind, and it seems to vaguely work. I have a grudge piece for this that makes it work in the distributed setting. We can worry about that if/when this is in.

Mainly submitting this for initial feedback; is this a good idea or a terrible one? @kaushikcfd, thoughts?

kaushikcfd commented 2 years ago

I agree with the idea here. Definitely better than inserting breakpoints in the developer-facing source code. LGTM, thanks!

inducer commented 2 years ago

OK, well... maybe we should just do this then. It's easy enough to remove if it's a bad idea, as I don't expect any long-lived code to actually depend on this.

inducer commented 2 years ago

cc @mtcam @anderson2981 Once this merges, you should be able to use code like this to get some insight into the compile process:

def compile_trace_callback(what, stage, ir):
    if "my_rhs" in str(what):
        print("--> TRACE", what, stage)

        if stage == "final":
            with open("rhs.cl", "w") as outf:
                outf.write(lp.generate_code_v2(ir.program).device_code())

You may want to wait for the grudge part of this in order to fully make use of this, but it should be mostly functional on its own.

FWIW, the main branches of arraycontext and grudge now support running @kaushikcfd's transform pipeline. (You'll still need his versions of pytato, loopy, and meshmode.)

kaushikcfd commented 2 years ago

(You'll still need his versions of pytato, loopy, and meshmode.)

I've rebased the fusion_actx branches to include this functionality.

inducer commented 2 years ago

The grudge part of this is at https://github.com/inducer/grudge/pull/262.