Floating-point exception scheduling reduction of zero-element tensor

[jacobhinkle]

Currently, scheduling a reduction with a tensor having zero size results in an unhandled SIGFPE due to 0/0 at https://github.com/NVIDIA/Fuser/blob/main/csrc/scheduler/reduction.cpp#L558 (n_elems=target_blocks=0).

Here is a minimal reproduction:

// Test that 0-dimensional tensors do not break reduction scheduler
TEST_F(NVFuserTest, FusionReduceZeroElementTensor_CUDA) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  std::vector<int64_t> input_shape{3, 4, 0, 5};
  auto tv0 = makeSymbolicTensor(4);
  fusion->addInput(tv0);
  auto tv1 = sum(tv0, {1});
  fusion->addOutput(tv1);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor at_x = at::randn(input_shape, options);
  FusionExecutorCache executor_cache(std::move(fusion));
  auto outputs = executor_cache.runFusionWithInputs({at_x});
  auto t2 = at_x.sum({2});

  auto reduction_params = getReductionHeuristics(fusion.get(), {at_x});
  TORCH_CHECK(reduction_params, "Reduction schedule was not generated!");
  scheduleReduction(fusion.get(), *reduction_params);

  testValidate(
      executor_cache.fusion(), outputs, {at_x}, {t2}, __LINE__, __FILE__);
}

Note that the reduction does not need to be along the 0 dimension to trigger the crash. We have a test for zero-dim tensors in the presence of reductions, but it is not actually reducing the zero-dim tensor:

Fuser/test/test_gpu2.cpp

Line 5984 in 13a90b6

TEST_F(NVFuserTest, FusionZeroSizeTensorReduction_CUDA) {

.

I tried guarding a few of the ceilDivs in outerReductionHeuristic which can avoid the SIGFPE but lead to a kernel segfault, so I'm posting this issue for a little discussion before I try a fix. We may just need to check numel at the beginning of that function and short-circuit there if it's zero.

[jacobhinkle]

If we have concrete input extents, we can handle reduction by rewriting. Letting n_red be the numel of the reduction axes, and n_it the numel of the non-reduction axes of the input, then:

1. If reducing over all dimensions, and n_red > 0 then we're reducing to a scalar. All output IDs are Reduction.
2. Else if n_it == 0: There are no elements in the input or output. This is a no-op and the output IDs have type Iteration.
3. Else if n_red==0. Output is equivalent to full using the reduction's starting value. All output IDs are Broadcast.
4. Else if n_red==1, this op is a reshape (squeeze) instead of a reduction.
5. Otherwise, the output is the usual mix of Iteration and Reduction, and the reduction scheduler proceeds as usual.

I think this could be handled with dynamic transforms similar to how we've handled reshape (#24) and resize (#256). Every reduction output would depend on the input extents, so we would need to set each output ID's IterType to Symbolic during definition. Since cases (1) and (2) result in removal of the reduction op, this may mean we need to use the pointwise scheduler instead of reduction and it may also impact segmentation.

Since 5 is probably overwhelmingly the most common case, we could also just try and make the reduction scheduler robust against all these cases instead.

[naoyam]

Else if n_red==0. Output is equivalent to full using the reduction's starting value. All output IDs are Broadcast.

This means that the input to the reduction has a size-0 dimension. I thought if a tensor has a size-0 dimension, it's an empty tensor, and reducing it should be no-op. Is this not true? Why is it equivalent to full?

[jacobhinkle]

In torch, reducing zero-size dimensions results in the starting value. For example:

[ins] In [1]: import torch
[ins] In [2]: x = torch.randn(3, 0, 2)
[ins] In [3]: y = x.sum(1)
[ins] In [4]: y
Out[4]: 
tensor([[0., 0.],
        [0., 0.],
        [0., 0.]])
[ins] In [5]: z = x.prod(1)
[ins] In [6]: z
Out[6]: 
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]])

[jacobhinkle]

In these cases, note that the output does not have any size-0 dimensions. Otherwise, you're right that it's a no-op (case 2)

[naoyam]

Wow. Do you know this behavior is specific to PyTorch? Is it also the same with numpy?

[jacobhinkle]

Yeah, I think this is identical to numpy. Also in pure python if you do sum([]) the result is 0.

[naoyam]

So,

Every reduction output would depend on the input extents, so we would need to set each output ID's IterType to Symbolic during definition.

This seems to be the case, which means Symbolic domains would be much more common. We could do this right now, but we would really need to pay much more attention to the latency overhead.

[jacobhinkle]

Zero-element tensors do seem like a corner case. Still, I will start looking more closely at #244

[jacobhinkle]

This seems to be the case, which means Symbolic domains would be much more common. We could do this right now, but we would really need to pay much more attention to the latency overhead.

I am coming around to this approach. I have been handling zero-size in the schedulers recently for #269 and it does not scale very well; for example, the NoOp scheduler needs to process and replace resizes at runtime with full ops, and do so in its own way in order to avoid dynamic allocations in local memory (even though those will be zero size!). That is different from how pointwise works, which can simiply replace with full() and the scheduler takes care of the rest. These issues can be overcome, but it feels brittle as there is always the possibility that we will find more ops that require this kind of multiple checks.

If most Fusions are dynamic, it doesn't necessarily introduce a lot of additional runtime. The first level of caching is done with KernelArgumentHolder's cache ID. We can use a minimal set of values for that hash so that we have the largest possible chance at a cache hit right there. Most of the time, from iteration to iteration, these vals won't change so that cache ID will provide us with a new kernel for which we don't even have to check heuristics. So that's the fast path that won't change much at all. In the event that we have a miss, we're back in the situation we are currently, but the logic for the schedulers can be simpler since we can assume that reduction outputs are not actually fulls or squeezes.

[naoyam]

the NoOp scheduler needs to process and replace resizes at runtime with full ops,

Did you mean replacing reductions?