[spirv] Allow return ops to be in control flow ops

Use `getParentOfType<FunctionOp>()` instead of `cast<FuncOp>(getParentOp())`
to avoid crash when return ops are used inside spv.selection/spv.loop.

PiperOrigin-RevId: 273006041
This commit is contained in:
Lei Zhang 2019-10-04 20:08:05 -07:00 committed by TensorFlower Gardener
parent 38a28a71c6
commit c5c69cd6a8

View File

@ -1745,7 +1745,7 @@ static LogicalResult verify(spirv::ReferenceOfOp referenceOfOp) {
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
static LogicalResult verify(spirv::ReturnOp returnOp) { static LogicalResult verify(spirv::ReturnOp returnOp) {
auto funcOp = cast<FuncOp>(returnOp.getParentOp()); auto funcOp = returnOp.getParentOfType<FuncOp>();
auto numOutputs = funcOp.getType().getNumResults(); auto numOutputs = funcOp.getType().getNumResults();
if (numOutputs != 0) if (numOutputs != 0)
return returnOp.emitOpError("cannot be used in functions returning value") return returnOp.emitOpError("cannot be used in functions returning value")
@ -1774,7 +1774,7 @@ static void print(spirv::ReturnValueOp retValOp, OpAsmPrinter &printer) {
} }
static LogicalResult verify(spirv::ReturnValueOp retValOp) { static LogicalResult verify(spirv::ReturnValueOp retValOp) {
auto funcOp = cast<FuncOp>(retValOp.getParentOp()); auto funcOp = retValOp.getParentOfType<FuncOp>();
auto numFnResults = funcOp.getType().getNumResults(); auto numFnResults = funcOp.getType().getNumResults();
if (numFnResults != 1) if (numFnResults != 1)
return retValOp.emitOpError( return retValOp.emitOpError(