llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/lib/Python/MlirIr.cpp

988 lines
38 KiB
C++
Raw Blame History

//===- MlirIr.cpp - MLIR IR Bindings --------------------------------------===//
//
// 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 "npcomp/Python/MlirIr.h"
#include "npcomp/Python/MlirInit.h"
#include "npcomp/Python/NpcompModule.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Location.h"
#include "mlir/Parser.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
namespace mlir {
//===----------------------------------------------------------------------===//
// Forward declarations
//===----------------------------------------------------------------------===//
struct PyContext;
/// Parses an MLIR module from a string.
/// For maximum efficiency, the `contents` should be zero terminated.
static OwningModuleRef parseMLIRModuleFromString(StringRef contents,
MLIRContext *context);
//===----------------------------------------------------------------------===//
// Direct type bindings
//===----------------------------------------------------------------------===//
static void bindInsertPoint(py::module m) {
py::class_<OpBuilder::InsertPoint>(m, "InsertPoint");
}
//===----------------------------------------------------------------------===//
// Internal only template definitions
// Since it is only legal to use explicit instantiations of templates in
// mlir_ir.h, implementations are kept in this module to keep things scoped
// well for the compiler.
//===----------------------------------------------------------------------===//
template <typename ListTy, typename ItemWrapperTy>
void PyIpListWrapper<ListTy, ItemWrapperTy>::bind(py::module m,
const char *className) {
struct PyItemIterator : public llvm::iterator_adaptor_base<
PyItemIterator, typename ListTy::iterator,
typename std::iterator_traits<
typename ListTy::iterator>::iterator_category,
typename ListTy::value_type> {
PyItemIterator() = default;
PyItemIterator(typename ListTy::iterator &&other)
: PyItemIterator::iterator_adaptor_base(std::move(other)) {}
ItemWrapperTy operator*() const { return ItemWrapperTy(*this->I); }
};
py::class_<ThisTy>(m, className)
.def_property_readonly(
"front",
[](ThisTy &self) { return ItemWrapperTy(self.list.front()); })
.def("__len__", [](ThisTy &self) { return self.list.size(); })
.def(
"__iter__",
[](ThisTy &self) {
PyItemIterator begin(self.list.begin());
PyItemIterator end(self.list.end());
return py::make_iterator(begin, end);
},
py::keep_alive<0, 1>());
}
//===----------------------------------------------------------------------===//
// Explicit template instantiations
//===----------------------------------------------------------------------===//
template class PyIpListWrapper<Region::BlockListType, PyBlockRef>;
using PyBlockList = PyIpListWrapper<Region::BlockListType, PyBlockRef>;
template class PyIpListWrapper<Block::OpListType, PyOperationRef>;
using PyOperationList = PyIpListWrapper<Block::OpListType, PyOperationRef>;
//===----------------------------------------------------------------------===//
// Conversions
//===----------------------------------------------------------------------===//
Type mapBufferFormatToType(MLIRContext *context, const std::string &format,
py::ssize_t itemSize) {
// Floating point formats.
if (format == "f")
return FloatType::getF32(context);
if (format == "d")
return FloatType::getF64(context);
if (format == "D")
return ComplexType::get(FloatType::getF64(context));
// Signed integer formats.
if (format == "b" || format == "h" || format == "i" || format == "l" ||
format == "L") {
unsigned width = itemSize * 8;
return IntegerType::get(context, width,
IntegerType::SignednessSemantics::Signed);
}
// Unsigned integer format.
if (format == "B" || format == "H" || format == "I" || format == "k" ||
format == "K") {
unsigned width = itemSize * 8;
return IntegerType::get(context, width,
IntegerType::SignednessSemantics::Unsigned);
}
return Type();
}
/// Creates a DenseElementsAttr from a python buffer which must have been
/// requested to be C-Contiguous.
Attribute createDenseElementsAttrFromBuffer(MLIRContext *context,
py::buffer_info &array) {
Type elementType =
mapBufferFormatToType(context, array.format, array.itemsize);
if (!elementType) {
throw py::raiseValueError(
"Unsupported buffer/array type for conversion to DenseElementsAttr");
}
SmallVector<int64_t, 4> shape(array.shape.begin(),
array.shape.begin() + array.ndim);
RankedTensorType type = RankedTensorType::get(shape, elementType);
const char *rawBufferPtr = reinterpret_cast<const char *>(array.ptr);
ArrayRef<char> rawBuffer(rawBufferPtr, array.size * array.itemsize);
return DenseElementsAttr::getFromRawBuffer(type, rawBuffer, false);
}
//===----------------------------------------------------------------------===//
// Diagnostics
//===----------------------------------------------------------------------===//
/// RAII class to capture diagnostics for later reporting back to the python
/// layer.
class DiagnosticCapture {
public:
DiagnosticCapture(mlir::MLIRContext *mlir_context)
: mlir_context(mlir_context) {
handler_id = mlir_context->getDiagEngine().registerHandler(
[&](Diagnostic &d) -> LogicalResult {
diagnostics.push_back(std::move(d));
return success();
});
}
~DiagnosticCapture() {
if (mlir_context) {
mlir_context->getDiagEngine().eraseHandler(handler_id);
}
}
DiagnosticCapture(DiagnosticCapture &&other) {
mlir_context = other.mlir_context;
diagnostics.swap(other.diagnostics);
handler_id = other.handler_id;
other.mlir_context = nullptr;
}
std::vector<mlir::Diagnostic> &getDiagnostics() { return diagnostics; }
// Consumes/clears diagnostics.
std::string consumeDiagnosticsAsString(const char *error_message);
void clearDiagnostics() { diagnostics.clear(); }
private:
MLIRContext *mlir_context;
std::vector<mlir::Diagnostic> diagnostics;
mlir::DiagnosticEngine::HandlerID handler_id;
};
//===----------------------------------------------------------------------===//
// PyDialectHelper
//===----------------------------------------------------------------------===//
void PyDialectHelper::bind(py::module m) {
py::class_<PyDialectHelper>(m, "DialectHelper")
.def(py::init<PyContext &, PyOpBuilder &>(), py::keep_alive<1, 2>(),
py::keep_alive<1, 3>())
.def_property_readonly("builder",
[](PyDialectHelper &self) -> PyBaseOpBuilder & {
return self.pyOpBuilder;
})
.def_property_readonly(
"context",
[](PyDialectHelper &self) -> std::shared_ptr<PyContext> {
return self.context.shared_from_this();
})
.def(
"op",
[](PyDialectHelper &self, const std::string &opNameStr,
std::vector<PyType> pyResultTypes, std::vector<PyValue> pyOperands,
llvm::Optional<PyAttribute> attrs) -> PyOperationRef {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(false);
Location loc = self.pyOpBuilder.getCurrentLoc();
OperationName opName(opNameStr, opBuilder.getContext());
SmallVector<Type, 4> types(pyResultTypes.begin(),
pyResultTypes.end());
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
DictionaryAttr attrList;
if (attrs) {
auto dictAttrs = attrs->attr.dyn_cast<DictionaryAttr>();
if (!dictAttrs) {
throw py::raiseValueError(
"Expected `attrs` to be a DictionaryAttr");
}
attrList = dictAttrs;
} else {
attrList = DictionaryAttr::get({}, self.getContext());
}
Operation *op =
Operation::create(loc, opName, types, operands, attrList);
opBuilder.insert(op);
return op;
},
py::arg("op_name"), py::arg("result_types"), py::arg("operands"),
py::arg("attrs") = llvm::Optional<PyAttribute>())
.def(
"func_op",
[](PyDialectHelper &self, const std::string &name, PyType type,
bool createEntryBlock, llvm::Optional<PyAttribute> attrs) {
auto functionType = type.type.dyn_cast_or_null<FunctionType>();
if (!functionType) {
throw py::raiseValueError("Illegal function type");
}
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc();
// TODO: Dedup attr creation from op().
DictionaryAttr attrList;
if (attrs) {
auto dictAttrs = attrs->attr.dyn_cast<DictionaryAttr>();
if (!dictAttrs) {
throw py::raiseValueError(
"Expected `attrs` to be a DictionaryAttr");
}
attrList = dictAttrs;
} else {
attrList = DictionaryAttr::get({}, self.getContext());
}
FuncOp op =
opBuilder.create<FuncOp>(loc, StringRef(name), functionType,
/*attrs=*/attrList.getValue());
if (createEntryBlock) {
Block *entryBlock = new Block();
entryBlock->addArguments(functionType.getInputs());
op.getBody().push_back(entryBlock);
opBuilder.setInsertionPointToStart(entryBlock);
}
return PyOperationRef(op);
},
py::arg("name"), py::arg("type"),
py::arg("create_entry_block") = false,
py::arg("attrs") = llvm::Optional<PyAttribute>(),
R"(Creates a new `func` op, optionally creating an entry block.
If an entry block is created, the builder will be positioned
to its start.)")
.def(
"select_op",
[](PyDialectHelper &self, PyValue conditionValue, PyValue trueValue,
PyValue falseValue) -> PyOperationRef {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc();
return PyOperationRef(opBuilder.create<SelectOp>(
loc, conditionValue, trueValue, falseValue));
},
py::arg("condition"), py::arg("true_value"), py::arg("false_value"))
.def("return_op",
[](PyDialectHelper &self, std::vector<PyValue> pyOperands) {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc();
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
return PyOperationRef(opBuilder.create<ReturnOp>(loc, operands));
})
.def("constant_op",
[](PyDialectHelper &self, PyType type, PyAttribute value) {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = self.pyOpBuilder.getCurrentLoc();
return PyOperationRef(
opBuilder.create<ConstantOp>(loc, type.type, value.attr));
})
// Types.
.def_property_readonly("index_type",
[](PyDialectHelper &self) -> PyType {
return IndexType::get(self.getContext());
})
.def(
"integer_type",
[](PyDialectHelper &self, unsigned width) -> PyType {
return IntegerType::get(self.getContext(), width);
},
py::arg("width") = 32)
.def_property_readonly("i1_type",
[](PyDialectHelper &self) -> PyType {
return IntegerType::get(self.getContext(), 1);
})
.def_property_readonly("i16_type",
[](PyDialectHelper &self) -> PyType {
return IntegerType::get(self.getContext(), 32);
})
.def_property_readonly("i32_type",
[](PyDialectHelper &self) -> PyType {
return IntegerType::get(self.getContext(), 32);
})
.def_property_readonly("i64_type",
[](PyDialectHelper &self) -> PyType {
return IntegerType::get(self.getContext(), 64);
})
.def_property_readonly("f32_type",
[](PyDialectHelper &self) -> PyType {
return FloatType::getF32(self.getContext());
})
.def_property_readonly("f64_type",
[](PyDialectHelper &self) -> PyType {
return FloatType::getF64(self.getContext());
})
.def(
"tensor_type",
[](PyDialectHelper &self, PyType elementType,
llvm::Optional<std::vector<int64_t>> shape) -> PyType {
if (!elementType.type) {
throw py::raiseValueError("Null element type");
}
if (shape) {
return RankedTensorType::get(*shape, elementType.type);
} else {
return UnrankedTensorType::get(elementType.type);
}
},
py::arg("element_type"),
py::arg("shape") = llvm::Optional<std::vector<int64_t>>())
.def("function_type",
[](PyDialectHelper &self, std::vector<PyType> inputs,
std::vector<PyType> results) -> PyType {
llvm::SmallVector<Type, 4> inputTypes;
llvm::SmallVector<Type, 1> resultTypes;
for (auto input : inputs) {
inputTypes.push_back(input.type);
}
for (auto result : results) {
resultTypes.push_back(result.type);
}
return FunctionType::get(self.getContext(), inputTypes,
resultTypes);
});
}
//===----------------------------------------------------------------------===//
// Module initialization
//===----------------------------------------------------------------------===//
static void emitDiagnostic(DiagnosticSeverity severity, PyAttribute loc,
std::string &message) {
auto locAttr = loc.attr.dyn_cast_or_null<LocationAttr>();
if (!locAttr) {
throw py::raiseValueError("Expected a LocationAttr");
}
auto &diagEngine = locAttr.getContext()->getDiagEngine();
diagEngine.emit(Location(locAttr), severity) << message;
}
void defineMlirIrModule(py::module m) {
m.doc() = "Python bindings for constructs in the mlir/IR library";
// Globals.
m.def(
"emit_error",
[](PyAttribute loc, std::string message) {
emitDiagnostic(DiagnosticSeverity::Error, loc, message);
},
py::arg("loc"), py::arg("message"));
m.def(
"emit_warning",
[](PyAttribute loc, std::string message) {
emitDiagnostic(DiagnosticSeverity::Warning, loc, message);
},
py::arg("loc"), py::arg("message"));
m.def(
"emit_remark",
[](PyAttribute loc, std::string message) {
emitDiagnostic(DiagnosticSeverity::Remark, loc, message);
},
py::arg("loc"), py::arg("message"));
// Python only types.
PyDialectHelper::bind(m);
// Utility types.
PyBlockList::bind(m, "BlockList");
PyOperationList::bind(m, "OperationList");
// Wrapper types.
PyAttribute::bind(m);
PyBaseOperation::bind(m);
PyBaseOpBuilder::bind(m);
PyBlockRef::bind(m);
PyContext::bind(m);
PyIdentifier::bind(m);
PyModuleOp::bind(m);
PyOperationRef::bind(m);
PyOpBuilder::bind(m);
PyRegionRef::bind(m);
PySymbolTable::bind(m);
PyType::bind(m);
PyValue::bind(m);
// Direct wrappings.
bindInsertPoint(m);
}
//===----------------------------------------------------------------------===//
// PyContext
//===----------------------------------------------------------------------===//
PyContext::PyContext() {
mlir::npcomp::python::loadGlobalDialectsIntoContext(&context);
}
void PyContext::bind(py::module m) {
py::class_<PyContext, std::shared_ptr<PyContext>>(m, "MLIRContext")
.def(py::init<>([]() {
// Need explicit make_shared to avoid UB with enable_shared_from_this.
return std::make_shared<PyContext>();
}))
.def("new_module",
[&](PyContext &self) -> PyModuleOp {
Location loc = UnknownLoc::get(&self.context);
auto m = ModuleOp::create(loc);
return PyModuleOp(self.shared_from_this(), m);
})
.def("parse_asm", &PyContext::parseAsm)
.def(
"new_builder",
[](PyContext &self) {
// Note: we collapse the Builder and OpBuilder into one because
// there is little reason to expose the inheritance hierarchy to
// Python.
return PyOpBuilder(self);
},
py::keep_alive<0, 1>())
.def("identifier",
[](PyContext &self, std::string s) -> PyIdentifier {
return Identifier::get(s, &self.context);
})
.def(
"file_line_col_loc_attr",
[](PyContext &self, PyIdentifier filename, unsigned line,
unsigned column) -> PyAttribute {
return static_cast<LocationAttr>(FileLineColLoc::get(
filename.identifier, line, column, &self.context));
},
py::arg("filename"), py::arg("line"), py::arg("column"))
// Salient functions from Builder.
.def("parse_type",
[](PyContext &self, const std::string &asmText) {
Type t = parseType(asmText, &self.context);
if (!t) {
std::string message = "Unable to parse MLIR type: ";
message.append(asmText);
throw py::raiseValueError(message);
}
return PyType(t);
})
.def(
"integer_attr",
[](PyContext &self, PyType type, int64_t value) -> PyAttribute {
if (!type.type.isa<IntegerType>()) {
throw py::raiseValueError("Expected IntegerType");
}
return IntegerAttr::get(type.type, value);
},
py::arg("type"), py::arg("value"))
.def("float_attr",
[](PyContext &self, PyType type, double value) -> PyAttribute {
if (!type.type.isa<FloatType>()) {
throw py::raiseValueError("Expected FloatType");
}
return FloatAttr::get(type.type, value);
})
.def("index_attr",
[](PyContext &self, int64_t indexValue) -> PyAttribute {
return IntegerAttr::get(IndexType::get(&self.context), indexValue);
})
.def("string_attr",
[](PyContext &self, const std::string &s) -> PyAttribute {
return StringAttr::get(s, &self.context);
})
.def("bytes_attr",
[](PyContext &self, py::bytes bytes) -> PyAttribute {
char *buffer;
ssize_t length;
if (PYBIND11_BYTES_AS_STRING_AND_SIZE(bytes.ptr(), &buffer,
&length)) {
throw py::raiseValueError("Cannot extract bytes");
}
return StringAttr::get(StringRef(buffer, length), &self.context);
})
.def("flat_symbol_ref_attr",
[](PyContext &self, const std::string &s) -> PyAttribute {
return FlatSymbolRefAttr::get(s, &self.context);
})
.def("dictionary_attr",
[](PyContext &self, py::dict d) -> PyAttribute {
SmallVector<NamedAttribute, 4> attrs;
for (auto &it : d) {
auto key = it.first.cast<std::string>();
auto value = it.second.cast<PyAttribute>();
auto keyIdent = Identifier::get(key, &self.context);
attrs.emplace_back(keyIdent, value.attr);
}
return DictionaryAttr::get(attrs, &self.context);
})
.def("array_attr",
[](PyContext &self, py::list l) -> PyAttribute {
SmallVector<Attribute, 4> attrs;
for (auto &it : l) {
attrs.push_back(it.cast<PyAttribute>().attr);
}
return ArrayAttr::get(attrs, &self.context);
})
.def(
"dense_elements_attr",
[](PyContext &self, py::buffer array) -> PyAttribute {
// Request a contiguous view.
int flags = PyBUF_C_CONTIGUOUS | PyBUF_FORMAT;
Py_buffer *view = new Py_buffer();
if (PyObject_GetBuffer(array.ptr(), view, flags) != 0) {
delete view;
throw py::error_already_set();
}
py::buffer_info array_info(view);
return createDenseElementsAttrFromBuffer(&self.context, array_info);
},
py::arg("array"))
.def_property_readonly("unit_attr", [](PyContext &self) -> PyAttribute {
return UnitAttr::get(&self.context);
});
}
PyModuleOp PyContext::parseAsm(const std::string &asm_text) {
// Arrange to get a view that includes a terminating null to avoid
// additional copy.
// TODO: Consider using the buffer protocol to access and avoid more copies.
const char *asm_chars = asm_text.c_str();
StringRef asm_sr(asm_chars, asm_text.size() + 1);
// TODO: Output non failure diagnostics (somewhere)
DiagnosticCapture diag_capture(&context);
auto module_ref = parseMLIRModuleFromString(asm_sr, &context);
if (!module_ref) {
throw py::raiseValueError(
diag_capture.consumeDiagnosticsAsString("Error parsing ASM"));
}
return PyModuleOp{shared_from_this(), module_ref.release()};
}
//===----------------------------------------------------------------------===//
// PyBaseOperation
//===----------------------------------------------------------------------===//
PyBaseOperation::~PyBaseOperation() = default;
void PyBaseOperation::bind(py::module m) {
py::class_<PyBaseOperation>(m, "BaseOperation")
.def_property_readonly(
"name",
[](PyBaseOperation &self) {
return std::string(self.getOperation()->getName().getStringRef());
})
.def_property_readonly("is_registered",
[](PyBaseOperation &self) {
return self.getOperation()->isRegistered();
})
.def_property_readonly("num_regions",
[](PyBaseOperation &self) {
return self.getOperation()->getNumRegions();
})
.def_property_readonly("results",
[](PyBaseOperation &self) {
auto *op = self.getOperation();
std::vector<PyValue> results(op->result_begin(),
op->result_end());
return results;
})
.def_property_readonly("result",
[](PyBaseOperation &self) -> PyValue {
auto *op = self.getOperation();
if (op->getNumResults() != 1) {
throw py::raiseValueError(
"Operation does not have 1 result");
}
return op->getOpResult(0);
})
.def("region",
[](PyBaseOperation &self, int index) {
auto *op = self.getOperation();
if (index < 0 || index >= static_cast<int>(op->getNumRegions())) {
throw py::raisePyError(PyExc_IndexError,
"Region index out of bounds");
}
return PyRegionRef(op->getRegion(index));
})
.def_property_readonly("first_block", [](PyBaseOperation &self) {
Operation *op = self.getOperation();
assert(op);
if (op->getNumRegions() == 0) {
throw py::raiseValueError("Op has no regions");
}
auto &region = op->getRegion(0);
if (region.empty()) {
throw py::raiseValueError("Op has no blocks");
}
return PyBlockRef(region.front());
});
}
//===----------------------------------------------------------------------===//
// PyOperationRef
//===----------------------------------------------------------------------===//
PyOperationRef::~PyOperationRef() = default;
void PyOperationRef::bind(py::module m) {
py::class_<PyOperationRef, PyBaseOperation>(m, "OperationRef");
}
Operation *PyOperationRef::getOperation() { return operation; }
//===----------------------------------------------------------------------===//
// PyModuleOp
//===----------------------------------------------------------------------===//
PyModuleOp::~PyModuleOp() = default;
void PyModuleOp::bind(py::module m) {
py::class_<PyModuleOp, PyBaseOperation>(m, "ModuleOp")
.def_property_readonly("context",
[](PyModuleOp &self) { return self.context; })
.def("to_asm", &PyModuleOp::toAsm, py::arg("debug_info") = false,
py::arg("pretty") = false, py::arg("large_element_limit") = -1);
}
Operation *PyModuleOp::getOperation() { return moduleOp; }
std::string PyModuleOp::toAsm(bool enableDebugInfo, bool prettyForm,
int64_t largeElementLimit) {
// Print to asm.
std::string asmOutput;
llvm::raw_string_ostream sout(asmOutput);
OpPrintingFlags printFlags;
if (enableDebugInfo) {
printFlags.enableDebugInfo(prettyForm);
}
if (largeElementLimit >= 0) {
printFlags.elideLargeElementsAttrs(largeElementLimit);
}
moduleOp.print(sout, printFlags);
return sout.str();
}
static OwningModuleRef parseMLIRModuleFromString(StringRef contents,
MLIRContext *context) {
std::unique_ptr<llvm::MemoryBuffer> contents_buffer;
if (contents.back() == 0) {
// If it has a nul terminator, just use as-is.
contents_buffer = llvm::MemoryBuffer::getMemBuffer(contents.drop_back());
} else {
// Otherwise, make a copy.
contents_buffer = llvm::MemoryBuffer::getMemBufferCopy(contents, "EMBED");
}
llvm::SourceMgr source_mgr;
source_mgr.AddNewSourceBuffer(std::move(contents_buffer), llvm::SMLoc());
OwningModuleRef mlir_module = parseSourceFile(source_mgr, context);
return mlir_module;
}
// Custom location printer that prints prettier, multi-line file output
// suitable for human readable error messages. The standard printer just prints
// a long nested expression not particularly human friendly). Note that there
// is a location pretty printer in the MLIR AsmPrinter. It is private and
// doesn't do any path shortening, which seems to make long Python stack traces
// a bit easier to scan.
// TODO: Upstream this.
void printLocation(Location loc, raw_ostream &out) {
TypeSwitch<Location>(loc)
.Case<OpaqueLoc>(
[&](OpaqueLoc loc) { printLocation(loc.getFallbackLocation(), out); })
.Case<UnknownLoc>([&](Location) { out << " [unknown location]\n"; })
.Case<FileLineColLoc>([&](FileLineColLoc line_col_loc) {
StringRef this_filename = line_col_loc.getFilename();
auto slash_pos = this_filename.find_last_of("/\\");
// We print both the basename and extended names with a structure like
// `foo.py:35:4`. Even though technically the line/col
// information is redundant to include in both names, having it on both
// makes it easier to paste the paths into an editor and jump to the
// exact location.
std::string line_col_suffix =
":" + std::to_string(line_col_loc.getLine()) + ":" +
std::to_string(line_col_loc.getColumn());
bool has_basename = false;
StringRef basename = this_filename;
if (slash_pos != StringRef::npos) {
has_basename = true;
basename = this_filename.substr(slash_pos + 1);
}
out << " at: " << basename << line_col_suffix;
if (has_basename) {
StringRef extended_name = this_filename;
// Print out two tabs, as basenames usually vary in length by more
// than one tab width.
out << "\t\t( " << extended_name << line_col_suffix << " )";
}
out << "\n";
})
.Case<NameLoc>([&](NameLoc nameLoc) {
out << " @'" << nameLoc.getName() << "':\n";
auto childLoc = nameLoc.getChildLoc();
if (!childLoc.isa<UnknownLoc>()) {
out << "(...\n";
printLocation(childLoc, out);
out << ")\n";
}
})
.Case<CallSiteLoc>([&](CallSiteLoc callSite) {
printLocation(callSite.getCaller(), out);
printLocation(callSite.getCallee(), out);
});
}
//===----------------------------------------------------------------------===//
// PySymbolTable
//===----------------------------------------------------------------------===//
void PySymbolTable::bind(py::module m) {
py::class_<PySymbolTable>(m, "SymbolTable")
.def_property_readonly_static("symbol_attr_name",
[](const py::object &) {
auto sr =
SymbolTable::getSymbolAttrName();
return py::str(sr.data(), sr.size());
})
.def_property_readonly_static(
"visibility_attr_name", [](const py::object &) {
auto sr = SymbolTable::getVisibilityAttrName();
return py::str(sr.data(), sr.size());
});
}
//===----------------------------------------------------------------------===//
// DiagnosticCapture
//===----------------------------------------------------------------------===//
std::string
DiagnosticCapture::consumeDiagnosticsAsString(const char *error_message) {
std::string s;
llvm::raw_string_ostream sout(s);
bool first = true;
if (error_message) {
sout << error_message;
first = false;
}
for (auto &d : diagnostics) {
if (!first) {
sout << "\n\n";
} else {
first = false;
}
switch (d.getSeverity()) {
case DiagnosticSeverity::Note:
sout << "[NOTE]";
break;
case DiagnosticSeverity::Warning:
sout << "[WARNING]";
break;
case DiagnosticSeverity::Error:
sout << "[ERROR]";
break;
case DiagnosticSeverity::Remark:
sout << "[REMARK]";
break;
}
// Message.
sout << ": " << d << "\n";
printLocation(d.getLocation(), sout);
}
diagnostics.clear();
return sout.str();
}
//===----------------------------------------------------------------------===//
// PyBlockRef
//===----------------------------------------------------------------------===//
void PyBlockRef::bind(py::module m) {
py::class_<PyBlockRef>(m, "BlockRef")
.def_property_readonly("operations",
[](PyBlockRef &self) {
return PyOperationList(
self.block.getOperations());
})
.def_property_readonly("args", [](PyBlockRef &self) {
return std::vector<PyValue>(self.block.args_begin(),
self.block.args_end());
});
}
//===----------------------------------------------------------------------===//
// PyRegionRef
//===----------------------------------------------------------------------===//
void PyRegionRef::bind(py::module m) {
py::class_<PyRegionRef>(m, "RegionRef")
.def_property_readonly("blocks", [](PyRegionRef &self) {
return PyBlockList(self.region.getBlocks());
});
}
//===----------------------------------------------------------------------===//
// PyType
//===----------------------------------------------------------------------===//
void PyType::bind(py::module m) {
py::class_<PyType>(m, "Type").def("__repr__",
[](PyType &self) -> std::string {
if (!self.type)
return "<undefined type>";
std::string res;
llvm::raw_string_ostream os(res);
self.type.print(os);
return res;
});
}
//===----------------------------------------------------------------------===//
// PyIdentifier
//===----------------------------------------------------------------------===//
void PyIdentifier::bind(py::module m) {
py::class_<PyIdentifier>(m, "Identifier")
.def("__str__", [](PyIdentifier &self) { return self.identifier.str(); })
.def("__repr__", [](PyIdentifier &self) {
std::string s("<Identifier \"");
s.append(self.identifier.str());
s.append("\">");
return s;
});
}
//===----------------------------------------------------------------------===//
// PyValue
//===----------------------------------------------------------------------===//
void PyValue::bind(py::module m) {
py::class_<PyValue>(m, "Value")
.def_property_readonly(
"type", [](PyValue &self) -> PyType { return self.value.getType(); })
.def("__repr__", [](PyValue &self) {
std::string res;
llvm::raw_string_ostream os(res);
os << self.value;
return res;
});
}
//===----------------------------------------------------------------------===//
// PyAttribute
//===----------------------------------------------------------------------===//
void PyAttribute::bind(py::module m) {
py::class_<PyAttribute>(m, "Attribute")
.def_property_readonly(
"type",
[](PyAttribute &self) -> PyType { return self.attr.getType(); })
.def("__repr__", [](PyAttribute &self) {
std::string res;
llvm::raw_string_ostream os(res);
os << self.attr;
return res;
});
}
//===----------------------------------------------------------------------===//
// OpBuilder implementations
//===----------------------------------------------------------------------===//
PyBaseOpBuilder::~PyBaseOpBuilder() = default;
PyOpBuilder::~PyOpBuilder() = default;
OpBuilder &PyOpBuilder::getBuilder(bool requirePosition) {
if (requirePosition && !builder.getBlock()) {
throw py::raisePyError(PyExc_IndexError, "Insertion point not set");
}
return builder;
}
void PyBaseOpBuilder::bind(py::module m) {
py::class_<PyBaseOpBuilder>(m, "BaseOpBuilder");
}
void PyOpBuilder::bind(py::module m) {
py::class_<PyOpBuilder, PyBaseOpBuilder>(m, "OpBuilder")
.def(py::init<PyContext &>(), py::keep_alive<1, 2>())
.def_property(
"current_loc",
[](PyOpBuilder &self) -> PyAttribute {
return static_cast<Attribute>(self.getCurrentLoc());
},
[](PyOpBuilder &self, PyAttribute attr) {
auto loc_attr = attr.attr.dyn_cast_or_null<LocationAttr>();
if (!loc_attr) {
throw py::raiseValueError("Expected a LocationAttr");
}
self.setCurrentLoc(Location(loc_attr));
})
.def_property(
"insertion_point",
[](PyOpBuilder &self) {
return self.getBuilder(true).saveInsertionPoint();
},
[](PyOpBuilder &self, OpBuilder::InsertPoint ip) {
self.getBuilder(false).restoreInsertionPoint(ip);
})
.def(
"set_file_line_col",
[](PyOpBuilder &self, PyIdentifier filename, unsigned line,
unsigned column) {
Location loc = FileLineColLoc::get(filename.identifier, line,
column, self.getContext());
self.setCurrentLoc(loc);
},
py::arg("filename"), py::arg("line"), py::arg("column"),
"Shortcut to set a FileLineCol current location")
.def("clear_insertion_point",
[](PyOpBuilder &self) { self.builder.clearInsertionPoint(); })
.def(
"insert_op_before",
[](PyOpBuilder &self, PyBaseOperation &pyOp) {
Operation *op = pyOp.getOperation();
self.builder.setInsertionPoint(op);
},
"Sets the insertion point to just before the specified op.")
.def(
"insert_op_after",
[](PyOpBuilder &self, PyBaseOperation &pyOp) {
Operation *op = pyOp.getOperation();
self.builder.setInsertionPointAfter(op);
},
"Sets the insertion point to just after the specified op.")
.def(
"insert_block_start",
[](PyOpBuilder &self, PyBlockRef block) {
self.builder.setInsertionPointToStart(&block.block);
},
"Sets the insertion point to the start of the block.")
.def(
"insert_block_end",
[](PyOpBuilder &self, PyBlockRef block) {
self.builder.setInsertionPointToEnd(&block.block);
},
"Sets the insertion point to the end of the block.")
.def(
"insert_before_terminator",
[](PyOpBuilder &self, PyBlockRef block) {
auto *terminator = block.block.getTerminator();
if (!terminator) {
throw py::raiseValueError("Block has no terminator");
}
self.builder.setInsertionPoint(terminator);
},
"Sets the insertion point to just before the block terminator.");
}
} // namespace mlir
Reference in New Issue View Git Blame Copy Permalink