First commit

pkgs/triton/debugger/memory_map.py

@@ -0,0 +1,100 @@
+import dataclasses
+
+from triton.debugger import torch_wrapper
+
+torch = torch_wrapper.torch
+
+
+@dataclasses.dataclass
+class RegisteredStorage:
+    storage: torch.Storage
+    dtype: torch.dtype
+    size: int
+    ptr: int
+
+    @property
+    def end_ptr(self) -> int:
+        return self.ptr + self.size
+
+    @property
+    def access_tensor(self) -> torch.Tensor:
+        return torch.tensor(self.storage, dtype=self.dtype, device=self.storage.device)
+
+    def ensure_immutable(self):
+        assert self.storage.data_ptr() == self.ptr and self.storage.size() == self.size
+
+
+class MemoryMap:
+    storages: [RegisteredStorage]
+
+    def __init__(self):
+        self.storages = []
+
+    def _get_registered_storage(self, pointer: torch.Tensor):
+        max_pointer = torch.max(pointer).item()
+        min_pointer = torch.min(pointer).item()
+
+        registered_storage = next(
+            filter(
+                lambda registered: min_pointer >= registered.ptr and max_pointer < registered.end_ptr, self.storages
+            ),
+            None,
+        )
+        if registered_storage is None:
+            raise Exception("Storage not found or pointers spanning multiple tensors")
+        registered_storage.ensure_immutable()
+        return registered_storage
+
+    def add_tensor(self, t: torch.Tensor):
+        storage = t.untyped_storage()
+        self.storages.append(RegisteredStorage(storage, t.dtype, storage.size(), storage.data_ptr()))
+        return t.data_ptr()
+
+    def load(
+        self,
+        pointer: torch.Tensor,
+        mask: torch.Tensor = None,
+        other=0.0,
+    ):
+        assert pointer.is_cuda
+        assert 0 < pointer.dim() < 3
+        assert pointer.dtype == torch.int64
+
+        if mask is None:
+            mask = torch.ones_like(pointer).bool()
+        assert mask.is_cuda
+        assert 0 < mask.dim() < 3
+        assert mask.dtype == torch.bool
+        mask = mask.expand(pointer.size())
+
+        if torch.all(~mask):
+            # Todo: The type is wrong here, we can't determine the correct type
+            return torch.full_like(pointer, fill_value=other, dtype=torch.float16, device="cuda")
+
+        registered_storage = self._get_registered_storage(pointer[mask])
+        access_tensor = registered_storage.access_tensor
+
+        index_tensor = pointer - registered_storage.ptr
+
+        block = torch.full_like(pointer, fill_value=other, dtype=access_tensor.dtype, device="cuda")
+        block[mask] = access_tensor[index_tensor[mask]]
+        return block
+
+    def store(self, pointer: torch.Tensor, value: torch.Tensor, mask=None):
+        assert 0 < pointer.dim() < 3
+        assert pointer.dtype == torch.int64
+
+        if mask is None:
+            mask = torch.ones_like(pointer).bool()
+        assert 0 < mask.dim() < 3
+        assert mask.dtype == torch.bool
+        mask = mask.expand(pointer.size())
+
+        if torch.all(~mask):
+            return
+
+        registered_storage = self._get_registered_storage(pointer[mask])
+        access_tensor = registered_storage.access_tensor
+
+        index_tensor = pointer - registered_storage.ptr
+        access_tensor[index_tensor[mask]] = value[mask].to(access_tensor.dtype)