progress

README.md

@@ -11,6 +11,7 @@ A new Python DSL frontend for LuisaCompute. Will be integrated into LuisaCompute
     - [Functions](#functions)
     - [User-defined Structs](#user-defined-structs)
     - [Control Flow](#control-flow)
+    - [Define DSL Operation for Non-DSL Types](#define-dsl-operation-for-non-dsl-types)
 - [Advanced Usage](#advanced-syntax)
     - [Generics](#generics)
     - [Metaprogramming](#metaprogramming)
@@ -104,6 +105,35 @@ class Sphere:
     radius: lc.float
 ```
 
+### Define DSL Operation for Non-DSL Types
+Sometimes we want to use a non-DSL type in our DSL code. Such type could be imported from a third-party library or a built-in Python type. As long as we know the object layout, we can define the DSL operation for it by first defining a proxy struct that mirrors the object layout, and then define the operation for the proxy struct.
+
+```python
+# Assume we have a third-party library that defines a Vec3 class
+class Vec3:
+    def __init__(self, x, y, z):
+        self.x = x
+        self.y = y
+        self.z = z
+
+@lc.struct
+class Vec3Proxy:
+    x: lc.float
+    y: lc.float
+    z: lc.float
+
+    # write DSL operations here
+
+lc.register_dsl_type_alias(Vec3, Vec3Proxy)
+
+@lc.func
+def use_vec3(v: Vec3): # Vec3 is now treated as Vec3Proxy internally
+    v.x = 1.0
+    v.y = 2.0
+    v.z = 3.0
+
+```
+
 ### Generics
 ```python
 T = TypeVar('T', bound=Any)

luisa_lang/hir.py

@@ -394,6 +394,8 @@ def member(self, field: Any) -> Optional['Type']:
     def __len__(self) -> int:
         return self.count
 
+class MatrixType(Type):
+    pass
 
 class ArrayType(Type):
     element: Type

luisa_lang/lang_builtins.py

@@ -208,6 +208,14 @@ def _instantiate_array_type(args: List[Any]) -> hir.Type:
     _instantiate_array_type
 )
 class Array(Generic[T, N]):
+    """
+    An array is a fixed-size collection of elements of type T. N need to be a Literal type.
+    Example:
+    ```python
+    arr = Array[int, Literal[10]]()
+    ```
+    """
+
     def __init__(self) -> None:
         self = intrinsic("init.array", Array[T, N])
 
@@ -236,6 +244,10 @@ def __len__(self) -> u64:
 
 @opaque("Buffer")
 class Buffer(Generic[T]):
+    """
+    A buffer is a contiguuos memory of elements of type T.
+    """
+
     def __getitem__(self, index: int | u32 | i64 | u64) -> T:
         return intrinsic("buffer.ref", T, self, index)  # type: ignore
 
@@ -262,13 +274,30 @@ def _inst_pointer_type(args: List[Any]) -> hir.Type:
     _inst_pointer_type
 )
 class Pointer(Generic[T]):
+    """
+    A physical pointer (just like C pointer) to a memory location of type T.
+    Note that pointers might not be available in all backends.
+
+    ```python
+    p = Pointer[int](123456798) # 123456798 is the address of the memory location
+    i = p[0] # read the value at the memory location
+    p[0] = 10 # write the value at the memory location
+    # alternatively
+    i = p.read()
+    p.write(10)
+    # offset the pointer
+    x = p[1]
+    y = (p + 1).read()
+    ```
+    """
+
     def __init__(self, addr: u64) -> None:
         self = intrinsic("init.pointer", Pointer[T], addr)
 
-    def __getitem__(self, index: int | i32 | i64 | u32 | u64) -> T:
+    def __getitem__(self, index:  i32 | i64 | u32 | u64) -> T:
         return intrinsic("pointer.read", T, self, index)  # type: ignore
 
-    def __setitem__(self, index: int | i32 | i64 | u32 | u64, value: T) -> None:
+    def __setitem__(self, index: i32 | i64 | u32 | u64, value: T) -> None:
         pass
 
     def read(self) -> T:
@@ -277,6 +306,12 @@ def read(self) -> T:
     def write(self, value: T) -> None:
         intrinsic("pointer.write", None, self, value)  # type: ignore
 
+    def __add__(self, offset: i32 | i64 | u32 | u64) -> 'Pointer[T]':
+        return intrinsic("pointer.add", Pointer[T], self, offset)
+
+    def __sub__(self, offset: i32 | i64 | u32 | u64) -> 'Pointer[T]':
+        return intrinsic("pointer.sub", Pointer[T], self, offset)
+
 
 __all__: List[str] = [
     # 'Pointer',

luisa_lang/utils.py

@@ -41,7 +41,7 @@ def increment_lineno_and_col_offset(
     return node
 
 
-def dedent_and_retrieve_indentation(lines: str) -> Tuple[str, int]:
+def dedent_and_retrieve_indentation(lines: List[str]) -> Tuple[str, int]:
     """
     Dedent the lines and return the indentation level of the first line.
     """

scripts/gen_math_types.py

@@ -212,6 +212,22 @@ class {ty}{inherits_str}:
             )
             print("")
 
+        def gen_matrix_type(ty:str, vector_ty:str, scalar_ty:str, literal_scalar_ty:str, dim:int):
+            nonlocal exports
+            exports.append(ty)
+            comps = "xyzw"[:dim]
+            fields_def = "".join([f"    {comp}: {vector_ty}\n" for comp in comps])
+            inherits:List[str] = []
+            # if kind == Kind.FLOAT:
+            #     inherits.append(f"FloatBuiltin['{ty}']")
+            inherits_str = "" if len(inherits) == 0 else f"({', '.join(inherits)})"
+            print(
+                f"""@builtin_type(_hir.MatrixType({dim}))
+class {ty}{inherits_str}:
+{fields_def}
+    def __init__(self) -> None: self = intrinsic("init.{ty}", {ty})
+""")
+
         float_types = ["f32", "f64"]
         for size in [2, 3, 4]:
             float_types.append(f"float{size}")