Add variance_mean composite function using Welford

[IvanYashchuk]

Fixes #1801.

Currently complex, half, bfloat16 types give compilation errors with Welford, so these types raise an explicit error for now.

I'm not sure how to include this function in parser.cpp to be tied with aten::var_mean.correction correctly, so integration with TorchScript is omitted from this PR.

[jjsjann123]

LGTM. Don't worry about TS support at this moment. I can patch a PR for that after this one is merged.

[jjsjann123]

Should we consider just promote math to float for reduced precision then? that seems to be an universal way used by all others computing variance.

[IvanYashchuk]

I handle type promotion on the Python side.

[jjsjann123]

Doesn't sound very right to me. cc'ing @zasdfgbnm

[zasdfgbnm]

This is definitely not the correct behavior, but I am not surprised at all😜. I got moved to more important tasks before finishing complex support, so currently, complex support in nvfuser is very broken. Please consider complex as not supported when doing other works.

[jjsjann123]

Sorry I forgot to stamp earlier.

[jjsjann123]

I would hope that we can have variance to also use welford, so we can unify the interface for variance and variance_mean.
But WelfordOp appears to be a single op in codegen, does this means the generated kernel would still write mean to a register somewhere, even though they are not used at all? That we probably not want.

cc'ing @naoyam @shmsong on this.

[IvanYashchuk]

We can of course try using

TensorView* variance(
    TensorView* x,
    const std::vector<int>& dims,
    int64_t correction,
    bool keepdim) {
  auto var_mean = variance_mean(x, dims, correction, keepdim);
  return var_mean.var;
}

If it positively impacts the performance. That's how ATen implements this function and it would justify var being a prim in PrimTorch.

[IvanYashchuk]

What's the procedure for moving this further? Who runs the tests and clicks "merge"?

[jjsjann123]

What's the procedure for moving this further? Who runs the tests and clicks "merge"?

We don't have CI, so you'll need to run tests locally.

./build/bin/test_jit --gtest_filter="*NVFuser*"
python test/test_jit_cuda_fuser.py

^^^ These two are the usual ones that we run before merging each PR. You've already got my stamp, once test is verified to pass, you can click the squash and merge.

[IvanYashchuk]

I see two test failures that seem to be unrelated:

[ RUN      ] NVFuserTest.FusionBiasGeluBwd_CUDA
unknown file: Failure
C++ exception with description "aten_output_tensor.allclose( fusion_output_tensor.to(aten_output_tensor.dtype()), tolerance_values.second, tolerance_values.first) INTERNAL ASSERT FAILED at "/home/iyashchuk/dev/pytorch/jit-devel/torch/csrc/jit/codegen/cuda/test/test_gpu_validator.h":416, please report a bug to PyTorch.

Validation error in output 0 on line 11691 in file /home/iyashchuk/dev/pytorch/jit-devel/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp.
  Detected abs error of: 1.66893e-06
    absolute tolerance was set to 1.51992e-06
    and relative tolerance set to 2.23704e-06
Exception raised from testValidate at /home/iyashchuk/dev/pytorch/jit-devel/torch/csrc/jit/codegen/cuda/test/test_gpu_validator.h:416 (most recent call first):

[ RUN      ] NVFuserTest.FusionScheduleTransposeMultipleInputOutput_CUDA
unknown file: Failure
C++ exception with description "CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 11.77 GiB total capacity; 8.25 GiB already allocated; 800.44 MiB free; 9.77 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Exception raised from malloc at /home/iyashchuk/dev/pytorch/jit-devel/c10/cuda/CUDACachingAllocator.cpp:639 (most recent call first):