Fix tv parallelization

[naoyam]

Fixes #757

[naoyam]

It turned out this is actually not the right thing to do. The parallel map contains mappings created with the the forward-bcast-mismatch enabled, so for example, it would create a mapping between I2*I1 of t4 and I1 of t2:

t0 = makeSymbolicTensor(1); // t0: [I1]
t1 = makeSymbolicTensor(2); // t1: [I2, I1]
t2 = t0 + 1; // t2: [I1]
t3 = broadcast(t2, {true, false}); // t3: [B1, I1]
t4 = t1 + t3; // t4: [I2, I1]
t4->merge(0, 1); // t4: [I2*I1]
t2->computeAt(t4, -1);
t4->axis(0)->parallelize(TIDx);

This PR would do t2->axis(0)->parallelize(TIDx) as the axis is mapped with the t4 axis. However, this is a problem since that would mean I2*I1 == blockDim.x and I1 == blockDim.x.

[naoyam]

Ended up adding ParallelTypeBitmap to kir::BroadcastOp. See 9f9f902. This is the only way I can think of for properly finding the parallelism of kir::BroadcastOp.

More broadly, I think a bigger problem is that we don't have straightforward way to know parallelism of kir:TensorView and kir::IterDomain. The truth only lies in the ComputeAt parallel map, which only exists at the lowering time. Referring to kir::IterDomain::parallelType() is not robust as it may not be the same as the real parallel type. We use that for codegen of kir::ReductionOp, where it is fortunately safe as there should be no reduction in the CA axes.

[csarofeen]

I2*I1 == blockDim.x and I1 == blockDim.x actually seems fine to me in this case.

blockDim.x should be set to the maximum size. I think the issue is that we substitute values in the IR for blockDim.x where we shouldn't. I think the original issue stems from the fact that a tensor size gets substituted for this value, and thread bindings can be ambiguous in the presence of broadcst. I'm uncertain we can resolve all cases.

[csarofeen]

I thought that was how we started down this path: #622

[csarofeen]

Clearing approval as looks like we should discuss/think about this more. Broadcast parallelism is definitely a complex topioc.

[naoyam]

The substitution is done here:

https://github.com/csarofeen/pytorch/blob/20_12_3_devel/torch/csrc/jit/codegen/cuda/kernel_ir.cpp#L100-L109

If we allow I2*I1 == blockDim.x and I1 == blockDim.x, the substitution is no longer valid.

I see two options:

1. Parallelize all kir::IterDomain when inferred from the computeAt map and stop the substitution. (ab83e3e)
2. Annotate kir::BroadcastOp (9f9f902)

I'd say the second option is more conservative. It only solves this particular problem with BroadcastOp. The first option may be more preferable as it conveys the parallelism information to KIR. However, there may be side effects as the substitution problem.

The first option seems better to me from long-term perspectives.

[csarofeen]

Can you try 1 and run the reduction benchmark suite before and after to see if you find any serious perf regressions. We use a lot of tensor size information in our kernels anyways, I think removing the substitution shouldn't be significant, though I could definitely be wrong.

[naoyam]

@csarofeen Done. I changed kir::IterDomain::extent() to just return extent_, which is just the same as rawExtent(). All tests are working fine.