Repository: pytorch
Issue: [NNC] loop vectorization fails, `Ramp` and `Broadcast` undefined
Labels: triaged,NNC

### 🐛 Describe the bug

Loop vectorization introduces `Ramp` and `Broadcast`, which are both undefined.

As an aside, is this supposed to be the case that the loops are still explicit in the generated code or could that be the reason why the code (without vectorisation) is two orders of magnitude slower than plain Pytorch? 

```python
import torch
import torch.utils.benchmark as benchmark
import torch._C._te as te

def construct_gather(n: int, backend, dtype=torch.float32):
    A = te.BufHandle("A", [n, n], dtype)
    INDEX = te.BufHandle("INDEX", [n, n], torch.long)
    B = te.BufHandle("B", [n, n], dtype)
    i = te.VarHandle("i", torch.int)
    j = te.VarHandle("j", torch.int)

    store = te.Store.make(A, [i, j], B.load([INDEX.load([j + i*n]), j])) # flatten index manually for now
    for_j = te.For.make(j, te.ExprHandle.int(0), te.ExprHandle.int(n), store)
    for_i = te.For.make(i, te.ExprHandle.int(0), te.ExprHandle.int(n), for_j)

    loopnest = te.LoopNest(te.Block([for_i]), [A, B, INDEX])

    loopnest.vectorize_inner_loops() #this line causes the issue

    loopnest.prepare_for_codegen()
    stmt = te.simplify(loopnest.root_stmt())

    return te.construct_codegen(backend, stmt, [A, B, INDEX])

if __name__ == "__main__":
    n = 2000
    torch.manual_seed(42)
    d = torch.randn((n, n), device=0)
    index = (torch.rand((n, n), device=0) * (n-2)).long()
    to = torch.zeros((n, n), device=0)

    print("Pytorch GPU")
    t1 = benchmark.Timer("torch.gather(d, 0, index, out=to)", globals={"d": d, "index": index, "to": to})
    print(t1.timeit(2))

    print("NNC/CUDA GPU")
    gather_cuda = construct_gather(n, "cuda")
    t0 = benchmark.Timer("gather.call([to, d, index])", globals={"d": d, "index": index, "to": to, "gather":gather_cuda})
    print(t0.timeit(2))
```
Gives the following error message:
```
RuntimeError                              Traceback (most recent call last)
[<ipython-input-12-d7d2ee163d59>](https://localhost:8080/#) in <module>
     35 
     36     print("NNC/CUDA GPU")
---> 37     gather_cuda = construct_gather(n, "cuda")
     38     t0 = benchmark.Timer("gather.call([to, d, index])", globals={"d": d, "index": index, "to": to, "gather":gather_cuda})
     39     print(t0.timeit(2))

[<ipython-input-12-d7d2ee163d59>](https://localhost:8080/#) in construct_gather(n, backend, dtype)
     21     stmt = te.simplify(loopnest.root_stmt())
     22 
---> 23     return te.construct_codegen(backend, stmt, [A, B, INDEX])
     24 
     25 if __name__ == "__main__":

RuntimeError: default_program(22): error: identifier "Ramp" is undefined

default_program(23): error: identifier "Broadcast" is undefined

2 errors detected in the compilation of "default_program".

nvrtc compilation failed: 

#define NAN __int_as_float(0x7fffffff)
#define POS_INFINITY __int_as_float(0x7f800000)
#define NEG_INFINITY __int_as_float(0xff800000)


template<typename T>
__device__ T maximum(T a, T b) {
  return isnan(a) ? a : (a > b ? a : b);
}

template<typename T>
__device__ T minimum(T a, T b) {
  return isnan(a) ? a : (a < b ? a : b);
}

extern "C" __global__
void func(float* A, float* B, long long* INDEX) {
{
for (int i = 0; i < 2000; i++) {
for (int j_outer = 0; j_outer < 250; j_outer++) {
      long long v = __ldg(INDEX + Ramp(8 * (j_outer + 250 * i), 1, 8));
      float v_1 = __ldg(B + ((Broadcast(0ll, 8)) + (Broadcast(1ll, 8)) * (long long)(Ramp(8 * j_outer, 1, 8))) + (Broadcast(2000ll, 8)) * v);
      A[Ramp(8 * (j_outer + 250 * i), 1, 8)] = v_1;
    }  }}
}
```

### Versions

I ran this on Google colab:

```
Collecting environment information...
PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.6
Libc version: glibc-2.26

Python version: 3.7.13 (default, Apr 24 2022, 01:04:09)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 460.32.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.13.1
[pip3] torchvision==0.13.1+cu113
[conda] Could not collect
```