Remove Allocate, Free, and ModeLiteral from IR (#72)

src/tensora/codegen/_ir_to_c.py

@@ -2,11 +2,8 @@
 
 from functools import singledispatch
 
-from ..format import Mode
 from ..ir.ast import (
     Add,
-    Address,
-    Allocate,
     And,
     ArrayAllocate,
     ArrayIndex,
@@ -22,7 +19,6 @@
     Equal,
     Expression,
     FloatLiteral,
-    Free,
     FunctionDefinition,
     GreaterThan,
     GreaterThanOrEqual,
@@ -32,7 +28,6 @@
     Loop,
     Max,
     Min,
-    ModeLiteral,
     Module,
     Multiply,
     NotEqual,
@@ -91,15 +86,6 @@ def ir_to_c_boolean_literal(self: BooleanLiteral) -> str:
     return str(int(self.value))
 
 
-@ir_to_c_expression.register(ModeLiteral)
-def ir_to_c_mode_literal(self: ModeLiteral) -> str:
-    match self.value:
-        case Mode.dense:
-            return "taco_mode_dense"
-        case Mode.compressed:
-            return "taco_mode_sparse"
-
-
 @ir_to_c_expression.register(Add)
 def ir_to_c_add(self: Add) -> str:
     return f"{ir_to_c_expression(self.left)} + {ir_to_c_expression(self.right)}"
@@ -167,21 +153,11 @@ def ir_to_c_min(self: Min) -> str:
     return f"TACO_MIN({ir_to_c_expression(self.left)}, {ir_to_c_expression(self.right)})"
 
 
-@ir_to_c_expression.register(Address)
-def ir_to_c_address(self: Address) -> str:
-    return f"&{ir_to_c_expression(self.target)}"
-
-
 @ir_to_c_expression.register(BooleanToInteger)
 def ir_to_c_boolean_to_integer(self: BooleanToInteger) -> str:
     return f"(int32_t)({ir_to_c_expression(self.expression)})"
 
 
-@ir_to_c_expression.register(Allocate)
-def ir_to_c_allocate(self: Allocate) -> str:
-    return f"malloc(sizeof({type_to_c(self.type)}))"
-
-
 @ir_to_c_expression.register(ArrayAllocate)
 def ir_to_c_array_allocate(self: ArrayAllocate) -> str:
     return f"malloc(sizeof({type_to_c(self.type)}) * {parens(self.n_elements, (Add, Subtract))})"
@@ -208,11 +184,6 @@ def convert_expression_to_statement(self: Expression) -> list[str]:
     return [ir_to_c_expression(self) + ";"]
 
 
-@ir_to_c_statement.register(Free)
-def ir_to_c_free(self: Free) -> list[str]:
-    return [f"free({ir_to_c_expression(self.target)});"]
-
-
 @ir_to_c_statement.register(Declaration)
 def convert_declaration_to_statement(self: Declaration) -> list[str]:
     return [ir_to_c_declaration(self) + ";"]

src/tensora/ir/_peephole.py

@@ -26,8 +26,6 @@
 
 from .ast import (
     Add,
-    Address,
-    Allocate,
     And,
     ArrayAllocate,
     ArrayIndex,
@@ -44,7 +42,6 @@
     Equal,
     Expression,
     FloatLiteral,
-    Free,
     FunctionDefinition,
     GreaterThan,
     GreaterThanOrEqual,
@@ -54,7 +51,6 @@
     Loop,
     Max,
     Min,
-    ModeLiteral,
     Module,
     Multiply,
     NotEqual,
@@ -100,9 +96,6 @@ def peephole_expression_assignable(self: Assignable) -> Assignable:
 @peephole_expression.register(IntegerLiteral)
 @peephole_expression.register(FloatLiteral)
 @peephole_expression.register(BooleanLiteral)
-@peephole_expression.register(ModeLiteral)
-@peephole_expression.register(Address)
-@peephole_expression.register(Allocate)
 def peephole_noop(self: Expression) -> Expression:
     return self
 
@@ -257,11 +250,6 @@ def peephole_declaration(self: Declaration) -> Declaration:
     return self
 
 
-@peephole_statement.register(Free)
-def peephole_free(self: Free) -> Statement:
-    return Free(peephole_assignable(self.target))
-
-
 @peephole_statement.register(Assignment)
 def peephole_assignment(self: Assignment) -> Statement:
     target = peephole_assignable(self.target)

src/tensora/ir/ast.py

@@ -10,7 +10,6 @@
     "IntegerLiteral",
     "FloatLiteral",
     "BooleanLiteral",
-    "ModeLiteral",
     "Add",
     "Subtract",
     "Multiply",
@@ -24,12 +23,9 @@
     "Or",
     "Max",
     "Min",
-    "Address",
     "BooleanToInteger",
-    "Allocate",
     "ArrayAllocate",
     "ArrayReallocate",
-    "Free",
     "Declaration",
     "Assignment",
     "DeclarationAssignment",
@@ -45,7 +41,6 @@
 from functools import reduce
 from typing import Sequence
 
-from ..format import Mode
 from .types import Type
 
 
@@ -133,11 +128,6 @@ class BooleanLiteral(Expression):
     value: bool
 
 
-@dataclass(frozen=True, slots=True)
-class ModeLiteral(Expression):
-    value: Mode
-
-
 @dataclass(frozen=True, slots=True)
 class Add(Expression):
     left: Expression
@@ -246,21 +236,11 @@ def join(operands: Sequence[Expression | int | str]) -> Expression:
         return reduce(Min, expression_operands)
 
 
-@dataclass(frozen=True, slots=True)
-class Address(Expression):
-    target: Assignable
-
-
 @dataclass(frozen=True, slots=True)
 class BooleanToInteger(Expression):
     expression: Expression
 
 
-@dataclass(frozen=True, slots=True)
-class Allocate(Expression):
-    type: Type
-
-
 @dataclass(frozen=True, slots=True)
 class ArrayAllocate(Expression):
     type: Type
@@ -274,11 +254,6 @@ class ArrayReallocate(Expression):
     n_elements: Expression
 
 
-@dataclass(frozen=True, slots=True)
-class Free(Statement):
-    target: Assignable
-
-
 @dataclass(frozen=True, slots=True)
 class Declaration(Statement):
     name: Variable

tests/codegen/test_ast_to_c.py

@@ -60,13 +60,9 @@ def clean(string: str) -> str:
     (Min(Variable("x"), Variable("y")), "TACO_MIN(x, y)"),
     (Max(Max(Variable("x"), Variable("y")), Variable("z")), "TACO_MAX(TACO_MAX(x, y), z)"),
     (Min(Min(Variable("x"), Variable("y")), Variable("z")), "TACO_MIN(TACO_MIN(x, y), z)"),
-    # Address
-    (Address(Variable("x")), "&x"),
-    (Address(ArrayIndex(Variable("x"), Variable("i"))), "&x[i]"),
     # Cast
     (BooleanToInteger(Equal(Variable("x"), Variable("y"))), "(int32_t)(x == y)"),
     # Allocate
-    (Allocate(tensor), "malloc(sizeof(taco_tensor_t))"),
     (ArrayAllocate(integer, Variable("capacity")), "malloc(sizeof(int32_t) * capacity)"),
     (
         ArrayAllocate(float, Add(Variable("previous"), Variable("new"))),
@@ -80,8 +76,6 @@ def clean(string: str) -> str:
         ArrayReallocate(Variable("old"), float, Add(Variable("previous"), Variable("new"))),
         "realloc(old, sizeof(double) * (previous + new))",
     ),
-    # Free
-    (Free(Variable("x")), "free(x)"),
     # Declaration and types
     (Declaration(Variable("x"), integer), "int32_t x"),
     (Declaration(Variable("x"), float), "double x"),

tests/ir/test_peephole.py

@@ -92,10 +92,6 @@
         ArrayReallocate(Variable("x"), float, Add(IntegerLiteral(0), Variable("y"))),
         ArrayReallocate(Variable("x"), float, Variable("y")),
     ),
-    (
-        Free(ArrayIndex(Variable("x"), Add(IntegerLiteral(0), Variable("y")))),
-        Free(ArrayIndex(Variable("x"), Variable("y"))),
-    ),
     (
         Assignment(Variable("x"), Add(IntegerLiteral(0), Variable("y"))),
         Assignment(Variable("x"), Variable("y")),