Handle empty tensors during definition of cat

[jacobhinkle]

In #3292, a PadOp is created for a pad of two inputs, one of which has zero size in the cat dimension. This caused an error when we replaced the empty input with a FullOp output. That is addressed in the remove_empty pass by PR #3301. This PR aims to additionally simplify the fusion definition when we have concrete sizes. In particular, arguments to pad and cat are inspected for empty dimensions and if found, we avoid using PadOp or CatOp unless necessary. The behavior is demonstrated in the included CatOfEmpty test:

Cat of tensors where some non-cat dimensions are zero results in a FullOp of the proper size
Cat of tensors that are all empty in cat dim results in LoadStoreOp of the first input
Cat of tensors, some empty and some not, results in ignoring the empty inputs. In this case that leaves only a set op
Cat of tensors, some empty and some not, results in ignoring the empty inputs. In this case that leaves us with a CatOp with two inputs

[jacobhinkle]

!build --diff

[naoyam]

Can we use Fusion::zeroVal(DatatType)?

[naoyam]

I'm not sure what the correct behavior should be here. Can we just use the first extent for a non-cat dimension? Suppose the first cat input has size zero for that dimension, the output of the cat would also have zero for the dimension. Would that be expected?

[jacobhinkle]

In that case it should be caught by the inputs.empty() condition at line 678, since all of the inputs should be empty. Actually though, I just realized that assumes that if one of the inputs has constant size 0 in one dimension that all the other inputs will have constant size zero in the same dimension. If some are symbolic this definition should be proof that they're equal by exact mapping, but they won't be detected at line 660. I guess what I should do instead is fire this condition if any input has a zero non-cat dimension.

[naoyam]

if one of the inputs has constant size 0 in one dimension that all the other inputs will have constant size zero in the same dimension.

Yeah, this was what I was thinking about. I'm not sure if that's actually allowed. Is it?

[jacobhinkle]

We could have something like this I think:

tv0[ 2, 0 ]
tv1[ 3, i0 ]
tv2 = cat({tv0, tv1}, /*axis=*/0)

In this case we would normally exact map tv0->axis(1) with tv1->axis(1) so that i0 must be 0 or else we'll hit an error in ExpressionEvaluator::propagateBoundValuesThroughExactMaps(). But now since we'll just translate this to full({5, 0}) I don't think there's any such constraint, so maybe it would be legal to pass in a tv1 with shape (3,2) and we would not hit an error.

[naoyam]

OK, makes sense.

[jacobhinkle]

The failure in jit_binary_tests_17_A100_1/3 is real. I have a fix but am waiting to push until the codediff is done.

[jacobhinkle]

!test --diff