[Primitive] Extensible primitives and transformations in Python (cornell-zhang#504)

* [Primitive] Refactor

* Decouple primitives

* Remove redundant code

* Fix pylint

* Move context outside

* Add user-defined primitives

* Fix typo

* Add create_buffer

* Fix format

* Fix linting

* Update transform

Author: chhzh123

heterocl/__init__.py

@@ -2,7 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=redefined-builtin
 
-from .schedule import Schedule, customize, create_schedule, Partition
+from .schedule import Schedule, customize, create_schedule
+from .primitives.base import Primitive, register_primitive
+from .primitives.partition import Partition
 from .scheme import Scheme, create_scheme, create_schedule_from_scheme
 from .build_module import lower, build
 from .operation import *

heterocl/ast/ir_builder.py

@@ -362,40 +362,10 @@ def build_visitor(self, op, ip):
             self.build_op_handle(op, ip)
         elif isinstance(op, ast.LoopHandle):
             self.build_loop_handle(op, ip)
-        elif isinstance(op, ast.ReuseAtOp):
-            self.build_reuse_at_op(op, ip)
-        elif isinstance(op, ast.PartitionOp):
-            self.build_partition_op(op, ip)
-        elif isinstance(op, ast.ReplaceOp):
-            self.build_replace_op(op, ip)
-        elif isinstance(op, ast.ReshapeOp):
-            self.build_reshape_op(op, ip)
-        elif isinstance(op, ast.ReformOp):
-            self.build_reform_op(op, ip)
-        elif isinstance(op, ast.BufferAtOp):
-            self.build_buffer_at_op(op, ip)
         elif isinstance(op, ast.InterKernelToOp):
             self.build_inter_kernel_to_op(op, ip)
         elif isinstance(op, ast.OutlineOp):
             self.build_outline_op(op, ip)
-        elif isinstance(op, ast.ReorderOp):
-            self.build_reorder_op(op, ip)
-        elif isinstance(op, ast.SplitOp):
-            self.build_split_op(op, ip)
-        elif isinstance(op, ast.TileOp):
-            self.build_tile_op(op, ip)
-        elif isinstance(op, ast.PipelineOp):
-            self.build_pipeline_op(op, ip)
-        elif isinstance(op, ast.UnrollOp):
-            self.build_unroll_op(op, ip)
-        elif isinstance(op, ast.ParallelOp):
-            self.build_parallel_op(op, ip)
-        elif isinstance(op, ast.FuseOp):
-            self.build_fuse_op(op, ip)
-        elif isinstance(op, ast.ComputeAtOp):
-            self.build_compute_at_op(op, ip)
-        elif isinstance(op, ast.SystolicOp):
-            self.build_systolic_op(op, ip)
         else:
             raise HCLNotImplementedError(
                 f"{type(op)}'s build visitor is not implemented yet."
@@ -469,6 +439,7 @@ def build_func_op(self, op: ast.FuncOp, ip):
         # as the same argument object may be refered in multiple functions
         # we need to make sure that the result is not reused
         for arg in op.args:
+            arg.prev_result = arg.result
             arg.result = None
 
     def build_call_op(self, op: ast.CallOp, ip):
@@ -1610,75 +1581,6 @@ def build_loop_handle(self, op: ast.LoopHandle, ip):
         op.ir_op = hdl_op
         op.result = hdl_op.result
 
-    def build_partition_op(self, op: ast.PartitionOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        i32 = IntegerType.get_signless(32)
-        ui32 = IntegerType.get_unsigned(32)
-        partition_type = IntegerAttr.get(i32, op.kind)
-        dim = IntegerAttr.get(ui32, op.dim)
-        factor = IntegerAttr.get(ui32, op.factor)
-        self.build_visitor(op.tensor, ip)
-        partition_op = hcl_d.PartitionOp(
-            op.tensor.result,
-            partition_kind=partition_type,
-            dim=dim,
-            factor=factor,
-            ip=ip,
-            loc=loc,
-        )
-        op.ir_op = partition_op
-
-    def build_replace_op(self, op: ast.ReplaceOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        self.build_visitor(op.src, ip)
-        replace_op = hcl_d.ReplaceOp(op.target.result, op.src.result, ip=ip, loc=loc)
-        op.ir_op = replace_op
-
-    def build_reshape_op(self, op: ast.ReshapeOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.tensor, ip)
-        eletype = hcl_dtype_to_mlir(op.tensor.dtype)
-        memref_type = MemRefType.get(op.shape, eletype, loc=loc)
-        reshape_op = hcl_d.ReshapeOp(memref_type, op.tensor.result, ip=ip, loc=loc)
-        op.ir_op = reshape_op
-
-    def build_reform_op(self, op: ast.ReformOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        if op.layout == "nhwc":
-            attr = AffineMap.get_permutation([0, 2, 3, 1])
-        else:
-            raise RuntimeError("Not supported layout")
-        memref_type = MemRefType.get(op.target.shape, op.target.ir_op.dtype)
-        reform_op = hcl_d.ReformOp(memref_type, op.target.result, ip=ip, loc=loc)
-        reform_op.attributes["layout"] = AffineMapAttr.get(attr)
-        op.ir_op = reform_op
-
-    def build_reuse_at_op(self, op: ast.ReuseAtOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        self.build_visitor(op.axis, ip)
-        f32 = F32Type.get()
-        memref_type = MemRefType.get((1,), f32, loc=loc)
-        reuse_at_op = hcl_d.ReuseAtOp(
-            memref_type, op.target.result, op.axis.result, ip=ip, loc=loc
-        )
-        op.ir_op = reuse_at_op
-        op.result = reuse_at_op.result
-
-    def build_buffer_at_op(self, op: ast.BufferAtOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        self.build_visitor(op.axis, ip)
-        f32 = F32Type.get()
-        memref_type = MemRefType.get((1,), f32, loc=loc)
-        buffer_at_op = hcl_d.BufferAtOp(
-            memref_type, op.target.result, op.axis.result, ip=ip, loc=loc
-        )
-        op.ir_op = buffer_at_op
-        op.result = buffer_at_op.result
-
     def build_inter_kernel_to_op(self, op: ast.InterKernelToOp, ip):
         loc = Location.file(op.loc.filename, op.loc.lineno, 0)
         self.build_visitor(op.tensor, ip)
@@ -1704,79 +1606,3 @@ def build_outline_op(self, op: ast.OutlineOp, ip):
         if op.axis is not None:
             outline_op.attributes["axis"] = StringAttr.get(op.axis)
         op.ir_op = outline_op
-
-    def build_reorder_op(self, op: ast.ReorderOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        for arg in op.args:
-            self.build_visitor(arg, ip)
-        arg_results = [arg.result for arg in op.args]
-        reorder_op = hcl_d.ReorderOp(arg_results, ip=ip, loc=loc)
-        op.ir_op = reorder_op
-
-    def build_split_op(self, op: ast.SplitOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.parent, ip)
-        i32 = IntegerType.get_unsigned(32)
-        factor = IntegerAttr.get(i32, op.factor)
-        split_op = hcl_d.SplitOp(op.parent.result, factor, ip=ip, loc=loc)
-        op.ir_op = split_op
-        for result_loop_hdl, hdl_result in zip(op.results, split_op.results):
-            result_loop_hdl.result = hdl_result
-
-    def build_tile_op(self, op: ast.TileOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        i32 = IntegerType.get_unsigned(32)
-        x_factor = IntegerAttr.get(i32, op.x_factor)
-        y_factor = IntegerAttr.get(i32, op.y_factor)
-        self.build_visitor(op.x_parent, ip)
-        self.build_visitor(op.y_parent, ip)
-        tile_op = hcl_d.TileOp(
-            op.x_parent.result, op.y_parent.result, x_factor, y_factor, ip=ip, loc=loc
-        )
-        op.ir_op = tile_op
-        for result_loop_hdl, hdl_result in zip(op.results, tile_op.results):
-            result_loop_hdl.result = hdl_result
-
-    def build_pipeline_op(self, op: ast.PipelineOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        i32 = IntegerType.get_unsigned(32)
-        ii = IntegerAttr.get(i32, op.ii)
-        pipeline_op = hcl_d.PipelineOp(op.target.result, ii=ii, ip=ip, loc=loc)
-        op.ir_op = pipeline_op
-
-    def build_unroll_op(self, op: ast.UnrollOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        i32 = IntegerType.get_unsigned(32)
-        factor = IntegerAttr.get(i32, op.factor)
-        unroll_op = hcl_d.UnrollOp(op.target.result, factor=factor, ip=ip, loc=loc)
-        op.ir_op = unroll_op
-
-    def build_parallel_op(self, op: ast.ParallelOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.target, ip)
-        parallel_op = hcl_d.ParallelOp(op.target.result, ip=ip, loc=loc)
-        op.ir_op = parallel_op
-
-    def build_fuse_op(self, op: ast.FuseOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        for arg in op.arg_list:
-            self.build_visitor(arg, ip)
-        arg_results = [arg.result for arg in op.arg_list]
-        fuse_op = hcl_d.FuseOp(arg_results, ip=ip, loc=loc)
-        op.ir_op = fuse_op
-        op.result = fuse_op.result
-
-    def build_compute_at_op(self, op: ast.ComputeAtOp, ip):
-        loc = Location.file(op.loc.filename, op.loc.lineno, 0)
-        self.build_visitor(op.stage, ip)
-        self.build_visitor(op.parent, ip)
-        self.build_visitor(op.axis, ip)
-        compute_at_op = hcl_d.ComputeAtOp(
-            op.stage.result, op.parent.result, op.axis.result, ip=ip, loc=loc
-        )
-        op.ir_op = compute_at_op
-
-    def build_systolic_op(self, op: ast.SystolicOp, ip):
-        op.target.ir_op.attributes["systolic"] = UnitAttr.get()

heterocl/build_module.py

@@ -24,9 +24,6 @@
 from .runtime import copy_build_files
 from .schedule import Schedule
 from .utils import hcl_dtype_to_mlir
-from .passes.pass_manager import PassManager as ast_pass_manager
-from .passes.nest_if import NestElseIf
-from .passes.promote_func import PromoteFunc
 from .ast.ir_builder import IRBuilder
 from .ast.build_cleaner import ASTCleaner
 from .ast import ast
@@ -51,26 +48,7 @@ def lower(
     """Lowering step before build into target
     by applying optimization pass
     """
-    if schedule.is_lowered():
-        raise APIError(
-            "The module has been lowered. Please apply schedule primitives before the lowering process."
-        )
-    # HeteroCL Transformation Pipeline
-    ast_pm = ast_pass_manager()
-    ast_pm.add_pass(NestElseIf)
-    ast_pm.add_pass(PromoteFunc)
-    device_agnostic_ast = ast_pm.run(schedule.ast)
-    schedule._ast = device_agnostic_ast
-
-    # Build MLIR IR
-    set_context()
-    agnostic_ir_builder = IRBuilder(device_agnostic_ast)
-    agnostic_ir_builder.build()
-    agnostic_module = agnostic_ir_builder.module
-    schedule._module = _mlir_lower_pipeline(agnostic_module)
-    schedule._top_func = agnostic_ir_builder.top_func
-    exit_context()
-
+    schedule._module = _mlir_lower_pipeline(schedule._module)
     schedule.set_lowered()
     return schedule.module
 

heterocl/ir/__init__.py

@@ -0,0 +1,2 @@
+# Copyright HeteroCL authors. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

heterocl/ir/transform.py

@@ -0,0 +1,64 @@
+# Copyright HeteroCL authors. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+# pylint: disable=no-name-in-module, no-value-for-parameter
+
+import hcl_mlir
+from hcl_mlir import UnitAttr, StringAttr, InsertionPoint, MemRefType
+from hcl_mlir.dialects import memref as memref_d
+
+
+def get_affine_loop_nests(func):
+    loops = hcl_mlir.get_affine_loop_nests(func)
+    res = []
+    for loop in loops:
+        res.append([(item["name"], item["body"]) for item in loop])
+    return res
+
+
+def annotate(op, name):
+    op.attributes[name] = UnitAttr.get()
+
+
+def build_for_loops(grid, ip, name="loop"):
+    for_loops = []
+    if isinstance(name, str):
+        names = [name + f"_l_{i}" for i in range(len(grid))]
+        stage_name = "S_" + name
+    else:  # list
+        names = name
+        stage_name = "S_" + "_".join(names)
+    assert len(grid) >= 1
+
+    def recursive_for(for_handle, idx):
+        if idx == len(grid):
+            return
+        with InsertionPoint(for_handle.body.operations[0]):
+            new_for = hcl_mlir.make_for(0, grid[idx], name=names[idx])
+            for_loops.append(new_for)
+            recursive_for(new_for, idx + 1)
+
+    if not isinstance(ip, InsertionPoint):
+        ip = InsertionPoint(ip)
+    with ip:
+        for_handle = hcl_mlir.make_for(0, grid[0], name=names[0], stage=stage_name)
+    for_loops.append(for_handle)
+    recursive_for(for_handle, 1)
+    return for_loops
+
+
+def create_buffer(tensor, name, ip):
+    with InsertionPoint(ip):
+        alloc_op = memref_d.AllocOp(tensor.type, [], [])
+        alloc_op.attributes["name"] = StringAttr.get(name)
+        shape = MemRefType(tensor.type).shape
+    for_loops = build_for_loops(shape, ip, name)
+    induction_vars = [for_loop.induction_variable for for_loop in for_loops]
+    with InsertionPoint(for_loops[-1].body.operations[0]):
+        load = memref_d.LoadOp(tensor, induction_vars)
+        memref_d.StoreOp(
+            load.result,
+            alloc_op.result,
+            induction_vars,
+        )
+    # TODO: Upgrade LLVM version and use the following code
+    # tensor.replace_all_uses_with(alloc_op.result)

heterocl/primitives/__init__.py

@@ -0,0 +1,21 @@
+# Copyright HeteroCL authors. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Schedule primitives."""
+
+import os
+from os.path import abspath, dirname, join, isfile
+from inspect import getsourcefile, getmembers
+from importlib import import_module
+
+from .base import register_primitive, PRIMITIVES, Primitive
+
+path = dirname(abspath(getsourcefile(lambda: 0)))
+files = [
+    f
+    for f in os.listdir(path)
+    if isfile(join(path, f)) and f not in {"__init__.py", "base.py"}
+]
+for file in files:
+    mod = import_module(f".{file.split('.')[0]}", package="heterocl.primitives")
+    # register the schedule primitive using the decorator
+    getmembers(mod)

heterocl/primitives/base.py

@@ -0,0 +1,43 @@
+# Copyright HeteroCL authors. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Schedule primitive base."""
+# pylint: disable=no-method-argument
+
+from __future__ import annotations
+from abc import ABCMeta, abstractmethod
+
+PRIMITIVES = {}
+STAGE_PRIMITIVES = {}
+
+
+def register_primitive():
+    """Register a primitive to the schedule."""
+
+    def dectorator(cls):
+        if cls.name in PRIMITIVES:
+            raise ValueError(f"Primitive {cls.name} already registered")
+        if not issubclass(cls, Primitive):
+            raise ValueError(f"Class {cls} is not a subclass of Primitive")
+        if hasattr(cls, "is_stage_primitive") and cls.is_stage_primitive:
+            STAGE_PRIMITIVES[cls.name] = cls
+        else:
+            PRIMITIVES[cls.name] = cls
+        return cls
+
+    return dectorator
+
+
+class Primitive(metaclass=ABCMeta):
+    """A base class of schedule primitives."""
+
+    @property
+    @abstractmethod
+    def name():
+        """The name of the primitive."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def apply(sch, *args, **kwargs):
+        """Apply the primitive to the schedule."""
+        raise NotImplementedError