[mlir] Fix a cast that should be a dyn_cast.

[jreiffers]

Summary

This fixes a crash for certain IR, see the new test case for an
example.

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

struct WhileOpInterface
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
                 const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     auto whileOp = cast<scf::WhileOp>(op);
     assert(getOwnerOfValue(value) == op && "invalid value");
     assert(value.getType().isa<TensorType>() && "expected tensor type");
 
     // Case 1: Block argument of the "before" region.
-    if (auto bbArg = value.cast<BlockArgument>()) {
+    if (auto bbArg = value.dyn_cast<BlockArgument>()) {
       if (bbArg.getOwner()->getParent() == &whileOp.getBefore()) {
         Value initArg = whileOp.getInits()[bbArg.getArgNumber()];
         auto yieldOp = whileOp.getYieldOp();
         Value yieldedValue = yieldOp.getOperand(bbArg.getArgNumber());
         return computeLoopRegionIterArgBufferType(bbArg, initArg, yieldedValue,
                                                   options, fixedTypes);
       }
     }

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

func.func @buffer_type_of_collapse_shape(%arg0: tensor<f64>) {
   } do {
   ^bb0(%_: tensor<f64>):
     %3 = bufferization.alloc_tensor() : tensor<1xf64>
     %16 = tensor.collapse_shape %3 [] : tensor<1xf64> into tensor<f64>
     scf.yield %16 : tensor<f64>
   }
   return
 }
+
+// -----
+
+// This is a regression test. Just check that the IR bufferizes.
+  
+// CHECK-LABEL: func @non_block_argument_yield
+func.func @non_block_argument_yield() {
+  %true = arith.constant true 
+  %0 = bufferization.alloc_tensor() : tensor<i32>
+  %1 = scf.while (%arg0 = %0) : (tensor<i32>) -> (tensor<i32>) {
+    scf.condition(%true) %arg0 : tensor<i32>
+  } do {
+  ^bb0(%arg0: tensor<i32>):
+    %ret = scf.while (%arg1 = %0) : (tensor<i32>) -> (tensor<i32>) {
+      scf.condition(%true) %arg1 : tensor<i32>
+    } do {
+    ^bb0(%arg7: tensor<i32>):
+      scf.yield %0 : tensor<i32>
+    }
+    scf.yield %ret : tensor<i32>
+  }
+  return
+}