Static pad to broadcast fails to schedule

[jacobhinkle]

This is a corner case for Resize ops that hasn't had a lot of attention yet, in which we use pad in a way that results in a broadcast that is then resolved.

TEST_F(NVFuserTest, PadToBroadcast) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto tv0 = makeConcreteTensor({2});
  auto tv1 = makeConcreteTensor({3});
  fusion->addInput(tv0);
  fusion->addInput(tv1);

  auto tv2 = pad(tv0, {fusion->zeroVal(), IrBuilder::create<Scalar>(-1)});
  auto tv3 = mul(tv1, tv2);
  fusion->addOutput(tv3);

  // Fusion is not dynamic
  EXPECT_FALSE(fusion->hasDynamicTransform());

  fusion->printMath();
  /*
Inputs:
  T0_g[ iS0{2} ], float
  T1_g[ iS1{3} ], float
Outputs:
  T3_g[ iS4{3} ], float                                                                                                                                                                                                                                           %kernel_math {
T2_l[ bS3{1}rf ]
   = pad( T0_g[ iS0{2} ], {0, -1} )
T3_g[ iS4{3} ]
   = T1_g[ iS1{3} ]
   * T2_l[ bS3{1}rf ];
}
  */
  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);

  auto t0 = at::randn({2}, options);
  auto t1 = at::randn({3}, options);
  std::vector<c10::IValue> aten_inputs({t0, t1});

  auto args = KernelArgumentHolder::createKernelArgumentHolder(aten_inputs);
  FusionKernelRuntime runtime(std::move(fusion), args);
  /*
terminate called after throwing an instance of 'c10::Error'
  what():  it != broadcast_origin_map_.end() INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/device_lower/analysis/trivial_broadcast.cpp":98, please report a bug to PyTorch.
Broadcast origin info not found for producer broadcast domain: bS3{1}rf of T2_l[ bS3{1}rf ]
Exception raised from handle at /opt/pytorch/nvfuser/csrc/device_lower/analysis/trivial_broadcast.cpp:98 (most recent call first)
  */

  runtime.compileFusionParallel(args);
  auto cg_outputs = runtime.runWithInputs(args);

  auto t2_padded = at::pad(t0, {0, -1});
  auto ref_t2 = t1 * t2_padded;

  testValidate(
      runtime.fusionSegments()->completeFusion(),
      cg_outputs,
      aten_inputs,
      {ref_t2},
      __LINE__,
      __FILE__);
}

Simply determining the output IterType is not sufficient to be able to actually resolve the broadcast. We instead need to have an actual tensor allocated to act as the broadcast origin. This can be done in this case by translating the pad to

  auto tv2 = broadcast(select(tv0, 0, fusion->zeroVal()), {true});

in which case select introduces the needed intermediate tensor which is the origin of the following broadcast.

Note that dynamic pad faces a similar issue but that we should probably fix this static case first. Static and dynamic slice will face a similar problem (related to #511).

[jacobhinkle]

Since this is a transform that should be applied the same way to static and dynamic fusions, it might be best to do it as a pre-segmentation pass. Currently, those passes are called "optimization passes", but only ConsecutiveCastPass is performance oriented. Both AddAxiomsPass and RemoveEmptyPass deal with correctness, and a new BroadcastPadPass would have a similar flavor: it takes an unambiguous Fusion and makes it much simpler to properly schedule/compile by translating to a more straightforward set of ops.

[naoyam]

The broadcast analysis fails because there's no BroadcastOp found for broadcast domains introduced by the pad op?

[jacobhinkle]

Yeah I believe so. It happens when constructing a ConcretizedBroadcastDomains which is a member of ComputeAtMap. One issue may be that ConcretizedBroadcastDomains only recognizes new broadcast origins by handling BroadcastOp. We could add specific behavior here for PadOp, which might be simpler. I'll give that a try but something else in codegen might break with this approach..

[jacobhinkle]

I gave it a quick try working around the limitation in ConcretizedBroadcastDomains. Then I hit an issue with ComputeAtRootDomainMapBuilder. I think the issue there is that there's just no machinery in place for introducing broadcasts in the rfactor when none exists in the root. I will keep trying to understand ComputeAtRootDomainMapBuilder::handle(TensorView* tv) to see if this can be added.

[jjsjann123]

Is this issue resolved?

I actually run into another corner case on this one, where we are trying to pad an unsqueezed size-0 tensor.

import torch
from nvfuser import FusionDefinition, DataType

def nvfuser_fusion_id11(fd : FusionDefinition) -> None :
    T0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[None, None, None], dtype=DataType.Double, is_cpu=False)
    S1 = fd.define_scalar(-3.70753, dtype=DataType.Double)
    T2 = fd.ops.pad(T0, [0, 0, 1, 1, 1, 0], S1)
    fd.add_output(T2)

with FusionDefinition() as fd:
    nvfuser_fusion_id11(fd)

inputs = [
    torch.randn((0,), dtype=torch.float64, device='cuda:0').as_strided((2, 0, 3), (0, 0, 0)),
]
fd.execute(inputs)

[jacobhinkle]

I actually run into another corner case on this one, where we are trying to pad an unsqueezed size-0 tensor.

This actually looks like a new issue. It concretizes to this:

Inputs:
  T0_g[ bS0{i0}, bS9{0}, bS2{i2} ], double
Outputs:
  T1_g[ bS11{( ( i0 + ( (nvfuser_index_t)(1) ) ) + ( (nvfuser_index_t)(0) ) )}rf, iS6{( ( 0 + ( (nvfuser_index_t)(1) ) ) + ( (nvfuser_index_t)(1) ) )}rf, bS7{i2}rf ], double

%kernel_math {
i12 = (nvfuser_index_t)(1);
i14 = (nvfuser_index_t)(0);
i16 = (nvfuser_index_t)(1);
i18 = (nvfuser_index_t)(1);
i20 = (nvfuser_index_t)(0);
i22 = (nvfuser_index_t)(0);
T1_g[ bS11{( ( i0 + ( (nvfuser_index_t)(1) ) ) + ( (nvfuser_index_t)(0) ) )}rf, iS6{( ( 0 + ( (nvfuser_index_t)(1) ) ) + ( (nvfuser_index_t)(1) ) )}rf, bS7{i2}rf ]
   = pad( T0_g[ bS0{i0}, bS9{0}, bS2{i2} ], {i12, i14, i16, i18, i20, i22} )
}

The input in this case should have expanded extent i0, 0, i2 = 2, 0, 3. It appears i0 and i2 are being recognized as 1, so the IterTypes are concretized improperly. I think this is an issue with binding expanded inputs.