Where is CPU_CAPABILITY_AVX2 defined? It is used in vec256* headers. (e.g https://github.com/pytorch/pytorch/blob/d1352192e21001521035dacc0eda574b58224198/aten/src/ATen/cpu/vec256/vec256.h#L53). Is it defined in system headers? I am trying understand if my ops are making use of AVX instructions?

These flags are generated when compiling the source files within aten/src/ATen/native/cpu.
The code is here,
cpu folder will be compiled with all these flags, pytorch automatically detect whether your CPU support those instruction sets.

I think if you want to see whether your op used the AVX, you don’t need to check these flags since as long as your cpu support AVX/2, pytorch will compile with CPU_CAPABILITY_AVX2/2 for sure. What you can do is debug your op using eg. dbx, gdb, setup some breakpoint and print out the
Assembly code, see whether your op really call the AVX/2 instructions.