better types; now int, float are supported dsl types; only float and …

…int types of other bit width needs to write lc.ixx, lc.fxx explicitly

luisa_lang/hir.py

@@ -1597,3 +1597,35 @@ def inline(func: Function, args: List[Value | Ref], body: BasicBlock, span: Opti
         inliner = FunctionInliner(func, args, body, span)
         assert inliner.ret
         return inliner.ret
+
+
+def register_dsl_type_alias(target: type, alias: type):
+    """
+    Allow a type to be remapped to another type within DSL code.
+    Parameters:
+    target (type): The type to be remapped.
+    alias (type): The type to which the target type will be remapped.
+    Example:
+
+    For example, 
+    ```python
+    @lc.struct
+    class Foo:
+        x: int
+        y: int
+
+    class SomeOtherFoo:
+        components: List[int]
+        
+    register_dsl_type_alias(SomeOtherFoo, Foo)
+    
+    @lc.func
+    def foo(f: SomeOtherFoo): # SomeOtherFoo is interpreted as Foo
+        ...
+
+    ```
+    """
+    ctx = GlobalContext.get()
+    alias_ty = get_dsl_type(alias)
+    assert alias_ty, f"alias type {alias} is not a DSL type"
+    ctx.types[target] = alias_ty

luisa_lang/lang_builtins.py

@@ -211,11 +211,11 @@ class Array(Generic[T, N]):
     def __init__(self) -> None:
         self = intrinsic("init.array", Array[T, N])
 
-    def __getitem__(self, index: int | i32 | u32 | i64 | u64) -> T:
+    def __getitem__(self, index: int | u32 | i64 | u64) -> T:
         return intrinsic("array.ref", T, byref(self), index)  # type: ignore
 
-    def __setitem__(self, index: int | i32 | u32 | i64 | u64, value: T | int | float) -> None:
-        """value: T | int | float annotation is to make mypy happy. this function is ignored by the compiler"""
+    def __setitem__(self, index: int | u32 | i64 | u64, value: T) -> None:
+        pass
 
     def __len__(self) -> u64:
         return intrinsic("array.size", u64, self)  # type: ignore
@@ -236,11 +236,10 @@ def __len__(self) -> u64:
 
 @opaque("Buffer")
 class Buffer(Generic[T]):
-    def __getitem__(self, index: int | i32 | u32 | i64 | u64) -> T:
+    def __getitem__(self, index: int | u32 | i64 | u64) -> T:
         return intrinsic("buffer.ref", T, self, index)  # type: ignore
 
-    def __setitem__(self,  index: int | i32 | u32 | i64 | u64,  value: T | int | float) -> None:
-        """value: T | int | float annotation is to make mypy happy. this function is ignored by the compiler"""
+    def __setitem__(self,  index: int | u32 | i64 | u64,  value: T) -> None:
         pass
 
     def __len__(self) -> u64: