Hack in `unshardedSizes` can cause inconsistent values for the same extent `Val*`

[zasdfgbnm]

As demonstrated in #5708, The same extent Val* could have different multiplier. This is a blocker for my PR #5082, because this could lead to PrecomputedValues binding the same Val* with different value for the same instance of FusionExecutorCache.

Any suggestion? cc: @wujingyue

[wujingyue]

I've yet to understand why

Fuser/csrc/multidevice/execution_utils.cpp

Line 73 in 878c3e2

// TODO(#5525): hack for MultiDeviceExecutor. MultiDeviceExecutor looks

can create problems.

I failed to reproduce this reliably. I suspect it has something to do with the use of static: https://github.com/NVIDIA/Fuser/pull/5708/files#diff-140da8af7df6a9ae6ddfbc972a5db1bfdd0efba4b6412902f285b32f3d131347R72. Some Val*s from a previous test might have stayed in the hashtable even though they've been deleted.

[wujingyue]

Repro

$ git fetch origin eval-segment
$ git checkout eval-segment
$ mpirun -np 1 pytest tests/python/multidevice/test_overlap.py -k test_row_parallel_linear_forward --only-mpi -vs
...
E       RuntimeError:  INTERNAL ASSERT FAILED at /opt/pytorch/nvfuser/csrc/evaluator_common.cpp:371, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. 
E       Expected isSame(values_[it.first], it.second) . Precomputed values failed to validate.
E       Something unexpected changed between the compilation and execution.
E       6 != 18
E       Exception raised from validate at /opt/pytorch/nvfuser/csrc/evaluator_common.cpp:371 (most recent call first):
...

Root cause

LinearOp::evaluate doesn't respect the stream parallel type in allocation. The LinearOp's output allocation ought to be stream-parallelized but LinearOp::evaluate still produces a tensor of shape [6, 2, 1]. Below is the preseg fusion IR. The problematic LinearOp is the local LinearOp before allreduce.

This problem doesn't show up today because, without #5082, we never use ExpressionEvaluator to evaluate a stream-parallel LinearOp.

Fusion IR after pre-segmenter optimization passes:
Inputs:
  T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0})
  T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})
Outputs:
  T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})

%kernel {
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0})
   = linear(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0}),
            T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})  )
T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})
   = reduction( T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0}), op = add, initial value = __bfloat(0), allreduce = false )

TransformPrinter : 
T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0})
 logical domain : (iS0{i0}, iS1{8})
 allocation domain : (iS0{i0}, ideviceIdx.x9{1}, iS10{8})
 contiguity: t t t
  Outer split: iS0{i0} by factor 3 -> istreamIdx7{3}, iS8{( ceilDiv(i0, 3) )}
  Outer split: iS1{8} by factor 1 -> ideviceIdx.x9{1}, iS10{8}
 loop domain : (istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8})
T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})
 logical domain : (iS2{2}, iS3{8})
 allocation domain : (iS2{2}, ideviceIdx.x11{1}, iS12{8})
 contiguity: t t t
  Outer split: iS3{8} by factor 1 -> ideviceIdx.x11{1}, iS12{8}
 loop domain : (ideviceIdx.x11{1}, iS2{2}, iS12{8})
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0})
 root domain : (iS17{i0}, iS18{2}, rS19{8}rf)
  Outer split: rS19{8}rf by factor 1 -> ideviceIdx.x22{1}rf, rS23{8}rf
 logical domain : (iS17{i0}, iS18{2}, ideviceIdx.x22{1}rf, rS23{8}rf)
 allocation domain : (istreamIdx20{3}, iS21{( ceilDiv(i0, 3) )}, iS18{2}, ideviceIdx.x22{1}rf, rS23{8}rf)
 contiguity: t t t t n
  Outer split: iS17{i0} by factor 3 -> istreamIdx20{3}, iS21{( ceilDiv(i0, 3) )}
  Outer split: rS19{8}rf by factor 1 -> ideviceIdx.x22{1}rf, rS23{8}rf
 loop domain : (istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf)
T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})
 logical domain : (iS24{i0}, iS25{2}, rdeviceIdx.x26{1})
 allocation domain : (iS24{i0}, iS25{2}, rdeviceIdx.x26{1})
 contiguity: t t n
  Outer split: iS24{i0} by factor 3 -> istreamIdx27{3}, iS28{( ceilDiv(i0, 3) )}
 loop domain : (istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2})
} // %kernel