1 //===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "mlir/EDSC/Builders.h"
10 #include "mlir/IR/AffineExpr.h"
11 #include "mlir/IR/AffineMap.h"
12
13 #include "llvm/ADT/Optional.h"
14
15 using namespace mlir;
16 using namespace mlir::edsc;
17
ScopedContext(OpBuilder & b)18 mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b)
19 : ScopedContext(b, b.getInsertionPoint()->getLoc()) {}
20
ScopedContext(OpBuilder & b,Location location)21 mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b, Location location)
22 : builder(b), guard(builder), location(location),
23 enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
24 getCurrentScopedContext() = this;
25 }
26
27 /// Sets the insertion point of the builder to 'newInsertPt' for the duration
28 /// of the scope. The existing insertion point of the builder is restored on
29 /// destruction.
ScopedContext(OpBuilder & b,OpBuilder::InsertPoint newInsertPt,Location location)30 mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b,
31 OpBuilder::InsertPoint newInsertPt,
32 Location location)
33 : builder(b), guard(builder), location(location),
34 enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
35 getCurrentScopedContext() = this;
36 builder.restoreInsertionPoint(newInsertPt);
37 }
38
~ScopedContext()39 mlir::edsc::ScopedContext::~ScopedContext() {
40 getCurrentScopedContext() = enclosingScopedContext;
41 }
42
getCurrentScopedContext()43 ScopedContext *&mlir::edsc::ScopedContext::getCurrentScopedContext() {
44 thread_local ScopedContext *context = nullptr;
45 return context;
46 }
47
getBuilderRef()48 OpBuilder &mlir::edsc::ScopedContext::getBuilderRef() {
49 assert(ScopedContext::getCurrentScopedContext() &&
50 "Unexpected Null ScopedContext");
51 return ScopedContext::getCurrentScopedContext()->builder;
52 }
53
getLocation()54 Location mlir::edsc::ScopedContext::getLocation() {
55 assert(ScopedContext::getCurrentScopedContext() &&
56 "Unexpected Null ScopedContext");
57 return ScopedContext::getCurrentScopedContext()->location;
58 }
59
getContext()60 MLIRContext *mlir::edsc::ScopedContext::getContext() {
61 return getBuilderRef().getContext();
62 }
63
createBlock(TypeRange argTypes)64 Block *mlir::edsc::createBlock(TypeRange argTypes) {
65 assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
66 OpBuilder &builder = ScopedContext::getBuilderRef();
67 Block *block = builder.getInsertionBlock();
68 assert(block != nullptr &&
69 "insertion point not set up in the builder within ScopedContext");
70
71 return createBlockInRegion(*block->getParent(), argTypes);
72 }
73
createBlockInRegion(Region & region,TypeRange argTypes)74 Block *mlir::edsc::createBlockInRegion(Region ®ion, TypeRange argTypes) {
75 assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
76 OpBuilder &builder = ScopedContext::getBuilderRef();
77
78 OpBuilder::InsertionGuard guard(builder);
79 return builder.createBlock(®ion, {}, argTypes);
80 }
81
appendToBlock(Block * block,function_ref<void (ValueRange)> builderFn)82 void mlir::edsc::appendToBlock(Block *block,
83 function_ref<void(ValueRange)> builderFn) {
84 assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
85 OpBuilder &builder = ScopedContext::getBuilderRef();
86
87 OpBuilder::InsertionGuard guard(builder);
88 if (block->empty() || block->back().isKnownNonTerminator())
89 builder.setInsertionPointToEnd(block);
90 else
91 builder.setInsertionPoint(&block->back());
92 builderFn(block->getArguments());
93 }
94
buildInNewBlock(TypeRange argTypes,function_ref<void (ValueRange)> builderFn)95 Block *mlir::edsc::buildInNewBlock(TypeRange argTypes,
96 function_ref<void(ValueRange)> builderFn) {
97 assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
98 OpBuilder &builder = ScopedContext::getBuilderRef();
99 Block *block = builder.getInsertionBlock();
100 assert(block != nullptr &&
101 "insertion point not set up in the builder within ScopedContext");
102 return buildInNewBlock(*block->getParent(), argTypes, builderFn);
103 }
104
buildInNewBlock(Region & region,TypeRange argTypes,function_ref<void (ValueRange)> builderFn)105 Block *mlir::edsc::buildInNewBlock(Region ®ion, TypeRange argTypes,
106 function_ref<void(ValueRange)> builderFn) {
107 assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
108 OpBuilder &builder = ScopedContext::getBuilderRef();
109
110 Block *block = createBlockInRegion(region, argTypes);
111 OpBuilder::InsertionGuard guard(builder);
112 builder.setInsertionPointToStart(block);
113 builderFn(block->getArguments());
114 return block;
115 }
116
117