[loopy-codegen, bug] Code generation for tagged `NamedArray`s is buggy

[kaushikcfd]

Tagging a named array and then accessing other results of the container loses the tag information as it might freshly create a "dirty" no-tagged variant, see

pytato/pytato/array.py

Lines 718 to 722 in 1c9ee78

	@memoize_method
	def __getitem__(self, name: str) -> NamedArray:
	if name not in self._data:
	raise KeyError(name)
	return NamedArray(self, name)

---

Loopy codegen is broken for tagged instances of NamedArray for the same reason. Here's an MWE:

class FooTag(Tag):
    pass

knl = lp.make_kernel(
    "{[i]: 0<=i<10}",
    """
    y[i] = i
    """)

y = call_loopy(knl, bindings={})["y"]
y = y.tagged(FooTag())
pt.generate_loopy(2 * y)

raises:

  File "/home/kgk2/projects/ceesd/pytato/pytato/target/loopy/codegen.py", line 414, in map_named_array
    assert expr in state.results
AssertionError

[inducer]

Hmm. On the one hand, for all other array nodes, "what you use counts", i.e. it's perfectly possible to have a tagged and an untagged version floating around, and only the version that you use actually matters. If you end up using the tagged one, great: it counts as tagged.

Maybe this isn't practical here? Could you explain?

[kaushikcfd]

i.e. it's perfectly possible to have a tagged and an untagged version floating around, and only the version that you use actually matters

I think this makes sense. I guess this is some error in the code-gen for NamedArrays, not too difficult to fix.

Thanks!

[nchristensen]

I'm guessing this issue is why my IsDOFArray tags don't survive the fusion process. I dug down a bit and I think the problem starts here

pytato/pytato/target/loopy/codegen.py

Line 416 in e88cac2

self.rec(expr._container, state)

While expr has tags, expr._container does not. When self.rec is called with expr._container as the argument, the tag information is lost.

[kaushikcfd]

I'm guessing this issue is why my IsDOFArray tags don't survive the fusion process

Could be.

But my first guess is that it might be because metadata propagation isn't done by pytato. The reason for this is outlined here: https://documen.tician.de/pytato/design.html#metadata-propagation. Essentially, there is no "correct" way to propagate tags so the user must provide the semantics by implementing their own Mapper.

[nchristensen]

I see. Thanks for the pointer.

[inducer]

Discussing with @kaushikcfd just now (as part of #364), we've come to the conclusion that the only well-formed way to tag NamedArrays is through their container.