Broadcast in dim with expand

[jjsjann123]

Fixes #1788

Added expand in broadcast_in_dim to support expanding to concrete size. Note that we are not supporting dynamic shape for concrete size at this moment.

[jjsjann123]

cc'ing @IvanYashchuk

[IvanYashchuk]

I tried to run the first code from the issue and I hit the same error:

Inputs:
  T0_g[ iS0{i0}, iS1{i2} ], float
  T1_g[ iS2{i3} ], float
Outputs:
  T6_g[ iS11{i0}, iS12{i2} ], float

%kernel_math {
T2_l[ iS3{i0}, iS4{i2} ]
   = T0_g[ iS0{i0}, iS1{i2} ];
T3_l[ iS5{i0}, iS6{i2} ]
   = T2_l[ iS3{i0}, iS4{i2} ];
T4_l[ bS7{1}, iS8{i3} ]
   = broadcast( T1_g[ iS2{i3} ] )
T5_l[ bS9{3}, iS10{i3} ] = expand( T4_l[ bS7{1}, iS8{i3} ], {3, i3} )
T6_g[ iS11{i0}, iS12{i2} ]
   = T3_l[ iS5{i0}, iS6{i2} ]
   + T5_l[ bS9{3}, iS10{i3} ];
}

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [1], in <cell line: 27>()
     24 input1 = torch.ones(3, 1, device='cuda')
     25 input2 = torch.ones(3, device='cuda')
---> 27 fusion.execute([input1, input2])

RuntimeError: Attempting to bind T0.size[1] to 3but it's already set to 1

[csarofeen]

Let me know if you need me to look into this.

[jjsjann123]

I tried to run the first code from the issue and I hit the same error:

Inputs:
  T0_g[ iS0{i0}, iS1{i2} ], float
  T1_g[ iS2{i3} ], float
Outputs:
  T6_g[ iS11{i0}, iS12{i2} ], float

%kernel_math {
T2_l[ iS3{i0}, iS4{i2} ]
   = T0_g[ iS0{i0}, iS1{i2} ];
T3_l[ iS5{i0}, iS6{i2} ]
   = T2_l[ iS3{i0}, iS4{i2} ];
T4_l[ bS7{1}, iS8{i3} ]
   = broadcast( T1_g[ iS2{i3} ] )
T5_l[ bS9{3}, iS10{i3} ] = expand( T4_l[ bS7{1}, iS8{i3} ], {3, i3} )
T6_g[ iS11{i0}, iS12{i2} ]
   = T3_l[ iS5{i0}, iS6{i2} ]
   + T5_l[ bS9{3}, iS10{i3} ];
}

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [1], in <cell line: 27>()
     24 input1 = torch.ones(3, 1, device='cuda')
     25 input2 = torch.ones(3, device='cuda')
---> 27 fusion.execute([input1, input2])

RuntimeError: Attempting to bind T0.size[1] to 3but it's already set to 1

I think you meant this example vvv

import torch

from torch._C._nvfuser import Fusion, FusionDefinition

# Construct and Define Fusion
fusion = Fusion()

with FusionDefinition(fusion) as fd :
    t0 = fd.define_tensor(2)
    t1 = fd.define_tensor(1)

    fd.add_input(t0)
    fd.add_input(t1)

    t0_b = fd.Ops.broadcast_in_dim(t0, [3, 3], [0, 1])
    t1_b = fd.Ops.broadcast_in_dim(t1, [3, 3], [1])
    t2 = fd.Ops.add(t0_b, t1_b)

    fd.add_output(t2)

fusion.print_ir()

# Execute Fusion
input1 = torch.ones(3, 1, device='cuda')
input2 = torch.ones(3, device='cuda')

fusion.execute([input1, input2])

I think that's not a valid python example. When you create t0, you need to tell fusion that its dimension-1 is of size-1 (a broadcasted dimension). so t0 = fd.define_tensor(2) need to be updated with t0 = fd.define_tensor([3, 1], [1, 1]). (sorry about 3 here, we don't have dynamic support in python API, we can change that tho).

Without knowing that, we won't know what to do with broadcast on t0.dim1 in the broadcast_in_dim.

[jjsjann123]

Note that in the above comment, I'm looking only at our current FusionDefinition API. It's likely going to change with Kevin's next iteration, which I believe is going to hide all the nasty tensor definition behind python API. 🤞
cc'ing @kevinstephano

[IvanYashchuk]

Alright using t0 = fd.define_tensor([3, 1], [1, 1]) works. This is what's currently used in primtorch's executor (concrete shape and concrete strides) and I discovered this error when trying to move the code to use the number of dimensions instead of concrete shape.

[kevinstephano]

Looks good to me, just wanted possibly 1 comment added.