Nvfuser code bump 12 5 (pytorch#69964)

Summary:
Pull Request resolved: pytorch#69964

Things added in this PR that requires review:
1. cuLaunchCooperativeKernel driver API added
aten/src/ATen/cuda/detail/LazyNVRTC.cpp
aten/src/ATen/cuda/nvrtc_stub/ATenNVRTC.h

nvfuser code update:
1. perf turning on codegen scheduler that improves performance.
2. permutation support has been extended beyond contiguous/channels-last. (The improvements could be observed on PW benchmark)

Things reverted from local changes:
1. aten::gelu with approximation
2. local changes that is upstreamed in PR pytorch#68804

Pull Request resolved: pytorch#69428

Reviewed By: ngimel

Differential Revision: D33073817

Pulled By: wconstab

fbshipit-source-id: e77d32e81d037d7370822b040456fd4c3bd68edb

Author: jjsjann123

aten/src/ATen/cuda/detail/LazyNVRTC.cpp

@@ -188,6 +188,36 @@ CUresult CUDAAPI cuLaunchKernel(CUfunction f,
             sharedMemBytes, hStream, kernelParams, extra);
 }
 
+// Irregularly shaped functions
+CUresult CUDAAPI cuLaunchCooperativeKernel(
+    CUfunction f,
+    unsigned int gridDimX,
+    unsigned int gridDimY,
+    unsigned int gridDimZ,
+    unsigned int blockDimX,
+    unsigned int blockDimY,
+    unsigned int blockDimZ,
+    unsigned int sharedMemBytes,
+    CUstream hStream,
+    void** kernelParams) {
+  auto fn = reinterpret_cast<decltype(&cuLaunchCooperativeKernel)>(
+      getCUDALibrary().sym(__func__));
+  if (!fn)
+    throw std::runtime_error("Can't get cuLaunchCooperativeKernel");
+  lazyNVRTC.cuLaunchCooperativeKernel = fn;
+  return fn(
+      f,
+      gridDimX,
+      gridDimY,
+      gridDimZ,
+      blockDimX,
+      blockDimY,
+      blockDimZ,
+      sharedMemBytes,
+      hStream,
+      kernelParams);
+}
+
 CUresult CUDAAPI cuModuleLoadDataEx(CUmodule *module,
                                     const void *image,
                                     unsigned int numOptions,

aten/src/ATen/cuda/nvrtc_stub/ATenNVRTC.h

@@ -49,6 +49,7 @@ namespace at { namespace cuda {
   _(cuOccupancyMaxActiveBlocksPerMultiprocessor) \
   _(cuGetErrorString)                            \
   _(cuLaunchKernel)                              \
+  _(cuLaunchCooperativeKernel)                   \
   _(cuCtxGetCurrent)                             \
   _(cuModuleUnload)                              \
   _(cuDevicePrimaryCtxGetState)                  \

caffe2/CMakeLists.txt

@@ -962,8 +962,11 @@ if(USE_CUDA OR USE_ROCM)
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/broadcast.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/fp16_support.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/bf16_support.cu
+    ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/grid_broadcast.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/grid_reduction.cu
+    ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/grid_sync.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/helpers.cu
+    ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/index_utils.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/random_numbers.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/tensor.cu
     ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/runtime/welford.cu