[mlir][python] Add basic python support for GPU dialect and passes

[nicolasvasilache]

step

[nicolasvasilache]

step

[gysit]

LGTM!

[nicolasvasilache]

rebase

mlir/include/mlir-c/Dialect/GPU.h

+//===-- 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

mlir/include/mlir/Dialect/GPU/CMakeLists.txt

 
 set(LLVM_TARGET_DEFINITIONS ParallelLoopMapperAttr.td)
 mlir_tablegen(ParallelLoopMapperEnums.h.inc -gen-enum-decls)
 mlir_tablegen(ParallelLoopMapperEnums.cpp.inc -gen-enum-defs)
 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)

mlir/lib/Bindings/Python/CMakeLists.txt

add_mlir_dialect_python_bindings(MLIRBindingsPythonAsyncOps
   DIALECT_NAME async_dialect)
 add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonAsyncOps)
 
 add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
   TD_FILE BuiltinOps.td
   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
   DEPENDS LinalgOdsGen)
 add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
 
 add_mlir_dialect_python_bindings(MLIRBindingsPythonMemRefOps
   TD_FILE MemRefOps.td

 add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
   INSTALL_DIR
     python
   SOURCES
     AsyncPasses.cpp
 )
 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
   SOURCES
     LinalgPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRLinalgPassesBindingsPythonExtension)

mlir/lib/Bindings/Python/GPUOps.td

+//===-- 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

mlir/lib/Bindings/Python/GPUPasses.cpp

+//===- 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 <pybind11/pybind11.h>
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirGPUPasses, m) {
+  m.doc() = "MLIR GPU Dialect Passes";
+
+  // Register all GPU passes on load.
+  mlirRegisterGPUPasses();
+}

mlir/lib/Bindings/Python/mlir/dialects/gpu/__init__.py

+#  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 *

mlir/lib/Bindings/Python/mlir/dialects/gpu/passes/__init__.py

+#  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")

mlir/lib/CAPI/Dialect/CMakeLists.txt

 # TODO: Make the check source feature optional as an argument on *_add_library.
 set(LLVM_OPTIONAL_SOURCES
   Async.cpp
   AsyncPasses.cpp
+  GPU.cpp
+  GPUPasses.cpp
   Linalg.cpp
   LinalgPasses.cpp
   SCF.cpp
   Shape.cpp
   Standard.cpp
   Tensor.cpp
 )
 
 add_mlir_public_c_api_library(MLIRCAPIAsync
   Async.cpp
   AsyncPasses.cpp
 
   DEPENDS
   MLIRAsyncPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRAsync
   MLIRAsyncTransforms
   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
 
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/CAPI/Dialect/GPU.cpp

+//===- 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)

mlir/test/Bindings/Python/dialects/gpu.py

+# 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)