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

Add DialectHelper for Basicpy dialect. 

* Involved native code for the types and slot_object_get ops.
pull/1/head
Stella Laurenzo 2020-05-06 18:24:51 -07:00
parent 4ebf972503
commit 6b7c913e0b
9 changed files with 355 additions and 169 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/npcomp/CMakeLists.txt
View File

@ -28,6 +28,7 @@ set(extension_pybind_sources
native.cpp
mlir_init.cpp
mlir_ir.cpp
npcomp_dialect.cpp
pybind_utils.cpp
)
set_source_files_properties(
@ -66,6 +67,7 @@ target_link_libraries(${extension_target}
pybind11::module
# Local depends
NPCOMPBasicpyDialect
NPCOMPNumpyDialect
# Upstream depends

65
python/npcomp/dialect/Basicpy.py 100644
View File

@ -0,0 +1,65 @@
# 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 ..native.dialect import BasicpyDialectHelper as _BaseDialectHelper
from ..native.mlir import ir
__all__ = [
"DialectHelper",
]
class DialectHelper(_BaseDialectHelper):
r"""Dialect helper for the Basicpy dialect.
>>> c = ir.MLIRContext()
>>> h = DialectHelper(c)
Dialect Types:
>>> h.basicpy_None_type
!basicpy.NoneType
>>> h.basicpy_Ellipsis_type
!basicpy.EllipsisType
>>> h.basicpy_SlotObject_type(
... "foobar", h.basicpy_None_type, h.basicpy_None_type)
!basicpy.SlotObject<foobar, !basicpy.NoneType, !basicpy.NoneType>
singleton op:
>>> m = c.new_module()
>>> h.builder.insert_block_start(m.first_block)
>>> _ = h.basicpy_singleton_op(h.basicpy_None_type)
>>> m.to_asm().strip()
'module {\n %0 = basicpy.singleton : !basicpy.NoneType\n}'
slot_object ops:
>>> m = c.new_module()
>>> h.builder.insert_block_start(m.first_block)
>>> v0 = h.basicpy_singleton_op(h.basicpy_None_type).result
>>> slot_object = h.basicpy_slot_object_make_op("foobar", v0, v0).result
>>> _ = h.basicpy_slot_object_get_op(slot_object, 0)
>>> print(m.to_asm().strip())
module {
%0 = basicpy.singleton : !basicpy.NoneType
%1 = basicpy.slot_object_make(%0, %0) -> !basicpy.SlotObject<foobar, !basicpy.NoneType, !basicpy.NoneType>
%2 = basicpy.slot_object_get %1[0] : !basicpy.SlotObject<foobar, !basicpy.NoneType, !basicpy.NoneType>
}
"""
def basicpy_singleton_op(self, singleton_type):
return self.op("basicpy.singleton", [singleton_type], [])
def basicpy_slot_object_make_op(self, class_name, *slot_values):
c = self.context
class_name_attr = c.string_attr(class_name)
object_type = self.basicpy_SlotObject_type(class_name,
*[v.type for v in slot_values])
attrs = c.dictionary_attr({"className": class_name_attr})
return self.op("basicpy.slot_object_make", [object_type], slot_values,
attrs)
if __name__ == "__main__":
import doctest
doctest.testmod()

3
python/npcomp/dialect/Numpy.py
View File

@ -2,6 +2,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from npcomp.dialect import Basicpy
from npcomp.native.mlir import ir
__all__ = [
@ -10,7 +11,7 @@ __all__ = [
]
class DialectHelper(ir.DialectHelper):
class DialectHelper(Basicpy.DialectHelper):
r"""Dialect helper.
>>> c = ir.MLIRContext()

317
python/npcomp/mlir_ir.cpp
View File

@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir_ir.h"
#include "native.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Attributes.h"
@ -117,163 +118,149 @@ private:
};
//===----------------------------------------------------------------------===//
// Python only classes
// PyDialectHelper
//===----------------------------------------------------------------------===//
/// Helper for creating (possibly dialect specific) IR objects. This class
/// is intended to be subclassed on the Python side (possibly with multiple
/// inheritance) to provide Python level APIs for custom dialects. The base
/// class contains helpers for std types and ops.
class PyDialectHelper {
public:
PyDialectHelper(std::shared_ptr<PyContext> context)
: pyOpBuilder(*context), context(std::move(context)) {
}
static void bind(py::module m) {
py::class_<PyDialectHelper>(m, "DialectHelper")
.def(py::init<std::shared_ptr<PyContext>>())
.def_property_readonly("builder",
[](PyDialectHelper &self) -> PyBaseOpBuilder & {
return self.pyOpBuilder;
})
.def_property_readonly(
"context",
[](PyDialectHelper &self) -> std::shared_ptr<PyContext> {
return self.context;
})
.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 = UnknownLoc::get(opBuilder.getContext());
OperationName opName(opNameStr, opBuilder.getContext());
SmallVector<Type, 4> types(pyResultTypes.begin(),
pyResultTypes.end());
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
MutableDictionaryAttr attrList;
if (attrs) {
auto dictAttrs = attrs->attr.dyn_cast<DictionaryAttr>();
if (!dictAttrs) {
throw py::raiseValueError(
"Expected `attrs` to be a DictionaryAttr");
}
attrList = MutableDictionaryAttr(dictAttrs);
void PyDialectHelper::bind(py::module m) {
py::class_<PyDialectHelper>(m, "DialectHelper")
.def(py::init<std::shared_ptr<PyContext>>())
.def_property_readonly("builder",
[](PyDialectHelper &self) -> PyBaseOpBuilder & {
return self.pyOpBuilder;
})
.def_property_readonly(
"context",
[](PyDialectHelper &self) -> std::shared_ptr<PyContext> {
return self.context;
})
.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 = UnknownLoc::get(opBuilder.getContext());
OperationName opName(opNameStr, opBuilder.getContext());
SmallVector<Type, 4> types(pyResultTypes.begin(),
pyResultTypes.end());
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
MutableDictionaryAttr attrList;
if (attrs) {
auto dictAttrs = attrs->attr.dyn_cast<DictionaryAttr>();
if (!dictAttrs) {
throw py::raiseValueError(
"Expected `attrs` to be a DictionaryAttr");
}
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) {
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 = UnknownLoc::get(opBuilder.getContext());
// TODO: Add function and arg/result attributes.
FuncOp op = opBuilder.create<FuncOp>(
loc, StringRef(name), functionType,
/*attrs=*/ArrayRef<NamedAttribute>());
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,
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("return_op",
[](PyDialectHelper &self, std::vector<PyValue> pyOperands) {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = UnknownLoc::get(opBuilder.getContext());
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
return PyOperationRef(opBuilder.create<ReturnOp>(loc, operands));
})
attrList = MutableDictionaryAttr(dictAttrs);
}
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) {
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 = UnknownLoc::get(opBuilder.getContext());
// TODO: Add function and arg/result attributes.
FuncOp op =
opBuilder.create<FuncOp>(loc, StringRef(name), functionType,
/*attrs=*/ArrayRef<NamedAttribute>());
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,
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("return_op",
[](PyDialectHelper &self, std::vector<PyValue> pyOperands) {
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = UnknownLoc::get(opBuilder.getContext());
SmallVector<Value, 4> operands(pyOperands.begin(),
pyOperands.end());
return PyOperationRef(opBuilder.create<ReturnOp>(loc, operands));
})
// Types.
.def("integer_type",
[](PyDialectHelper &self, unsigned width) {
return PyType(IntegerType::get(width, &self.context->context));
},
py::arg("width") = 32)
.def_property_readonly(
"i1_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(1, &self.context->context));
})
.def_property_readonly(
"i16_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(32, &self.context->context));
})
.def_property_readonly(
"i32_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(32, &self.context->context));
})
.def_property_readonly(
"i64_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(64, &self.context->context));
})
.def_property_readonly(
"f32_type",
[](PyDialectHelper &self) {
return PyType(
FloatType::get(StandardTypes::F32, &self.context->context));
})
.def_property_readonly(
"f64_type",
[](PyDialectHelper &self) {
return PyType(
FloatType::get(StandardTypes::F64, &self.context->context));
})
.def("tensor_type",
[](PyDialectHelper &self, PyType elementType,
llvm::Optional<std::vector<int64_t>> shape) {
if (!elementType.type) {
throw py::raiseValueError("Null element type");
}
if (shape) {
return PyType(RankedTensorType::get(*shape, elementType.type));
} else {
return PyType(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) {
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 PyType(FunctionType::get(inputTypes, resultTypes,
&self.context->context));
});
}
PyOpBuilder pyOpBuilder;
std::shared_ptr<PyContext> context;
};
// Types.
.def("integer_type",
[](PyDialectHelper &self, unsigned width) {
return PyType(IntegerType::get(width, &self.context->context));
},
py::arg("width") = 32)
.def_property_readonly("i1_type",
[](PyDialectHelper &self) {
return PyType(
IntegerType::get(1, &self.context->context));
})
.def_property_readonly(
"i16_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(32, &self.context->context));
})
.def_property_readonly(
"i32_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(32, &self.context->context));
})
.def_property_readonly(
"i64_type",
[](PyDialectHelper &self) {
return PyType(IntegerType::get(64, &self.context->context));
})
.def_property_readonly("f32_type",
[](PyDialectHelper &self) {
return PyType(FloatType::get(
StandardTypes::F32, &self.context->context));
})
.def_property_readonly("f64_type",
[](PyDialectHelper &self) {
return PyType(FloatType::get(
StandardTypes::F64, &self.context->context));
})
.def("tensor_type",
[](PyDialectHelper &self, PyType elementType,
llvm::Optional<std::vector<int64_t>> shape) {
if (!elementType.type) {
throw py::raiseValueError("Null element type");
}
if (shape) {
return PyType(RankedTensorType::get(*shape, elementType.type));
} else {
return PyType(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) {
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 PyType(FunctionType::get(inputTypes, resultTypes,
&self.context->context));
});
}
//===----------------------------------------------------------------------===//
// Module initialization
@ -414,6 +401,15 @@ void PyBaseOperation::bind(py::module m) {
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();
@ -672,12 +668,15 @@ void PyType::bind(py::module m) {
//===----------------------------------------------------------------------===//
void PyValue::bind(py::module m) {
py::class_<PyValue>(m, "Value").def("__repr__", [](PyValue &self) {
std::string res;
llvm::raw_string_ostream os(res);
os << self.value;
return res;
});
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;
});
}
//===----------------------------------------------------------------------===//

19
python/npcomp/mlir_ir.h
View File

@ -153,6 +153,25 @@ private:
OpBuilder builder;
};
//===----------------------------------------------------------------------===//
// Custom types
//===----------------------------------------------------------------------===//
/// Helper for creating (possibly dialect specific) IR objects. This class
/// is intended to be subclassed on the Python side (possibly with multiple
/// inheritance) to provide Python level APIs for custom dialects. The base
/// class contains helpers for std types and ops.
class PyDialectHelper {
public:
PyDialectHelper(std::shared_ptr<PyContext> context)
: pyOpBuilder(*context), context(std::move(context)) {}
static void bind(py::module m);
protected:
PyOpBuilder pyOpBuilder;
std::shared_ptr<PyContext> context;
};
} // namespace mlir
#endif // NPCOMP_PYTHON_MLIR_IR_H

9
python/npcomp/native.cpp
View File

@ -9,19 +9,15 @@
#include <cstddef>
#include <unordered_map>
#include "native.h"
#include "pybind_utils.h"
#include "llvm/Support/CommandLine.h"
namespace mlir {
void defineMlirIrModule(py::module m);
namespace npcomp {
namespace python {
// Externs
bool npcompMlirInitialize();
void defineLLVMModule(pybind11::module m) {
m.def("print_help_message", []() { llvm::cl::PrintHelpMessage(); });
m.def("add_option",
@ -66,6 +62,9 @@ PYBIND11_MODULE(native, m) {
auto mlir_m = m.def_submodule("mlir", "MLIR interop");
auto mlir_ir_m = mlir_m.def_submodule("ir");
defineMlirIrModule(mlir_ir_m);
auto npcomp_dialect = m.def_submodule("dialect", "NPComp custom dialects");
defineNpcompDialect(npcomp_dialect);
}
} // namespace python

27
python/npcomp/native.h 100644
View File

@ -0,0 +1,27 @@
//===- dialect.h - Module registrations -----------------------------------===//
//
// 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 NPCOMP_PYTHON_NATIVE_H
#define NPCOMP_PYTHON_NATIVE_H
#include "pybind_utils.h"
namespace mlir {
void defineMlirIrModule(py::module m);
namespace npcomp {
namespace python {
bool npcompMlirInitialize();
void defineNpcompDialect(py::module m);
} // namespace python
} // namespace npcomp
} // namespace mlir
#endif // NPCOMP_PYTHON_NATIVE_H

76
python/npcomp/npcomp_dialect.cpp 100644
View File

@ -0,0 +1,76 @@
//===- npcomp_dialect.cpp - Custom dialect classes ------------------------===//
//
// 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_ir.h"
#include "native.h"
#include "npcomp/Dialect/Basicpy/BasicpyDialect.h"
#include "npcomp/Dialect/Basicpy/BasicpyOps.h"
namespace mlir {
namespace NPCOMP {
class BasicpyDialectHelper : public PyDialectHelper {
public:
using PyDialectHelper::PyDialectHelper;
static void bind(py::module m) {
py::class_<BasicpyDialectHelper, PyDialectHelper>(m, "BasicpyDialectHelper")
.def(py::init<std::shared_ptr<PyContext>>())
.def_property_readonly("basicpy_None_type",
[](BasicpyDialectHelper &self) -> PyType {
return Basicpy::NoneType::get(
&self.context->context);
})
.def_property_readonly("basicpy_Ellipsis_type",
[](BasicpyDialectHelper &self) -> PyType {
return Basicpy::EllipsisType::get(
&self.context->context);
})
.def("basicpy_SlotObject_type",
[](BasicpyDialectHelper &self, std::string className,
py::args pySlotTypes) -> PyType {
SmallVector<Type, 4> slotTypes;
for (auto pySlotType : pySlotTypes) {
slotTypes.push_back(pySlotType.cast<PyType>());
}
auto classNameAttr =
StringAttr::get(className, &self.context->context);
return Basicpy::SlotObjectType::get(classNameAttr, slotTypes);
},
py::arg("className"))
.def("basicpy_slot_object_get_op",
[](BasicpyDialectHelper &self, PyValue slotObject,
unsigned index) -> PyOperationRef {
auto slotObjectType = slotObject.value.getType()
.dyn_cast<Basicpy::SlotObjectType>();
if (!slotObjectType) {
throw py::raiseValueError("Operand must be a SlotObject");
}
if (index >= slotObjectType.getSlotCount()) {
throw py::raiseValueError("Out of range slot index");
}
auto resultType = slotObjectType.getSlotTypes()[index];
auto indexAttr = IntegerAttr::get(
IndexType::get(&self.context->context), index);
OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true);
Location loc = UnknownLoc::get(opBuilder.getContext());
auto op = opBuilder.create<Basicpy::SlotObjectGetOp>(
loc, resultType, slotObject, indexAttr);
return op.getOperation();
});
}
};
} // namespace NPCOMP
} // namespace mlir
using namespace ::mlir::NPCOMP;
void mlir::npcomp::python::defineNpcompDialect(py::module m) {
BasicpyDialectHelper::bind(m);
}

6
python/npcomp/tracing/emitters.py
View File

@ -212,10 +212,8 @@ class GenericArrayFuncEmitter(FuncEmitter):
def emit(self, request: EmissionRequest):
h = request.dialect_helper
op_result_types = [h.tensor_type(h.numpy_any_dtype)
] * self._nresults
op = h.op(self._op_name, op_result_types,
request.input_ssa_values)
op_result_types = [h.tensor_type(h.numpy_any_dtype)] * self._nresults
op = h.op(self._op_name, op_result_types, request.input_ssa_values)
return op.results