Enable AIEX dialect bindings (#658)

* Enable AIEX dialect bindings * Replace 'Aie' prefix with 'AIE' in python cmake
Xilinx · Sep 27, 2023 · 57bd497 · 57bd497
1 parent 2a6a112
commit 57bd497
Show file tree

Hide file tree

Showing 8 changed files with 108 additions and 26 deletions.
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -15,34 +15,56 @@ add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=aie.mlir.")
 # Sources
 ################################################################################
 
-declare_mlir_python_sources(AiePythonSources
+declare_mlir_python_sources(AIEPythonSources
   ROOT_DIR "${AIE_PYTHON_ROOT_DIR}"
   SOURCES
     dialects/_ods_common.py)
 
-declare_mlir_python_sources(AiePythonExtensions)
+declare_mlir_python_sources(AIEPythonExtensions)
 
-declare_mlir_python_sources(AiePythonSources.Dialects
-  ADD_TO_PARENT AiePythonSources
+declare_mlir_python_sources(AIEPythonSources.Dialects
+  ADD_TO_PARENT AIEPythonSources
 )
 
 declare_mlir_dialect_python_bindings(
-  ADD_TO_PARENT AiePythonSources.Dialects
+  ADD_TO_PARENT AIEPythonSources.Dialects
   ROOT_DIR "${AIE_PYTHON_ROOT_DIR}"
-  TD_FILE dialects/AieBinding.td
+  TD_FILE dialects/AIEBinding.td
   SOURCES
     dialects/aie/__init__.py
     dialects/_AIE_ops_ext.py
   DIALECT_NAME AIE
 )
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT AIEPythonSources.Dialects
+  ROOT_DIR "${AIE_PYTHON_ROOT_DIR}"
+  TD_FILE dialects/AIEXBinding.td
+  SOURCES
+    dialects/aiex/__init__.py
+    dialects/_AIEX_ops_ext.py
+  DIALECT_NAME AIEX
+)
+
 ################################################################################
 # Extensions
 ################################################################################
 
-declare_mlir_python_extension(AiePythonExtensions.MLIR
+declare_mlir_python_extension(AIEPythonExtensions.MLIR
   MODULE_NAME _aieMlir
-  ADD_TO_PARENT AiePythonExtensions
+  ADD_TO_PARENT AIEPythonExtensions
+  ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}
+  SOURCES
+    AIEMLIRModule.cpp
+  EMBED_CAPI_LINK_LIBS
+    AIECAPI
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+)
+
+declare_mlir_python_extension(AIEXPythonExtensions.MLIR
+  MODULE_NAME _aiexMlir
+  ADD_TO_PARENT AIEPythonExtensions
   ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}
   SOURCES
     AIEMLIRModule.cpp
@@ -52,8 +74,8 @@ declare_mlir_python_extension(AiePythonExtensions.MLIR
     LLVMSupport
 )
 
-add_mlir_python_common_capi_library(AieAggregateCAPI
-  INSTALL_COMPONENT AiePythonModules
+add_mlir_python_common_capi_library(AIEAggregateCAPI
+  INSTALL_COMPONENT AIEPythonModules
   INSTALL_DESTINATION python/aie/mlir/_mlir_libs
   OUTPUT_DIRECTORY "${AIE_PYTHON_PACKAGES_DIR}/aie/mlir/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
@@ -62,32 +84,32 @@ add_mlir_python_common_capi_library(AieAggregateCAPI
     MLIRPythonExtension.Core
     MLIRPythonExtension.RegisterEverything
     MLIRPythonExtension.ExecutionEngine
-    AiePythonSources
-    AiePythonExtensions
+    AIEPythonSources
+    AIEPythonExtensions
 )
 
-add_mlir_python_modules(AieMLIRPythonModules
+add_mlir_python_modules(AIEMLIRPythonModules
   ROOT_PREFIX "${AIE_PYTHON_PACKAGES_DIR}/aie/mlir"
   INSTALL_PREFIX "python/aie/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.Core
     MLIRPythonExtension.RegisterEverything
     MLIRPythonExtension.ExecutionEngine
-    AiePythonExtensions
+    AIEPythonExtensions
   COMMON_CAPI_LINK_LIBS
-    AieAggregateCAPI
+    AIEAggregateCAPI
     AIECAPI
   )
 
-add_mlir_python_modules(AiePythonModules
+add_mlir_python_modules(AIEPythonModules
   ROOT_PREFIX "${AIE_PYTHON_PACKAGES_DIR}/aie"
   INSTALL_PREFIX "python/aie"
   DECLARED_SOURCES
-    AiePythonSources
+    AIEPythonSources
   COMMON_CAPI_LINK_LIBS
-    AieAggregateCAPI
+    AIEAggregateCAPI
 )
 
-add_dependencies(AiePythonModules AieCompilerPythonModules)
+add_dependencies(AIEPythonModules AIECompilerPythonModules)
 add_subdirectory(aie/compiler)
diff --git a/python/aie/compiler/CMakeLists.txt b/python/aie/compiler/CMakeLists.txt
@@ -1,18 +1,18 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-declare_mlir_python_sources(AieCompilerPythonSources)
+declare_mlir_python_sources(AIECompilerPythonSources)
 
-declare_mlir_python_sources(AieCompilerPythonSources.Core
+declare_mlir_python_sources(AIECompilerPythonSources.Core
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
-  ADD_TO_PARENT AieCompilerPythonSources
+  ADD_TO_PARENT AIECompilerPythonSources
   SOURCES_GLOB
     aiecc/*.py
     *.py
 )
 
-add_mlir_python_modules(AieCompilerPythonModules
+add_mlir_python_modules(AIECompilerPythonModules
   ROOT_PREFIX "${AIE_PYTHON_PACKAGES_DIR}/aie/compiler"
   INSTALL_PREFIX "${AIE_PYTHON_INSTALL_DIR}/aie/compiler"
-  DECLARED_SOURCES AieCompilerPythonSources
+  DECLARED_SOURCES AIECompilerPythonSources
   )
diff --git a/python/aie/dialects/AieBinding.td → python/aie/dialects/AIEBinding.td b/python/aie/dialects/AieBinding.td → python/aie/dialects/AIEBinding.td
@@ -12,4 +12,4 @@
 include "mlir/Bindings/Python/Attributes.td"
 include "aie/Dialect/AIE/IR/AIE.td"
 
-#endif // AIE_BINDING_TD
+#endif // AIE_BINDING_TD
diff --git a/python/aie/dialects/AIEXBinding.td b/python/aie/dialects/AIEXBinding.td
@@ -0,0 +1,15 @@
+//===- AiexBinding.td --------------------------------------*- tablegen -*-===//
+//
+// Copyright (C) 2023, Advanced Micro Devices, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef AIEX_BINDING_TD
+#define AIEX_BINDING_TD
+
+include "mlir/Bindings/Python/Attributes.td"
+include "aie/Dialect/AIEX/IR/AIEX.td"
+
+#endif // AIEX_BINDING_TD
diff --git a/python/aie/dialects/_AIEX_ops_ext.py b/python/aie/dialects/_AIEX_ops_ext.py
@@ -0,0 +1,9 @@
+# ./python/aie/dialects/_AIEX_ops_ext.py -*- Python -*-
+
+# Copyright (C) 2023, Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+try:
+  from ..mlir.ir import *
+except ImportError as e:
+  raise RuntimeError("Error loading imports from extension module") from e
diff --git a/python/aie/dialects/aie/__init__.py b/python/aie/dialects/aie/__init__.py
@@ -4,4 +4,4 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._AIE_ops_gen import *
-from ...mlir._mlir_libs._aieMlir import *
+from ...mlir._mlir_libs._aieMlir import *
diff --git a/python/aie/dialects/aiex/__init__.py b/python/aie/dialects/aiex/__init__.py
@@ -0,0 +1,7 @@
+# ./python/aie/dialects/aiex/__init__.py -*- Python -*-
+
+# Copyright (C) 2023, Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._AIEX_ops_gen import *
+from ...mlir._mlir_libs._aieMlir import *
diff --git a/test/python/aiex_ops.py b/test/python/aiex_ops.py
@@ -0,0 +1,29 @@
+# Copyright (C) 2023, Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# RUN: python3 %s | FileCheck %s
+
+import aie
+from aie.mlir.ir import *
+from aie.dialects.aie import *
+from aie.dialects.aiex import *
+from aie.mlir.dialects import arith
+
+def constructAndPrintInModule(f):
+    with Context() as ctx, Location.unknown():
+        aie.dialects.aiex.register_dialect(ctx)
+        module = Module.create()
+        print("\nTEST:", f.__name__)
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+
+
+# CHECK-LABEL: getTileOp
+# CHECK: AIEX.getTile
+@constructAndPrintInModule
+def getTileOp():
+    iTy = IndexType.get()
+    four = arith.ConstantOp(iTy, 4)
+    two = arith.ConstantOp(iTy, 2)
+    GetTileOp(IndexType.get(), four, two)