Add CUDA Examples/Improve CUDA robustness

aeon/__main__.py

@@ -72,6 +72,12 @@ def _parse_common_arguments(parser: ArgumentParser):
 
     parser.add_argument("-n", "--no-main", action="store_true", help="Disables introducing hole in main")
 
+    parser.add_argument(
+        "--skip-elaboration",
+        action="store_true",
+        help="Skip the elaboration step",
+    )
+
     parser.add_argument(
         "-s",
         "--synthesizer",
@@ -142,6 +148,7 @@ def main() -> None:
         timings=args.timings,
         no_main=args.no_main,
         synthesis_format=SynthesisFormat.from_string(args.synthesis_format),
+        skip_elaboration=args.skip_elaboration,
     )
     driver = AeonDriver(cfg)
 

aeon/backend/evaluator.py

@@ -2,7 +2,13 @@
 
 from typing import Any
 
-from aeon.core.terms import Abstraction, RefinementAbstraction, RefinementApplication, TypeAbstraction, TypeApplication
+from aeon.core.terms import (
+    Abstraction,
+    RefinementApplication,
+    TypeAbstraction,
+    TypeApplication,
+    RefinementAbstraction,
+)
 from aeon.core.terms import Annotation
 from aeon.core.terms import Application
 from aeon.core.terms import Hole
@@ -12,9 +18,9 @@
 from aeon.core.terms import Rec
 from aeon.core.terms import Term
 from aeon.core.terms import Var
-from aeon.utils.name import Name
 from aeon.decorators.api import Metadata
 from aeon.llvm.core import LLVMPipeline
+from aeon.utils.name import Name
 
 real_eval = eval
 
@@ -87,32 +93,34 @@ def eval(t: Term, ctx: EvaluationContext = EvaluationContext()) -> Any:
         case Let(var_name, var_value, body):
             return eval(body, ctx.with_var(var_name, eval(var_value, ctx)))
         case Rec(var_name, _, var_value, body, _, _):
-            found_llvm = False
-            if ctx.pipeline and ctx.metadata:
-                name_str = var_name.name
-                for k, v in ctx.metadata.items():
-                    k_name = k.name if isinstance(k, Name) else str(k)
-                    if k_name == name_str and v.get("llvm"):
-                        found_llvm = True
-                        break
-
-            if found_llvm:
-                try:
-                    v = ctx.pipeline.get_curried_function(var_name)
-                    if v is not None:
-                        return eval(body, ctx.with_var(var_name, v))
-                except Exception:
-                    pass
-
             if isinstance(var_value, Abstraction):
                 fun = var_value
 
-                def v(x):
+                def v_py(x):
                     return eval(
                         fun.body,
-                        ctx.with_var(var_name, v).with_var(fun.var_name, x),
+                        ctx.with_var(var_name, v_py).with_var(fun.var_name, x),
                     )
 
+                v = v_py
+
+                if ctx.pipeline and ctx.metadata:
+                    name_str = var_name.name
+                    found_llvm = False
+                    for k, meta in ctx.metadata.items():
+                        k_name = k.name if isinstance(k, Name) else str(k)
+                        if k_name == name_str and (meta.get("cpu") or meta.get("gpu")):
+                            found_llvm = True
+                            break
+
+                    if found_llvm:
+                        try:
+                            v_llvm = ctx.pipeline.get_curried_function(var_name, native_fallback=v_py)
+                            if v_llvm is not None:
+                                v = v_llvm
+                        except Exception:
+                            v = v_py
+                            pass
             else:
                 v = eval(var_value, ctx)
             return eval(body, ctx.with_var(var_name, v))

aeon/core/parser.py

@@ -63,7 +63,7 @@ def simple_t(self, args):
         n = str(args[0])
         if n == "Top":
             return top
-        elif n in ["Unit", "Int", "Bool", "Float", "String"]:
+        elif n in ["Unit", "Int", "Bool", "Float", "String", "Vector", "Set"]:
             return TypeConstructor(Name(n, 0))
         else:
             return TypeVar(Name(n))

aeon/core/types.py

@@ -173,11 +173,10 @@ def __hash__(self) -> int:
 t_int = TypeConstructor(Name("Int", 0), [])
 t_float = TypeConstructor(Name("Float", 0), [])
 t_string = TypeConstructor(Name("String", 0), [])
+t_vector = TypeConstructor(Name("Vector", 0), [])
 t_set = TypeConstructor(Name("Set", 0), [])
-t_tensor = TypeConstructor(Name("Tensor", 0), [])
-t_gpu_config = TypeConstructor(Name("GpuConfig", 0), [])
 
-builtin_core_types = [t_unit, t_bool, t_int, t_float, t_string, t_set, t_tensor, t_gpu_config]
+builtin_core_types = [t_unit, t_bool, t_int, t_float, t_string, t_vector, t_set]
 
 top = Top()
 

aeon/facade/driver.py

@@ -44,6 +44,7 @@ class AeonConfig:
     timings: bool = False
     no_main: bool = False
     synthesis_format: SynthesisFormat = SynthesisFormat.DEFAULT
+    skip_elaboration: bool = False
 
 
 class AeonDriver:
@@ -71,9 +72,18 @@ def parse(self, filename: str = None, aeon_code: str = None) -> Iterable[AeonErr
             )
             metadata: Metadata = desugared.metadata
 
+        from aeon.decorators.api import CORE_DECORATOR_QUEUE_META_KEY
+
+        auto_skip_elaboration = self.cfg.skip_elaboration or (
+            CORE_DECORATOR_QUEUE_META_KEY in metadata and len(metadata[CORE_DECORATOR_QUEUE_META_KEY]) > 0
+        )
+
         try:
-            with RecordTime("Elaboration"):
-                sterm: STerm = elaborate(desugared.elabcontext, desugared.program, st_top)
+            if auto_skip_elaboration:
+                sterm: STerm = desugared.program
+            else:
+                with RecordTime("Elaboration"):
+                    sterm = elaborate(desugared.elabcontext, desugared.program, st_top)
         except AeonError as e:
             return [e]  # TODO: Support multiple errors
 
@@ -85,11 +95,12 @@ def parse(self, filename: str = None, aeon_code: str = None) -> Iterable[AeonErr
         with RecordTime("ANF conversion"):
             core_ast_anf = ensure_anf(core_ast)
 
-        with RecordTime("TypeChecking"):
-            type_errors = check_type_errors(typing_ctx, core_ast_anf, top)
-            # TODO
-            if type_errors:
-                return type_errors
+        if not auto_skip_elaboration:
+            with RecordTime("TypeChecking"):
+                type_errors = check_type_errors(typing_ctx, core_ast_anf, top)
+                # TODO
+                if type_errors:
+                    return type_errors
 
         with RecordTime("CorePhaseDecorators"):
             metadata = apply_core_decorators_phase(typing_ctx, core_ast_anf, metadata)

aeon/llvm/constants.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Mapping
+
+
+class LLVMBackendName(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+
+
+class LLVMMetadataKey(str, Enum):
+    CPU_ENABLED = "cpu"
+    CPU_OPT_LEVEL = "cpu_opt_level"
+    GPU_ENABLED = "gpu"
+    GPU_DEVICE = "gpu_device"
+    GPU_OPT_LEVEL = "gpu_opt_level"
+    GPU_DEBUG = "gpu_debug"
+    GPU_BLOCK_SIZE = "gpu_block_size"
+    GPU_ARCH = "gpu_arch"
+    GPU_KERNEL_PLAN = "gpu_kernel_plan"
+    GPU_NO_DEVICE_ALLOC = "gpu_no_device_alloc"
+    SIZE = "size"
+
+
+class VectorOperation(str, Enum):
+    MAP = "map"
+    REDUCE = "reduce"
+    IMAP = "imap"
+    FILTER = "filter"
+    ZIP_WITH = "zipWith"
+    COUNT = "count"
+    SIZE = "size"
+
+
+VECTOR_OPERATION_NAMES: frozenset[str] = frozenset(op.value for op in VectorOperation)
+
+
+def metadata_get(metadata: Mapping[str, Any], key: LLVMMetadataKey, default: Any = None) -> Any:
+    return metadata.get(key.value, default)
+
+
+def metadata_has(metadata: Mapping[str, Any], key: LLVMMetadataKey) -> bool:
+    return key.value in metadata

aeon/llvm/core.py

@@ -27,6 +27,9 @@
         LLVMVectorFilter,
         LLVMVectorZipWith,
         LLVMVectorCount,
+        LLVMVectorGet,
+        LLVMVectorSet,
+        LLVMVectorSize,
     )
 else:
     LLVMTerm = Any
@@ -48,6 +51,9 @@
     LLVMVectorFilter = Any
     LLVMVectorZipWith = Any
     LLVMVectorCount = Any
+    LLVMVectorGet = Any
+    LLVMVectorSet = Any
+    LLVMVectorSize = Any
 
 
 class LLVMBackendError(Exception):
@@ -108,7 +114,13 @@ def optimize(self, llvm_ir: str, opt_level: int = 3) -> str:
 class LLVMExecutionEngine(ABC):
     @abstractmethod
     def execute(
-        self, llvm_ir: str, func_name: str, args: List[Any], arg_types: List[LLVMType], ret_type: LLVMType
+        self,
+        llvm_ir: str,
+        func_name: str,
+        args: List[Any],
+        arg_types: List[LLVMType],
+        ret_type: LLVMType,
+        metadata: dict[str, Any] | None = None,
     ) -> Any:
         pass
 
@@ -126,7 +138,7 @@ def compile(self, program: Term) -> None:
         pass
 
     @abstractmethod
-    def get_curried_function(self, name: Name) -> Any:
+    def get_curried_function(self, name: Name, native_fallback: Any = None) -> Any:
         pass
 
     @abstractmethod
@@ -206,3 +218,15 @@ def visit_vector_zipwith(self, node: LLVMVectorZipWith) -> Any:
     @abstractmethod
     def visit_vector_count(self, node: LLVMVectorCount) -> Any:
         pass
+
+    @abstractmethod
+    def visit_vector_get(self, node: LLVMVectorGet) -> Any:
+        pass
+
+    @abstractmethod
+    def visit_vector_set(self, node: LLVMVectorSet) -> Any:
+        pass
+
+    @abstractmethod
+    def visit_vector_size(self, node: LLVMVectorSize) -> Any:
+        pass