[GSan] Implement shadow memory allocator

[peterbell10]

Commits in this PR

1. [GSan] Implement shadow memory allocator

   This implements an allocator that hooks into PyTorch's memory allocation
   API to map tensors into a GSan-managed virtual address space. We also create
   a corresponding shadow memory region that is mapped into the lower half of the
   reserved address space.

   Usage is like:

   from triton.experimental import gsan
   allocator = gsan.get_allocator()
   pool = torch.cuda.MemPool(allocator.allocator())
   with torch.cuda.use_mem_pool(pool):
       t = torch.empty(4096, dtype=torch.uint8, device="cuda")

2. Misc cleanup/fixes

3. Sync stream before dealloc

4. More misc changes

PR chain

1. 👉 [GSan] Implement shadow memory allocator #9478 👈 YOU ARE HERE
2. [GSan] Instrument tl.{load,store} #9568
3. [GSan] Partially support TMA & cp.async ops #9699
4. [GSan] Add symmetric memory API #9493
5. [GSan] Support atomics #9700

[pawelszczerbuk]

What are the implications if we's bump it to 32?

[peterbell10]

It's more or less an arbitrary constant, I bump it to 32 in a later PR. Basically we reserve a memory space equal to kMaxGPUs * kPerDeviceStateStride and this is where all the non-shadow memory state lives. Because it's only virtual memory this is cheap.

[pawelszczerbuk]

Just curious, how is getGlobalState used, and why do we need the masking here?

[peterbell10]

The global state are actually all constants, and I have them in a common place so we don't have to store them duplicated in each thread state or keep them around in registers

triton/python/triton/experimental/gsan/src/GSan.h

Lines 50 to 63 in 661198b

	struct GlobalState {
	// Base address of gsan managed memory
	uintptr_t reserveBase;
	uintptr_t globalsBase;
	uint32_t rngSeed;
	thread_id_t numSms;
	thread_id_t numDevices;
	// numThreads = numSms * numDevices
	thread_id_t numThreads;
	uint16_t clockBufferSize;
	};

The masking is a bit of a trick with the memory layout. Each GPU holds

GlobalState | ThreadState0 | ThreadState1 | ThreadState2 | ...

with the global state aligned to kPerDeviceStateStride. This means you can have a single pointer to the thread state, and by masking down to an aligned pointer you get a pointer to the globals. This saves either carrying 2 pointers around, or doing an extra indirection.

[pawelszczerbuk]

Left some questions just for my education. Good stuff!