Skip to content

[mlir][tosa] Check for isolated regions in tosa.while_loop #144865

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[mlir][tosa] Check for isolated regions in tosa.while_loop #144865

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
25dc942
[mlir][tosa] Fix check for isolated regions in `tosa.cond_if`
lhutton1 Jun 11, 2025
7bd13d2
[mlir][tosa] Check for isolated regions in `tosa.while_loop`
lhutton1 Jun 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 75 additions & 37 deletions mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1193,61 +1193,99 @@ bool checkErrorIfPad(Operation *op) {
return true;
}

// Returns true if the operation takes no input operands, excluding attributes.
static bool isNullaryOperation(Operation *op) {
if (isa<tosa::ConstOp>(op) || isa<tosa::ConstShapeOp>(op) ||
isa<tosa::YieldOp>(op) || isa<tosa::VariableOp>(op))
return true;
return false;
static bool isOpIsolatedFromAbove(Operation *op, Region &region) {
Copy link
Contributor

@udaya-ranga udaya-ranga Jul 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will there just be 1 copy after merging this and the other review for cond.if or is this patch built on top of the other so that it appears here as well ? I guess the latter ?

return llvm::all_of(op->getOperands(), [&](auto operand) {
Region *operandRegion = operand.getParentRegion();
return region.isAncestor(operandRegion);
});
}

static bool isRegionIsolatedFromAbove(Region &region) {
bool noLiveInValue = true;
region.walk([&noLiveInValue, &region](Operation *op) {
if (!isOpIsolatedFromAbove(op, region)) {
noLiveInValue = false;
return WalkResult::interrupt();
}
return WalkResult::advance();
});
return noLiveInValue;
}

bool checkErrorIfCondIf(Operation *op) {
auto ifOp = dyn_cast<tosa::IfOp>(op);
if (!ifOp)
return true;

// Whether the types and shapes of operands between the input/output list and
// internal regions are validated by the operation verifier. However, with
// support for the simplified form - where redundant operand notations are
// omitted - is not conformant to the specification. According to the
// specification, all operands passed into an operation must be explicitly
// declared at each operation's structure. This code section verify that the
// operation's form complies with this requirement.

// Returns true if the region uses no external input operands.
auto isNullaryRegion = [](Region &region) -> bool {
bool noLiveInValue = true;
region.walk([&noLiveInValue](Operation *op) {
if (!isNullaryOperation(op)) {
noLiveInValue = false;
return WalkResult::interrupt();
}
return WalkResult::advance();
});
return noLiveInValue;
};
// Currently the dialect supports declaring cond_if operations that
// have then/else regions that reference values from outside these
// regions. According to the specification, all values used by the
// then/else regions must be explicitly declared within the regions.
// Therefore we must check that the then/else regions are
// "isolated from above", in order to be conformant to the
// specification.
//
// Note: the dialect currently supports two styles of syntax for
// declaring "cond_if" operations. We'll refer to these as follows:
//
// Generic:
// %0 = "tosa.cond_if"(%arg0, %arg1, %arg2) ({
// ^bb0(%arg3, %arg4):
// tosa.yield %arg3
// }, {
// ^bb0(%arg3, %arg4):
// tosa.yield %arg4
// })
//
// Simplified:
// %0 = tosa.cond_if %arg2 {
// tosa.yield %arg0
// } else {
// tosa.yield %arg1
// }
//
// Unfortunately, the simplified syntax does not encapsulate values
// used in then/else regions (see 'simplified' example above), so it
// must be rewritten to use the generic syntax in order to be conformant
// to the specification.
Region &thenGraph = ifOp.getThenGraph();
Region &elseGraph = ifOp.getElseGraph();
bool isThenGraphIsolatedRegion = isRegionIsolatedFromAbove(thenGraph);
bool isElseGraphIsolatedRegion = isRegionIsolatedFromAbove(elseGraph);

if (!isThenGraphIsolatedRegion || !isElseGraphIsolatedRegion) {
op->emitOpError()
<< "is not conformant to the TOSA specification. It requires the "
"then/else regions are isolated from above.\n";
return false;
}
return true;
}

mlir::Region &thenGraph = ifOp.getThenGraph();
mlir::Region &elseGraph = ifOp.getElseGraph();
bool isThenGraphNullaryRegion = isNullaryRegion(thenGraph);
bool isElseGraphNullaryRegion = isNullaryRegion(elseGraph);
bool isInputListEmpty = ifOp.getInputList().size() == 0;
bool checkErrorIfWhileLoop(Operation *op) {
auto whileOp = dyn_cast<tosa::WhileOp>(op);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this be const?

if (!whileOp)
return true;

if ((isInputListEmpty != isThenGraphNullaryRegion) ||
(isInputListEmpty != isElseGraphNullaryRegion)) {
Region &condGraph = whileOp.getCondGraph();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

const?

Region &bodyGraph = whileOp.getBodyGraph();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

const?

bool isCondGraphIsolatedRegion = isRegionIsolatedFromAbove(condGraph);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

const?

bool isBodyGraphIsolatedRegion = isRegionIsolatedFromAbove(bodyGraph);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

const?


if (!isCondGraphIsolatedRegion || !isBodyGraphIsolatedRegion) {
op->emitOpError()
<< "the current simplified form is not strictly conformant to the "
"spec, please use the generic format\n";
<< "is not conformant to the TOSA specification. It requires the "
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly to the if case, we could provide more information here if we were to split the checks into two. Only a week preference for this though.

"cond/body regions are isolated from above.\n";
return false;
}

return true;
}

LogicalResult TosaValidation::applyErrorIfCheck(Operation *op) {
if (!checkErrorIfResize(op) || !checkErrorIfMul(op) ||
!checkErrorIfTable(op) || !checkErrorIfRescale(op) ||
!checkErrorIfPad(op) || !checkErrorIfCondIf(op))
!checkErrorIfPad(op) || !checkErrorIfCondIf(op) ||
!checkErrorIfWhileLoop(op))
return failure();
return success();
}
Expand Down
78 changes: 71 additions & 7 deletions mlir/test/Dialect/Tosa/error_if_check.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,15 +227,79 @@ func.func @test_error_i32_unsigned_output(%arg0: tensor<1xi8>) -> tensor<1xi32>
}

// -----
// CHECK-LABEL: cond_if_simplified_form
func.func @test_cond_if_simplified_form(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
// expected-error@+1 {{'tosa.cond_if' op the current simplified form is not strictly conformant to the spec, please use the generic format}}

func.func @test_cond_if_not_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
// expected-error@+1 {{'tosa.cond_if' op is not conformant to the TOSA specification. It requires the then/else regions are isolated from above.}}
%0 = "tosa.cond_if"(%arg2) ({
^bb0():
tosa.yield %arg0 : tensor<f32>
}, {
^bb0():
tosa.yield %arg1 : tensor<f32>
}) : (tensor<i1>) -> tensor<f32>
return %0 : tensor<f32>
}

// -----

func.func @test_cond_if_simplified_form_not_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
// expected-error@+1 {{'tosa.cond_if' op is not conformant to the TOSA specification. It requires the then/else regions are isolated from above.}}
%0 = tosa.cond_if %arg2 -> (tensor<f32>) {
%1 = tosa.add %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
tosa.yield %1 : tensor<f32>
tosa.yield %arg0 : tensor<f32>
} else {
%1 = tosa.sub %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
tosa.yield %1 : tensor<f32>
tosa.yield %arg1 : tensor<f32>
}
return %0 : tensor<f32>
}

// -----

// COM: Check isolated cond_if's are valid
func.func @test_cond_if_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
%0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
tosa.yield %arg3 : tensor<f32>
}, {
^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
tosa.yield %arg4 : tensor<f32>
}) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
return %0 : tensor<f32>
}

// -----

func.func @test_while_loop_not_isolated_from_above(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>) {
%0 = "tosa.const"() {values = dense<0> : tensor<i32>} : () -> tensor<i32>
// expected-error@+1 {{'tosa.while_loop' op is not conformant to the TOSA specification. It requires the cond/body regions are isolated from above.}}
%1 = "tosa.while_loop"(%0) ({
^bb0(%arg3: tensor<i32>):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should the blocks have the same ID bb0?

%2 = "tosa.greater_equal"(%arg3, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
%3 = "tosa.logical_not"(%2) : (tensor<i1>) -> tensor<i1>
tosa.yield %3 : tensor<i1>
}, {
^bb0(%arg3: tensor<i32>):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here it looks like the body graph is isolated from above, and it is only the condition graph that isn't. Should we add another test to check the error is triggered when the body graph isn't isolated from above?

%2 = "tosa.const"() {values = dense<1> : tensor<i32>} : () -> tensor<i32>
%3 = "tosa.add"(%arg3, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
tosa.yield %3 : tensor<i32>
}) : (tensor<i32>) -> (tensor<i32>)
return
}

// -----

// COM: Check isolated while_loops are valid
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to use the COM directive here if there are no strings in the comment that could be interpreted as a FileCheck directive?

func.func @test_while_loop_isolated_from_above(%arg0: tensor<f32>, %arg1: tensor<i32>) {
%0 = "tosa.const"() {values = dense<0> : tensor<i32>} : () -> tensor<i32>
%1:3 = "tosa.while_loop"(%0, %arg0, %arg1) ({
^bb0(%arg3: tensor<i32>, %arg4: tensor<f32>, %arg5: tensor<i32>):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should the blocks have the same ID bb0?

%2 = "tosa.greater_equal"(%arg3, %arg5) : (tensor<i32>, tensor<i32>) -> tensor<i1>
%3 = "tosa.logical_not"(%2) : (tensor<i1>) -> tensor<i1>
"tosa.yield"(%3) : (tensor<i1>) -> ()
}, {
^bb0(%arg3: tensor<i32>, %arg4: tensor<f32>, %arg5: tensor<i32>):
%2 = "tosa.const"() {values = dense<1> : tensor<i32>} : () -> tensor<i32>
%3 = "tosa.add"(%arg3, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"tosa.yield"(%3, %arg4, %arg5) : (tensor<i32>, tensor<f32>, tensor<i32>) -> ()
}) : (tensor<i32>, tensor<f32>, tensor<i32>) -> (tensor<i32>, tensor<f32>, tensor<i32>)
return
}