[Frontend] Assert that TensorDescriptor fields have matching rank (#6911)

Currently a rank mismatch gives an error in the driver when creating the
tma descriptor, which makes it hard to track back to the original
argument. This just makes the stack trace more useful.

Also add some basic tensor descriptor stuff to the docs.

Author: peterbell10

docs/python-api/triton.language.rst

@@ -12,6 +12,7 @@ Programming Model
     :nosignatures:
 
     tensor
+    tensor_descriptor
     program_id
     num_programs
 
@@ -71,6 +72,9 @@ Memory/Pointer Ops
 
     load
     store
+    make_tensor_descriptor
+    load_tensor_descriptor
+    store_tensor_descriptor
     make_block_ptr
     advance
 

python/triton/tools/tensor_descriptor.py

@@ -9,6 +9,11 @@ class TensorDescriptor:
     strides: List[int]
     block_shape: List[int]
 
+    def __post_init__(self):
+        rank = len(self.shape)
+        assert len(self.strides) == rank, f"rank mismatch: {self}"
+        assert len(self.block_shape) == rank, f"rank mismatch: {self}"
+
     @staticmethod
     def from_tensor(tensor: Any, block_shape: List[int]):
         return TensorDescriptor(