e3d7bf290d
This PR adds some LLVM metadata attributes and an `llvm.named_metadata` container op (similar to `llvm.module_flags`) for those attributes. Summary: - Add MLIR attributes modeling LLVM IR metadata: `#llvm.md_string`, `#llvm.md_const`, `#llvm.md_func`, and `#llvm.md_node`; - Add `llvm.named_metadata` container op for module-level named metadata nodes; - Add MLIR-to-LLVM-IR translation for the new attributes and op; - Add C API functions (`mlirLLVMMDStringAttrGet`, `mlirLLVMMDNodeAttrGet`, etc.); - Add Python bindings (`llvm.MDStringAttr`, `llvm.MDConstantAttr`, `llvm.MDFuncAttr`, `llvm.MDNodeAttr`, `llvm.FunctionType`).
414 lines
15 KiB
C++
414 lines
15 KiB
C++
//===- DialectLLVM.cpp - Pybind module for LLVM dialect API support -------===//
|
|
//
|
|
// 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 <string>
|
|
#include <vector>
|
|
|
|
#include "mlir-c/Dialect/LLVM.h"
|
|
#include "mlir-c/IR.h"
|
|
#include "mlir-c/Support.h"
|
|
#include "mlir-c/Target/LLVMIR.h"
|
|
#include "mlir/Bindings/Python/Diagnostics.h"
|
|
#include "mlir/Bindings/Python/IRCore.h"
|
|
#include "mlir/Bindings/Python/Nanobind.h"
|
|
#include "mlir/Bindings/Python/NanobindAdaptors.h"
|
|
|
|
namespace nb = nanobind;
|
|
|
|
using namespace nanobind::literals;
|
|
using namespace mlir;
|
|
using namespace mlir::python::nanobind_adaptors;
|
|
|
|
namespace mlir {
|
|
namespace python {
|
|
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
|
|
namespace llvm {
|
|
//===--------------------------------------------------------------------===//
|
|
// StructType
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
struct StructType : PyConcreteType<StructType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirTypeIsALLVMStructType;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMStructTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "StructType";
|
|
static inline const MlirStringRef name = mlirLLVMStructTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get_literal",
|
|
[](const std::vector<PyType> &elements, bool packed,
|
|
DefaultingPyLocation loc, DefaultingPyMlirContext context) {
|
|
python::CollectDiagnosticsToStringScope scope(
|
|
mlirLocationGetContext(loc));
|
|
std::vector<MlirType> elements_(elements.size());
|
|
std::copy(elements.begin(), elements.end(), elements_.begin());
|
|
|
|
MlirType type = mlirLLVMStructTypeLiteralGetChecked(
|
|
loc, elements.size(), elements_.data(), packed);
|
|
if (mlirTypeIsNull(type)) {
|
|
throw nb::value_error(scope.takeMessage().c_str());
|
|
}
|
|
return StructType(context->getRef(), type);
|
|
},
|
|
"elements"_a, nb::kw_only(), "packed"_a = false, "loc"_a = nb::none(),
|
|
"context"_a = nb::none());
|
|
|
|
c.def_static(
|
|
"get_literal_unchecked",
|
|
[](const std::vector<PyType> &elements, bool packed,
|
|
DefaultingPyMlirContext context) {
|
|
python::CollectDiagnosticsToStringScope scope(context.get()->get());
|
|
|
|
std::vector<MlirType> elements_(elements.size());
|
|
std::copy(elements.begin(), elements.end(), elements_.begin());
|
|
|
|
MlirType type = mlirLLVMStructTypeLiteralGet(
|
|
context.get()->get(), elements.size(), elements_.data(), packed);
|
|
if (mlirTypeIsNull(type)) {
|
|
throw nb::value_error(scope.takeMessage().c_str());
|
|
}
|
|
return StructType(context->getRef(), type);
|
|
},
|
|
"elements"_a, nb::kw_only(), "packed"_a = false,
|
|
"context"_a = nb::none());
|
|
|
|
c.def_static(
|
|
"get_identified",
|
|
[](const std::string &name, DefaultingPyMlirContext context) {
|
|
return StructType(context->getRef(),
|
|
mlirLLVMStructTypeIdentifiedGet(
|
|
context.get()->get(),
|
|
mlirStringRefCreate(name.data(), name.size())));
|
|
},
|
|
"name"_a, nb::kw_only(), "context"_a = nb::none());
|
|
|
|
c.def_static(
|
|
"get_opaque",
|
|
[](const std::string &name, DefaultingPyMlirContext context) {
|
|
return StructType(context->getRef(),
|
|
mlirLLVMStructTypeOpaqueGet(
|
|
context.get()->get(),
|
|
mlirStringRefCreate(name.data(), name.size())));
|
|
},
|
|
"name"_a, "context"_a = nb::none());
|
|
|
|
c.def(
|
|
"set_body",
|
|
[](const StructType &self, const std::vector<PyType> &elements,
|
|
bool packed) {
|
|
std::vector<MlirType> elements_(elements.size());
|
|
std::copy(elements.begin(), elements.end(), elements_.begin());
|
|
MlirLogicalResult result = mlirLLVMStructTypeSetBody(
|
|
self, elements.size(), elements_.data(), packed);
|
|
if (!mlirLogicalResultIsSuccess(result)) {
|
|
throw nb::value_error(
|
|
"Struct body already set to different content.");
|
|
}
|
|
},
|
|
"elements"_a, nb::kw_only(), "packed"_a = false);
|
|
|
|
c.def_static(
|
|
"new_identified",
|
|
[](const std::string &name, const std::vector<PyType> &elements,
|
|
bool packed, DefaultingPyMlirContext context) {
|
|
std::vector<MlirType> elements_(elements.size());
|
|
std::copy(elements.begin(), elements.end(), elements_.begin());
|
|
return StructType(context->getRef(),
|
|
mlirLLVMStructTypeIdentifiedNewGet(
|
|
context.get()->get(),
|
|
mlirStringRefCreate(name.data(), name.length()),
|
|
elements.size(), elements_.data(), packed));
|
|
},
|
|
"name"_a, "elements"_a, nb::kw_only(), "packed"_a = false,
|
|
"context"_a = nb::none());
|
|
|
|
c.def_prop_ro("name",
|
|
[](const StructType &type) -> std::optional<MlirStringRef> {
|
|
if (mlirLLVMStructTypeIsLiteral(type))
|
|
return std::nullopt;
|
|
|
|
return mlirLLVMStructTypeGetIdentifier(type);
|
|
});
|
|
|
|
c.def_prop_ro("body", [](const StructType &type) -> nb::object {
|
|
// Don't crash in absence of a body.
|
|
if (mlirLLVMStructTypeIsOpaque(type))
|
|
return nb::none();
|
|
|
|
nb::list body;
|
|
for (intptr_t i = 0, e = mlirLLVMStructTypeGetNumElementTypes(type);
|
|
i < e; ++i) {
|
|
body.append(mlirLLVMStructTypeGetElementType(type, i));
|
|
}
|
|
return body;
|
|
});
|
|
|
|
c.def_prop_ro("packed", [](const StructType &type) {
|
|
return mlirLLVMStructTypeIsPacked(type);
|
|
});
|
|
|
|
c.def_prop_ro("opaque", [](const StructType &type) {
|
|
return mlirLLVMStructTypeIsOpaque(type);
|
|
});
|
|
}
|
|
};
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// ArrayType
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
struct ArrayType : PyConcreteType<ArrayType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirTypeIsALLVMArrayType;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMArrayTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "ArrayType";
|
|
static inline const MlirStringRef name = mlirLLVMArrayTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](PyType &elementType, unsigned numElements) {
|
|
return ArrayType(elementType.getContext(),
|
|
mlirLLVMArrayTypeGet(elementType, numElements));
|
|
},
|
|
"element_type"_a, "num_elements"_a);
|
|
c.def_prop_ro("element_type", [](const ArrayType &type) {
|
|
return mlirLLVMArrayTypeGetElementType(type);
|
|
});
|
|
c.def_prop_ro("num_elements", [](const ArrayType &type) {
|
|
return mlirLLVMArrayTypeGetNumElements(type);
|
|
});
|
|
}
|
|
};
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// PointerType
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
struct PointerType : PyConcreteType<PointerType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirTypeIsALLVMPointerType;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMPointerTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "PointerType";
|
|
static inline const MlirStringRef name = mlirLLVMPointerTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](std::optional<unsigned> addressSpace,
|
|
DefaultingPyMlirContext context) {
|
|
python::CollectDiagnosticsToStringScope scope(context.get()->get());
|
|
MlirType type = mlirLLVMPointerTypeGet(
|
|
context.get()->get(),
|
|
addressSpace.has_value() ? *addressSpace : 0);
|
|
if (mlirTypeIsNull(type)) {
|
|
throw nb::value_error(scope.takeMessage().c_str());
|
|
}
|
|
return PointerType(context->getRef(), type);
|
|
},
|
|
"address_space"_a = nb::none(), nb::kw_only(),
|
|
"context"_a = nb::none());
|
|
c.def_prop_ro("address_space", [](const PointerType &type) {
|
|
return mlirLLVMPointerTypeGetAddressSpace(type);
|
|
});
|
|
}
|
|
};
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// FunctionType
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
struct FunctionType : PyConcreteType<FunctionType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirTypeIsALLVMFunctionType;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMFunctionTypeGetTypeID;
|
|
static constexpr const char *pyClassName = "FunctionType";
|
|
static inline const MlirStringRef name = mlirLLVMFunctionTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](PyType &resultType, const std::vector<PyType> &argumentTypes,
|
|
bool isVarArg) {
|
|
std::vector<MlirType> argTypes(argumentTypes.size());
|
|
std::copy(argumentTypes.begin(), argumentTypes.end(),
|
|
argTypes.begin());
|
|
return FunctionType(
|
|
resultType.getContext(),
|
|
mlirLLVMFunctionTypeGet(resultType, argTypes.size(),
|
|
argTypes.data(), isVarArg));
|
|
},
|
|
"result_type"_a, "argument_types"_a, nb::kw_only(),
|
|
"is_var_arg"_a = false);
|
|
c.def_prop_ro("return_type", [](const FunctionType &type) {
|
|
return mlirLLVMFunctionTypeGetReturnType(type);
|
|
});
|
|
c.def_prop_ro("num_inputs", [](const FunctionType &type) {
|
|
return mlirLLVMFunctionTypeGetNumInputs(type);
|
|
});
|
|
c.def_prop_ro("inputs", [](const FunctionType &type) {
|
|
nb::list inputs;
|
|
for (intptr_t i = 0, e = mlirLLVMFunctionTypeGetNumInputs(type); i < e;
|
|
++i) {
|
|
inputs.append(mlirLLVMFunctionTypeGetInput(type, i));
|
|
}
|
|
return inputs;
|
|
});
|
|
c.def_prop_ro("is_var_arg", [](const FunctionType &type) {
|
|
return mlirLLVMFunctionTypeIsVarArg(type);
|
|
});
|
|
}
|
|
};
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Metadata Attributes
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
struct MDStringAttr : PyConcreteAttribute<MDStringAttr> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirLLVMAttrIsAMDStringAttr;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMMDStringAttrGetTypeID;
|
|
static constexpr const char *pyClassName = "MDStringAttr";
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](const std::string &value, DefaultingPyMlirContext context) {
|
|
return MDStringAttr(
|
|
context->getRef(),
|
|
mlirLLVMMDStringAttrGet(
|
|
context.get()->get(),
|
|
mlirStringRefCreate(value.data(), value.size())));
|
|
},
|
|
"value"_a, nb::kw_only(), "context"_a = nb::none());
|
|
c.def_prop_ro("value", [](const MDStringAttr &self) {
|
|
MlirStringRef ref = mlirLLVMMDStringAttrGetValue(self);
|
|
return nb::str(ref.data, ref.length);
|
|
});
|
|
}
|
|
};
|
|
|
|
struct MDConstantAttr : PyConcreteAttribute<MDConstantAttr> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirLLVMAttrIsAMDConstantAttr;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMMDConstantAttrGetTypeID;
|
|
static constexpr const char *pyClassName = "MDConstantAttr";
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](PyAttribute &valueAttr, DefaultingPyMlirContext context) {
|
|
return MDConstantAttr(
|
|
context->getRef(),
|
|
mlirLLVMMDConstantAttrGet(context.get()->get(), valueAttr));
|
|
},
|
|
"value"_a, nb::kw_only(), "context"_a = nb::none());
|
|
c.def_prop_ro("value", [](const MDConstantAttr &self) {
|
|
return mlirLLVMMDConstantAttrGetValue(self);
|
|
});
|
|
}
|
|
};
|
|
|
|
struct MDFuncAttr : PyConcreteAttribute<MDFuncAttr> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirLLVMAttrIsAMDFuncAttr;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMMDFuncAttrGetTypeID;
|
|
static constexpr const char *pyClassName = "MDFuncAttr";
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](const std::string &name, DefaultingPyMlirContext context) {
|
|
MlirAttribute symRef = mlirFlatSymbolRefAttrGet(
|
|
context.get()->get(),
|
|
mlirStringRefCreate(name.data(), name.size()));
|
|
return MDFuncAttr(
|
|
context->getRef(),
|
|
mlirLLVMMDFuncAttrGet(context.get()->get(), symRef));
|
|
},
|
|
"name"_a, nb::kw_only(), "context"_a = nb::none());
|
|
c.def_prop_ro("name", [](const MDFuncAttr &self) {
|
|
MlirAttribute symRef = mlirLLVMMDFuncAttrGetName(self);
|
|
MlirStringRef ref = mlirFlatSymbolRefAttrGetValue(symRef);
|
|
return nb::str(ref.data, ref.length);
|
|
});
|
|
}
|
|
};
|
|
|
|
struct MDNodeAttr : PyConcreteAttribute<MDNodeAttr> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirLLVMAttrIsAMDNodeAttr;
|
|
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
|
|
mlirLLVMMDNodeAttrGetTypeID;
|
|
static constexpr const char *pyClassName = "MDNodeAttr";
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](const std::vector<PyAttribute> &operands,
|
|
DefaultingPyMlirContext context) {
|
|
std::vector<MlirAttribute> operands_(operands.size());
|
|
std::copy(operands.begin(), operands.end(), operands_.begin());
|
|
return MDNodeAttr(context->getRef(),
|
|
mlirLLVMMDNodeAttrGet(context.get()->get(),
|
|
operands_.size(),
|
|
operands_.data()));
|
|
},
|
|
"operands"_a, nb::kw_only(), "context"_a = nb::none());
|
|
c.def_prop_ro("num_operands", [](const MDNodeAttr &self) {
|
|
return mlirLLVMMDNodeAttrGetNumOperands(self);
|
|
});
|
|
c.def("__getitem__", [](const MDNodeAttr &self, intptr_t index) {
|
|
intptr_t n = mlirLLVMMDNodeAttrGetNumOperands(self);
|
|
if (index < 0 || index >= n)
|
|
throw nb::index_error("MDNodeAttr operand index out of range");
|
|
return mlirLLVMMDNodeAttrGetOperand(self, index);
|
|
});
|
|
c.def("__len__", [](const MDNodeAttr &self) {
|
|
return mlirLLVMMDNodeAttrGetNumOperands(self);
|
|
});
|
|
}
|
|
};
|
|
|
|
static void populateDialectLLVMSubmodule(nanobind::module_ &m) {
|
|
StructType::bind(m);
|
|
ArrayType::bind(m);
|
|
PointerType::bind(m);
|
|
FunctionType::bind(m);
|
|
MDStringAttr::bind(m);
|
|
MDConstantAttr::bind(m);
|
|
MDFuncAttr::bind(m);
|
|
MDNodeAttr::bind(m);
|
|
|
|
m.def(
|
|
"translate_module_to_llvmir",
|
|
[](const PyOperation &module) {
|
|
return mlirTranslateModuleToLLVMIRToString(module);
|
|
},
|
|
"module"_a, nb::rv_policy::take_ownership);
|
|
}
|
|
} // namespace llvm
|
|
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
|
|
} // namespace python
|
|
} // namespace mlir
|
|
|
|
NB_MODULE(_mlirDialectsLLVM, m) {
|
|
m.doc() = "MLIR LLVM Dialect";
|
|
|
|
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::llvm::populateDialectLLVMSubmodule(
|
|
m);
|
|
}
|