(WIP) frontend: allow use of MLIR Types as "casts"

Author: AlexanderViand

frontend/heir/mlir/types.py

@@ -4,28 +4,64 @@
 from typing import Generic, Self, TypeVar, TypeVarTuple, get_args, get_origin
 from numba.core.types import Type as NumbaType
 from numba.core.types import boolean, int8, int16, int32, int64, float32, float64
+from numba.extending import typeof_impl, type_callable
 
 T = TypeVar("T")
 Ts = TypeVarTuple("Ts")
 
-operator_error_message = "MLIRType should only be used for annotations."
+# List of all MLIR types we define here, for use in other parts of the compiler
+MLIR_TYPES = []  # populated via MLIRType's __init_subclass__
+
+
+def check_for_value(a: "MLIRType"):
+  if not hasattr(a, "value"):
+    raise RuntimeError(
+        "Trying to use an operator on an MLIRType without a value."
+    )
 
 
 class MLIRType(ABC):
 
+  def __init__(self, value: int):
+    self.value = value
+
+  def __init_subclass__(cls, **kwargs):
+    super().__init_subclass__(**kwargs)
+    MLIR_TYPES.append(cls)
+
   @staticmethod
   @abstractmethod
   def numba_type() -> NumbaType:
     raise NotImplementedError("No numba type exists for a generic MLIRType")
 
-  def __add__(self, other) -> Self:
-    raise RuntimeError(operator_error_message)
+  @staticmethod
+  @abstractmethod
+  def mlir_type() -> str:
+    raise NotImplementedError("No mlir type exists for a generic MLIRType")
+
+  def __add__(self, other):
+    check_for_value(self)
+    return self.value + other
 
-  def __sub__(self, other) -> Self:
-    raise RuntimeError(operator_error_message)
+  def __radd__(self, other):
+    check_for_value(self)
+    return other + self.value
 
-  def __mul__(self, other) -> Self:
-    raise RuntimeError(operator_error_message)
+  def __sub__(self, other):
+    check_for_value(self)
+    return self.value - other
+
+  def __rsub__(self, other):
+    check_for_value(self)
+    return other - self.value
+
+  def __mul__(self, other):
+    check_for_value(self)
+    return self.value * other
+
+  def __rmul__(self, other):
+    check_for_value(self)
+    return other * self.value
 
 
 class Secret(Generic[T], MLIRType):
@@ -34,62 +70,106 @@ class Secret(Generic[T], MLIRType):
   def numba_type() -> NumbaType:
     raise NotImplementedError("No numba type exists for a generic Secret")
 
+  @staticmethod
+  def mlir_type() -> str:
+    raise NotImplementedError("No mlir type exists for a generic Secret")
+
 
 class Tensor(Generic[*Ts], MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     raise NotImplementedError("No numba type exists for a generic Tensor")
 
+  @staticmethod
+  def mlir_type() -> str:
+    raise NotImplementedError("No mlir type exists for a generic Tensor")
+
 
 class F32(MLIRType):
   # TODO (#1162): For CKKS/Float: allow specifying actual intended precision/scale and warn/error if not achievable  @staticmethod
   @staticmethod
   def numba_type() -> NumbaType:
     return float32
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "f32"
+
 
 class F64(MLIRType):
   # TODO (#1162): For CKKS/Float: allow specifying actual intended precision/scale and warn/error if not achievable  @staticmethod
   @staticmethod
   def numba_type() -> NumbaType:
     return float64
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "f64"
+
 
 class I1(MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     return boolean
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "i1"
+
 
 class I8(MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     return int8
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "i8"
+
 
 class I16(MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     return int16
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "i16"
+
 
 class I32(MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     return int32
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "i32"
+
 
 class I64(MLIRType):
 
   @staticmethod
   def numba_type() -> NumbaType:
     return int64
 
+  @staticmethod
+  def mlir_type() -> str:
+    return "i64"
+
+
+# Register the types defined above with Numba
+for typ in [I8, I16, I32, I64, I1, F32, F64]:
+
+  @type_callable(typ)
+  def build_typer_function(context, typ=typ):
+    return lambda value: typ.numba_type()
+
 
 # Helper functions
 

frontend/heir/mlir_emitter.py

@@ -12,7 +12,8 @@
 from numba.core import controlflow
 from numba.core.types import Type as NumbaType
 
-from heir.interfaces import InternalCompilerError
+from heir.mlir.types import MLIRType, MLIR_TYPES, I1, I8, I16, I32, I64, F32, F64
+from heir.interfaces import CompilerError, DebugMessage, InternalCompilerError
 
 
 def mlirType(numba_type: NumbaType) -> str:
@@ -36,6 +37,61 @@ def mlirType(numba_type: NumbaType) -> str:
   raise InternalCompilerError("Unsupported type: " + str(numba_type))
 
 
+def isIntegerLike(typ: NumbaType | MLIRType) -> bool:
+  if isinstance(typ, type) and issubclass(typ, MLIRType):
+    return typ in {I1, I8, I16, I32, I64}
+  if isinstance(typ, NumbaType):
+    return isinstance(typ, types.Integer) or isinstance(typ, types.Boolean)
+  raise InternalCompilerError(f"Encountered unexpected type {typ}")
+
+
+def isFloatLike(typ: NumbaType | MLIRType) -> bool:
+  if isinstance(typ, type) and issubclass(type, MLIRType):
+    return typ in {F32, F64}
+  if isinstance(typ, NumbaType):
+    return isinstance(typ, types.Float)
+  raise InternalCompilerError(f"Encountered unexpected type {type}")
+
+
+def mlirCastOp(
+    from_type: NumbaType, to_type: MLIRType, value: str, loc: ir.Loc
+) -> str:
+  if isIntegerLike(from_type) and isIntegerLike(to_type):
+    if from_type.bitwidth == to_type.numba_type().bitwidth:
+      raise CompilerError(
+          f"Cannot create cast of {value} from {from_type} to {to_type} as they"
+          " have the same bitwidth",
+          loc,
+      )
+    if from_type.bitwidth > to_type.numba_type().bitwidth:
+      return (
+          f"arith.trunci {value} : {mlirType(from_type)} to"
+          f" {to_type.mlir_type()} {mlirLoc(loc)}"
+      )
+    if from_type.bitwidth < to_type.numba_type().bitwidth:
+      # FIXME: signedness for extensions?
+      return (
+          f"arith.extui {value} : {mlirType(from_type)} to"
+          f" {to_type.mlir_type()} {mlirLoc(loc)}"
+      )
+  if isFloatLike(from_type) and isIntegerLike(to_type):
+    # FIXME: signedness?
+    return (
+        f"arith.fptoui {value} : {mlirType(from_type)} to"
+        f" {mlirType(to_type)} {mlirLoc(loc)}"
+    )
+  if isIntegerLike(from_type) and isFloatLike(to_type):
+    # FIXME: signendess?
+    return (
+        f"arith.uitofp {value} : {mlirType(from_type)} to"
+        f" {mlirType(to_type)} {mlirLoc(loc)}"
+    )
+  raise CompilerError(
+      f"Encountered unsupported cast of {value} from {from_type} to {to_type}",
+      loc,
+  )
+
+
 def mlirLoc(loc: ir.Loc) -> str:
   return (
       f"loc(\"{loc.filename or '<unknown>'}\":{loc.line or 0}:{loc.col or 0})"
@@ -419,12 +475,39 @@ def emit_assign(self, assign):
         func = assign.value.func
         # if assert fails, variable was undefined
         assert func.name in self.globals_map
-        if self.globals_map[func.name] == "bool":
+        name, global_ = self.globals_map[func.name]
+        if name == "bool":
           # nothing to do, forward the name to the arg of bool()
           self.forward_name(from_var=assign.target, to_var=assign.value.args[0])
           return ""
+        if global_ in MLIR_TYPES:
+          if len(assign.value.args) != 1:
+            raise CompilerError(
+                "MLIR type cast requires exactly one argument", assign.value.loc
+            )
+          value = assign.value.args[0].name
+          if (
+              mlirType(self.typemap.get(assign.target.name))
+              != global_.mlir_type()
+          ):
+            raise InternalCompilerError(
+                f"MLIR type cast of {value} from"
+                f" {mlirType(self.typemap.get(value))} to"
+                f" {global_.mlir_type()} is not correctly reflected in types"
+                " inferred for the assignment, which expects"
+                f" {mlirType(self.typemap.get(assign.target.name))}"
+            )
+          target_ssa = self.get_or_create_name(assign.target)
+          ssa_id = self.get_or_create_name(assign.value.args[0])
+          cast = mlirCastOp(
+              self.typemap.get(value),
+              global_,
+              ssa_id,
+              assign.loc,
+          )
+          return f"{target_ssa} = {cast}"
         else:
-          raise InternalCompilerError("Unknown global " + func.name)
+          raise InternalCompilerError("Call to unknown function " + name)
       case ir.Expr(op="cast"):
         # not sure what to do here. maybe will be needed for type conversions
         # when interfacing with C
@@ -446,7 +529,10 @@ def emit_assign(self, assign):
           self.forward_name_to_id(assign.target, name.strip("%"))
         return const_str
       case ir.Global():
-        self.globals_map[assign.target.name] = assign.value.name
+        self.globals_map[assign.target.name] = (
+            assign.value.name,
+            assign.value.value,
+        )
         return ""
       case ir.Var():
         # Sometimes we need this to be assigned?
@@ -469,6 +555,7 @@ def emit_ext_if_needed(self, lhs, rhs):
       raise InternalCompilerError(
           "Extension handling for non-integer (e.g., floats, tensors) types"
           " is not yet supported. Please ensure (inferred) bit-widths match."
+          f" Failed to extend {lhs_type} and {rhs_type} types."
       )
       # TODO (#1162): Support bitwidth extension for float types
       #      (this probably requires adding support for local variable type hints,