Pass immutable arg_id_to_dtype to InKernelCallables

[alexfikl]

This uses immutabledict whenever possible. It's not exactly clear to me if some of these should be passed in as immutable higher up the stack, so someone more knowledgeable should have a serious look.

Fixes #905.

[alexfikl]

@inducer I can't say I'm a fan of this. Maybe a nicer workaround would be to have InKernelCallable.copy automatically cast to immutabledict?

[inducer]

I'm torn. On the one hand, I agree it's a mess. On the other hand, I dislike all the implicit casting in constructor-like interfaces (copy counts as such, to me). While it seems convenient at first, I worry about cost, and dataclasses have a more-than-implicit nudge towards constructors without intelligence. So I don't hate having to cast to the right type where these things are being made from scratch.

A few practical concerns next up:

• Batch mutation should likely happen via an immutables-style mutation context manager. immutabledict doesn't have this, but we're switching to constandict anyway, which does. (cc @matthiasdiener, x-ref: replace pyrsistent.PMap, immutables.Map, immutabledict with constantdict #884 (comment))
• This currently breaks Firedrake, on a mutability assumption. @connorjward, could you chime in? IIUC, this should be a simple patch on your end. (Cast to dict and back.)

[alexfikl]

I'm torn. On the one hand, I agree it's a mess. On the other hand, I dislike all the implicit casting in constructor-like interfaces (copy counts as such, to me). While it seems convenient at first, I worry about cost, and dataclasses have a more-than-implicit nudge towards constructors without intelligence. So I don't hate having to cast to the right type where these things are being made from scratch.

Yeah, agreed. My main worry was that it's a pain to find all the places that should be wrapped, but it's slowly getting there by rummaging through the warnings.

A few practical concerns next up:

• Batch mutation should likely happen via an immutables-style mutation context manager. immutabledict doesn't have this, but we're switching to constandict anyway, which does. (cc @matthiasdiener, x-ref: replace pyrsistent.PMap, immutables.Map with immutabledict #884 (comment))

Definitely fine with waiting for that to get in first.

• This currently breaks Firedrake, on a mutability assumption. @connorjward, could you chime in? IIUC, this should be a simple patch on your end. (Cast to dict and back.)

Maybe it can be fixed on this side? I haven't looked too much into it..

[connorjward]

This currently breaks Firedrake, on a mutability assumption. @connorjward, could you chime in? IIUC, this should be a simple patch on your end. (Cast to dict and back.)

I'm sure the fix would be straightforward, which I'm happy to do.

[inducer]

Definitely fine with waiting for that to get in first.

It's in!

[connorjward]

I have a branch that I am currently testing where I think this is fixed. I am happy for you to merge this when ready and then I will update our fork and get my things merged.

[alexfikl]

Batch mutation should likely happen via an immutables-style mutation context manager. immutabledict doesn't have this, but we're switching to constandict anyway, which does.

I'm note quite sure where to use batch mutation stuff in here (if any?). Most places that cast to constantdict are in with_dtypes and they all cast to dict and then mutate that. Is that ok?

[matthiasdiener]

This could probably be smth like:

        new_arg_id_to_dtype = constantdict(arg_id_to_dtype).mutate()
        for i in range(len(arg_id_to_dtype)):
            if i in arg_id_to_dtype and arg_id_to_dtype[i] is not None:
                new_arg_id_to_dtype[-i-1] = new_arg_id_to_dtype[i]

        return (self.copy(arg_id_to_dtype=new_arg_id_to_dtype.finish()),

This avoids a second copy.

[alexfikl]

Ah, I see! Thanks a lot for the example. I'll make those changes.

[alexfikl]

I think I got all the places where this made sense (at least in the parts that were already modified). Let me know if you spot anything else!

[alexfikl]

@inducer This should be ready for a look now.

@connorjward It seems like the Firedrake failure was on this side.. I prematurely cast something to constantdict. It all works now!

[inducer]

Thanks! In it goes.