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