From 57bd497f83e91d288e670a4157cc2b0dbb4c2e18 Mon Sep 17 00:00:00 2001 From: Jeff Fifield Date: Wed, 27 Sep 2023 10:44:54 -0600 Subject: [PATCH] Enable AIEX dialect bindings (#658) * Enable AIEX dialect bindings * Replace 'Aie' prefix with 'AIE' in python cmake --- python/CMakeLists.txt | 60 +++++++++++++------ python/aie/compiler/CMakeLists.txt | 10 ++-- .../dialects/{AieBinding.td => AIEBinding.td} | 2 +- python/aie/dialects/AIEXBinding.td | 15 +++++ python/aie/dialects/_AIEX_ops_ext.py | 9 +++ python/aie/dialects/aie/__init__.py | 2 +- python/aie/dialects/aiex/__init__.py | 7 +++ test/python/aiex_ops.py | 29 +++++++++ 8 files changed, 108 insertions(+), 26 deletions(-) rename python/aie/dialects/{AieBinding.td => AIEBinding.td} (94%) create mode 100644 python/aie/dialects/AIEXBinding.td create mode 100644 python/aie/dialects/_AIEX_ops_ext.py create mode 100644 python/aie/dialects/aiex/__init__.py create mode 100644 test/python/aiex_ops.py diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 54391d74e0..4e6cc33fc1 100644 --- 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,11 +84,11 @@ 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 @@ -74,20 +96,20 @@ add_mlir_python_modules(AieMLIRPythonModules 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 index 73bfd25e85..1947116fbc 100644 --- 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 b/python/aie/dialects/AIEBinding.td similarity index 94% rename from python/aie/dialects/AieBinding.td rename to python/aie/dialects/AIEBinding.td index 5501e2fda7..b749f7aafd 100644 --- a/python/aie/dialects/AieBinding.td +++ b/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 \ No newline at end of file +#endif // AIE_BINDING_TD diff --git a/python/aie/dialects/AIEXBinding.td b/python/aie/dialects/AIEXBinding.td new file mode 100644 index 0000000000..d3dfca8ea8 --- /dev/null +++ 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 new file mode 100644 index 0000000000..d3060c7d5c --- /dev/null +++ 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 index ed4fcd8496..8402fe555d 100644 --- 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 * \ No newline at end of file +from ...mlir._mlir_libs._aieMlir import * diff --git a/python/aie/dialects/aiex/__init__.py b/python/aie/dialects/aiex/__init__.py new file mode 100644 index 0000000000..cee5235b7f --- /dev/null +++ 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 new file mode 100644 index 0000000000..fc6148cdd8 --- /dev/null +++ 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)