setup.py fails to pass USE_ROCM to CAFFE2 build 
Labels: module: rocm, triaged

### 🐛 Describe the bug

When building pytorch main branch (978a2f2b276b51f615aa860d47fadd16a284b2f6) with: 

```
python tools/amd_build/build_amd.py
export USE_ROCM=1
export BUILD_CAFFE2=1
python setup.py develop
```

Following compiler error is produced:

```
[1/10] Building CXX object caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o
FAILED: caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o 
/usr/bin/c++ -DAT_PER_OPERATOR_HEADERS -DBUILD_ONEDNN_GRAPH -DCAFFE2_BUILD_MAIN_LIB -DCPUINFO_SUPPORTED_PLATFORM=1 -DFMT_HEADER_ONLY=1 -DFXDIV_USE_INLINE_ASSEMBLY=0 -DHAVE_MALLOC_USABLE_SIZE=1 -DHAVE_MMAP=1 -DHAVE_SHM_OPEN=1 -DHAVE_SHM_UNLINK=1 -DMINIZ_DISABLE_ZIP_READER_CRC32_CHECKS -DNNP_CONVOLUTION_ONLY=0 -DNNP_INFERENCE_ONLY=0 -DONNXIFI_ENABLE_EXT=1 -DONNX_ML=1 -DONNX_NAMESPACE=onnx_torch -DUSE_C10D_GLOO -DUSE_C10D_MPI -DUSE_DISTRIBUTED -DUSE_EXTERNAL_MZCRC -DUSE_RPC -DUSE_TENSORPIPE -D_FILE_OFFSET_BITS=64 -Dtorch_cpu_EXPORTS -Iaten/src -I../aten/src -I. -I../ -I../cmake/../third_party/benchmark/include -Icaffe2/contrib/aten -I../third_party/onnx -Ithird_party/onnx -I../third_party/foxi -Ithird_party/foxi -I../torch/csrc/api -I../torch/csrc/api/include -I../caffe2/aten/src/TH -Icaffe2/aten/src/TH -Icaffe2/aten/src -I../caffe2/../third_party -Icaffe2/../aten/src -I../torch/csrc -I../third_party/miniz-2.1.0 -I../third_party/kineto/libkineto/include -I../third_party/kineto/libkineto/src -I../aten/src/ATen/.. -I../caffe2/core/nomnigraph/include -I../third_party/FXdiv/include -I../c10/.. -I../third_party/pthreadpool/include -I../third_party/cpuinfo/include -I../third_party/QNNPACK/include -I../aten/src/ATen/native/quantized/cpu/qnnpack/include -I../aten/src/ATen/native/quantized/cpu/qnnpack/src -I../third_party/cpuinfo/deps/clog/include -I../third_party/NNPACK/include -I../third_party/fbgemm/include -I../third_party/fbgemm -I../third_party/fbgemm/third_party/asmjit/src -I../third_party/ittapi/src/ittnotify -I../third_party/FP16/include -I../third_party/tensorpipe -Ithird_party/tensorpipe -I../third_party/tensorpipe/third_party/libnop/include -Ithird_party/ideep/mkl-dnn/third_party/oneDNN/include -I../third_party/ideep/mkl-dnn/third_party/oneDNN/src/../include -I../third_party/fmt/include -I../third_party/flatbuffers/include -isystem third_party/gloo -isystem ../cmake/../third_party/gloo -isystem ../cmake/../third_party/googletest/googlemock/include -isystem ../cmake/../third_party/googletest/googletest/include -isystem ../third_party/protobuf/src -isystem ../third_party/gemmlowp -isystem ../third_party/neon2sse -isystem ../third_party/XNNPACK/include -isystem ../third_party/ittapi/include -isystem ../cmake/../third_party/eigen -isystem /opt/cray/pe/mpich/8.1.21/ofi/cray/10.0/include -isystem ../third_party/ideep/mkl-dnn/third_party/oneDNN/include -isystem ../third_party/ideep/include -isystem ../third_party/ideep/mkl-dnn/include -isystem include -D_GLIBCXX_USE_CXX11_ABI=1 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=old-style-cast -Wno-invalid-partial-specialization -Wno-unused-private-field -Wno-aligned-allocation-unavailable -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow -DHAVE_AVX2_CPU_DEFINITION -O3 -DNDEBUG -DNDEBUG -fPIC -DCAFFE2_USE_GLOO -DTH_HAVE_THREAD -Wall -Wextra -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-missing-field-initializers -Wno-unknown-pragmas -Wno-type-limits -Wno-array-bounds -Wno-strict-overflow -Wno-strict-aliasing -Wno-missing-braces -Wno-maybe-uninitialized -fvisibility=hidden -O2 -Wno-sign-compare -pthread -DASMJIT_STATIC -fopenmp -std=gnu++17 -MD -MT caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o -MF caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o.d -o caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o -c ../torch/csrc/jit/ir/ir.cpp
../torch/csrc/jit/ir/ir.cpp: In member function ‘bool torch::jit::Node::hasSideEffects() const’:
../torch/csrc/jit/ir/ir.cpp:1191:10: error: ‘hip’ has not been declared
     case hip::set_stream:
          ^~~
../torch/csrc/jit/ir/ir.cpp:1192:10: error: ‘hip’ has not been declared
     case hip::_set_device:
          ^~~
../torch/csrc/jit/ir/ir.cpp:1193:10: error: ‘hip’ has not been declared
     case hip::_current_device:
          ^~~
../torch/csrc/jit/ir/ir.cpp:1194:10: error: ‘hip’ has not been declared
     case hip::synchronize:
          ^~~
At global scope:
cc1plus: warning: unrecognized command line option ‘-Wno-aligned-allocation-unavailable’
cc1plus: warning: unrecognized command line option ‘-Wno-unused-private-field’
cc1plus: warning: unrecognized command line option ‘-Wno-invalid-partial-specialization’
ninja: build stopped: subcommand failed.
```

And after lines in torch/csrc/jit/ir/ir.cpp:1191-1194 from:

```
#if !defined(USE_ROCM)
    case hip::set_stream:
    case hip::_set_device:
    case hip::_current_device:
    case hip::synchronize:
#endif
```

to:

```
#if !defined(USE_ROCM)
    case c10::cuda::set_stream:
    case c10::cuda::_set_device:
    case c10::cuda::_current_device:
    case c10::cuda::synchronize:
#endif
```

it manages to build.

It seems that when building `caffe2/CMakeFiles/torch_cpu.dir/__/torch/csrc/jit/ir/ir.cpp.o` cmake does not pass `USE_ROCM` to preprocessor.

 


### Versions

Collecting environment information...
PyTorch version: 2.1.0a0+git978a2f2
Is debug build: False
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: 0.0.0

OS: Red Hat Enterprise Linux 8.6 (Ootpa) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-10)
Clang version: 15.0.0  (324a8e7de6a18594c06a0ee5d8c0eda2109c6ac6)
CMake version: version 3.20.2
Libc version: glibc-2.28

Python version: 3.9.13 (main, Aug  2 2022, 03:25:18)  [GCC 9.3.0 20200312 (Cray Inc.)] (64-bit runtime)
Python platform: Linux-4.18.0-372.9.1.el8.x86_64-x86_64-with-glibc2.28
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              128
On-line CPU(s) list: 0-127
Thread(s) per core:  2
Core(s) per socket:  32
Socket(s):           2
NUMA node(s):        4
Vendor ID:           AuthenticAMD
CPU family:          23
Model:               49
Model name:          AMD EPYC 7542 32-Core Processor
Stepping:            0
CPU MHz:             2900.000
CPU max MHz:         2900.0000
CPU min MHz:         1500.0000
BogoMIPS:            5789.48
Virtualization:      AMD-V
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            16384K
NUMA node0 CPU(s):   0-15,64-79
NUMA node1 CPU(s):   16-31,80-95
NUMA node2 CPU(s):   32-47,96-111
NUMA node3 CPU(s):   48-63,112-127
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.1.0a0+git978a2f2
[conda] Could not collect


cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo