From eb19b4eefceea2f07abd4d2be4d43ec47d06b938 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Fri, 8 Mar 2019 16:41:25 -0800 Subject: [PATCH] Add support for custom ops in declarative builders. This CL adds support for named custom instructions in declarative builders. To allow this, it introduces a templated `CustomInstruction` class. This CL also splits ValueHandle which can capture only the **value** in single-valued instructions from InstructionHandle which can capture any instruction but provide no typing and sugaring to extract the potential Value*. PiperOrigin-RevId: 237543222 --- mlir/include/mlir/EDSC/Builders.h | 90 ++++++++++++++++++++++++++--- mlir/include/mlir/EDSC/Helpers.h | 18 +++--- mlir/include/mlir/EDSC/Intrinsics.h | 32 +++++----- mlir/lib/EDSC/Builders.cpp | 39 ++++++++++++- mlir/lib/EDSC/Helpers.cpp | 8 +-- mlir/lib/EDSC/Intrinsics.cpp | 34 ++++++----- mlir/test/EDSC/builder-api-test.cpp | 33 +++++++++++ 7 files changed, 198 insertions(+), 56 deletions(-) diff --git a/mlir/include/mlir/EDSC/Builders.h b/mlir/include/mlir/EDSC/Builders.h index f0132703fd9b..c569f577b996 100644 --- a/mlir/include/mlir/EDSC/Builders.h +++ b/mlir/include/mlir/EDSC/Builders.h @@ -38,6 +38,7 @@ struct index_t { }; class BlockHandle; +class CapturableHandle; class NestedBuilder; class ValueHandle; @@ -162,7 +163,7 @@ public: /// In order to be admissible in a nested ArrayRef, operator() /// returns a ValueHandle::null() that cannot be captured. // TODO(ntv): when loops return escaping ssa-values, this should be adapted. - ValueHandle operator()(ArrayRef stmts); + ValueHandle operator()(ArrayRef stmts); }; /// Explicit nested LoopBuilder. Offers a compressed multi-loop builder to avoid @@ -192,7 +193,7 @@ public: ArrayRef ubs, ArrayRef steps); // TODO(ntv): when loops return escaping ssa-values, this should be adapted. - ValueHandle operator()(ArrayRef stmts); + ValueHandle operator()(ArrayRef stmts); private: SmallVector loops; @@ -225,13 +226,20 @@ public: /// The only purpose of this operator is to serve as a sequence point so that /// the evaluation of `stmts` (which build IR snippets in a scoped fashion) is /// sequenced strictly after the constructor of BlockBuilder. - void operator()(ArrayRef stmts); + void operator()(ArrayRef stmts); private: BlockBuilder(const BlockBuilder &) = delete; BlockBuilder &operator=(const BlockBuilder &other) = delete; }; +/// Base class for Handles that cannot be constructed explicitly by a user of +/// the API. +struct CapturableHandle { +protected: + CapturableHandle() = default; +}; + /// ValueHandle implements a (potentially "delayed") typed Value abstraction. /// ValueHandle should be captured by pointer but otherwise passed by Value /// everywhere. @@ -245,7 +253,13 @@ private: /// 2. delayed state (empty value), in which case it represents an eagerly /// typed "delayed" value that can be hold a Value in the future; /// 3. constructed state,in which case it holds a Value. -class ValueHandle { +/// +/// A ValueHandle is meant to capture a single Value* and should be used for +/// instructions that have a single result. For convenience of use, we also +/// include AffineForOp in this category although it does not return a value. +/// In the case of AffineForOp, the captured Value* is the loop induction +/// variable. +class ValueHandle : public CapturableHandle { public: /// A ValueHandle in a null state can never be captured; static ValueHandle null() { return ValueHandle(); } @@ -275,14 +289,13 @@ public: /// ValueHandle is a value type, the assignment operator typechecks before /// assigning. - /// ``` ValueHandle &operator=(const ValueHandle &other); /// Implicit conversion useful for automatic conversion to Container. operator Value *() const { return getValue(); } /// Generic mlir::Op create. This is the key to being extensible to the whole - /// of MLIR without duplicating the type system or the AST. + /// of MLIR without duplicating the type system or the op definitions. template static ValueHandle create(Args... args); @@ -291,6 +304,11 @@ public: static ValueHandle createComposedAffineApply(AffineMap map, ArrayRef operands); + /// Generic create for a named instruction producing a single value. + static ValueHandle create(StringRef name, ArrayRef operands, + ArrayRef resultTypes, + ArrayRef attributes = {}); + bool hasValue() const { return v != nullptr; } Value *getValue() const { return v; } bool hasType() const { return t != Type(); } @@ -303,12 +321,59 @@ private: Value *v; }; +/// An InstructionHandle can be used in lieu of ValueHandle to capture the +/// instruction in cases when one does not care about, or cannot extract, a +/// unique Value* from the instruction. +/// This can be used for capturing zero result instructions as well as +/// multi-result instructions that are not supported by ValueHandle. +/// We do not distinguish further between zero and multi-result instructions at +/// this time. +struct InstructionHandle : public CapturableHandle { + InstructionHandle() : inst(nullptr) {} + InstructionHandle(Instruction *inst) : inst(inst) {} + + InstructionHandle(const InstructionHandle &) = default; + InstructionHandle &operator=(const InstructionHandle &) = default; + + /// Generic mlir::Op create. This is the key to being extensible to the whole + /// of MLIR without duplicating the type system or the op definitions. + template + static InstructionHandle create(Args... args); + + /// Generic create for a named instruction. + static InstructionHandle create(StringRef name, + ArrayRef operands, + ArrayRef resultTypes, + ArrayRef attributes = {}); + + operator Instruction *() { return inst; } + +private: + Instruction *inst; +}; + +/// Simple wrapper to build a generic instruction without successor blocks. +template struct CustomInstruction { + CustomInstruction(StringRef name) : name(name) { + static_assert(std::is_same() || + std::is_same(), + "Only CustomInstruction or " + "CustomInstruction can be constructed."); + } + HandleType operator()(ArrayRef operands = {}, + ArrayRef resultTypes = {}, + ArrayRef attributes = {}) { + return HandleType::create(name, operands, resultTypes, attributes); + } + std::string name; +}; + /// A BlockHandle represents a (potentially "delayed") Block abstraction. /// This extra abstraction is necessary because an mlir::Block is not an /// mlir::Value. /// A BlockHandle should be captured by pointer but otherwise passed by Value /// everywhere. -class BlockHandle { +class BlockHandle : public CapturableHandle { public: /// A BlockHandle constructed without an mlir::Block* represents a "delayed" /// Block. A delayed Block represents the declaration (in the PL sense) of a @@ -338,6 +403,14 @@ private: mlir::Block *block; }; +template +InstructionHandle InstructionHandle::create(Args... args) { + return InstructionHandle( + ScopedContext::getBuilder() + ->create(ScopedContext::getLocation(), args...) + ->getInstruction()); +} + template ValueHandle ValueHandle::create(Args... args) { Instruction *inst = ScopedContext::getBuilder() @@ -350,9 +423,8 @@ ValueHandle ValueHandle::create(Args... args) { f->createBody(); return ValueHandle(f->getInductionVar()); } - return ValueHandle(); } - llvm_unreachable("unsupported inst with > 1 results"); + llvm_unreachable("unsupported instruction, use an InstructionHandle instead"); } namespace op { diff --git a/mlir/include/mlir/EDSC/Helpers.h b/mlir/include/mlir/EDSC/Helpers.h index 9c4618211e74..5974439a2660 100644 --- a/mlir/include/mlir/EDSC/Helpers.h +++ b/mlir/include/mlir/EDSC/Helpers.h @@ -106,7 +106,7 @@ struct IndexedValue { /// Emits a `store`. // NOLINTNEXTLINE: unconventional-assign-operator - ValueHandle operator=(ValueHandle rhs) { + InstructionHandle operator=(ValueHandle rhs) { return intrinsics::STORE(rhs, getBase(), indices); } @@ -122,10 +122,10 @@ struct IndexedValue { ValueHandle operator-(ValueHandle e); ValueHandle operator*(ValueHandle e); ValueHandle operator/(ValueHandle e); - ValueHandle operator+=(ValueHandle e); - ValueHandle operator-=(ValueHandle e); - ValueHandle operator*=(ValueHandle e); - ValueHandle operator/=(ValueHandle e); + InstructionHandle operator+=(ValueHandle e); + InstructionHandle operator-=(ValueHandle e); + InstructionHandle operator*=(ValueHandle e); + InstructionHandle operator/=(ValueHandle e); ValueHandle operator+(IndexedValue e) { return *this + static_cast(e); } @@ -138,16 +138,16 @@ struct IndexedValue { ValueHandle operator/(IndexedValue e) { return *this / static_cast(e); } - ValueHandle operator+=(IndexedValue e) { + InstructionHandle operator+=(IndexedValue e) { return this->operator+=(static_cast(e)); } - ValueHandle operator-=(IndexedValue e) { + InstructionHandle operator-=(IndexedValue e) { return this->operator-=(static_cast(e)); } - ValueHandle operator*=(IndexedValue e) { + InstructionHandle operator*=(IndexedValue e) { return this->operator*=(static_cast(e)); } - ValueHandle operator/=(IndexedValue e) { + InstructionHandle operator/=(IndexedValue e) { return this->operator/=(static_cast(e)); } diff --git a/mlir/include/mlir/EDSC/Intrinsics.h b/mlir/include/mlir/EDSC/Intrinsics.h index 6e69506fb213..cbfe43efc422 100644 --- a/mlir/include/mlir/EDSC/Intrinsics.h +++ b/mlir/include/mlir/EDSC/Intrinsics.h @@ -30,6 +30,7 @@ namespace mlir { namespace edsc { class BlockHandle; +class InstructionHandle; class ValueHandle; /// Provides a set of first class intrinsics. @@ -41,7 +42,7 @@ namespace intrinsics { /// /// Prerequisites: /// All Handles have already captured previously constructed IR objects. -ValueHandle BR(BlockHandle bh, ArrayRef operands); +InstructionHandle BR(BlockHandle bh, ArrayRef operands); /// Creates a new mlir::Block* and branches to it from the current block. /// Argument types are specified by `operands`. @@ -56,8 +57,8 @@ ValueHandle BR(BlockHandle bh, ArrayRef operands); /// All `operands` have already captured an mlir::Value* /// captures.size() == operands.size() /// captures and operands are pairwise of the same type. -ValueHandle BR(BlockHandle *bh, ArrayRef captures, - ArrayRef operands); +InstructionHandle BR(BlockHandle *bh, ArrayRef captures, + ArrayRef operands); /// Branches into the mlir::Block* captured by BlockHandle `trueBranch` with /// `trueOperands` if `cond` evaluates to `true` (resp. `falseBranch` and @@ -65,9 +66,10 @@ ValueHandle BR(BlockHandle *bh, ArrayRef captures, /// /// Prerequisites: /// All Handles have captured previouly constructed IR objects. -ValueHandle COND_BR(ValueHandle cond, BlockHandle trueBranch, - ArrayRef trueOperands, BlockHandle falseBranch, - ArrayRef falseOperands); +InstructionHandle COND_BR(ValueHandle cond, BlockHandle trueBranch, + ArrayRef trueOperands, + BlockHandle falseBranch, + ArrayRef falseOperands); /// Eagerly creates new mlir::Block* with argument types specified by /// `trueOperands`/`falseOperands`. @@ -85,12 +87,12 @@ ValueHandle COND_BR(ValueHandle cond, BlockHandle trueBranch, /// `falseCaptures`.size() == `falseOperands`.size() /// `trueCaptures` and `trueOperands` are pairwise of the same type /// `falseCaptures` and `falseOperands` are pairwise of the same type. -ValueHandle COND_BR(ValueHandle cond, BlockHandle *trueBranch, - ArrayRef trueCaptures, - ArrayRef trueOperands, - BlockHandle *falseBranch, - ArrayRef falseCaptures, - ArrayRef falseOperands); +InstructionHandle COND_BR(ValueHandle cond, BlockHandle *trueBranch, + ArrayRef trueCaptures, + ArrayRef trueOperands, + BlockHandle *falseBranch, + ArrayRef falseCaptures, + ArrayRef falseOperands); //////////////////////////////////////////////////////////////////////////////// // TODO(ntv): Intrinsics below this line should be TableGen'd. @@ -103,13 +105,13 @@ ValueHandle LOAD(ValueHandle base, llvm::ArrayRef indices); /// Builds an mlir::ReturnOp with the proper `operands` that each must have /// captured an mlir::Value*. /// Returns an empty ValueHandle. -ValueHandle RETURN(llvm::ArrayRef operands); +InstructionHandle RETURN(llvm::ArrayRef operands); /// Builds an mlir::StoreOp with the proper `operands` that each must have /// captured an mlir::Value*. /// Returns an empty ValueHandle. -ValueHandle STORE(ValueHandle value, ValueHandle base, - llvm::ArrayRef indices); +InstructionHandle STORE(ValueHandle value, ValueHandle base, + llvm::ArrayRef indices); } // namespace intrinsics diff --git a/mlir/lib/EDSC/Builders.cpp b/mlir/lib/EDSC/Builders.cpp index 7a5da8ab0dda..ffe2dda82c3a 100644 --- a/mlir/lib/EDSC/Builders.cpp +++ b/mlir/lib/EDSC/Builders.cpp @@ -92,6 +92,38 @@ mlir::edsc::ValueHandle::createComposedAffineApply(AffineMap map, return ValueHandle(inst->getResult(0)); } +ValueHandle ValueHandle::create(StringRef name, ArrayRef operands, + ArrayRef resultTypes, + ArrayRef attributes) { + Instruction *inst = + InstructionHandle::create(name, operands, resultTypes, attributes); + if (auto f = inst->dyn_cast()) { + // Immediately create the loop body so we can just insert instructions right + // away. + f->createBody(); + return ValueHandle(f->getInductionVar()); + } + if (inst->getNumResults() == 1) { + return ValueHandle(inst->getResult(0)); + } + llvm_unreachable("unsupported instruction, use an InstructionHandle instead"); +} + +InstructionHandle +InstructionHandle::create(StringRef name, ArrayRef operands, + ArrayRef resultTypes, + ArrayRef attributes) { + OperationState state(ScopedContext::getContext(), + ScopedContext::getLocation(), name); + SmallVector ops(operands.begin(), operands.end()); + state.addOperands(ops); + state.addTypes(resultTypes); + for (const auto &attr : attributes) { + state.addAttribute(attr.first, attr.second); + } + return InstructionHandle(ScopedContext::getBuilder()->createOperation(state)); +} + BlockHandle mlir::edsc::BlockHandle::create(ArrayRef argTypes) { BlockHandle res; res.block = ScopedContext::getBuilder()->createBlock(); @@ -139,7 +171,8 @@ mlir::edsc::LoopBuilder::LoopBuilder(ValueHandle *iv, enter(body); } -ValueHandle mlir::edsc::LoopBuilder::operator()(ArrayRef stmts) { +ValueHandle +mlir::edsc::LoopBuilder::operator()(ArrayRef stmts) { // Call to `exit` must be explicit and asymmetric (cannot happen in the // destructor) because of ordering wrt comma operator. /// The particular use case concerns nested blocks: @@ -176,7 +209,7 @@ mlir::edsc::LoopNestBuilder::LoopNestBuilder(ArrayRef ivs, } ValueHandle -mlir::edsc::LoopNestBuilder::operator()(ArrayRef stmts) { +mlir::edsc::LoopNestBuilder::operator()(ArrayRef stmts) { // Iterate on the calling operator() on all the loops in the nest. // The iteration order is from innermost to outermost because enter/exit needs // to be asymmetric (i.e. enter() occurs on LoopBuilder construction, exit() @@ -212,7 +245,7 @@ mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle *bh, /// Only serves as an ordering point between entering nested block and creating /// stmts. -void mlir::edsc::BlockBuilder::operator()(ArrayRef stmts) { +void mlir::edsc::BlockBuilder::operator()(ArrayRef stmts) { // Call to `exit` must be explicit and asymmetric (cannot happen in the // destructor) because of ordering wrt comma operator. exit(); diff --git a/mlir/lib/EDSC/Helpers.cpp b/mlir/lib/EDSC/Helpers.cpp index 9df9a97b155a..6400c73e2295 100644 --- a/mlir/lib/EDSC/Helpers.cpp +++ b/mlir/lib/EDSC/Helpers.cpp @@ -71,19 +71,19 @@ ValueHandle mlir::edsc::IndexedValue::operator/(ValueHandle e) { return static_cast(*this) / e; } -ValueHandle mlir::edsc::IndexedValue::operator+=(ValueHandle e) { +InstructionHandle mlir::edsc::IndexedValue::operator+=(ValueHandle e) { using op::operator+; return intrinsics::STORE(*this + e, getBase(), indices); } -ValueHandle mlir::edsc::IndexedValue::operator-=(ValueHandle e) { +InstructionHandle mlir::edsc::IndexedValue::operator-=(ValueHandle e) { using op::operator-; return intrinsics::STORE(*this - e, getBase(), indices); } -ValueHandle mlir::edsc::IndexedValue::operator*=(ValueHandle e) { +InstructionHandle mlir::edsc::IndexedValue::operator*=(ValueHandle e) { using op::operator*; return intrinsics::STORE(*this * e, getBase(), indices); } -ValueHandle mlir::edsc::IndexedValue::operator/=(ValueHandle e) { +InstructionHandle mlir::edsc::IndexedValue::operator/=(ValueHandle e) { using op::operator/; return intrinsics::STORE(*this / e, getBase(), indices); } diff --git a/mlir/lib/EDSC/Intrinsics.cpp b/mlir/lib/EDSC/Intrinsics.cpp index 87c9fec5bdc7..5afdb4cd364c 100644 --- a/mlir/lib/EDSC/Intrinsics.cpp +++ b/mlir/lib/EDSC/Intrinsics.cpp @@ -22,15 +22,15 @@ using namespace mlir; using namespace mlir::edsc; -ValueHandle mlir::edsc::intrinsics::BR(BlockHandle bh, - ArrayRef operands) { +InstructionHandle mlir::edsc::intrinsics::BR(BlockHandle bh, + ArrayRef operands) { assert(bh && "Expected already captured BlockHandle"); for (auto &o : operands) { (void)o; assert(o && "Expected already captured ValueHandle"); } SmallVector ops(operands.begin(), operands.end()); - return ValueHandle::create(bh.getBlock(), ops); + return InstructionHandle::create(bh.getBlock(), ops); } static void enforceEmptyCapturesMatchOperands(ArrayRef captures, ArrayRef operands) { @@ -46,9 +46,9 @@ static void enforceEmptyCapturesMatchOperands(ArrayRef captures, } } -ValueHandle mlir::edsc::intrinsics::BR(BlockHandle *bh, - ArrayRef captures, - ArrayRef operands) { +InstructionHandle mlir::edsc::intrinsics::BR(BlockHandle *bh, + ArrayRef captures, + ArrayRef operands) { assert(!*bh && "Unexpected already captured BlockHandle"); enforceEmptyCapturesMatchOperands(captures, operands); { // Clone the scope explicitly to avoid modifying the insertion point in the @@ -60,21 +60,21 @@ ValueHandle mlir::edsc::intrinsics::BR(BlockHandle *bh, BlockBuilder(bh, captures)({/* no body */}); } // Release before adding the branch to the eagerly created block. SmallVector ops(operands.begin(), operands.end()); - return ValueHandle::create(bh->getBlock(), ops); + return InstructionHandle::create(bh->getBlock(), ops); } -ValueHandle +InstructionHandle mlir::edsc::intrinsics::COND_BR(ValueHandle cond, BlockHandle trueBranch, ArrayRef trueOperands, BlockHandle falseBranch, ArrayRef falseOperands) { SmallVector trueOps(trueOperands.begin(), trueOperands.end()); SmallVector falseOps(falseOperands.begin(), falseOperands.end()); - return ValueHandle::create(cond, trueBranch.getBlock(), trueOps, - falseBranch.getBlock(), falseOps); + return InstructionHandle::create( + cond, trueBranch.getBlock(), trueOps, falseBranch.getBlock(), falseOps); } -ValueHandle mlir::edsc::intrinsics::COND_BR( +InstructionHandle mlir::edsc::intrinsics::COND_BR( ValueHandle cond, BlockHandle *trueBranch, ArrayRef trueCaptures, ArrayRef trueOperands, BlockHandle *falseBranch, ArrayRef falseCaptures, @@ -93,7 +93,7 @@ ValueHandle mlir::edsc::intrinsics::COND_BR( } // Release before adding the branch to the eagerly created block. SmallVector trueOps(trueOperands.begin(), trueOperands.end()); SmallVector falseOps(falseOperands.begin(), falseOperands.end()); - return ValueHandle::create( + return InstructionHandle::create( cond, trueBranch->getBlock(), trueOps, falseBranch->getBlock(), falseOps); } @@ -107,14 +107,16 @@ mlir::edsc::intrinsics::LOAD(ValueHandle base, return ValueHandle::create(base.getValue(), ops); } -ValueHandle mlir::edsc::intrinsics::RETURN(ArrayRef operands) { +InstructionHandle +mlir::edsc::intrinsics::RETURN(ArrayRef operands) { SmallVector ops(operands.begin(), operands.end()); - return ValueHandle::create(ops); + return InstructionHandle::create(ops); } -ValueHandle +InstructionHandle mlir::edsc::intrinsics::STORE(ValueHandle value, ValueHandle base, llvm::ArrayRef indices = {}) { SmallVector ops(indices.begin(), indices.end()); - return ValueHandle::create(value.getValue(), base.getValue(), ops); + return InstructionHandle::create(value.getValue(), base.getValue(), + ops); } diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp index f58f15abcbed..0a4c4b99e34d 100644 --- a/mlir/test/EDSC/builder-api-test.cpp +++ b/mlir/test/EDSC/builder-api-test.cpp @@ -363,6 +363,39 @@ TEST_FUNC(builder_helpers) { f->print(llvm::outs()); } +TEST_FUNC(custom_ops) { + using namespace edsc; + using namespace edsc::intrinsics; + using namespace edsc::op; + auto indexType = IndexType::get(&globalContext()); + auto f = makeFunction("custom_ops", {}, {indexType, indexType}); + + ScopedContext scope(f.get()); + CustomInstruction MY_CUSTOM_OP("my_custom_op"); + CustomInstruction MY_CUSTOM_INST_0("my_custom_inst_0"); + CustomInstruction MY_CUSTOM_INST_2("my_custom_inst_2"); + + // clang-format off + ValueHandle vh(indexType); + InstructionHandle ih0, ih2; + IndexHandle m, n, M(f->getArgument(0)), N(f->getArgument(1)); + IndexHandle ten(index_t(10)), twenty(index_t(20)); + LoopNestBuilder({&m, &n}, {M, N}, {M + ten, N + twenty}, {1, 1})({ + vh = MY_CUSTOM_OP({m, m + n}, {indexType}, {}), + ih0 = MY_CUSTOM_INST_0({m, m + n}, {}), + ih2 = MY_CUSTOM_INST_2({m, m + n}, {indexType, indexType}), + }); + + // CHECK-LABEL: @custom_ops + // CHECK: for %i0 {{.*}} + // CHECK: for %i1 {{.*}} + // CHECK: {{.*}} = "my_custom_op"{{.*}} : (index, index) -> index + // CHECK: "my_custom_inst_0"{{.*}} : (index, index) -> () + // CHECK: {{.*}} = "my_custom_inst_2"{{.*}} : (index, index) -> (index, index) + // clang-format on + f->print(llvm::outs()); +} + int main() { RUN_TESTS(); return 0;