Make apply-target extensible with new targets

ftn/tools/ftn_opt.py

@@ -3,10 +3,15 @@
 
 from xdsl.dialects.builtin import ModuleOp
 
+from typing import IO
+
 from ftn.transforms.rewrite_fir_to_core import RewriteFIRToCore
 from ftn.transforms.merge_memref_deref import MergeMemRefDeref
+from ftn.transforms.extract_target import ExtractTarget
+from ftn.transforms.fpga.target_to_hls import TargetToHLSPass
 from ftn.transforms.lower_omp_target_data import LowerOmpTargetDataPass
-# from ftn.transforms.extract_target import ExtractTarget
+from ftn.transforms.apply_target_config import ApplyTargetConfig
+from ftn.transforms.omp_target_to_kernel import OmpTargetToKernelPass
 # from ftn.transforms.isolate_target import IsolateTarget
 # from psy.extract_stencil import ExtractStencil
 # from ftn.transforms.tenstorrent.convert_to_tt import ConvertToTT
@@ -18,6 +23,7 @@
 
 from xdsl.xdsl_opt_main import xDSLOptMain
 from ftn.dialects import ftn_relative_cf
+from ftn.dialects import device
 
 import traceback
 
@@ -27,13 +33,23 @@ def register_all_passes(self):
         super().register_all_passes()
         self.register_pass("rewrite-fir-to-core", lambda: RewriteFIRToCore)
         self.register_pass("merge-memref-deref", lambda: MergeMemRefDeref)
+        self.register_pass("extract-target", lambda: ExtractTarget)
+        self.register_pass("target-to-hls", lambda: TargetToHLSPass)
         self.register_pass("lower-omp-target-data", lambda: LowerOmpTargetDataPass)
-        # self.register_pass("extract-target", lambda: ExtractTarget)
+        self.register_pass("apply-target", lambda: ApplyTargetConfig)
+        self.register_pass("omp-target-to-kernel", lambda: OmpTargetToKernelPass)
         # self.register_pass("isolate-target", lambda: IsolateTarget)
         # self.register_pass("convert-to-tt", lambda: ConvertToTT)
 
     def register_all_targets(self):
+        def _output_fpga_host(prog: ModuleOp, output: IO[str]):
+            from ftn.dialects.fpga.host_printer import HostPrinter
+
+            printer = HostPrinter(stream=output)
+            printer.print(prog)
+
         super().register_all_targets()
+        self.available_targets["fpga-host"] = _output_fpga_host
 
     def setup_pipeline(self):
         super().setup_pipeline()
@@ -50,6 +66,7 @@ def register_all_arguments(self, arg_parser: argparse.ArgumentParser):
     def register_all_dialects(self):
         super().register_all_dialects()
         self.ctx.load_dialect(ftn_relative_cf.Ftn_relative_cf)
+        self.ctx.load_dialect(device.Device)
 
     @staticmethod
     def get_passes_as_dict() -> Dict[str, Callable[[ModuleOp], None]]:

ftn/transforms/apply_target_config.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from abc import ABC, abstractmethod
 
 from xdsl.context import Context
 from xdsl.dialects import builtin, dlti
@@ -7,16 +8,32 @@
 from ftn.dialects import device
 
 
-class TenstorrentConfiguration:
-    def get():
+class TargetConfiguration(ABC):
+    @classmethod
+    @abstractmethod
+    def get(cls) -> dlti.TargetDeviceSpecAttr: ...
+
+    @classmethod
+    @abstractmethod
+    def _memory_subsystem(cls) -> dlti.MapAttr: ...
+
+    @classmethod
+    @abstractmethod
+    def _compute_subsystem(cls) -> dlti.MapAttr: ...
+
+
+class TenstorrentConfiguration(TargetConfiguration):
+    @classmethod
+    def get(cls):
         return dlti.TargetDeviceSpecAttr(
             {
-                "memory": TenstorrentConfiguration._memory_subsystem(),
-                "compute": TenstorrentConfiguration._compute_subsystem(),
+                "memory": cls._memory_subsystem(),
+                "compute": cls._compute_subsystem(),
             }
         )
 
-    def _memory_subsystem():
+    @classmethod
+    def _memory_subsystem(cls):
         config = {
             "DRAM": {
                 "kind": device.MemoryKindAttr(device.MemoryKind.DDR),
@@ -25,7 +42,8 @@ def _memory_subsystem():
         }
         return dlti.MapAttr(config)
 
-    def _compute_subsystem():
+    @classmethod
+    def _compute_subsystem(cls):
         config = {
             "architecture_type": device.ArchitectureKindAttr(
                 device.ArchitectureKind.MANYCORE
@@ -53,16 +71,18 @@ def _compute_subsystem():
         return dlti.MapAttr(config)
 
 
-class U280Configuration:
-    def get():
+class U280Configuration(TargetConfiguration):
+    @classmethod
+    def get(cls):
         return dlti.TargetDeviceSpecAttr(
             {
-                "memory": U280Configuration._memory_subsystem(),
-                "compute": U280Configuration._compute_subsystem(),
+                "memory": cls._memory_subsystem(),
+                "compute": cls._compute_subsystem(),
             }
         )
 
-    def _memory_subsystem():
+    @classmethod
+    def _memory_subsystem(cls):
         config = {
             "DRAM": {
                 "kind": device.MemoryKindAttr(device.MemoryKind.DDR),
@@ -76,7 +96,8 @@ def _memory_subsystem():
             }
         return dlti.MapAttr(config)
 
-    def _compute_subsystem():
+    @classmethod
+    def _compute_subsystem(cls):
         config = {
             "architecture_type": device.ArchitectureKindAttr(
                 device.ArchitectureKind.FPGA
@@ -153,14 +174,22 @@ def generate_system_config(self, accelerator_name, accelerator_config):
             }
         )
 
+    def _get_config(self) -> dlti.TargetDeviceSpecAttr:
+        """
+        Get the device spec for the current `self.taregt`
+
+        If overriding this function, make sure to *not* specify `name` field again
+        """
+        if config := SYSTEM_CONFIGURATIONS.get(self.target):
+            return config.get()
+        raise ValueError(f"No such target configuration {self.target}")
+
     def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
         op.attributes["omp.target_triples"] = builtin.ArrayAttr(
             [builtin.StringAttr(self.target)]
         )
 
-        assert self.target in SYSTEM_CONFIGURATIONS.keys()
-
-        config = SYSTEM_CONFIGURATIONS[self.target].get()
+        config = self._get_config()
 
         op.attributes["dlti.target_system_spec"] = self.generate_system_config(
             self.target, config