From c0d99c98d1c4a18ed0548b1e510252a614fa9733 Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Wed, 16 Oct 2024 20:52:56 -0400 Subject: [PATCH] [WIP] iree_gpu Python bindings --- compiler/bindings/c/BUILD.bazel | 1 + compiler/bindings/c/CMakeLists.txt | 1 + .../c/iree/compiler/dialects/iree_gpu.h | 37 ++++++++++++++++ compiler/bindings/python/CMakeLists.txt | 20 +++++++++ .../python/IREECompilerDialectsModule.cpp | 43 +++++++++++++++++++ .../iree/compiler/dialects/IREEGPUBinding.td | 12 ++++++ .../python/iree/compiler/dialects/iree_gpu.py | 9 ++++ .../bindings/python/test/ir/dialects_test.py | 27 +++++++++--- .../Codegen/Dialect/GPU/IR/BUILD.bazel | 1 + .../Codegen/Dialect/GPU/IR/CMakeLists.txt | 1 + .../Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp | 36 ++++++++++++++++ .../Codegen/Dialect/GPU/IR/IREEGPUAttrs.td | 2 +- 12 files changed, 182 insertions(+), 8 deletions(-) create mode 100644 compiler/bindings/c/iree/compiler/dialects/iree_gpu.h create mode 100644 compiler/bindings/python/IREECompilerDialectsModule.cpp create mode 100644 compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td create mode 100644 compiler/bindings/python/iree/compiler/dialects/iree_gpu.py diff --git a/compiler/bindings/c/BUILD.bazel b/compiler/bindings/c/BUILD.bazel index 21857d035db0..01c939c528f7 100644 --- a/compiler/bindings/c/BUILD.bazel +++ b/compiler/bindings/c/BUILD.bazel @@ -14,6 +14,7 @@ cc_library( name = "headers", hdrs = [ "iree/compiler/api_support.h", + "iree/compiler/dialects/iree_gpu.h", "iree/compiler/embedding_api.h", "iree/compiler/loader.h", "iree/compiler/mlir_interop.h", diff --git a/compiler/bindings/c/CMakeLists.txt b/compiler/bindings/c/CMakeLists.txt index 2ac9f5e9356c..f13bf426f6c8 100644 --- a/compiler/bindings/c/CMakeLists.txt +++ b/compiler/bindings/c/CMakeLists.txt @@ -22,6 +22,7 @@ iree_cc_library( headers HDRS "iree/compiler/api_support.h" + "iree/compiler/dialects/iree_gpu.h" "iree/compiler/embedding_api.h" "iree/compiler/loader.h" "iree/compiler/mlir_interop.h" diff --git a/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h new file mode 100644 index 000000000000..93c1e087b912 --- /dev/null +++ b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h @@ -0,0 +1,37 @@ +// Copyright 2024 The IREE Authors +// +// Licensed 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 IREE_COMPILER_DIALECTS_IREE_GPU_H +#define IREE_COMPILER_DIALECTS_IREE_GPU_H + +#include "mlir-c/IR.h" +#include "mlir-c/Support.h" +#include "mlir/CAPI/Wrap.h" + +#ifdef __cplusplus +extern "C" { +#endif + +MLIR_CAPI_EXPORTED bool +ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirAttribute +ireeGPUObjectAttrGet(MlirContext mlirCtx, bool prefetchSharedMemory, + bool noReduceSharedMemoryBankConflicts); + +MLIR_CAPI_EXPORTED bool +ireeGPUObjectAttrGetPrefetchSharedMemory(MlirAttribute attr); + +MLIR_CAPI_EXPORTED bool +ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirTypeID ireeGPUObjectAttrGetTypeID(void); + +#ifdef __cplusplus +} +#endif + +#endif // IREE_COMPILER_DIALECTS_IREE_GPU_H diff --git a/compiler/bindings/python/CMakeLists.txt b/compiler/bindings/python/CMakeLists.txt index e27e8e526216..0a0b36474885 100644 --- a/compiler/bindings/python/CMakeLists.txt +++ b/compiler/bindings/python/CMakeLists.txt @@ -100,6 +100,15 @@ declare_mlir_dialect_python_bindings( DIALECT_NAME vm ) +declare_mlir_dialect_python_bindings( + ADD_TO_PARENT IREEPythonSources.Dialects + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler" + TD_FILE dialects/IREEGPUBinding.td + GEN_ENUM_BINDINGS + SOURCES dialects/iree_gpu.py + DIALECT_NAME iree_gpu +) + declare_mlir_python_sources(IREECompilerAPIPythonCore ADD_TO_PARENT IREEPythonSources ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler" @@ -165,6 +174,17 @@ declare_mlir_python_extension(IREECompilerPythonExtensions.Registration LLVMSupport ) +declare_mlir_python_extension(IREECompilerPythonExtensions.CompilerDialects + MODULE_NAME _ireeCompilerDialects + ADD_TO_PARENT IREECompilerPythonExtensions + SOURCES + IREECompilerDialectsModule.cpp + EMBED_CAPI_LINK_LIBS + iree_compiler_API_SharedImpl + PRIVATE_LINK_LIBS + LLVMSupport +) + ################################################################################ # Generate packages and shared library ################################################################################ diff --git a/compiler/bindings/python/IREECompilerDialectsModule.cpp b/compiler/bindings/python/IREECompilerDialectsModule.cpp new file mode 100644 index 000000000000..754837c4bf12 --- /dev/null +++ b/compiler/bindings/python/IREECompilerDialectsModule.cpp @@ -0,0 +1,43 @@ +// Copyright 2024 The IREE Authors +// +// Licensed 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 "iree/compiler/dialects/iree_gpu.h" +#include "mlir-c/IR.h" +#include "mlir/Bindings/Python/PybindAdaptors.h" + +namespace py = pybind11; +using namespace mlir::python::adaptors; + +PYBIND11_MODULE(_ireeCompilerDialects, m) { + m.doc() = "iree-compiler dialects python extension"; + + //===-------------------------------------------------------------------===// + // GPUPipelineOptionsAttr + //===-------------------------------------------------------------------===// + + mlir_attribute_subclass(m, "GPUPipelineOptionsAttr", + ireeAttributeIsAGPUPipelineOptionsAttr, + ireeGPUObjectAttrGetTypeID) + .def_classmethod( + "get", + [](const py::object &, bool prefetchSharedMemory, + bool noReduceSharedMemoryBankConflicts, MlirContext ctx) { + return ireeGPUObjectAttrGet(ctx, prefetchSharedMemory, + noReduceSharedMemoryBankConflicts); + }, + "cls"_a, "prefetch_shared_memory"_a, + "no_reduce_shared_memory_bank_conflicts"_a, "ctx"_a = py::none(), + "Gets a gpu.pipeline_options from parameters.") + .def_property_readonly("prefetch_shared_memory", + [](MlirAttribute self) { + return ireeGPUObjectAttrGetPrefetchSharedMemory( + self); + }) + .def_property_readonly( + "no_reduce_shared_memory_bank_conflicts", [](MlirAttribute self) { + return ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(self); + }); +} diff --git a/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td b/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td new file mode 100644 index 000000000000..7674967caafd --- /dev/null +++ b/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td @@ -0,0 +1,12 @@ +// Copyright 2024 The IREE Authors +// +// Licensed 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_IREEGPU_OPS +#define PYTHON_BINDINGS_IREEGPU_OPS + +include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td" + +#endif // PYTHON_BINDINGS_IREEGPU_OPS diff --git a/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py b/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py new file mode 100644 index 000000000000..71c5d1b4a638 --- /dev/null +++ b/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py @@ -0,0 +1,9 @@ +# Copyright 2024 The IREE Authors +# +# Licensed 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 ._iree_gpu_ops_gen import * +from ._iree_gpu_enum_gen import * +from .._mlir_libs._ireeCompilerDialects import * diff --git a/compiler/bindings/python/test/ir/dialects_test.py b/compiler/bindings/python/test/ir/dialects_test.py index 63a960360e41..7c8f452c72d5 100644 --- a/compiler/bindings/python/test/ir/dialects_test.py +++ b/compiler/bindings/python/test/ir/dialects_test.py @@ -7,10 +7,23 @@ from iree.compiler import ir # Make sure that our dialects import. -from iree.compiler.dialects import ( - flow, - hal, - stream, - vm, - util, -) +from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu + +with ir.Context() as ctx, ir.Location.unknown(): + module = ir.Module.create() + with ir.InsertionPoint(module.body): + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True, False) + assert ( + repr(gpu_attr) + == "GPUPipelineOptionsAttr(#iree_gpu.pipeline_options)" + ) + assert ( + repr(type(gpu_attr)) + == "" + ) + assert gpu_attr.prefetch_shared_memory + assert not gpu_attr.no_reduce_shared_memory_bank_conflicts + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(False, True) + assert not gpu_attr.prefetch_shared_memory + assert gpu_attr.no_reduce_shared_memory_bank_conflicts diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel index 2e9d283a7756..961091ec43f0 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel @@ -87,6 +87,7 @@ iree_compiler_cc_library( "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect", "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils", "//compiler/src/iree/compiler/Dialect/LinalgExt/IR", + "//compiler/src/iree/compiler/bindings/c:headers", "@llvm-project//llvm:Support", "@llvm-project//mlir:AMDGPUDialect", "@llvm-project//mlir:AffineDialect", diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt index 4c160b495976..0a80fd68e71a 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt @@ -67,6 +67,7 @@ iree_cc_library( iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect iree::compiler::Codegen::Utils::VectorOpUtils iree::compiler::Dialect::LinalgExt::IR + iree::compiler::bindings::c::headers PUBLIC ) diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp index 471ba8b0c70c..40a47928712c 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp @@ -14,12 +14,16 @@ #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h" #include "iree/compiler/Codegen/Utils/VectorOpUtils.h" +#include "iree/compiler/dialects/iree_gpu.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLForwardCompat.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" #include "llvm/Support/ErrorHandling.h" +#include "mlir-c/IR.h" +#include "mlir/CAPI/IR.h" +#include "mlir/CAPI/Support.h" #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -1687,3 +1691,35 @@ void IREEGPUDialect::registerAttributes() { } } // namespace mlir::iree_compiler::IREE::GPU + +bool ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr) { + return llvm::isa( + unwrap(attr)); +} + +MlirAttribute ireeGPUObjectAttrGet(MlirContext mlirCtx, + bool prefetchSharedMemory, + bool noReduceSharedMemoryBankConflicts) { + mlir::MLIRContext *ctx = unwrap(mlirCtx); + return wrap(mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::get( + ctx, prefetchSharedMemory, noReduceSharedMemoryBankConflicts)); +} + +bool ireeGPUObjectAttrGetPrefetchSharedMemory(MlirAttribute attr) { + auto gpuAttr = + llvm::cast( + unwrap(attr)); + return gpuAttr.getPrefetchSharedMemory().getValue(); +} + +bool ireeGPUObjectAttrGetNoReduceSharedMemoryBankConflicts(MlirAttribute attr) { + auto gpuAttr = + llvm::cast( + unwrap(attr)); + return gpuAttr.getNoReduceSharedMemoryBankConflicts().getValue(); +} + +MlirTypeID ireeGPUObjectAttrGetTypeID(void) { + return wrap( + mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::getTypeID()); +} diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td index 138128988fb1..e0de6bb8c85b 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td @@ -482,7 +482,7 @@ def IREEGPU_LaneIdAttr : AttrDef { - let assemblyFormat = "``$value"; + let assemblyFormat = "$value"; let cppNamespace = "::mlir::iree_compiler::IREE::GPU"; }