Continue broadcast support

[csarofeen]

Rebased #56

• Rework computeAt so we do less forward propagation, as that can move our computeAt position and broadcast dims in computeAt will be lost and won't be recoverable.
• Add braodcast type to iterDomain
• Add broadcast IR node
• Add broadcast op to Arith
• Work broadcast through transform ops and transform replay.
• Work broadcast through rfactor
• Figure out if we can support merging broadcast dims with non-broadcast dims
• Add broadcast op to lowering logic: lower2gpu, lower loops, and unroll_loops
• Add broadcast logic to index compute and predicate compute
• Test pointwise and broadcast operations
• Test broadcast operations before and after reduction
• Test softmax
• Test Normalization
• Test GEMM

[jjsjann123]

Nitpcking: any reason for a nested block here?

[csarofeen]

It's for my own safety as I often reuse index names and I have size_t it_c = 0, it_p = 0; surrounding the while loop. I could probably get rid of them, but I don't think it harms much. thoughts?

[jjsjann123]

Perfectly fine to me.

I was suggesting similar things as well for running_producer below.

[tlemo]

constexpr?

[naoyam]

Probably not. We should want to verify the tests work with dynamically sized tensors.

[csarofeen]

I don't think constexpr buys us anything here.

[tlemo]

These are compile time constants, right? If that's the case constexpr is the most explict language construct to express it (it would also prevent accidental value updates)

[tlemo]

can you expand the comment to list what has to be fixed / what's not quite right ?

[naoyam]

I believe this has the issue that only some threads have valid results after reductions. I'm working on broadcasting within a block, which may solve some of the problems.

[csarofeen]

added context, thanks.

[tlemo]

c_inds can be const

[csarofeen]

thanks, added.

[tlemo]

cleanup the iteration style since it's touched?

[csarofeen]

Thanks, changed all instances in this file as well.

[tlemo]

I like this pattern :)

[tlemo]

these can all be const, right?

[csarofeen]

yes, thanks.

[tlemo]

nit - no need for double parantheses

[csarofeen]

thanks, too much copy-pasting

[tlemo]

const?

[csarofeen]

oh cool, I thought it wouldn't work in leaf_ids_.find(...) but does fine. Thanks.

[tlemo]

why the leading underscore?

[csarofeen]

just a habit, I've been doing this on constructors when I copy it directly to a member named with a trailing underscore, is it bothersome?

[tlemo]

naming patterns are a powerful tool/hint, best when used sparingly (otherwise you end up with https://en.wikipedia.org/wiki/Hungarian_notation)

If we have a good reason to adopt a _ prefix we could, but then we'd need to use it consistently.

[kevinstephano]

I personally liked this for Constructor parameters. I thought this was being used consistently.

[csarofeen]

I have been trying to be consistent about this and really like it as clang tidy tends to force us to use std::move which helps me remember not to use things with a prefix except in initialization. We can discuss further but thinking to leave it for this PR.

[tlemo]

reformat this comment to make the individual notes more explicit? (just split it into multiple paragraphs should do the trick)

[csarofeen]

Done, lets see if it gets through clang_format

[jjsjann123]

LGTM

[naoyam]

I'm having a problem to make broadcast work correctly in a fusion. Here's the fusion:

{
    torch::jit::fuser::cuda::CudaKernel prog;
    Fusion& fusion = *prog.fusion_;
    FusionGuard fg(&fusion);

    // Set up your input tensor views
    TensorView* tv0 = makeDummyTensor(1);
    TensorView* tv1 = makeDummyTensor(2);
    fusion.addInput(tv0);
    fusion.addInput(tv1);

    auto tv2 = broadcast(tv0, {false, true});
    auto tv3 = add(tv2, tv1);
    fusion.addOutput(tv3);

    tv2->computeAt(tv3, 1);

    GPULower gpulw(&fusion);
    gpulw.printKernel(std::cout);
  }

It's available here too: https://github.com/naoyam/pytorch/blob/bcast_bug/test/cpp/jit/test_gpu.cpp#L2677

And here's the generated code using the broadcast_rebased branch:

__global__ void CUDAGeneratedKernel(Tensor<float, 1> T0, Tensor<float, 2> T1, Tensor<float, 2> T3){
  for(size_t i12 = 0; i12 < T3.size[0]; ++i12 ) {
    float T2[1];
    T2[ 0 ]
       = T0[ ( i12 * T0.stride[0] ) ];
    for(size_t i14 = 0; i14 < T3.size[1]; ++i14 ) {
      T3[ ( i12 * T3.stride[0] ) + ( i14 * T3.stride[1] ) ]
         = T2[ i14 ]
         + T1[ ( i12 * T1.stride[0] ) + ( i14 * T1.stride[1] ) ];
    }
  }
}

The problem is that T2 is of size 1, but it's accessed with subscript i14. It seems that allocating T2 as a size-1 array makes sense as its 2nd dimension is broadcast, but accessing it with the index of that dimension seems wrong.

Is this because I have something wrong in the fusion?

[csarofeen]

Issue Naoya found fixed in #98