Allow vectorization of ternaries, indices

[nkoskelo]

knl = lp.make_kernel( " { [i] : 0 <= i  < 8}"," out[i] = a if i == 0 else (b if i == 1 else c)")
knl = lp.tag_inames(knl, {"i": "vec"})
from loopy.kernel.array import VectorArrayDimTag
knl = lp.tag_array_axes(knl, "out", [VectorArrayDimTag()])
knl = lp.add_and_infer_dtypes(knl, {"a": np.float32, "b": np.float32, "c": np.float32})
lp.generate_code_v2(knl).device_code()

fails due to there being a dependence on i within the store operation due to the presence of an if statement.

Removing the tag_array_axes call and trying again will result in a warning but device code will be generated.

knl = lp.make_kernel( " { [i] : 0 <= i  < 8}"," out[i] = a if i == 0 else (b if i == 1 else c)")

knl = lp.tag_inames(knl, {"i": "vec"})
from loopy.kernel.array import VectorArrayDimTag
lp.generate_code_v2(knl).device_code()

will generate the following device code:

__kernel void __attribute__  (( reqn_work_group_size(1,1,1))) loopy_kernel(float const a, float const b, float const c, __global float *__restrict__ out)
{
out[0] = ((0 == 0) ? a : ((0 == 1) ? b : c));
out[1] = ((1 == 0) ? a : ((1 == 1) ? b : c));
out[2] = ((2 == 0) ? a : ((2 == 1) ? b : c));
out[3] = ((3 == 0) ? a : ((3 == 1) ? b : c));
out[4] = ((4 == 0) ? a : ((4 == 1) ? b : c));
out[5] = ((5 == 0) ? a : ((5 == 1) ? b : c));
out[6] = ((6 == 0) ? a : ((6 == 1) ? b : c));
out[7] = ((7 == 0) ? a : ((7 == 1) ? b : c));
}

[inducer]

As discussed a few times, while the lack of simplification is a potential problem, I think it's probably more useful to allow this to vectorize. There are two missing pieces AFAICT:

• The ternary operator cleanly vectorizes according to the OpenCL spec. Loopy just needs to be told that this is so.
• The loop iname for a vec loop currently isn't available in vector form. But this is straightforward, it's just a constant that needs to be generated:

  const int8 lp_idx = (int8) (0, 1, 2, 3, 4, 5, 6, 7);