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
114
115
116
117
|
//===- RewriteInsertsPass.cpp - MLIR conversion pass ----------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to rewrite sequential chains of
// `spirv::CompositeInsert` operations into `spirv::CompositeConstruct`
// operations.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
namespace mlir {
namespace spirv {
#define GEN_PASS_DEF_SPIRVREWRITEINSERTSPASS
#include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
} // namespace spirv
} // namespace mlir
using namespace mlir;
namespace {
/// Replaces sequential chains of `spirv::CompositeInsertOp` operation into
/// `spirv::CompositeConstructOp` operation if possible.
class RewriteInsertsPass
: public spirv::impl::SPIRVRewriteInsertsPassBase<RewriteInsertsPass> {
public:
void runOnOperation() override;
private:
/// Collects a sequential insertion chain by the given
/// `spirv::CompositeInsertOp` operation, if the given operation is the last
/// in the chain.
LogicalResult
collectInsertionChain(spirv::CompositeInsertOp op,
SmallVectorImpl<spirv::CompositeInsertOp> &insertions);
};
} // namespace
void RewriteInsertsPass::runOnOperation() {
SmallVector<SmallVector<spirv::CompositeInsertOp, 4>, 4> workList;
getOperation().walk([this, &workList](spirv::CompositeInsertOp op) {
SmallVector<spirv::CompositeInsertOp, 4> insertions;
if (succeeded(collectInsertionChain(op, insertions)))
workList.push_back(insertions);
});
for (const auto &insertions : workList) {
auto lastCompositeInsertOp = insertions.back();
auto compositeType = lastCompositeInsertOp.getType();
auto location = lastCompositeInsertOp.getLoc();
SmallVector<Value, 4> operands;
// Collect inserted objects.
for (auto insertionOp : insertions)
operands.push_back(insertionOp.getObject());
OpBuilder builder(lastCompositeInsertOp);
auto compositeConstructOp = builder.create<spirv::CompositeConstructOp>(
location, compositeType, operands);
lastCompositeInsertOp.replaceAllUsesWith(
compositeConstructOp->getResult(0));
// Erase ops.
for (auto insertOp : llvm::reverse(insertions)) {
auto *op = insertOp.getOperation();
if (op->use_empty())
insertOp.erase();
}
}
}
LogicalResult RewriteInsertsPass::collectInsertionChain(
spirv::CompositeInsertOp op,
SmallVectorImpl<spirv::CompositeInsertOp> &insertions) {
auto indicesArrayAttr = cast<ArrayAttr>(op.getIndices());
// TODO: handle nested composite object.
if (indicesArrayAttr.size() == 1) {
auto numElements = cast<spirv::CompositeType>(op.getComposite().getType())
.getNumElements();
auto index = cast<IntegerAttr>(indicesArrayAttr[0]).getInt();
// Need a last index to collect a sequential chain.
if (index + 1 != numElements)
return failure();
insertions.resize(numElements);
while (true) {
insertions[index] = op;
if (index == 0)
return success();
op = op.getComposite().getDefiningOp<spirv::CompositeInsertOp>();
if (!op)
return failure();
--index;
indicesArrayAttr = cast<ArrayAttr>(op.getIndices());
if ((indicesArrayAttr.size() != 1) ||
(cast<IntegerAttr>(indicesArrayAttr[0]).getInt() != index))
return failure();
}
}
return failure();
}
|
