[python][tuner] Add bindings for `iree_codegen.compilation_info` (#19129)
diff --git a/compiler/bindings/c/iree/compiler/dialects/iree_codegen.h b/compiler/bindings/c/iree/compiler/dialects/iree_codegen.h
index 357ac87..029e672 100644
--- a/compiler/bindings/c/iree/compiler/dialects/iree_codegen.h
+++ b/compiler/bindings/c/iree/compiler/dialects/iree_codegen.h
@@ -50,6 +50,22 @@
MLIR_CAPI_EXPORTED ireeCodegenTranslationInfoParameters
ireeCodegenTranslationInfoAttrGetParameters(MlirAttribute attr);
+MLIR_CAPI_EXPORTED bool
+ireeAttributeIsACodegenCompilationInfoAttr(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirTypeID ireeCodegenCompilationInfoAttrGetTypeID(void);
+
+struct ireeCodegenCompilationInfoParameters {
+ MlirAttribute loweringConfig;
+ MlirAttribute translationInfo;
+};
+
+MLIR_CAPI_EXPORTED MlirAttribute ireeCodegenCompilationInfoAttrGet(
+ MlirContext mlirCtx, ireeCodegenCompilationInfoParameters parameters);
+
+MLIR_CAPI_EXPORTED ireeCodegenCompilationInfoParameters
+ireeCodegenCompilationInfoAttrGetParameters(MlirAttribute attr);
+
#ifdef __cplusplus
}
#endif
diff --git a/compiler/bindings/python/IREECompilerDialectsModule.cpp b/compiler/bindings/python/IREECompilerDialectsModule.cpp
index 1f8be3f..7ece224 100644
--- a/compiler/bindings/python/IREECompilerDialectsModule.cpp
+++ b/compiler/bindings/python/IREECompilerDialectsModule.cpp
@@ -83,8 +83,7 @@
"cls"_a, "pass_pipeline"_a, "codegen_spec"_a = py::none(),
"workgroup_size"_a = py::none(), "subgroup_size"_a = py::none(),
"configuration"_a = py::none(), py::kw_only(), "ctx"_a = py::none(),
- "Gets an #iree_codegen.translation_info from "
- "parameters.")
+ "Gets an #iree_codegen.translation_info from parameters.")
.def_property_readonly(
"pass_pipeline",
[](MlirAttribute self) -> MlirAttribute {
@@ -124,6 +123,37 @@
return parameters.configuration;
});
+ //===-------------------------------------------------------------------===//
+ // CodegenCompilationInfoAttr
+ //===-------------------------------------------------------------------===//
+
+ mlir_attribute_subclass(iree_codegen_module, "CompilationInfoAttr",
+ ireeAttributeIsACodegenCompilationInfoAttr,
+ ireeCodegenCompilationInfoAttrGetTypeID)
+ .def_classmethod(
+ "get",
+ [](const py::object &, MlirAttribute loweringConfig,
+ MlirAttribute translationInfo, MlirContext ctx) {
+ ireeCodegenCompilationInfoParameters parameters = {};
+ parameters.loweringConfig = loweringConfig;
+ parameters.translationInfo = translationInfo;
+ return ireeCodegenCompilationInfoAttrGet(ctx, parameters);
+ },
+ "cls"_a, "lowering_config"_a, "translation_info"_a,
+ "ctx"_a = py::none(),
+ "Gets an #iree_codegen.compilation_info from parameters.")
+ .def_property_readonly(
+ "lowering_config",
+ [](MlirAttribute self) -> MlirAttribute {
+ auto parameters = ireeCodegenCompilationInfoAttrGetParameters(self);
+ return parameters.loweringConfig;
+ })
+ .def_property_readonly(
+ "translation_info", [](MlirAttribute self) -> MlirAttribute {
+ auto parameters = ireeCodegenCompilationInfoAttrGetParameters(self);
+ return parameters.translationInfo;
+ });
+
//===--------------------------------------------------------------------===//
auto iree_gpu_module =
diff --git a/compiler/bindings/python/test/ir/dialects_test.py b/compiler/bindings/python/test/ir/dialects_test.py
index 378d6ca..381ecea 100644
--- a/compiler/bindings/python/test/ir/dialects_test.py
+++ b/compiler/bindings/python/test/ir/dialects_test.py
@@ -215,3 +215,20 @@
assert lowering_config is not None
assert lowering_config.attributes == attributes
+
+
+@run
+def compilation_info():
+ attributes = ir.DictAttr.get({"reduction": ir.ArrayAttr.get([])})
+ lowering_config = iree_gpu.LoweringConfigAttr.get(attributes)
+ pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
+ iree_codegen.DispatchLoweringPassPipeline.None_
+ )
+ translation_info = iree_codegen.TranslationInfoAttr.get(pipeline_attr)
+
+ compilation_info = iree_codegen.CompilationInfoAttr.get(
+ lowering_config, translation_info
+ )
+ assert compilation_info is not None
+ assert compilation_info.lowering_config == lowering_config
+ assert compilation_info.translation_info == translation_info
diff --git a/compiler/src/iree/compiler/API/Internal/IREECodegenDialectCAPI.cpp b/compiler/src/iree/compiler/API/Internal/IREECodegenDialectCAPI.cpp
index 13f2257..c295d48 100644
--- a/compiler/src/iree/compiler/API/Internal/IREECodegenDialectCAPI.cpp
+++ b/compiler/src/iree/compiler/API/Internal/IREECodegenDialectCAPI.cpp
@@ -9,15 +9,20 @@
#include <optional>
#include <type_traits>
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/dialects/iree_codegen.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/IR.h"
#include "mlir/CAPI/IR.h"
#include "mlir/CAPI/Support.h"
+#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/MLIRContext.h"
+using mlir::iree_compiler::IREE::Codegen::CompilationInfoAttr;
using mlir::iree_compiler::IREE::Codegen::DispatchLoweringPassPipeline;
using mlir::iree_compiler::IREE::Codegen::DispatchLoweringPassPipelineAttr;
+using mlir::iree_compiler::IREE::Codegen::LoweringConfigAttrInterface;
using mlir::iree_compiler::IREE::Codegen::TranslationInfoAttr;
bool ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr(
@@ -109,3 +114,38 @@
return parameters;
}
+
+bool ireeAttributeIsACodegenCompilationInfoAttr(MlirAttribute attr) {
+ return llvm::isa<CompilationInfoAttr>(unwrap(attr));
+}
+
+MlirTypeID ireeCodegenCompilationInfoAttrGetTypeID() {
+ return wrap(CompilationInfoAttr::getTypeID());
+}
+
+MlirAttribute ireeCodegenCompilationInfoAttrGet(
+ MlirContext mlirCtx, ireeCodegenCompilationInfoParameters parameters) {
+ assert(!mlirAttributeIsNull(parameters.loweringConfig) &&
+ "Invalid lowering config attr");
+ assert(
+ !mlirAttributeIsNull(parameters.translationInfo) &&
+ ireeAttributeIsACodegenTranslationInfoAttr(parameters.translationInfo) &&
+ "Invalid translation info attr");
+
+ auto loweringConfig = llvm::cast<LoweringConfigAttrInterface>(
+ unwrap(parameters.loweringConfig));
+ auto translationInfo =
+ llvm::cast<TranslationInfoAttr>(unwrap(parameters.translationInfo));
+
+ mlir::MLIRContext *ctx = unwrap(mlirCtx);
+ return wrap(CompilationInfoAttr::get(ctx, loweringConfig, translationInfo));
+}
+
+ireeCodegenCompilationInfoParameters
+ireeCodegenCompilationInfoAttrGetParameters(MlirAttribute attr) {
+ auto compilationInfo = llvm::cast<CompilationInfoAttr>(unwrap(attr));
+ ireeCodegenCompilationInfoParameters parameters = {};
+ parameters.loweringConfig = wrap(compilationInfo.getLoweringConfig());
+ parameters.translationInfo = wrap(compilationInfo.getTranslationInfo());
+ return parameters;
+}
diff --git a/compiler/src/iree/compiler/API/api_exports.c b/compiler/src/iree/compiler/API/api_exports.c
index 9105596..ffb8f08 100644
--- a/compiler/src/iree/compiler/API/api_exports.c
+++ b/compiler/src/iree/compiler/API/api_exports.c
@@ -10,6 +10,7 @@
#include <stdint.h>
+extern void ireeAttributeIsACodegenCompilationInfoAttr();
extern void ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr();
extern void ireeAttributeIsACodegenTranslationInfoAttr();
extern void ireeAttributeIsAGPULoweringConfigAttr();
@@ -17,6 +18,9 @@
extern void ireeAttributeIsAGPUMMAIntrinsicAttr();
extern void ireeAttributeIsAGPUPipelineOptionsAttr();
extern void ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr();
+extern void ireeCodegenCompilationInfoAttrGet();
+extern void ireeCodegenCompilationInfoAttrGetParameters();
+extern void ireeCodegenCompilationInfoAttrGetTypeID();
extern void ireeCodegenDispatchLoweringPassPipelineAttrGet();
extern void ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID();
extern void ireeCodegenDispatchLoweringPassPipelineAttrGetValue();
@@ -868,6 +872,7 @@
uintptr_t __iree_compiler_hidden_force_extern() {
uintptr_t x = 0;
+ x += (uintptr_t)&ireeAttributeIsACodegenCompilationInfoAttr;
x += (uintptr_t)&ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr;
x += (uintptr_t)&ireeAttributeIsACodegenTranslationInfoAttr;
x += (uintptr_t)&ireeAttributeIsAGPULoweringConfigAttr;
@@ -875,6 +880,9 @@
x += (uintptr_t)&ireeAttributeIsAGPUMMAIntrinsicAttr;
x += (uintptr_t)&ireeAttributeIsAGPUPipelineOptionsAttr;
x += (uintptr_t)&ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr;
+ x += (uintptr_t)&ireeCodegenCompilationInfoAttrGet;
+ x += (uintptr_t)&ireeCodegenCompilationInfoAttrGetParameters;
+ x += (uintptr_t)&ireeCodegenCompilationInfoAttrGetTypeID;
x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGet;
x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID;
x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGetValue;
diff --git a/compiler/src/iree/compiler/API/api_exports.def b/compiler/src/iree/compiler/API/api_exports.def
index 9844a9f..ed5e12c 100644
--- a/compiler/src/iree/compiler/API/api_exports.def
+++ b/compiler/src/iree/compiler/API/api_exports.def
@@ -1,5 +1,6 @@
; Generated by generate_exports.py: Do not edit.
EXPORTS
+ ireeAttributeIsACodegenCompilationInfoAttr
ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr
ireeAttributeIsACodegenTranslationInfoAttr
ireeAttributeIsAGPULoweringConfigAttr
@@ -7,6 +8,9 @@
ireeAttributeIsAGPUMMAIntrinsicAttr
ireeAttributeIsAGPUPipelineOptionsAttr
ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr
+ ireeCodegenCompilationInfoAttrGet
+ ireeCodegenCompilationInfoAttrGetParameters
+ ireeCodegenCompilationInfoAttrGetTypeID
ireeCodegenDispatchLoweringPassPipelineAttrGet
ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID
ireeCodegenDispatchLoweringPassPipelineAttrGetValue
diff --git a/compiler/src/iree/compiler/API/api_exports.ld b/compiler/src/iree/compiler/API/api_exports.ld
index 0c36027..0808927 100644
--- a/compiler/src/iree/compiler/API/api_exports.ld
+++ b/compiler/src/iree/compiler/API/api_exports.ld
@@ -1,6 +1,7 @@
# Generated by generate_exports.py: Do not edit.
VER_0 {
global:
+ ireeAttributeIsACodegenCompilationInfoAttr;
ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr;
ireeAttributeIsACodegenTranslationInfoAttr;
ireeAttributeIsAGPULoweringConfigAttr;
@@ -8,6 +9,9 @@
ireeAttributeIsAGPUMMAIntrinsicAttr;
ireeAttributeIsAGPUPipelineOptionsAttr;
ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr;
+ ireeCodegenCompilationInfoAttrGet;
+ ireeCodegenCompilationInfoAttrGetParameters;
+ ireeCodegenCompilationInfoAttrGetTypeID;
ireeCodegenDispatchLoweringPassPipelineAttrGet;
ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID;
ireeCodegenDispatchLoweringPassPipelineAttrGetValue;
diff --git a/compiler/src/iree/compiler/API/api_exports.macos.lst b/compiler/src/iree/compiler/API/api_exports.macos.lst
index d0683e8..11169bf 100644
--- a/compiler/src/iree/compiler/API/api_exports.macos.lst
+++ b/compiler/src/iree/compiler/API/api_exports.macos.lst
@@ -1,4 +1,5 @@
# Generated by generate_exports.py: Do not edit.
+_ireeAttributeIsACodegenCompilationInfoAttr
_ireeAttributeIsACodegenDispatchLoweringPassPipelineAttr
_ireeAttributeIsACodegenTranslationInfoAttr
_ireeAttributeIsAGPULoweringConfigAttr
@@ -6,6 +7,9 @@
_ireeAttributeIsAGPUMMAIntrinsicAttr
_ireeAttributeIsAGPUPipelineOptionsAttr
_ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr
+_ireeCodegenCompilationInfoAttrGet
+_ireeCodegenCompilationInfoAttrGetParameters
+_ireeCodegenCompilationInfoAttrGetTypeID
_ireeCodegenDispatchLoweringPassPipelineAttrGet
_ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID
_ireeCodegenDispatchLoweringPassPipelineAttrGetValue
