Broadcasting rules distinguish between host scalar and device scalars

[kaushikcfd]

Consider this script:

actx = _acf()

@with_container_arithmetic(
 bcast_obj_array=True,
 bcast_numpy_array=True,
 rel_comparison=True,
 _cls_has_array_context_attr=True)
@dataclass_array_container
@dataclass(frozen=True)
class Foo:
 u: DOFArray

foo = Foo(DOFArray(actx, (actx.zeros(3, dtype=np.float64) + 41, )))

print(foo) # prints: Foo(u=DOFArray((cl.Array([41., 41., 41.]),)))
print(foo + 1) # prints: Foo(u=DOFArray((cl.Array([42., 42., 42.]),)))
print(foo + actx.from_numpy(np.ones(()))) # ERROR! 😲

I don't see a good reason why we don't support this.

I think we should also go a step ahead and make sure that we broadcast every cl.Array to be added (if legal) with all the leaf arrays of the array container. One way I could see this working is if every ArrayContext comes with a ARRAY_TYPE static, frozen attribute which gives us a hint to broadcast to all successive arrays.

[inducer]

I agree that this is inconsistent, however there's an elephant of a practicality in the room: How do you propose we recognize scalars?

>>> import pyopencl as cl
>>> ctx = cl.create_some_context()
>>> queue = cl.CommandQueue(ctx)
>>> import pyopencl.array as clas
>>> import pyopencl.array as cla
>>> import numpy as np

>>> one = cla.to_device(queue, np.ones(()))
>>> np.isscalar(one)
False
>>> from numbers import Number
>>> isinstance(one, Number)
False

A shot-in-the-dark .shape == ()?

[kaushikcfd]

How do you propose we recognize scalars?

I think every ArrayContext providing a type for a leaf array's type isn't too bad. And we can accordingly broadcast the operation with all the other leaf arrays in the array container.

I.e. I believe even the below operation could also work with the approach.

foo = Foo(DOFArray(actx, (actx.zeros(3, dtype=np.float64) + 41, )))
foo + cla.zeros(cq, (3,))

[inducer]

Arithmetic is a pretty important code path, eager+strong-scaling-wise. (Hence all that generated code.) I'm hesitant about doing something expensive there.

Finding the array context to go with a container (so that we can ask it what its array class is) seems kind of difficult. There currently isn't a great way to do so, and what little way there is would require recursion, i.e. it's quite expensive.

We could also do a global registry of array classes... 🤷 not that I'm a fan of that solution either.

[kaushikcfd]

How about we take this code path iff getting the array context is cheap like here:

arraycontext/test/test_arraycontext.py

Lines 506 to 508 in c006285

	@property
	def array_context(self):
	return self.mass.array_context

[inducer]

That wouldn't fix your example though (unless you add an .array_context attribute). with_arithmetic already has _cls_has_array_context_attr (which, it occurs to me, could have a smarter default):

arraycontext/arraycontext/container/arithmetic.py

Line 136 in c006285

_cls_has_array_context_attr: bool = False,

I guess I'm not super opposed to this, though certainly not delighted. Want to put together a PR so we can see what it might look like?

[kaushikcfd]

Want to put together a PR so we can see what it might look like?

Yup. Thanks for considering this.

[alexfikl]

+1 for something like this on my side too. I was just playing with getting pytential working with force_device_scalars=True and it died pretty quicky because of expressions like

sym.Ones() * sym.mean(ambient_dim, dim, expr)

as it didn't know how to multiply a DOFArray (from Ones) with a device scalar (from mean).

[kaushikcfd]

There's actually another issue hiding there:

>>> a = cla.to_device(cq, np.ones(3))
>>> b = cla.to_device(cq, np.ones(()))
>>> a
cl.Array([1., 1., 1.])
>>> b
cl.Array(1.)
>>> a * b
Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
 File "/shared/home/kgk2/pack/virtual_envs/emirge/pyopencl/pyopencl/array.py", line 1161, in __mul__
 self._elwise_multiply(result, self, other))
 File "/shared/home/kgk2/pack/virtual_envs/emirge/pyopencl/pyopencl/array.py", line 183, in kernel_runner
 knl = kernel_getter(*args, **kwargs)
 File "/shared/home/kgk2/pack/virtual_envs/emirge/pyopencl/pyopencl/array.py", line 897, in _elwise_multiply
 assert out.shape == b.shape

I.e. PyOpenCL doesn't support broadcasting.

[inducer]

Let's take the PyOpenCL discussion here: inducer/pyopencl#498

[inducer]

See #51 for work by @kaushikcfd towards solving this.