More lost scalars during segmentation for dynamic reshape

[jacobhinkle]

// Test concretizing a pad that follows a reshape. This requires the
// ExpressionEvaluator used in concretization to propagate shapes properly
// across symbolic reshapes in order to infer the size of the downstream pad.
TEST_F(NVFuserTest, DynamicReshapePad_CUDA) {
  std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
  Fusion& fusion = *fusion_ptr.get();
  FusionGuard fg(&fusion);

  TensorView* tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto s0 = IrBuilder::create<Val>(DataType::Int);
  auto s1 = IrBuilder::create<Val>(DataType::Int);
  auto s2 = IrBuilder::create<Val>(DataType::Int);
  auto s3 = IrBuilder::create<Val>(DataType::Int);
  fusion.addInput(s0);
  fusion.addInput(s1);
  fusion.addInput(s2);
  fusion.addInput(s3);

  auto tv1 = reshape(tv0, {s2, s3});
  auto tv2 = pad(tv1, {fusion.zeroVal(), s0, fusion.zeroVal(), s1});
  fusion.addOutput(tv2);

  /*
Symbolic Fusion:
Inputs:
  T0_g[ iS0{i0}, iS1{i2} ], float
  i3, int64_t
  i4, int64_t
  i5, int64_t
  i6, int64_t
Outputs:
  T2_g[ ?S7{( i5 + i4 )}rf, ?S9{( i6 + i3 )}rf ], float

%kernel_math {
T1_l[ ?S4{i5}rf, ?S5{i6}rf ] = view( T0_g[ iS0{i0}, iS1{i2} ] )
T2_g[ ?S7{( i5 + i4 )}rf, ?S9{( i6 + i3 )}rf ]
   = pad( T1_l[ ?S4{i5}rf, ?S5{i6}rf ], {0, i4, 0, i3} )
}

Concretized Fusion:
Inputs:
  T0_g[ iS0{i0}, iS1{i2} ], float
  i3, int64_t
  i4, int64_t
  i5, int64_t
  i6, int64_t
Outputs:
  T2_g[ iS21{( i5 + i4 )}rf, iS22{( i6 + i3 )}rf ], float

%kernel_math {
T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ] = view( T0_g[ iS0{i0}, iS1{i2} ] )
T2_g[ iS21{( i5 + i4 )}rf, iS22{( i6 + i3 )}rf ]
   = pad( T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ], {0, i4, 0, i3} )
}

T0_g[ iS0{i0}, iS1{i2} ]
 root domain : (iS0{i0}, iS1{i2})
 contiguity: f f
 leaf domain : (iS0{i0}, iS1{i2})
T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ]
 root domain : (iS12{i0}rf, iS13{i2}rf)
  Merge: iS12{i0}rf and iS13{i2}rf -> iS14{( i0 * i2 )}rf
  Outer split: iS14{( i0 * i2 )}rf by factor 3 -> iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf, start offset: 0, stop offset: 0
 rfactor domain : (iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf)
 contiguity: t t
 leaf domain : (iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf)
T2_g[ iS21{( i5 + i4 )}rf, iS22{( i6 + i3 )}rf ]
 root domain : (iS19{i5}rf, iS20{i6}rf)
  Resize: iS19{i5}rf by 0 and i4 -> iS21{( i5 + i4 )}rf
  Resize: iS20{i6}rf by 0 and i3 -> iS22{( i6 + i3 )}rf
 rfactor domain : (iS21{( i5 + i4 )}rf, iS22{( i6 + i3 )}rf)
 contiguity: t t
 leaf domain : (iS21{( i5 + i4 )}rf, iS22{( i6 + i3 )}rf)
  */

  FusionExecutorCache fusion_executor_cache(std::move(fusion_ptr));

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor at_x = at::randn({4, 3}, options);
  std::vector<c10::IValue> aten_inputs = {at_x, 1, 1, 3, 4};
  auto at_y = at::pad(at_x.reshape({3, 4}), {0, 1, 0, 1});

  auto outputs = fusion_executor_cache.runFusionWithInputs(aten_inputs);
  /*
terminate called after throwing an instance of 'c10::Error'
  what():  val.hasValue() INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/executor.cpp":1039, please report a bug to
PyTorch. Tried to evaluate the extent, ( ceilDiv(( ( ( ( i5 * i6 ) + ( i5 * i3 ) ) + ( i4 * i6 ) ) + ( i4 * i3 ) ), 128) )
 for the ptype: blockIdx.x to set launch bounds but could not.
  */

  testValidate(
      fusion_executor_cache.fusion(),
      outputs,
      aten_inputs,
      {at_y},
      __LINE__,
      __FILE__);
}

The fusion gets segmented at T3 and the second segment cannot evaluate the output extent, just like in #629. The fix in that case (#630) was to propagate the replaced extents during concretization. In this case that approach does not resolve the issue, since the downstream pad op builds an extent expression using the fusion inputs i5 and i6. Although we register those values for mutation, the actual rfactor extents in T2 are i5 + i4 and i6 + i3 so they are not replaced. It is not necessarily clear that replacing them before segmentation is always advantageous either: before segmentation we don't know if the original expression with fusion inputs, e.g. i6 + i3, will be easier to compute than the concretized one: ceilDiv(( i0 * i2), 3) + i3.

[jacobhinkle]

There might be a simpler solution to this issue that I haven't noticed yet, but if not then it might be time to consider using a union-find to organize these collections of equivalent scalars, and using those at segmentation to select representatives to bind. There is a data structure available in #382 we can use for this. This is a sketch of how we would use it:

• Introduce a UnionFind object that tracks scalar ints and provides two key functions: merge(a, b) which marks two scalars as always equal and bool equiv(a, b) which tests transitively whether they have been merged.
• Add a function that sets vals to be equivalent in the UnionFind whenever they are exact-mapped extents, similar to ExpressionEvaluator::propagateBoundValuesThroughExactMaps().
• Update ExpressionEvaluator to optionally try and use the UnionFind to substitute scalars if it fails to directly evaluate them. This gives us a way to test which scalars are computable from the bound inputs and which ones will require substitution.
• Add an OptOutMutator subclass that traverses the Fusion and uses an ExpressionEvaluator with bound inputs along with the UnionFind to swap scalars for equivalent ones which are directly computable from the bound inputs. This resembles the replaceSymbolicSizes pass and can similarly just use ir_utils::replaceValue.

In the example above, for the second segment's kernel we would bind the shape of T3 as {3, 4} and the traversal would replace instances of i5 with 3 and i6 with ceilDiv(( i0 * i2), 3) which is bound to 4. In this case, ceilDiv(( i0 * i2), 3) is actually converted to a new root scalar during convertInputRfactorsToRoots(), so we need to ensure that that function is also able to merge scalars in the UnionFind. Since these types of operations are somewhat disperse throughout the code, I think the UnionFind should be a member of Fusion called int_equiv_ or something like that.

[jacobhinkle]

Alternatives to the union-find approach:

• Instead of just registering the reshaped extents for mutation when concretized, actually replace them in all uses. This would mean that if an input scalar is used as a reshape, it would get replaced everywhere with an expression like ceilDiv(( i0 * i2), i3).
• Always set the output extents of a dynamic reshape to placeholder Vals, and store the requested shapes as inputs. Then, during concretization we could replace the placeholder Vals in all uses. We would need to store the requested shapes as inputs since currently we use the output's rfactor extents directly to determine the reshaped size. I have tried this method and the problem is that we need to propagate bound values through exact maps in many more places, including FusionExecutor; even then, I tend to get missing scalars since the exact maps don't necessarily cover every expression we need. For example, if we bind through exact maps a ceilDiv(a, 1) that actually does not make a available for evaluation, but this type of simplification is done at some points it seems.

[jacobhinkle]

I tried this example with slice instead of pad:

TEST_F(NVFuserTest, DynamicReshapeSlice_CUDA) {
  std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
  Fusion& fusion = *fusion_ptr.get();
  FusionGuard fg(&fusion);

  TensorView* tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto s0 = IrBuilder::create<Val>(DataType::Int);
  auto s1 = IrBuilder::create<Val>(DataType::Int);
  auto s2 = IrBuilder::create<Val>(DataType::Int);
  auto s3 = IrBuilder::create<Val>(DataType::Int);
  fusion.addInput(s0);
  fusion.addInput(s1);
  fusion.addInput(s2);
  fusion.addInput(s3);

  auto tv1 = reshape(tv0, {s2, s3});
  auto tv2 = slice(tv1, {{fusion.zeroVal(), s0}, {fusion.zeroVal(), s1}});
  fusion.addOutput(tv2);

  FusionExecutorCache fusion_executor_cache(std::move(fusion_ptr));

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor at_x = at::randn({4, 3}, options);
  std::vector<c10::IValue> aten_inputs = {at_x, 3, 2, 3, 4};
  auto at_y = at::slice(at::slice(at_x.reshape({3, 4}), 0, 0, 3), 1, 0, 2);

  auto outputs = fusion_executor_cache.runFusionWithInputs(aten_inputs);
  /*
...
__global__ void kernel2(int64_t i0, int64_t i1, Tensor<float, 2, 2> T3, Tensor<float, 2, 2> T2) {
  nvfuser_index_t i2;
  i2 = ((nvfuser_index_t)threadIdx.x) + (128 * ((nvfuser_index_t)blockIdx.x));
  nvfuser_index_t i3;
  i3 = (nvfuser_index_t)(i1);
  nvfuser_index_t i4;
  i4 = (i3 - (nvfuser_index_t)(i5)) + T3.logical_size[1];
  nvfuser_index_t i6;
  i6 = (nvfuser_index_t)(i0);
  if ((i2 < (i6 * i3))) {
    float T5[1];
    T5[0]
       = T3[((T3.logical_size[1] * (i2 / i4)) + (i2 % i4))];
    T2[i2]
       = T5[0];
  }
}
}

CUDA NVRTC compile error: __tmp_kernel2.cu(9566): error: identifier "i5" is undefined
    i4 = (i3 - (nvfuser_index_t)(i5)) + T3.logical_size[1];
                                 ^

1 error detected in the compilation of "__tmp_kernel2.cu".

Exception raised from invoke at /opt/pytorch/nvfuser/csrc/executor_utils.cpp:857 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xae (0x7ffff5c2c55e in /opt/pytorch/pytorch/torch/lib/libc10.so)
  */

  testValidate(
      fusion_executor_cache.fusion(),
      outputs,
      aten_inputs,
      {at_y},
      __LINE__,
      __FILE__);
}

I think this is the same scalar getting lost, but it the extent is simplified such that the error is not encountered until compilation. The complete fusion in this case is:

Fusion IR after pre-segmenter optimization passes:
Inputs:
  T0_g[ iS0{i0}, iS1{i2} ], float
  i3, int64_t
  i4, int64_t
  i5, int64_t
  i6, int64_t
Outputs:
  T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ], float

%kernel_math {
T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ] = view( T0_g[ iS0{i0}, iS1{i2} ] )
i16 = (nvfuser_index_t)(i3);
i20 = (nvfuser_index_t)(i4);
T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ]
   = slice( T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ], { {0, i16, 1} {0, i20, 1} } )
}

T0_g[ iS0{i0}, iS1{i2} ]
 root domain : (iS0{i0}, iS1{i2})
 contiguity: f f
 leaf domain : (iS0{i0}, iS1{i2})
T3_l[ iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf ]
 root domain : (iS12{i0}rf, iS13{i2}rf)
  Merge: iS12{i0}rf and iS13{i2}rf -> iS14{( i0 * i2 )}rf
  Outer split: iS14{( i0 * i2 )}rf by factor 3 -> iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf, start offset: 0, stop offset: 0
 rfactor domain : (iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf)
 contiguity: t t
 leaf domain : (iS15{3}rf, iS16{( ceilDiv(( i0 * i2 ), 3) )}rf)
T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ]
 root domain : (iS21{3}rf, iS22{( ceilDiv(( i0 * i2 ), 3) )}rf)
  Resize: iS21{3}rf by 0 and ( ( (nvfuser_index_t)(i3) ) - ( (nvfuser_index_t)(i5) ) ) -> iS23{( (nvfuser_index_t)(i3) )}rf
  Resize: iS22{( ceilDiv(( i0 * i2 ), 3) )}rf by 0 and ( ( (nvfuser_index_t)(i4) ) - ( (nvfuser_index_t)(i6) ) ) -> iS24{( (nvfuser_index_t)(i4) )}rf
 rfactor domain : (iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf)
 contiguity: t t
 leaf domain : (iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf)

After segmentation, the second segment in this case is:

Inputs:
  i3, int64_t
  i4, int64_t
  T3_g[ iS25{i124}, iS26{i125} ], float
Outputs:
  T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ], float

%kernel_math {
i16 = (nvfuser_index_t)(i3);
i20 = (nvfuser_index_t)(i4);
T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ]
   = slice( T3_g[ iS25{i124}, iS26{i125} ], { {0, i16, 1} {0, i20, 1} } )
}

T3_g[ iS25{i124}, iS26{i125} ]
 root domain : (iS25{i124}, iS26{i125})
 contiguity: t t
 leaf domain : (iS25{i124}, iS26{i125})
T2_g[ iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf ]
 root domain : (iS29{i124}rf, iS30{i125}rf)
  Resize: iS29{i124}rf by 0 and ( ( (nvfuser_index_t)(i3) ) - ( (nvfuser_index_t)(i5) ) ) -> iS23{( (nvfuser_index_t)(i3) )}rf
  Resize: iS30{i125}rf by 0 and ( ( (nvfuser_index_t)(i4) ) - ( (nvfuser_index_t)(i6) ) ) -> iS24{( (nvfuser_index_t)(i4) )}rf
 rfactor domain : (iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf)
 contiguity: t t
 leaf domain : (iS23{( (nvfuser_index_t)(i3) )}rf, iS24{( (nvfuser_index_t)(i4) )}rf)

In this case the fusion math looks fine, but the Resize ops still hold expansions with i5 and i6 in them.

[jacobhinkle]

• Instead of just registering the reshaped extents for mutation when concretized, actually replace them in all uses. This would mean that if an input scalar is used as a reshape, it would get replaced everywhere with an expression like ceilDiv(( i0 * i2), i3).

I'm going to try this approach out. I have another branch exploring more placeholders https://github.com/NVIDIA/Fuser/tree/dynamic_extent_placeholders, which has some other disadvantages, and another branch implementing a union-find for scalar propagation https://github.com/NVIDIA/Fuser/tree/unionfind_exactmapping. Those two will require a bit more work probably, so I'll leave those as backup plans.