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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// These BufferizableOpInterface implementations provide analysis-related
// interface methods only. They are getting bufferized by the
// SparseTensorConversion pass.
#include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
using namespace mlir::bufferization;
using namespace mlir::sparse_tensor;
namespace mlir {
namespace sparse_tensor {
namespace {
struct ConvertOpInterface
: public BufferizableOpInterface::ExternalModel<ConvertOpInterface,
sparse_tensor::ConvertOp> {
bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
// ConvertOps may allocate. (Unless they convert between two identical
// types, then they fold away.)
return true;
}
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return false;
}
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return {};
}
bool isWritable(Operation *op, Value value,
const AnalysisState &state) const {
return true;
}
};
struct LoadOpInterface
: public BufferizableOpInterface::ExternalModel<LoadOpInterface,
sparse_tensor::LoadOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return false;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return false;
}
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return {op->getOpResult(0)};
}
BufferRelation bufferRelation(Operation *op, OpResult opResult,
const AnalysisState &state) const {
return BufferRelation::Equivalent;
}
};
struct NewOpInterface
: public BufferizableOpInterface::ExternalModel<NewOpInterface,
sparse_tensor::NewOp> {
bool isMemoryWrite(Operation *op, OpResult opResult,
const AnalysisState &state) const {
// NewOps allocate but do not write.
return false;
}
bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
return true;
}
};
struct InsertOpInterface
: public BufferizableOpInterface::ExternalModel<InsertOpInterface,
sparse_tensor::InsertOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
// InsertOp writes to memory.
return true;
}
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
// InsertOp returns an alias of its operand.
assert(op->getNumResults() == 1);
return op->getResults();
}
BufferRelation bufferRelation(Operation *oo, OpResult opResult,
const AnalysisState &state) const {
// InsertOp returns the same object (realloc should not invalidate
// aliases).
return BufferRelation::Equivalent;
}
};
} // namespace
} // namespace sparse_tensor
} // namespace mlir
void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(
+[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) {
sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
sparse_tensor::InsertOp::attachInterface<InsertOpInterface>(*ctx);
});
}
|
