slow test failure: test_nvfuser_extremal_values_masked_amin_cuda_float32

[zasdfgbnm]

🐛 Describe the bug

To reproduce:

PYTORCH_TEST_WITH_SLOW=1 python test/test_jit_cuda_fuser.py -v -k test_nvfuser_extremal_values_masked_amin_cuda_float32

Fusion math is:

Inputs:
  T0_g[ iS0{i1}, iS1{i2}, bS2{1}, iS3{i4} ], float
  T1_g[ 0 ], float
  T2_g[ iS4{i5}, iS5{i6}, bS6{1}, iS7{i7} ], bool
Outputs:
  T9_g[ 0 ], float

%kernel_math {
T3_l[ iS8{i5}, iS9{i6}, bS10{1}, iS11{i7} ]
   = T2_g[ iS4{i5}, iS5{i6}, bS6{1}, iS7{i7} ];
T4_l[ iS12{i5}, iS13{i6}, bS14{1}, iS15{i7} ]
   = T3_l[ iS8{i5}, iS9{i6}, bS10{1}, iS11{i7} ];
T5_l[ bS16{1}, bS17{1}, bS18{1}, bS19{1} ]
   = broadcast( T1_g[ 0 ] )
T6_l[ iS20{i5}, iS21{i6}, bS22{1}, iS23{i7} ]
   = where(T4_l[ iS12{i5}, iS13{i6}, bS14{1}, iS15{i7} ]
  , T0_g[ iS0{i1}, iS1{i2}, bS2{1}, iS3{i4} ]
  , T5_l[ bS16{1}, bS17{1}, bS18{1}, bS19{1} ]);
T7_l[ iS24{i5}, iS25{i6}, iS26{i7} ]
   = squeeze( T6_l[ iS20{i5}, iS21{i6}, bS22{1}, iS23{i7} ] )
T8_l[ rS27{i5}, rS28{i6}, rS29{i7} ]
   = reduction( T7_l[ iS24{i5}, iS25{i6}, iS26{i7} ], op = fmin, initial value = double(inf), allreduce = false )
T9_g[ 0 ]
   = T8_l[ rS27{i5}, rS28{i6}, rS29{i7} ];
}

but the given input T2 has shape (3, 2, 1, 1). The last dim was not correctly marked as broadcast, so I think during codegen, our system is assuming i4 == i7 and generating code based on that.

Versions

TOT devel

[zasdfgbnm]

cc @jjsjann123

[jjsjann123]

Thx for the debugging and repro instructions~~~

[jjsjann123]

hmmm. Started picking up this one. My current run fails on the same kernel, but input2 is still with shape (3, 2, 1, 2), stride (2, 1, 1, 0).

So it's indeed a broadcast dimension, but that's an explicit broadcast, which is viewed by us us as a non-broadcasted dimension. So guard logic also seems to be right here. (something might have changed?!?!)
https://github.com/csarofeen/pytorch/blob/devel/torch/csrc/jit/codegen/cuda/interface.cpp#L264-L269
https://github.com/csarofeen/pytorch/blob/devel/torch/csrc/jit/codegen/cuda/interface.cpp#L364-L369

I'll try to get a standalone repro on this one.

[jjsjann123]

Ha, I see where the confusion you had is... Changing input tensor shape isn't affecting our guard logic, since it gets expanded to a broadcast dimension before passing to nvfuser.

I believe this is actually a legit codegen issue.

A simpler repro:

import torch

def fn(t0, t1, mask):
    out = torch.where(mask, t0, t1)
    out = torch.amin(out, [0, 1, 2, 3])
    return out

t0 = torch.rand(3, 2, 1, 2, device="cuda", dtype=torch.float32)
t1 = torch.tensor(float('inf'), device="cuda")
mask = torch.rand(3, 2, 1, 1, device="cuda").bool().fill_(False)
mask[1, 0, :, :]=True
mask[2, 1, :, :]=True
mask = mask.expand(3, 2, 1, 2)

jit_fn = torch.jit.script(fn)
for i in range(5):
    o = jit_fn(t0, t1, mask)
    print(i, " val: ", o)
ref = fn(t0, t1, mask)
print("ref: ", ref)
assert(ref == o)

Pay attention to the mask input here. which is broadcasted on its last dimension (with stride 0).

Which is input T0_g in the fusion definition.

inputs:
T0_g[ iS0{i0}, iS1{i2}, bS2{1}, iS3{i4} ] bool
T1_g[ iS4{i5}, iS5{i6}, bS6{1}, iS7{i7} ] float
T2_g[ 0 ] float
outputs:
T6_g[ rS19{i0}, rS20{i2}, rS21{i4} ] float


T3_l[ bS8{1}, bS9{1}, bS10{1}, bS11{1} ]
   = broadcast( T2_g[ 0 ] )
T4_l[ iS12{i0}, iS13{i2}, bS14{1}, iS15{i4} ]
   = where(T0_g[ iS0{i0}, iS1{i2}, bS2{1}, iS3{i4} ]
  , T1_g[ iS4{i5}, iS5{i6}, bS6{1}, iS7{i7} ]
  , T3_l[ bS8{1}, bS9{1}, bS10{1}, bS11{1} ]);
T5_l[ iS16{i0}, iS17{i2}, iS18{i4} ]
   = squeeze( T4_l[ iS12{i0}, iS13{i2}, bS14{1}, iS15{i4} ] )
T6_g[ rS19{i0}, rS20{i2}, rS21{i4} ]
   = reduction( T5_l[ iS16{i0}, iS17{i2}, iS18{i4} ], op = fmin, initial value = double(inf), allreduce = false )
}

Printed out the transformation, I saw

T0_g[ iS104{( ceilDiv(( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) ), 1) )}, iS103{blockDim.x}, iS105{1}, iS101{4}, bS2{1} ]
 root domain : (iS0{i0},iS1{i2},bS2{1},iS3{i4})
 contiguity: 1 1 0 0
  Merge: iS1{i2} and iS3{i4} -> iS98{( i2 * i4 )}
  Merge: iS0{i0} and iS98{( i2 * i4 )} -> iS99{( i0 * ( i2 * i4 ) )}
  Split: iS99{( i0 * ( i2 * i4 ) )} by factor 4 -> iS100{( ceilDiv(( i0 * ( i2 * i4 ) ), 4) )}, iS101{4}, start offset: 0, stop offset: 0
  Split: iS100{( ceilDiv(( i0 * ( i2 * i4 ) ), 4) )} by factor blockDim.x -> iS102{( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) )}, iS103{blockDim.x}, start offset: 0, stop offset: 0
  Split: iS102{( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) )} by factor 1 -> iS104{( ceilDiv(( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) ), 1) )}, iS105{1}, start offset: 0, stop offset: 0

T1_g[ iS88{( ceilDiv(( ceilDiv(( ceilDiv(( i5 * ( i6 * i7 ) ), 4) ), blockDim.x) ), 1) )}, iS87{blockDim.x}, iS89{1}, iS85{4}, bS6{1} ]
 root domain : (iS4{i5},iS5{i6},bS6{1},iS7{i7})
 contiguity: 1 1 1 1
  Merge: iS5{i6} and iS7{i7} -> iS82{( i6 * i7 )}
  Merge: iS4{i5} and iS82{( i6 * i7 )} -> iS83{( i5 * ( i6 * i7 ) )}
  Split: iS83{( i5 * ( i6 * i7 ) )} by factor 4 -> iS84{( ceilDiv(( i5 * ( i6 * i7 ) ), 4) )}, iS85{4}, start offset: 0, stop offset: 0
  Split: iS84{( ceilDiv(( i5 * ( i6 * i7 ) ), 4) )} by factor blockDim.x -> iS86{( ceilDiv(( ceilDiv(( i5 * ( i6 * i7 ) ), 4) ), blockDim.x) )}, iS87{blockDim.x}, start offset: 0, stop offset: 0
  Split: iS86{( ceilDiv(( ceilDiv(( i5 * ( i6 * i7 ) ), 4) ), blockDim.x) )} by factor 1 -> iS88{( ceilDiv(( ceilDiv(( ceilDiv(( i5 * ( i6 * i7 ) ), 4) ), blockDim.x) ), 1) )}, iS89{1}, start offset: 0, stop offset: 0

Where the contiguity on T0_g is marked correctly (1, 1, 0, 0)
However, in the generated kernel, we seems to be indexing into T0 incorrectly (by collapsing everything to its fastest stride)

      for(nvfuser_index_t i77 = 0; i77 < 4; ++i77) {
        bool T7[1];
        T7[0] = 0;
        T7[0]
           = T0[(((((i75 * ((nvfuser_index_t)blockDim.x)) + ((nvfuser_index_t)threadIdx.x)) * 4) + (i77 + nvfuser_zero)) * T0.stride[3])];
        float T4[1];
        T4[0]
           = where(T7[0], T8[i77], T3[0]);
        float T5[1];
        T5[0]
           = T4[0];
        T11[0] = fmin(
          T11[0],
          T5[0]);
      }

Full transformation & kernel attached in text.

issue_2169.txt

[naoyam]

With this:

T0_g[ iS104{( ceilDiv(( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) ), 1) )}, iS103{blockDim.x}, iS105{1}, iS101{4}, bS2{1} ]
 root domain : (iS0{i0},iS1{i2},bS2{1},iS3{i4})
 contiguity: 1 1 0 0

The T0 index seems to make sense since the innermost dimension isn't recognized as a broadcast. We are just ignoring the only broadcast (bS2), and since the outer two dimensions are contiguous, the contiguous index should be valid.

If the innermost dimension is a broadcast, why isn't it created as a broadcast domain?

[jjsjann123]

With this:

T0_g[ iS104{( ceilDiv(( ceilDiv(( ceilDiv(( i0 * ( i2 * i4 ) ), 4) ), blockDim.x) ), 1) )}, iS103{blockDim.x}, iS105{1}, iS101{4}, bS2{1} ]
 root domain : (iS0{i0},iS1{i2},bS2{1},iS3{i4})
 contiguity: 1 1 0 0

The T0 index seems to make sense since the innermost dimension isn't recognized as a broadcast. We are just ignoring the only broadcast (bS2), and since the outer two dimensions are contiguous, the contiguous index should be valid.

If the innermost dimension is a broadcast, why isn't it created as a broadcast domain?

Feels like there are two parts here:

1. If the innermost dimension is a broadcast, why isn't it created as a broadcast domain?

Your accusation here is legit. We can refactor the guard we have for broadcast dimension, where we have concrete sizes and 0 stride. That is currently not how we parse things.

2. However, the indexing logic is still wrong.. We shouldn't be able to collapse i2 to i4 in any circumstance here. In a counter example, we could have the innermost dimension being non-broadcast with an arbitrary stride (that is not collapsible to i2 here).
   More concretely, if we look at T0 with size(3, 2, 1, 2) and stride(8, 4, 4, 1).
   It would still satisfy

 root domain : (iS0{i0},iS1{i2},bS2{1},iS3{i4})
 contiguity: 1 1 0 0

But we can't collapse the index to only use the innermost stride, which would give us the wrong result here.

[naoyam]

2. However, the indexing logic is still wrong.. We shouldn't be able to collapse i2 to i4 in any circumstance here. In a counter example, we could have the innermost dimension being non-broadcast with an arbitrary stride (that is not collapsible to i2 here).
   More concretely, if we look at T0 with size(3, 2, 1, 2) and stride(8, 4, 4, 1).
   It would still satisfy

 root domain : (iS0{i0},iS1{i2},bS2{1},iS3{i4})
 contiguity: 1 1 0 0

But we can't collapse the index to only use the innermost stride, which would give us the wrong result here.

OK, I didn't think about strided broadcast domains. There seems like something we need to fix.

[zasdfgbnm]

I think it has something to do with the definition of contiguity. For example, if we have IterDomains [I1, I2, b, I3] and contiguity [1, 1, 0, 0], currently I2 being contiguous means it could collapse to b instead of I3. In indexing, we must not be honoring this definition and skipping the b. I believe we had some discussion before, and we believe that the contiguity should ignore the broadcast, so that I2 being contiguous should mean it can be collapsed to I3. Self-assigning, and I will try to make a fix.

[jjsjann123]

Periodic bump on this issue 😝

[jjsjann123]

FYI, I noticed that this test has been disabled in both upstream & devel, Since we pulled upstream into devel with the disablement.

But the issue is still around. 😢