From b87219f77e8b946c03c1aac8c357244ce9e262c6 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Wed, 28 Apr 2021 13:22:34 +0000 Subject: [PATCH] [mlir][python] Add basic python support for GPU dialect and passes Differential Revision: https://reviews.llvm.org/D101449 --- mlir/include/mlir-c/Dialect/GPU.h | 28 +++++++++++++++++++ mlir/include/mlir/Dialect/GPU/CMakeLists.txt | 2 ++ mlir/lib/Bindings/Python/CMakeLists.txt | 13 +++++++++ mlir/lib/Bindings/Python/GPUOps.td | 15 ++++++++++ mlir/lib/Bindings/Python/GPUPasses.cpp | 22 +++++++++++++++ .../Python/mlir/dialects/gpu/__init__.py | 5 ++++ .../mlir/dialects/gpu/passes/__init__.py | 6 ++++ mlir/lib/CAPI/Dialect/CMakeLists.txt | 15 ++++++++++ mlir/lib/CAPI/Dialect/GPU.cpp | 13 +++++++++ mlir/lib/CAPI/Dialect/GPUPasses.cpp | 26 +++++++++++++++++ mlir/test/Bindings/Python/dialects/gpu.py | 19 +++++++++++++ 11 files changed, 164 insertions(+) create mode 100644 mlir/include/mlir-c/Dialect/GPU.h create mode 100644 mlir/lib/Bindings/Python/GPUOps.td create mode 100644 mlir/lib/Bindings/Python/GPUPasses.cpp create mode 100644 mlir/lib/Bindings/Python/mlir/dialects/gpu/__init__.py create mode 100644 mlir/lib/Bindings/Python/mlir/dialects/gpu/passes/__init__.py create mode 100644 mlir/lib/CAPI/Dialect/GPU.cpp create mode 100644 mlir/lib/CAPI/Dialect/GPUPasses.cpp create mode 100644 mlir/test/Bindings/Python/dialects/gpu.py diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h new file mode 100644 index 000000000000..e4797a7ee4e8 --- /dev/null +++ b/mlir/include/mlir-c/Dialect/GPU.h @@ -0,0 +1,28 @@ +//===-- mlir-c/Dialect/GPU.h - C API for GPU dialect -------------*- C -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM +// Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===---------------------------------------------------------------------===// + +#ifndef MLIR_C_DIALECT_GPU_H +#define MLIR_C_DIALECT_GPU_H + +#include "mlir-c/Registration.h" +#include "mlir-c/Support.h" + +#ifdef __cplusplus +extern "C" { +#endif + +MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(GPU, gpu); + +#ifdef __cplusplus +} +#endif + +#include "mlir/Dialect/GPU/Passes.capi.h.inc" + +#endif // MLIR_C_DIALECT_GPU_H diff --git a/mlir/include/mlir/Dialect/GPU/CMakeLists.txt b/mlir/include/mlir/Dialect/GPU/CMakeLists.txt index 7db7596bdf10..73aa1d92ffc1 100644 --- a/mlir/include/mlir/Dialect/GPU/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/GPU/CMakeLists.txt @@ -18,6 +18,8 @@ add_public_tablegen_target(MLIRParallelLoopMapperEnumsGen) set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name GPU) +mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix GPU) +mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix GPU) add_public_tablegen_target(MLIRGPUPassIncGen) add_mlir_doc(Passes GPUPasses ./ -gen-pass-doc) diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt index eba4d2886ec2..bbccea63c461 100644 --- a/mlir/lib/Bindings/Python/CMakeLists.txt +++ b/mlir/lib/Bindings/Python/CMakeLists.txt @@ -41,6 +41,11 @@ add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps DIALECT_NAME builtin) add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonBuiltinOps) +add_mlir_dialect_python_bindings(MLIRBindingsPythonGPUOps + TD_FILE GPUOps.td + DIALECT_NAME gpu) +add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonGPUOps) + add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps TD_FILE LinalgOps.td DIALECT_NAME linalg @@ -133,6 +138,14 @@ add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasse ) add_dependencies(MLIRBindingsPythonExtension MLIRAsyncPassesBindingsPythonExtension) +add_mlir_python_extension(MLIRGPUPassesBindingsPythonExtension _mlirGPUPasses + INSTALL_DIR + python + SOURCES + GPUPasses.cpp +) +add_dependencies(MLIRBindingsPythonExtension MLIRGPUPassesBindingsPythonExtension) + add_mlir_python_extension(MLIRLinalgPassesBindingsPythonExtension _mlirLinalgPasses INSTALL_DIR python diff --git a/mlir/lib/Bindings/Python/GPUOps.td b/mlir/lib/Bindings/Python/GPUOps.td new file mode 100644 index 000000000000..bf0980f2930d --- /dev/null +++ b/mlir/lib/Bindings/Python/GPUOps.td @@ -0,0 +1,15 @@ +//===-- GPUOps.td - Entry point GPU_dialect bindings ------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===---------------------------------------------------------------------===// + +#ifndef PYTHON_BINDINGS_GPU_OPS +#define PYTHON_BINDINGS_GPU_OPS + +include "mlir/Bindings/Python/Attributes.td" +include "mlir/Dialect/GPU/GPUOps.td" + +#endif diff --git a/mlir/lib/Bindings/Python/GPUPasses.cpp b/mlir/lib/Bindings/Python/GPUPasses.cpp new file mode 100644 index 000000000000..cb623a11b6c4 --- /dev/null +++ b/mlir/lib/Bindings/Python/GPUPasses.cpp @@ -0,0 +1,22 @@ +//===- GPUPasses.cpp - Pybind module for the GPU passes ------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===---------------------------------------------------------------------===// + +#include "mlir-c/Dialect/GPU.h" + +#include + +// ----------------------------------------------------------------------------- +// Module initialization. +// ----------------------------------------------------------------------------- + +PYBIND11_MODULE(_mlirGPUPasses, m) { + m.doc() = "MLIR GPU Dialect Passes"; + + // Register all GPU passes on load. + mlirRegisterGPUPasses(); +} diff --git a/mlir/lib/Bindings/Python/mlir/dialects/gpu/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/gpu/__init__.py new file mode 100644 index 000000000000..67bf7bd854e1 --- /dev/null +++ b/mlir/lib/Bindings/Python/mlir/dialects/gpu/__init__.py @@ -0,0 +1,5 @@ +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from .._gpu_ops_gen import * diff --git a/mlir/lib/Bindings/Python/mlir/dialects/gpu/passes/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/gpu/passes/__init__.py new file mode 100644 index 000000000000..dd28e91a4646 --- /dev/null +++ b/mlir/lib/Bindings/Python/mlir/dialects/gpu/passes/__init__.py @@ -0,0 +1,6 @@ +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from ...._cext_loader import _load_extension +_cextGPUPasses = _load_extension("_mlirGPUPasses") diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt index 3f6265e8a2fe..dd9bd6f67881 100644 --- a/mlir/lib/CAPI/Dialect/CMakeLists.txt +++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt @@ -2,6 +2,8 @@ set(LLVM_OPTIONAL_SOURCES Async.cpp AsyncPasses.cpp + GPU.cpp + GPUPasses.cpp Linalg.cpp LinalgPasses.cpp SCF.cpp @@ -24,6 +26,19 @@ add_mlir_public_c_api_library(MLIRCAPIAsync MLIRPass ) +add_mlir_public_c_api_library(MLIRCAPIGPU + GPU.cpp + GPUPasses.cpp + + DEPENDS + MLIRGPUPassIncGen + + LINK_LIBS PUBLIC + MLIRCAPIIR + MLIRGPU + MLIRPass +) + add_mlir_public_c_api_library(MLIRCAPILinalg Linalg.cpp LinalgPasses.cpp diff --git a/mlir/lib/CAPI/Dialect/GPU.cpp b/mlir/lib/CAPI/Dialect/GPU.cpp new file mode 100644 index 000000000000..0de2cfa330ef --- /dev/null +++ b/mlir/lib/CAPI/Dialect/GPU.cpp @@ -0,0 +1,13 @@ +//===- GPUc.cpp - C Interface for GPU dialect ----------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir-c/Dialect/GPU.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Dialect/GPU/GPUDialect.h" + +MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(GPU, gpu, mlir::gpu::GPUDialect) diff --git a/mlir/lib/CAPI/Dialect/GPUPasses.cpp b/mlir/lib/CAPI/Dialect/GPUPasses.cpp new file mode 100644 index 000000000000..4ec167f8854f --- /dev/null +++ b/mlir/lib/CAPI/Dialect/GPUPasses.cpp @@ -0,0 +1,26 @@ +//===- GPUPasses.cpp - C API for GPU Dialect Passes ----------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/CAPI/Pass.h" +#include "mlir/Dialect/GPU/Passes.h" +#include "mlir/Pass/Pass.h" + +// Must include the declarations as they carry important visibility attributes. +#include "mlir/Dialect/GPU/Passes.capi.h.inc" + +using namespace mlir; + +#ifdef __cplusplus +extern "C" { +#endif + +#include "mlir/Dialect/GPU/Passes.capi.cpp.inc" + +#ifdef __cplusplus +} +#endif diff --git a/mlir/test/Bindings/Python/dialects/gpu.py b/mlir/test/Bindings/Python/dialects/gpu.py new file mode 100644 index 000000000000..edf59dfc9c8f --- /dev/null +++ b/mlir/test/Bindings/Python/dialects/gpu.py @@ -0,0 +1,19 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir.ir import * +import mlir.dialects.gpu +import mlir.dialects.gpu.passes +from mlir.passmanager import * + +def run(f): + print("\nTEST:", f.__name__) + f() + +def testGPUPass(): + with Context() as context: + PassManager.parse('gpu-kernel-outlining') + print('SUCCESS') + +# CHECK-LABEL: testGPUPass +# CHECK: SUCCESS +run(testGPUPass)