Enable use of mlir pass manager in aiecc (#628)

* Enable use of mlir pass manager in aiecc * clang-format * limit scope of mlir context, rebase * fixup
Xilinx · Sep 25, 2023 · 346c222 · 346c222
1 parent cc42e2b
commit 346c222
Show file tree

Hide file tree

Showing 27 changed files with 114 additions and 50 deletions.
diff --git a/include/aie-c/Dialects.h b/include/aie-c/Dialects.h
@@ -15,6 +15,8 @@ extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(AIE, aie);
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(AIEVec, aievec);
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(AIEX, aiex);
 
 //===---------------------------------------------------------------------===//
 // ObjectFifoType

diff --git a/include/aie/InitialAllDialect.h b/include/aie/InitialAllDialect.h
@@ -17,6 +17,8 @@
 #include "aie/Dialect/ADF/ADFDialect.h"
 #include "aie/Dialect/AIE/IR/AIEDialect.h"
 #include "aie/Dialect/AIEVec/IR/AIEVecDialect.h"
+#include "aie/Dialect/AIEX/IR/AIEXDialect.h"
+
 #include "mlir/IR/Dialect.h"
 
 namespace xilinx {
@@ -27,7 +29,8 @@ inline void registerAllDialects(mlir::DialectRegistry &registry) {
   registry.insert<
     ADF::ADFDialect,
     aievec::AIEVecDialect,
-    AIE::AIEDialect
+    AIE::AIEDialect,
+    AIEX::AIEXDialect
   >();
   // clang-format on
 }

diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt
@@ -8,5 +8,12 @@ add_mlir_public_c_api_library(
 
   LINK_LIBS PUBLIC
   AIE
+  AIEX
   ADF
-  MLIRAIEVec)
+  MLIRAIEVec
+  AIETransforms
+  AIEUtils
+  MLIRAIEVecTransforms
+  MLIRAIEVecUtils
+  AIEXTransforms
+  AIEXUtils)
diff --git a/lib/CAPI/Dialects.cpp b/lib/CAPI/Dialects.cpp
@@ -8,9 +8,15 @@
 #include "aie-c/Dialects.h"
 
 #include "aie/Dialect/AIE/IR/AIEDialect.h"
+#include "aie/Dialect/AIEVec/IR/AIEVecDialect.h"
+#include "aie/Dialect/AIEX/IR/AIEXDialect.h"
+
 #include "mlir/CAPI/Registration.h"
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(AIE, aie, xilinx::AIE::AIEDialect)
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(AIEX, aiex, xilinx::AIEX::AIEXDialect)
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(AIEVec, aievec,
+                                      xilinx::aievec::AIEVecDialect)
 
 //===---------------------------------------------------------------------===//
 // ObjectFifoType

diff --git a/lib/CAPI/Registration.cpp b/lib/CAPI/Registration.cpp
@@ -6,16 +6,24 @@
 //===----------------------------------------------------------------------===//
 
 #include "aie-c/Registration.h"
+
 #include "aie/InitialAllDialect.h"
 
+#include "aie/Dialect/AIEVec/Analysis/Passes.h"
+#include "aie/Dialect/AIEVec/Pipelines/Passes.h"
+#include "aie/Dialect/AIEVec/Transforms/Passes.h"
+
 #include "mlir/CAPI/IR.h"
-// #include "mlir/InitAllPasses.h"
 
 void aieRegisterAllDialects(MlirContext context) {
   mlir::DialectRegistry registry;
   xilinx::registerAllDialects(registry);
 }
 
 void aieRegisterAllPasses() {
-  // xilinx::AIE::registerAllPasses();
+  xilinx::AIE::registerAIEPasses();
+  xilinx::AIEX::registerAIEXPasses();
+  xilinx::aievec::registerAIEVecAnalysisPasses();
+  xilinx::aievec::registerAIEVecPasses();
+  xilinx::aievec::registerAIEVecPipelines();
 }
diff --git a/python/AIEMLIRModule.cpp b/python/AIEMLIRModule.cpp
@@ -15,7 +15,7 @@ using namespace mlir::python::adaptors;
 
 PYBIND11_MODULE(_aieMlir, m) {
 
-  //::aieRegisterAllPasses();
+  ::aieRegisterAllPasses();
 
   m.doc() = R"pbdoc(
     AIE MLIR Python bindings
@@ -30,10 +30,16 @@ PYBIND11_MODULE(_aieMlir, m) {
   m.def(
       "register_dialect",
       [](MlirContext context, bool load) {
-        MlirDialectHandle handle = mlirGetDialectHandle__aie__();
-        mlirDialectHandleRegisterDialect(handle, context);
+        MlirDialectHandle aie_handle = mlirGetDialectHandle__aie__();
+        mlirDialectHandleRegisterDialect(aie_handle, context);
+        MlirDialectHandle aiex_handle = mlirGetDialectHandle__aiex__();
+        mlirDialectHandleRegisterDialect(aiex_handle, context);
+        MlirDialectHandle aievec_handle = mlirGetDialectHandle__aievec__();
+        mlirDialectHandleRegisterDialect(aievec_handle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(aie_handle, context);
+          mlirDialectHandleLoadDialect(aiex_handle, context);
+          mlirDialectHandleLoadDialect(aievec_handle, context);
         }
       },
       py::arg("context"), py::arg("load") = true);
@@ -48,7 +54,5 @@ PYBIND11_MODULE(_aieMlir, m) {
           "Get an instance of ObjectFifoType with given element type.",
           py::arg("self"), py::arg("type") = py::none());
 
-  // m.def("_register_all_passes", ::aieRegisterAllPasses);
-
   m.attr("__version__") = "dev";
 }
diff --git a/python/aie/compiler/aiecc/main.py b/python/aie/compiler/aiecc/main.py
@@ -15,12 +15,14 @@
 import platform
 import sys
 import time
-from subprocess import PIPE, run, call
-import tempfile
+import subprocess
 import shutil
-import timeit
 import asyncio
 
+from aie.mlir.passmanager import PassManager
+from aie.mlir.ir import Module, Context, Location
+from aie.dialects import aie as aiedialect
+
 import aie.compiler.aiecc.cl_arguments
 import aie.compiler.aiecc.configure
 
@@ -42,7 +44,8 @@
                   '--cse']
 
 class flow_runner:
-  def __init__(self, opts, tmpdirname):
+  def __init__(self, mlir_module_str, opts, tmpdirname):
+      self.mlir_module_str = mlir_module_str
       self.opts = opts
       self.tmpdirname = tmpdirname
       self.runtimes = dict()
@@ -83,9 +86,23 @@ async def do_call(self, task, command, force=False):
   def do_run(self, command):
       if(self.opts.verbose):
           print(" ".join(command))
-      ret = run(command, stdout=PIPE, stderr=PIPE, universal_newlines=True)
+      m = subprocess.PIPE
+      ret = subprocess.run(command, stdout=m, stderr=m, universal_newlines=True)
       return ret
 
+  def run_passes(self, pass_pipeline, mlir_module_str, outputfile=None):
+      if self.opts.verbose:
+        print("Running:", pass_pipeline)
+      with Context() as ctx, Location.unknown():
+        aiedialect.register_dialect(ctx)
+        module = Module.parse(mlir_module_str)
+        PassManager.parse(pass_pipeline).run(module.operation)
+        mlir_module_str = str(module)
+        if outputfile:
+          with open(outputfile, 'w') as g:
+            g.write(mlir_module_str)
+      return mlir_module_str
+
   def corefile(self, dirname, core, ext):
       (corecol, corerow, _) = core
       return os.path.join(dirname, 'core_%d_%d.%s' % (corecol, corerow, ext))
@@ -418,17 +435,19 @@ async def run_flow(self):
         progress_bar.task = progress_bar.add_task("[green] MLIR compilation:", total=1, command="1 Worker")
 
         self.file_with_addresses = os.path.join(self.tmpdirname, 'input_with_addresses.mlir')
-        await self.do_call(progress_bar.task, ['aie-opt',
-                                          '--lower-affine',
-                                          '--aie-canonicalize-device',
-                                          '--aie-assign-lock-ids',
-                                          '--aie-register-objectFifos',
-                                          '--aie-objectFifo-stateful-transform',
-                                          '--aie-lower-broadcast-packet',
-                                          '--aie-create-packet-flows',
-                                          '--aie-lower-multicast',
-                                          '--aie-assign-buffer-addresses',
-                                          '-convert-scf-to-cf', opts.filename, '-o', self.file_with_addresses], True)
+        pass_pipeline = ','.join(['lower-affine',
+                                  'aie-canonicalize-device',
+                                  'AIE.device('+
+                                    'aie-assign-lock-ids',
+                                    'aie-register-objectFifos',
+                                    'aie-objectFifo-stateful-transform',
+                                    'aie-lower-broadcast-packet',
+                                    'aie-create-packet-flows',
+                                    'aie-lower-multicast',
+                                    'aie-assign-buffer-addresses)',
+                                  'convert-scf-to-cf'])
+        self.run_passes('builtin.module('+pass_pipeline+')', self.mlir_module_str, self.file_with_addresses)
+
         t = self.do_run(['aie-translate', '--aie-generate-corelist', self.file_with_addresses])
         cores = eval(t.stdout)
         t = self.do_run(['aie-translate', '--aie-generate-target-arch', self.file_with_addresses])
@@ -478,10 +497,10 @@ def dumpprofile(self):
         if(i < len(sortedruntimes)):
           print("%.4f sec: %s" % (sortedruntimes[i][1], sortedruntimes[i][0]))
 
-
-def main(builtin_params={}):
+def run(mlir_module, args=None):
     global opts
-    opts = aie.compiler.aiecc.cl_arguments.parse_args()
+    if args is not None:
+      opts = aie.compiler.aiecc.cl_arguments.parse_args(args)
 
     is_windows = platform.system() == 'Windows'
 
@@ -539,8 +558,23 @@ def main(builtin_params={}):
     if(opts.verbose):
       print('created temporary directory', tmpdirname)
 
-    runner = flow_runner(opts, tmpdirname)
+    runner = flow_runner(mlir_module, opts, tmpdirname)
     asyncio.run(runner.run_flow())
 
     if(opts.profiling):
       runner.dumpprofile()
+
+def main():
+  global opts
+  opts = aie.compiler.aiecc.cl_arguments.parse_args()
+
+  try:
+    with Context() as ctx, Location.unknown():
+      aiedialect.register_dialect(ctx)
+      with open(opts.filename, 'r') as f:
+        module = Module.parse(f.read())
+      module_str = str(module)
+  except Exception as e:
+    print(e)
+    sys.exit(1)
+  run(module_str)
diff --git a/test/dialect/AIE/badbuffer-ve2802.mlir b/test/dialect/AIE/badbuffer-ve2802.mlir
@@ -10,7 +10,7 @@
 
 // RUN: not aiecc.py %s |& FileCheck %s
 // Row 2 is a memtile, not a coretile.
-// CHECK: error: 'AIE.buffer' op in Column 1 and Row 2 is accessed from an unreachable tile in Column 1 and Row 3
+// CHECK: error{{.*}}'AIE.buffer' op in Column 1 and Row 2 is accessed from an unreachable tile in Column 1 and Row 3
 
 module @test {
  AIE.device(xcve2802) {

diff --git a/test/dialect/AIE/badbuffer.mlir b/test/dialect/AIE/badbuffer.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.buffer' op in Column 1 and Row 1 is accessed from an unreachable tile in Column 4 and Row 4
+// CHECK: error{{.*}}'AIE.buffer' op in Column 1 and Row 1 is accessed from an unreachable tile in Column 4 and Row 4
 
 module @test {
   %t1 = AIE.tile(1, 1)

diff --git a/test/dialect/AIE/badconnect.mlir b/test/dialect/AIE/badconnect.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.connect' op source index cannot be less than zero
+// CHECK: error{{.*}} 'AIE.connect' op source index cannot be less than zero
 
 module {
   %20 = AIE.tile(2, 0)

diff --git a/test/dialect/AIE/badcore.mlir b/test/dialect/AIE/badcore.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.core' op failed to verify that op exists in a core tile
+// CHECK: error{{.*}}'AIE.core' op failed to verify that op exists in a core tile
 
 module @test {
   %t1 = AIE.tile(4, 0)

diff --git a/test/dialect/AIE/badcore2.mlir b/test/dialect/AIE/badcore2.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.core' op failed to verify that op exists in a core tile
+// CHECK: error{{.*}}'AIE.core' op failed to verify that op exists in a core tile
 
 module @test {
   AIE.device(xcve2802) {

diff --git a/test/dialect/AIE/badlock-vc1902.mlir b/test/dialect/AIE/badlock-vc1902.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.lock' op lock assigned invalid id (maximum is 15)
+// CHECK: error{{.*}}'AIE.lock' op lock assigned invalid id (maximum is 15)
 module @test {
   %t1 = AIE.tile(1, 1)
   %lock = AIE.lock(%t1, 16) { sym_name = "lock1" }

diff --git a/test/dialect/AIE/badlock-ve2802.mlir b/test/dialect/AIE/badlock-ve2802.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.lock' op lock assigned invalid id (maximum is 63)
+// CHECK: error{{.*}}'AIE.lock' op lock assigned invalid id (maximum is 63)
 module @test {
   AIE.device(xcve2802) {
     %t1 = AIE.tile(1, 1)

diff --git a/test/dialect/AIE/badlockdma.mlir b/test/dialect/AIE/badlockdma.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.lock' op in Column 4 and Row 4 is accessed from an unreachable tile in Column 1 and Row 1
+// CHECK: error{{.*}}'AIE.lock' op in Column 4 and Row 4 is accessed from an unreachable tile in Column 1 and Row 1
 module @test {
   %t1 = AIE.tile(1, 1)
   %t2 = AIE.tile(4, 4)

diff --git a/test/dialect/AIE/badlockfunc.mlir b/test/dialect/AIE/badlockfunc.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.lock' op is accessed outside of a tile
+// CHECK: error{{.*}}'AIE.lock' op is accessed outside of a tile
 module @test {
   %t1 = AIE.tile(1, 1)
   %t2 = AIE.tile(4, 4)

diff --git a/test/dialect/AIE/badmemtiledma_channel4buffer.mlir b/test/dialect/AIE/badmemtiledma_channel4buffer.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aie-opt --canonicalize %s |& FileCheck %s
-// CHECK:  error: 'AIE.dmaBd' op is reachable from DMA channel 4 and attempts to access a non-local buffer
+// CHECK:  error{{.*}}'AIE.dmaBd' op is reachable from DMA channel 4 and attempts to access a non-local buffer
 // CHECK: note: channel
 // CHECK:     AIE.dmaStart("MM2S", 4, ^bd1, ^dma1)
 // CHECK: note: buffer

diff --git a/test/dialect/AIE/badmemtiledma_channel4lock.mlir b/test/dialect/AIE/badmemtiledma_channel4lock.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aie-opt --canonicalize %s |& FileCheck %s
-// CHECK: error: 'AIE.useLock' op is reachable from DMA channel 4 and attempts to access a non-local lock
+// CHECK: error{{.*}}'AIE.useLock' op is reachable from DMA channel 4 and attempts to access a non-local lock
 // CHECK: note: channel
 // CHECK:     AIE.dmaStart("MM2S", 4, ^bd0, ^dma1)
 // CHECK: note: lock

diff --git a/test/dialect/AIE/badmemtiledma_neighboraccess.mlir b/test/dialect/AIE/badmemtiledma_neighboraccess.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.buffer' op in Column 3 and Row 1 is accessed from an unreachable tile in Column 1 and Row 1
+// CHECK: error{{.*}}'AIE.buffer' op in Column 3 and Row 1 is accessed from an unreachable tile in Column 1 and Row 1
 
 // memtiles can only access neighboring memtiles
 

diff --git a/test/dialect/AIE/badswitchbox_memtile_nofifo-ve2802.mlir b/test/dialect/AIE/badswitchbox_memtile_nofifo-ve2802.mlir
@@ -10,7 +10,7 @@
 
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.connect' op source bundle FIFO not supported
+// CHECK: error{{.*}}'AIE.connect' op source bundle FIFO not supported
 
 module {
   AIE.device(xcve2802) {

diff --git a/test/dialect/AIE/badswitchbox_shimtile_nodma-ve2802.mlir b/test/dialect/AIE/badswitchbox_shimtile_nodma-ve2802.mlir
@@ -10,7 +10,7 @@
 
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.connect' op source bundle DMA not supported
+// CHECK: error{{.*}}'AIE.connect' op source bundle DMA not supported
 
 module {
   AIE.device(xcve2802) {

diff --git a/test/dialect/AIE/badtile-ve2802.mlir b/test/dialect/AIE/badtile-ve2802.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.tile' op column index (50) must be less than the number of columns in the device (38)
+// CHECK: error{{.*}}'AIE.tile' op column index (50) must be less than the number of columns in the device (38)
 
 module @test {
  AIE.device(xcve2802) {

diff --git a/test/dialect/AIE/badtile.mlir b/test/dialect/AIE/badtile.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.tile' op attribute 'col' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0
+// CHECK: error{{.*}}'AIE.tile' op attribute 'col' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 0
 
 module @test {
   %t1 = AIE.tile(-1, -1)

diff --git a/test/dialect/AIE/badtile2.mlir b/test/dialect/AIE/badtile2.mlir
@@ -9,7 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: not aiecc.py %s |& FileCheck %s
-// CHECK: error: 'AIE.tile' op column index (50) must be less than the number of columns in the device (50)
+// CHECK: error{{.*}}'AIE.tile' op column index (50) must be less than the number of columns in the device (50)
 
 module @test {
   %t1 = AIE.tile(50, 50)