Skip to content

Commit 78c8db8

Browse files
Jerry-GeTai78641
Jerry-Ge
authored and
Tai78641
committed
Add StatefulOps to TOSA Dialect
This patch adds tosa.variable, tosa.variable.read and tosa.variable.write operators and tests. Signed-off-by: Jerry Ge <[email protected]> Change-Id: I647e2e5c3762d7890b03f6aa7c09a29198b7d355
1 parent 750c8e3 commit 78c8db8

File tree

8 files changed

+290
-5
lines changed

8 files changed

+290
-5
lines changed

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ class PatternRewriter;
3434

3535
namespace tosa {
3636

37+
ParseResult parseTypeOrAttr(OpAsmParser &parser, TypeAttr &typeAttr,
38+
Attribute &attr);
39+
void printTypeOrAttr(OpAsmPrinter &p, Operation *op, TypeAttr type,
40+
Attribute attr);
41+
3742
#include "mlir/Dialect/Tosa/IR/TosaInterfaces.h.inc"
3843

3944
} // namespace tosa

‎mlir/include/mlir/Dialect/Tosa/IR/TosaUtilOps.td

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,71 @@ def Tosa_YieldOp : Tosa_Op<"yield", [
7979
let assemblyFormat = "$inputs attr-dict `:` type($inputs)";
8080
}
8181

82+
//===----------------------------------------------------------------------===//
83+
// Operator: variable
84+
//===----------------------------------------------------------------------===//
85+
def Tosa_VariableOp : Tosa_Op<"variable", []> {
86+
let summary = "Defines a variable";
87+
88+
let description = [{
89+
Defines a new TOSA variable. This is a mutable value.
90+
Modifications are expressed using read/write semantics.
91+
}];
92+
93+
let arguments = (ins
94+
SymbolNameAttr:$name,
95+
TypeAttr:$type,
96+
OptionalAttr<AnyAttr>:$initial_value
97+
);
98+
99+
let assemblyFormat = [{
100+
$name
101+
attr-dict
102+
custom<TypeOrAttr>($type, $initial_value)
103+
}];
104+
}
105+
106+
//===----------------------------------------------------------------------===//
107+
// Operator: variable.write
108+
//===----------------------------------------------------------------------===//
109+
def Tosa_VariableWriteOp : Tosa_Op<"variable.write", []> {
110+
let summary = "write_buffer operator";
111+
112+
let description = [{
113+
Assigns a value to pseudo-buffer resource holding a mutable tensor.
114+
}];
115+
116+
let arguments = (ins
117+
FlatSymbolRefAttr:$name,
118+
AnyType:$value
119+
);
120+
121+
let assemblyFormat = [{
122+
$name attr-dict `,` $value `:` type($value)
123+
}];
124+
}
125+
126+
//===----------------------------------------------------------------------===//
127+
// Operator: variable.read
128+
//===----------------------------------------------------------------------===//
129+
def Tosa_VariableReadOp : Tosa_Op<"variable.read", []> {
130+
let summary = "read_buffer operator";
131+
132+
let description = [{
133+
Reads the value from a pseudo-buffer resource holding a mutable tensor.
134+
}];
135+
136+
let arguments = (ins
137+
FlatSymbolRefAttr:$name
138+
);
139+
140+
let results = (outs
141+
AnyType:$value
142+
);
143+
144+
let assemblyFormat = [{
145+
$name attr-dict `:` type($value)
146+
}];
147+
}
148+
82149
#endif // TOSA_UTIL_OPS

‎mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ struct ValidationOptions {
6868
}
6969
};
7070

71-
std::unique_ptr<Pass> createTosaValidationPass(
71+
std::unique_ptr<OperationPass<ModuleOp>> createTosaValidationPass(
7272
ValidationOptions const &options = ValidationOptions());
7373

7474
#define GEN_PASS_REGISTRATION

‎mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ def TosaLevelType : I32EnumAttr<"TosaLevelEnum", "Tosa level",
8989
let cppNamespace = "mlir::tosa";
9090
}
9191

92-
def TosaValidation : Pass<"tosa-validate", "func::FuncOp"> {
92+
def TosaValidation : Pass<"tosa-validate", "mlir::ModuleOp"> {
9393
let summary = "Validates TOSA dialect";
9494
let description = [{
9595
This pass validates if input TOSA operations match the specification for given

‎mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,49 @@ Operation *TosaDialect::materializeConstant(OpBuilder &builder, Attribute value,
146146
return nullptr;
147147
}
148148

149+
//===----------------------------------------------------------------------===//
150+
// Parsers and printers
151+
//===----------------------------------------------------------------------===//
152+
153+
ParseResult mlir::tosa::parseTypeOrAttr(OpAsmParser &parser, TypeAttr &typeAttr,
154+
Attribute &attr) {
155+
if (succeeded(parser.parseOptionalEqual())) {
156+
if (failed(parser.parseAttribute(attr))) {
157+
return parser.emitError(parser.getCurrentLocation())
158+
<< "expected attribute";
159+
}
160+
if (auto typedAttr = attr.dyn_cast<TypedAttr>()) {
161+
typeAttr = TypeAttr::get(typedAttr.getType());
162+
}
163+
return success();
164+
}
165+
166+
Type type;
167+
if (failed(parser.parseColonType(type))) {
168+
return parser.emitError(parser.getCurrentLocation()) << "expected type";
169+
}
170+
typeAttr = TypeAttr::get(type);
171+
172+
return success();
173+
}
174+
175+
void mlir::tosa::printTypeOrAttr(OpAsmPrinter &p, Operation *op, TypeAttr type,
176+
Attribute attr) {
177+
bool needsSpace = false;
178+
auto typedAttr = attr.dyn_cast_or_null<TypedAttr>();
179+
if (!typedAttr || typedAttr.getType() != type.getValue()) {
180+
p << ": ";
181+
p.printAttribute(type);
182+
needsSpace = true; // subsequent attr value needs a space separator
183+
}
184+
if (attr) {
185+
if (needsSpace)
186+
p << ' ';
187+
p << "= ";
188+
p.printAttribute(attr);
189+
}
190+
}
191+
149192
//===----------------------------------------------------------------------===//
150193
// TOSA Operator Verifiers.
151194
//===----------------------------------------------------------------------===//

‎mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp

Lines changed: 95 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
1515
#include "mlir/Dialect/Tosa/Transforms/PassesEnums.cpp.inc"
1616

17+
#include <string>
18+
#include <unordered_map>
19+
1720
#include "mlir/Dialect/Func/IR/FuncOps.h"
1821
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
1922
#include "mlir/IR/Builders.h"
@@ -101,7 +104,7 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
101104
this->StrictOperationSpecAlignment = options.strictOperationSpecAlignment;
102105
this->level = options.level;
103106
}
104-
void runOnOperation() override;
107+
void runOnOperation() final;
105108

106109
LogicalResult applyConstantOperandCheck(Operation *op) {
107110
for (auto &checker : const_checkers) {
@@ -113,6 +116,9 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
113116

114117
LogicalResult applyLevelCheck(Operation *op);
115118

119+
// check variable read/write data types against variable declarations
120+
LogicalResult applyVariableCheck(Operation *op);
121+
116122
private:
117123
void populateConstantOperandChecks() {
118124
const_checkers.emplace_back(checkConstantOperandPad);
@@ -398,8 +404,12 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
398404
}
399405
}
400406

407+
bool CheckVariable(Operation *op);
408+
bool CheckVariableReadOrWrite(Operation *op);
409+
401410
SmallVector<std::function<LogicalResult(Operation *)>> const_checkers;
402411
tosa_level_t tosa_level;
412+
std::unordered_map<std::string, mlir::Type> variables_map;
403413
};
404414

405415
LogicalResult TosaValidation::applyLevelCheck(Operation *op) {
@@ -427,6 +437,83 @@ LogicalResult TosaValidation::applyLevelCheck(Operation *op) {
427437
return success();
428438
}
429439

440+
inline bool CompatibleTypes(const mlir::Type &type,
441+
const mlir::Type &declared_type) {
442+
// for now, simply use type equality comparison
443+
return type == declared_type;
444+
}
445+
446+
bool TosaValidation::CheckVariable(Operation *op) {
447+
if (isa<mlir::tosa::VariableOp>(op)) {
448+
auto name_attr = dyn_cast<mlir::StringAttr>(op->getAttr("name"));
449+
if (!name_attr) {
450+
op->emitOpError() << "Name attribute is not StringAttr";
451+
return false;
452+
}
453+
std::string name = name_attr.getValue().str();
454+
455+
if (variables_map.count(name)) {
456+
op->emitOpError() << "name has already been declared";
457+
return false;
458+
}
459+
460+
auto type_attr = dyn_cast<mlir::TypeAttr>(op->getAttr("type"));
461+
if (!type_attr) {
462+
op->emitOpError() << "type attribute is not TypeAttr";
463+
return false;
464+
}
465+
mlir::Type type = type_attr.getValue();
466+
467+
variables_map[name] = type;
468+
}
469+
470+
return true;
471+
}
472+
473+
bool TosaValidation::CheckVariableReadOrWrite(Operation *op) {
474+
if (isa<mlir::tosa::VariableReadOp>(op) ||
475+
isa<mlir::tosa::VariableWriteOp>(op)) {
476+
auto name_attr = dyn_cast<mlir::FlatSymbolRefAttr>(op->getAttr("name"));
477+
if (!name_attr) {
478+
op->emitOpError() << "name attribute is not FlatSymbolRefAttr";
479+
return false;
480+
}
481+
std::string name = name_attr.getValue().str();
482+
483+
if (!variables_map.count(name)) {
484+
op->emitOpError() << "name has not been declared";
485+
return false;
486+
}
487+
488+
auto var_type = variables_map[name];
489+
490+
for (auto v : op->getOperands()) {
491+
auto type = v.getType();
492+
if (!CompatibleTypes(type, var_type)) {
493+
op->emitOpError() << "operand type does not equal variable type";
494+
return false;
495+
}
496+
}
497+
498+
for (auto v : op->getResults()) {
499+
auto type = v.getType();
500+
if (!CompatibleTypes(type, var_type)) {
501+
op->emitOpError() << "result type does not equal variable type";
502+
return false;
503+
}
504+
}
505+
}
506+
507+
return true;
508+
}
509+
510+
LogicalResult TosaValidation::applyVariableCheck(Operation *op) {
511+
if (!CheckVariable(op) || !CheckVariableReadOrWrite(op)) {
512+
return failure();
513+
}
514+
return success();
515+
}
516+
430517
void TosaValidation::runOnOperation() {
431518
configLevelAndProfile();
432519
getOperation().walk([&](Operation *op) {
@@ -440,18 +527,23 @@ void TosaValidation::runOnOperation() {
440527
}
441528
}
442529

443-
// Some uses of TOSA rely on the constant operands of particular operations.
530+
// Some uses of TOSA rely on the constant operands of particular
531+
// operations.
444532
if (StrictOperationSpecAlignment && failed(applyConstantOperandCheck(op)))
445533
signalPassFailure();
446534

447535
// do level checks
448536
if (failed(applyLevelCheck(op)))
449537
signalPassFailure();
538+
539+
// do variable type checks
540+
if (failed(applyVariableCheck(op)))
541+
signalPassFailure();
450542
});
451543
}
452544
} // namespace
453545

454-
std::unique_ptr<Pass>
546+
std::unique_ptr<OperationPass<ModuleOp>>
455547
mlir::tosa::createTosaValidationPass(ValidationOptions const &options) {
456548
return std::make_unique<TosaValidation>(options);
457549
}

‎mlir/test/Dialect/Tosa/invalid.mlir

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,3 +203,48 @@ func.func @test_avg_pool2d_zero_dim_input(%arg0: tensor<1x0x?x9xf32>) -> tensor<
203203
: (tensor<1x0x?x9xf32>) -> tensor<1x7x7x9xf32>
204204
return %0 : tensor<1x7x7x9xf32>
205205
}
206+
207+
// -----
208+
209+
func.func @test_variable_duplicates(%arg0: tensor<2x4x8xi32>) -> () {
210+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
211+
// expected-error@+1 {{'tosa.variable' op name has already been declared}}
212+
tosa.variable @stored_var : tensor<1x4x8xi32>
213+
return
214+
}
215+
216+
// -----
217+
218+
func.func @test_variable_read_type(%arg0: tensor<2x4x8xi32>) -> () {
219+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
220+
// expected-error@+1 {{'tosa.variable.read' op result type does not equal variable type}}
221+
%0 = tosa.variable.read @stored_var : tensor<2x4x8xi16>
222+
return
223+
}
224+
225+
// -----
226+
227+
func.func @test_variable_read_shape(%arg0: tensor<2x4x8xi32>) -> () {
228+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
229+
// expected-error@+1 {{'tosa.variable.read' op result type does not equal variable type}}
230+
%0 = tosa.variable.read @stored_var : tensor<1x4x8xi32>
231+
return
232+
}
233+
234+
// -----
235+
236+
func.func @test_variable_write_type(%arg0: tensor<2x4x8xi16>) -> () {
237+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
238+
// expected-error@+1 {{'tosa.variable.write' op operand type does not equal variable type}}
239+
tosa.variable.write @stored_var, %arg0 : tensor<2x4x8xi16>
240+
return
241+
}
242+
243+
// -----
244+
245+
func.func @test_variable_write_shape(%arg0: tensor<1x4x8xi32>) -> () {
246+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
247+
// expected-error@+1 {{'tosa.variable.write' op operand type does not equal variable type}}
248+
tosa.variable.write @stored_var, %arg0 : tensor<1x4x8xi32>
249+
return
250+
}

‎mlir/test/Dialect/Tosa/variables.mlir

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
// RUN: mlir-opt %s | mlir-opt | FileCheck %s
2+
// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
3+
4+
5+
// -----
6+
// CHECK-LABEL: @test_variable_scalar(
7+
// CHECK-SAME: %[[ADD_VAL:.*]]: tensor<f32>) {
8+
func.func @test_variable_scalar(%arg0: tensor<f32>) -> () {
9+
// CHECK: tosa.variable @stored_var = dense<3.140000e+00> : tensor<f32>
10+
tosa.variable @stored_var = dense<3.14> : tensor<f32>
11+
// CHECK: %[[STORED_VAL:.*]] = tosa.variable.read @stored_var : tensor<f32>
12+
%0 = tosa.variable.read @stored_var : tensor<f32>
13+
// CHECK: %[[RESULT_ADD:.*]] = tosa.add %[[ADD_VAL]], %[[STORED_VAL]] : (tensor<f32>, tensor<f32>) -> tensor<f32>
14+
%1 = "tosa.add"(%arg0, %0) : (tensor<f32>, tensor<f32>) -> tensor<f32>
15+
// CHECK: tosa.variable.write @stored_var, %[[RESULT_ADD]] : tensor<f32>
16+
tosa.variable.write @stored_var, %1 : tensor<f32>
17+
return
18+
}
19+
20+
// -----
21+
// CHECK-LABEL: @test_variable_tensor(
22+
// CHECK-SAME: %[[ADD_VAL:.*]]: tensor<2x4x8xi32>) {
23+
func.func @test_variable_tensor(%arg0: tensor<2x4x8xi32>) -> () {
24+
// CHECK: tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
25+
tosa.variable @stored_var = dense<-1> : tensor<2x4x8xi32>
26+
// CHECK: %[[STORED_VAL:.*]] = tosa.variable.read @stored_var : tensor<2x4x8xi32>
27+
%0 = tosa.variable.read @stored_var : tensor<2x4x8xi32>
28+
// CHECK: %[[RESULT_ADD:.*]] = tosa.add %[[ADD_VAL]], %[[STORED_VAL]] : (tensor<2x4x8xi32>, tensor<2x4x8xi32>) -> tensor<2x4x8xi32>
29+
%1 = "tosa.add"(%arg0, %0) : (tensor<2x4x8xi32>, tensor<2x4x8xi32>) -> tensor<2x4x8xi32>
30+
// CHECK: tosa.variable.write @stored_var, %[[RESULT_ADD]] : tensor<2x4x8xi32>
31+
tosa.variable.write @stored_var, %1 : tensor<2x4x8xi32>
32+
return
33+
}

0 commit comments

Comments
 (0)