[mlir][spirv] Fix scf.yield pattern conversion

[kuhar]

Summary

Only rewrite scf.yield when the parent op is supported by
scf-to-spirv.

Fixes: #61380, #61107, #61148

[kuhar]

I discovered I skipped some existing tests, this doesn't always work. Working on a fix.

[kuhar]

Fix test failures. Add a regression test.

[antiagainst]

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

297

We don't need to check the dialect here given we are explicitly checking the ops?

[kuhar]

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

297

scf.yield can be inside an already converted op like spirv::LoopOp or other ops we don't know about. Not sure if there's a better way to guard against this. Alternatively, we could list all scf. ops we definitely don't support.

[antiagainst]

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

297

I see. Makes sense.

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

 public:
   using SCFToSPIRVPattern::SCFToSPIRVPattern;
 
   LogicalResult
   matchAndRewrite(scf::YieldOp terminatorOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     ValueRange operands = adaptor.getOperands();
 
-    // If the region is return values, store each value into the associated
+    Operation *parent = terminatorOp->getParentOp();
+    if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp>(parent))
+      return rewriter.notifyMatchFailure(
+          terminatorOp,
+          llvm::formatv("conversion not supported for parent op: '{0}'",
+                        parent->getName()));
+
+    // If the region return values, store each value into the associated
     // VariableOp created during lowering of the parent region.
     if (!operands.empty()) {
-      auto &allocas =
-          scfToSPIRVContext->outputVars[terminatorOp->getParentOp()];
+      auto &allocas = scfToSPIRVContext->outputVars[parent];
       if (allocas.size() != operands.size())
         return failure();
 
       auto loc = terminatorOp.getLoc();
       for (unsigned i = 0, e = operands.size(); i < e; i++)
         rewriter.create<spirv::StoreOp>(loc, allocas[i], operands[i]);
-      if (isa<spirv::LoopOp>(terminatorOp->getParentOp())) {
+      if (isa<spirv::LoopOp>(parent)) {
         // For loops we also need to update the branch jumping back to the
         // header.
         auto br = cast<spirv::BranchOp>(
             rewriter.getInsertionBlock()->getTerminator());
         SmallVector<Value, 8> args(br.getBlockArguments());
         args.append(operands.begin(), operands.end());
         rewriter.setInsertionPoint(br);
         rewriter.create<spirv::BranchOp>(terminatorOp.getLoc(), br.getTarget(),