[mlir][spirv] Account for type conversion failures in scf-to-spirv

[kuhar]

Summary

Fixes: https://github.com/llvm/llvm-project/issues/59136

[antiagainst]

mlir/test/Conversion/SCFToSPIRV/if.mlir

156

typo: class

[kuhar]

fix typo

mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp

 
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Context
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
IfOpConversion::matchAndRewrite(scf::IfOp ifOp, OpAdaptor adaptor,
   rewriter.setInsertionPointToEnd(selectionHeaderBlock);
   rewriter.create<spirv::BranchConditionalOp>(loc, adaptor.getCondition(),
                                               thenBlock, ArrayRef<Value>(),
                                               elseBlock, ArrayRef<Value>());
 
   SmallVector<Type, 8> returnTypes;
   for (auto result : ifOp.getResults()) {
     auto convertedType = typeConverter.convertType(result.getType());
+    if (!convertedType)
+      return rewriter.notifyMatchFailure(
+          loc, llvm::formatv("failed to convert type '{0}'", result.getType()));
+
     returnTypes.push_back(convertedType);
   }
   replaceSCFOutputValue(ifOp, selectionOp, rewriter, scfToSPIRVContext,
                         returnTypes);
   return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/SCFToSPIRV/if.mlir

%0 = scf.if %arg4 -> (memref<10xf32, #spirv.storage_class<StorageBuffer>>) {
     scf.yield %arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>
   } else {
     scf.yield %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>
   }
   memref.store %value, %0[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>
   return
 }
 
+// Memrefs without a spirv storage class are not supported. The conversion
+// should preserve the `scf.if` and not crash.
+func.func @unsupported_yield_type(%arg0 : memref<8xi32>, %arg1 : memref<8xi32>, %c : i1) {
+// CHECK-LABEL: @unsupported_yield_type
+// CHECK-NEXT:    scf.if
+// CHECK:         spirv.Return
+  %r = scf.if %c -> (memref<8xi32>) {
+    scf.yield %arg0 : memref<8xi32>
+  } else {
+    scf.yield %arg1 : memref<8xi32>
+  }
+  return
+}
+
 } // end module