diff --git a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h --- a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h +++ b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h @@ -52,11 +52,14 @@ using std_addf = ValueBuilder<AddFOp>; using std_alloc = ValueBuilder<AllocOp>; using std_call = OperationBuilder<CallOp>; +using std_constant = ValueBuilder<ConstantOp>; using std_constant_float = ValueBuilder<ConstantFloatOp>; using std_constant_index = ValueBuilder<ConstantIndexOp>; using std_constant_int = ValueBuilder<ConstantIntOp>; using std_dealloc = OperationBuilder<DeallocOp>; using std_dim = ValueBuilder<DimOp>; +using std_extract_element = ValueBuilder<ExtractElementOp>; +using std_index_cast = ValueBuilder<IndexCastOp>; using std_muli = ValueBuilder<MulIOp>; using std_mulf = ValueBuilder<MulFOp>; using std_memref_cast = ValueBuilder<MemRefCastOp>; @@ -65,7 +68,10 @@ using std_load = ValueBuilder<LoadOp>; using std_store = OperationBuilder<StoreOp>; using std_subi = ValueBuilder<SubIOp>; +using std_sub_view = ValueBuilder<SubViewOp>; using std_tanh = ValueBuilder<TanhOp>; +using std_tensor_load = ValueBuilder<TensorLoadOp>; +using std_tensor_store = OperationBuilder<TensorStoreOp>; using std_view = ValueBuilder<ViewOp>; using std_zero_extendi = ValueBuilder<ZeroExtendIOp>; using std_sign_extendi = ValueBuilder<SignExtendIOp>; @@ -74,7 +80,7 @@ /// /// Prerequisites: /// All Handles have already captured previously constructed IR objects. -OperationHandle br(BlockHandle bh, ArrayRef<ValueHandle> operands); +OperationHandle std_br(BlockHandle bh, ArrayRef<ValueHandle> operands); /// Creates a new mlir::Block* and branches to it from the current block. /// Argument types are specified by `operands`. @@ -89,8 +95,8 @@ /// All `operands` have already captured an mlir::Value /// captures.size() == operands.size() /// captures and operands are pairwise of the same type. -OperationHandle br(BlockHandle *bh, ArrayRef<ValueHandle *> captures, - ArrayRef<ValueHandle> operands); +OperationHandle std_br(BlockHandle *bh, ArrayRef<ValueHandle *> captures, + ArrayRef<ValueHandle> operands); /// Branches into the mlir::Block* captured by BlockHandle `trueBranch` with /// `trueOperands` if `cond` evaluates to `true` (resp. `falseBranch` and @@ -98,10 +104,10 @@ /// /// Prerequisites: /// All Handles have captured previously constructed IR objects. -OperationHandle cond_br(ValueHandle cond, BlockHandle trueBranch, - ArrayRef<ValueHandle> trueOperands, - BlockHandle falseBranch, - ArrayRef<ValueHandle> falseOperands); +OperationHandle std_cond_br(ValueHandle cond, BlockHandle trueBranch, + ArrayRef<ValueHandle> trueOperands, + BlockHandle falseBranch, + ArrayRef<ValueHandle> falseOperands); /// Eagerly creates new mlir::Block* with argument types specified by /// `trueOperands`/`falseOperands`. @@ -119,12 +125,12 @@ /// `falseCaptures`.size() == `falseOperands`.size() /// `trueCaptures` and `trueOperands` are pairwise of the same type /// `falseCaptures` and `falseOperands` are pairwise of the same type. -OperationHandle cond_br(ValueHandle cond, BlockHandle *trueBranch, - ArrayRef<ValueHandle *> trueCaptures, - ArrayRef<ValueHandle> trueOperands, - BlockHandle *falseBranch, - ArrayRef<ValueHandle *> falseCaptures, - ArrayRef<ValueHandle> falseOperands); +OperationHandle std_cond_br(ValueHandle cond, BlockHandle *trueBranch, + ArrayRef<ValueHandle *> trueCaptures, + ArrayRef<ValueHandle> trueOperands, + BlockHandle *falseBranch, + ArrayRef<ValueHandle *> falseCaptures, + ArrayRef<ValueHandle> falseOperands); /// Provide an index notation around sdt_load and std_store. using StdIndexedValue = diff --git a/mlir/include/mlir/EDSC/Builders.h b/mlir/include/mlir/EDSC/Builders.h --- a/mlir/include/mlir/EDSC/Builders.h +++ b/mlir/include/mlir/EDSC/Builders.h @@ -340,7 +340,8 @@ /// of MLIR without duplicating the type system or the op definitions. template <typename Op, typename... Args> static OperationHandle create(Args... args); - template <typename Op, typename... Args> static Op createOp(Args... args); + template <typename Op, typename... Args> + static Op createOp(Args... args); /// Generic create for a named operation. static OperationHandle create(StringRef name, ArrayRef<ValueHandle> operands, @@ -355,7 +356,8 @@ }; /// Simple wrapper to build a generic operation without successor blocks. -template <typename HandleType> struct CustomOperation { +template <typename HandleType> +struct CustomOperation { CustomOperation(StringRef name) : name(name) { static_assert(std::is_same<HandleType, ValueHandle>() || std::is_same<HandleType, OperationHandle>(), @@ -490,7 +492,8 @@ /// Store parameters. Assigning to an IndexedValue emits an actual `Store` /// operation, while converting an IndexedValue to a ValueHandle emits an actual /// `Load` operation. -template <typename Load, typename Store> class TemplatedIndexedValue { +template <typename Load, typename Store> +class TemplatedIndexedValue { public: explicit TemplatedIndexedValue(Type t) : base(t) {} explicit TemplatedIndexedValue(Value v) @@ -529,7 +532,7 @@ } /// Emits a `load` when converting to a Value. - Value operator*(void) const { + Value operator*(void)const { return Load(getBase(), {indices.begin(), indices.end()}).getValue(); } diff --git a/mlir/include/mlir/EDSC/Intrinsics.h b/mlir/include/mlir/EDSC/Intrinsics.h --- a/mlir/include/mlir/EDSC/Intrinsics.h +++ b/mlir/include/mlir/EDSC/Intrinsics.h @@ -55,7 +55,10 @@ SmallVector<Value, 8> values; }; -template <typename T> inline T unpack(T value) { return value; } +template <typename T> +inline T unpack(T value) { + return value; +} inline detail::ValueHandleArray unpack(ArrayRef<ValueHandle> values) { return detail::ValueHandleArray(values); @@ -70,7 +73,8 @@ /// Implementing it as a subclass allows it to compose all the way to Value. /// Without subclassing, implicit conversion to Value would fail when composing /// in patterns such as: `select(a, b, select(c, d, e))`. -template <typename Op> struct ValueBuilder : public ValueHandle { +template <typename Op> +struct ValueBuilder : public ValueHandle { // Builder-based template <typename... Args> ValueBuilder(Args... args) @@ -94,7 +98,8 @@ ValueBuilder() : ValueHandle(ValueHandle::create<Op>()) {} }; -template <typename Op> struct OperationBuilder : public OperationHandle { +template <typename Op> +struct OperationBuilder : public OperationHandle { template <typename... Args> OperationBuilder(Args... args) : OperationHandle(OperationHandle::create<Op>(detail::unpack(args)...)) {} diff --git a/mlir/lib/Dialect/StandardOps/EDSC/Intrinsics.cpp b/mlir/lib/Dialect/StandardOps/EDSC/Intrinsics.cpp --- a/mlir/lib/Dialect/StandardOps/EDSC/Intrinsics.cpp +++ b/mlir/lib/Dialect/StandardOps/EDSC/Intrinsics.cpp @@ -12,8 +12,8 @@ using namespace mlir; using namespace mlir::edsc; -OperationHandle mlir::edsc::intrinsics::br(BlockHandle bh, - ArrayRef<ValueHandle> operands) { +OperationHandle mlir::edsc::intrinsics::std_br(BlockHandle bh, + ArrayRef<ValueHandle> operands) { assert(bh && "Expected already captured BlockHandle"); for (auto &o : operands) { (void)o; @@ -22,6 +22,7 @@ SmallVector<Value, 4> ops(operands.begin(), operands.end()); return OperationHandle::create<BranchOp>(bh.getBlock(), ops); } + static void enforceEmptyCapturesMatchOperands(ArrayRef<ValueHandle *> captures, ArrayRef<ValueHandle> operands) { assert(captures.size() == operands.size() && @@ -36,9 +37,9 @@ } } -OperationHandle mlir::edsc::intrinsics::br(BlockHandle *bh, - ArrayRef<ValueHandle *> captures, - ArrayRef<ValueHandle> operands) { +OperationHandle mlir::edsc::intrinsics::std_br(BlockHandle *bh, + ArrayRef<ValueHandle *> captures, + ArrayRef<ValueHandle> operands) { assert(!*bh && "Unexpected already captured BlockHandle"); enforceEmptyCapturesMatchOperands(captures, operands); BlockBuilder(bh, captures)(/* no body */); @@ -47,17 +48,17 @@ } OperationHandle -mlir::edsc::intrinsics::cond_br(ValueHandle cond, BlockHandle trueBranch, - ArrayRef<ValueHandle> trueOperands, - BlockHandle falseBranch, - ArrayRef<ValueHandle> falseOperands) { +mlir::edsc::intrinsics::std_cond_br(ValueHandle cond, BlockHandle trueBranch, + ArrayRef<ValueHandle> trueOperands, + BlockHandle falseBranch, + ArrayRef<ValueHandle> falseOperands) { SmallVector<Value, 4> trueOps(trueOperands.begin(), trueOperands.end()); SmallVector<Value, 4> falseOps(falseOperands.begin(), falseOperands.end()); return OperationHandle::create<CondBranchOp>( cond, trueBranch.getBlock(), trueOps, falseBranch.getBlock(), falseOps); } -OperationHandle mlir::edsc::intrinsics::cond_br( +OperationHandle mlir::edsc::intrinsics::std_cond_br( ValueHandle cond, BlockHandle *trueBranch, ArrayRef<ValueHandle *> trueCaptures, ArrayRef<ValueHandle> trueOperands, BlockHandle *falseBranch, ArrayRef<ValueHandle *> falseCaptures, diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp --- a/mlir/test/EDSC/builder-api-test.cpp +++ b/mlir/test/EDSC/builder-api-test.cpp @@ -192,16 +192,16 @@ // b2 has not yet been constructed, need to come back later. // This is a byproduct of non-structured control-flow. ); - BlockBuilder(&b2, {&arg3, &arg4})([&] { br(b1, {arg3, arg4}); }); + BlockBuilder(&b2, {&arg3, &arg4})([&] { std_br(b1, {arg3, arg4}); }); // The insertion point within the toplevel function is now past b2, we will // need to get back the entry block. // This is what happens with unstructured control-flow.. BlockBuilder(b1, Append())([&] { r = arg1 + arg2; - br(b2, {arg1, r}); + std_br(b2, {arg1, r}); }); // Get back to entry block and add a branch into b1 - BlockBuilder(functionBlock, Append())([&] { br(b1, {c1, c2}); }); + BlockBuilder(functionBlock, Append())([&] { std_br(b1, {c1, c2}); }); // clang-format off // CHECK-LABEL: @builder_blocks @@ -234,15 +234,15 @@ BlockHandle b1, b2; { // Toplevel function scope. // Build a new block for b1 eagerly. - br(&b1, {&arg1, &arg2}, {c1, c2}); + std_br(&b1, {&arg1, &arg2}, {c1, c2}); // Construct a new block b2 explicitly with a branch into b1. BlockBuilder(&b2, {&arg3, &arg4})([&]{ - br(b1, {arg3, arg4}); + std_br(b1, {arg3, arg4}); }); /// And come back to append into b1 once b2 exists. BlockBuilder(b1, Append())([&]{ r = arg1 + arg2; - br(b2, {arg1, r}); + std_br(b2, {arg1, r}); }); } @@ -278,7 +278,7 @@ BlockBuilder(&b2, {&arg2, &arg3})([&] { std_ret(); }); // Get back to entry block and add a conditional branch BlockBuilder(functionBlock, Append())([&] { - cond_br(funcArg, b1, {c32}, b2, {c64, c42}); + std_cond_br(funcArg, b1, {c32}, b2, {c64, c42}); }); // clang-format off @@ -311,7 +311,7 @@ // clang-format off BlockHandle b1, b2; - cond_br(funcArg, &b1, {&arg1}, {c32}, &b2, {&arg2, &arg3}, {c64, c42}); + std_cond_br(funcArg, &b1, {&arg1}, {c32}, &b2, {&arg2, &arg3}, {c64, c42}); BlockBuilder(b1, Append())([]{ std_ret(); });