fix #55: support pos only and kw only args

Author: Ritsuka314

src/pydsl/analysis/names.py

@@ -110,7 +110,12 @@ def analyze(func: ast.FunctionDef | ast.Lambda) -> set[str]:
         ba = BoundAnalysis()
 
         # names of arguments are also bound in the function
-        ba.bound.update([a.arg for a in func.args.args])
+        ba.bound.update([
+            a.arg
+            for a in func.args.args
+            + func.args.posonlyargs
+            + func.args.kwonlyargs
+        ])
 
         body = func.body if isinstance(func.body, Iterable) else [func.body]
 

src/pydsl/frontend.py

@@ -2,6 +2,7 @@
 import ctypes
 import dataclasses
 import inspect
+from inspect import BoundArguments
 import logging
 import re
 import subprocess
@@ -762,8 +763,30 @@ def temp_file_to_path(tempf):
             self.ll_to_so,
         ])(file)
 
+    def parse_args(self, signature, *args, **kwargs) -> BoundArguments:
+        """
+        Try to bind arguments, apply defaults, and type cast. Arguments should
+        already be compiled to PyDSL types by this point.
+        """
+        params = signature.parameters
+
+        try:
+            # Associate each argument value passed into the call with
+            # the parameter of the function
+            bound_args = signature.bind(*args, **kwargs)
+            binding = bound_args.arguments
+        except TypeError as e:
+            raise TypeError(
+                f"couldn't bind arguments when calling an inline function: {e}"
+            ) from e
+
+        # Apply defaults for unfilled arguments
+        bound_args.apply_defaults()
+
+        return bound_args
+
     # TODO: this function is too large. Should break it down a bit
-    def call_function(self, fname: str, *args) -> Any:
+    def call_function(self, fname: str, *args, **kwargs) -> Any:
         if not hasattr(self, "_so"):
             raise RuntimeError(
                 f"function {fname} is called before it is compiled"
@@ -772,22 +795,26 @@ def call_function(self, fname: str, *args) -> Any:
         f = self.get_func(fname)
         sig = f.signature
         so_f = self.load_function(f)
-        if not len(sig.parameters) == len(args):
+        if not len(sig.parameters) == len(args) + len(kwargs):
             raise TypeError(
-                f"{f.name} takes {len(sig.parameters)} positional "
+                f"{f.name} takes {len(sig.parameters)} "
                 f"argument{"s" if len(sig.parameters) > 1 else ""} "
-                f"but {len(args)} were given"
+                f"but {len(args) + len(kwargs)} were given"
             )
 
         arg_cont = ArgContainer()
 
+        bound_args = self.parse_args(sig, *args, **kwargs)
+
         mapped_args_ct = [
             (
                 ct,
                 self.val_to_CType(arg_cont, sig.parameters[key].annotation, a),
             )
-            for ct, key, a in zip(
-                self.get_args_ctypes(f), sig.parameters, args, strict=False
+            for ct, (key, a) in zip(
+                self.get_args_ctypes(f),
+                bound_args.arguments.items(),
+                strict=False,
             )
         ]
 
@@ -1324,8 +1351,8 @@ def emit_mlir(self) -> str:
 
 
 class CompiledFunction(CompiledObject[Callable[..., Any]]):
-    def __call__(self, *args) -> Any:
-        return self._target.call_function(self._o.__name__, *args)
+    def __call__(self, *args, **kwargs) -> Any:
+        return self._target.call_function(self._o.__name__, *args, **kwargs)
 
 
 class CompiledClass(CompiledObject[type]):
@@ -1416,8 +1443,8 @@ def selfattr(name: str):
 
         if name in _target.funcs():
 
-            def make_call(*args) -> Any:
-                return _target.call_function(name, *args)
+            def make_call(*args, **kwargs) -> Any:
+                return _target.call_function(name, *args, **kwargs)
 
             return make_call
 

src/pydsl/func.py

@@ -616,7 +616,10 @@ def on_Call(
     ) -> SubtreeOut:
         self = attr_chain[-1]
         prefix_args = [visitor.visit(a) for a in prefix_args]
-        args = [visitor.visit(a) for a in node.args]
+        args = [
+            visitor.visit(a)
+            for a in node.args + [kw.value for kw in node.keywords]
+        ]
         args = [
             a if isinstance(a, t) else t(a) for t, a in zip(self.argst, args)
         ]

tests/e2e/test_syntax.py

@@ -130,6 +130,55 @@ def g(m1: MemRef[UInt8, 2, 3], m2: MemRef[Index, 1]) -> None:
     assert (m1 == expected).all()
 
 
+def test_pos_only_args():
+    @compile()
+    def positional_only(
+        a: UInt8, /, b: UInt8, c: UInt8
+    ) -> Tuple[UInt8, UInt8, UInt8]:
+        return a, b, c
+
+    assert positional_only(2, c=4, b=3) == (2, 3, 4)
+
+
+def test_kw_only_args():
+    @compile()
+    def keyword_only(
+        a: UInt8, *, b: UInt8, c: UInt8
+    ) -> Tuple[UInt8, UInt8, UInt8]:
+        return a, b, c
+
+    assert keyword_only(2, c=4, b=3) == (2, 3, 4)
+
+
+def test_pos_or_kw_args():
+    @compile()
+    def pos_or_kw(a: UInt8, b: UInt8, c: UInt8) -> Tuple[UInt8, UInt8, UInt8]:
+        return a, b, c
+
+    assert pos_or_kw(2, c=4, b=3) == (2, 3, 4)
+
+
+def test_fun_args():
+    @compile()
+    class Mod:
+        def pos_or_kw(
+            a: UInt8, b: UInt8, c: UInt8
+        ) -> Tuple[UInt8, UInt8, UInt8]:
+            return a, b, c
+
+        def keyword_only(
+            a: UInt8, *, b: UInt8, c: UInt8
+        ) -> Tuple[UInt8, UInt8, UInt8]:
+            return pos_or_kw(a, c=c, b=b)
+
+        def positional_only(
+            a: UInt8, /, b: UInt8, c: UInt8
+        ) -> Tuple[UInt8, UInt8, UInt8]:
+            return keyword_only(2, c=4, b=3)
+
+    assert Mod.positional_only(2, c=4, b=3) == (2, 3, 4)
+
+
 if __name__ == "__main__":
     run(test_annassign)
     run(test_illegal_annassign)
@@ -142,3 +191,6 @@ def g(m1: MemRef[UInt8, 2, 3], m2: MemRef[Index, 1]) -> None:
     run(test_minus_eq)
     run(test_plus_eq_memref)
     run(test_plus_eq_side_effect)
+    run(test_pos_only_args)
+    run(test_kw_only_args)
+    run(test_pos_or_kw_args)