Add Backend LLVM Code Gen Support

aeon/backend/evaluator.py

@@ -13,24 +13,35 @@
 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
 
 real_eval = eval
 
 
 class EvaluationContext:
     variables: dict[Name, Any]
-
-    def __init__(self, prev: dict[Name, Any] | None = None):
+    metadata: Metadata | None
+    pipeline: LLVMPipeline | None
+
+    def __init__(
+        self,
+        prev: dict[Name, Any] | None = None,
+        metadata: Metadata | None = None,
+        pipeline: LLVMPipeline | None = None,
+    ):
         if prev:
             self.variables = {k: v for (k, v) in prev.items()}
         else:
             self.variables = {}
+        self.metadata = metadata
+        self.pipeline = pipeline
 
     def with_var(self, name: Name, value: Any):
         assert isinstance(name, Name)
         v = self.variables.copy()
         v.update({name: value})
-        return EvaluationContext(v)
+        return EvaluationContext(v, metadata=self.metadata, pipeline=self.pipeline)
 
     def get(self, name: Name):
         return self.variables[name]
@@ -76,6 +87,23 @@ 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
 
@@ -87,7 +115,7 @@ def v(x):
 
             else:
                 v = eval(var_value, ctx)
-            return eval(t.body, ctx.with_var(t.var_name, v))
+            return eval(body, ctx.with_var(var_name, v))
         case If(cond, then, otherwise):
             c = eval(cond, ctx)
             if c:

aeon/core/types.py

@@ -159,8 +159,10 @@ def __hash__(self) -> int:
 t_int = TypeConstructor(Name("Int", 0), [])
 t_float = TypeConstructor(Name("Float", 0), [])
 t_string = TypeConstructor(Name("String", 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]
+builtin_core_types = [t_unit, t_bool, t_int, t_float, t_string, t_tensor, t_gpu_config]
 
 top = Top()
 

aeon/decorators/__init__.py

@@ -9,8 +9,9 @@ def fun(...) { ... }
 eventual complementary definitions.
 """
 
-from aeon.decorators.api import DecoratorType
-from aeon.decorators.api import Metadata
+from aeon.decorators.api import Metadata, DecoratorType
+from aeon.llvm.decorators.gpu import gpu
+from aeon.llvm.decorators.llvm import llvm
 from aeon.sugar.program import Definition
 from aeon.synthesis.decorators import (
     minimize_int,
@@ -40,6 +41,8 @@ def fun(...) { ... }
     "error_fitness": error_fitness,
     "disable_control_flow": disable_control_flow,
     "prompt": prompt,
+    "llvm": llvm,
+    "gpu": gpu,
     "csv_data": csv_data,
     "csv_file": csv_file,
 }
@@ -53,7 +56,7 @@ def apply_decorators(fun: Definition, metadata: Metadata) -> tuple[Definition, l
     for decorator in fun.decorators:
         dname = decorator.name.name
         if dname not in decorators_environment:
-            raise Exception(f"Unknown decorator named {dname}, in function {fun.name}.")
+            raise Exception(f"Unknown decorator named {dname}, in function {fun.name.pretty()}.")
         decorator_processor = decorators_environment[dname]
         (fun, extra, metadata) = decorator_processor(decorator.macro_args, fun, metadata)
         total_extra.extend(extra)

aeon/facade/driver.py

@@ -32,6 +32,10 @@
 from aeon.utils.name import Name
 from aeon.utils.pprint import pretty_print_node
 from aeon.utils.time_utils import RecordTime
+from aeon.llvm.cpu.pipeline import CPULLVMPipeline
+from aeon.llvm.cpu.executor import CPULLVMExecutionEngine
+from aeon.llvm.cpu.converter import CPULLVMIRGenerator
+from aeon.llvm.cpu.lowerer import CPULLVMLowerer
 
 
 def read_file(filename: str) -> str:
@@ -107,17 +111,25 @@ def parse(self, filename: str = None, aeon_code: str = None) -> Iterable[AeonErr
                 return type_errors
 
         with RecordTime("Preparing execution env"):
-            evaluation_ctx = EvaluationContext(evaluation_vars)
+            executor = CPULLVMExecutionEngine()
+            generator = CPULLVMIRGenerator()
+            lowerer = CPULLVMLowerer()
+            pipeline = CPULLVMPipeline(executor, generator, lowerer, metadata=metadata)
+            evaluation_ctx = EvaluationContext(evaluation_vars, metadata=metadata, pipeline=pipeline)
 
         self.metadata = metadata
         self.core = core_ast_anf
         self.typing_ctx = typing_ctx
         self.evaluation_ctx = evaluation_ctx
+
+        with RecordTime("LLVM compilation"):
+            pipeline.compile(self.core)
+
         return []
 
-    def run(self) -> None:
+    def run(self) -> Any:
         with RecordTime("Evaluation"):
-            eval(self.core, self.evaluation_ctx)
+            return eval(self.core, self.evaluation_ctx)
 
     def has_synth(self) -> bool:
         with RecordTime("DetectSynthesis"):

aeon/llvm/core.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List
+from aeon.core.terms import Term
+from aeon.utils.name import Name
+from aeon.llvm.llvm_ast import LLVMTerm, LLVMType
+
+
+class LLVMBackendError(Exception):
+    pass
+
+
+class LLVMValidationError(LLVMBackendError):
+    pass
+
+
+@dataclass(frozen=True)
+class ValidationContext:
+    pass
+
+
+class ValidationStep(ABC):
+    @abstractmethod
+    def validate(self, t: Term, ctx: ValidationContext) -> None:
+        pass
+
+
+class LLVMLowerer(ABC):
+    def validate(self, t: Term, ctx: ValidationContext) -> None:
+        for step in self.get_validation_steps():
+            step.validate(t, ctx)
+
+    @abstractmethod
+    def get_validation_steps(self) -> List[ValidationStep]:
+        pass
+
+    @abstractmethod
+    def lower(
+        self,
+        t: Term,
+        expected_type: LLVMType = None,
+        type_env: Dict[Name, LLVMType] = None,
+        env: Dict[Name, LLVMTerm] = None,
+    ) -> LLVMTerm:
+        pass
+
+    @abstractmethod
+    def get_signature(self, ty: LLVMType) -> tuple[List[LLVMType], LLVMType]:
+        pass
+
+
+class LLVMIRGenerator(ABC):
+    @abstractmethod
+    def generate_ir(self, definitions: List[LLVMTerm]) -> str:
+        pass
+
+
+class LLVMOptimizer(ABC):
+    @abstractmethod
+    def optimize(self, llvm_ir: str, opt_level: int = 3) -> str:
+        pass
+
+
+class LLVMExecutionEngine(ABC):
+    @abstractmethod
+    def execute(
+        self, llvm_ir: str, func_name: str, args: List[Any], arg_types: List[LLVMType], ret_type: LLVMType
+    ) -> Any:
+        pass
+
+
+class LLVMPipeline(ABC):
+    @abstractmethod
+    def compile(self, program: Term) -> None:
+        pass
+
+    @abstractmethod
+    def get_curried_function(self, name: Name) -> Any:
+        pass
+
+    @abstractmethod
+    def invoke(self, name_id: Name, arguments: List[Any]) -> Any:
+        pass