[mlir] Fix vectorize transform crashing on none-op operand

[asaadaldien]

Summary

Not all operands have a defining operator. without this lowering the following will crash:

func @fill_func(%arg0 : memref<2x2xf32>, %arg1: f32) {
  linalg.fill(%arg0, %arg1) : memref<2x2xf32>, f32
  return
}

[rriddle]

mlir/lib/Transforms/Vectorize.cpp

1009

You can just change the dyn_cast to dyn_cast_or_null instead.

[merge_guards_bot]

Unit tests: pass. 62114 tests passed, 0 failed and 808 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[asaadaldien]

use dyn_cast_or_null

[merge_guards_bot]

Unit tests: pass. 62114 tests passed, 0 failed and 808 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nicolasvasilache]

Note that for the purpose of vectorizing linalg.fill I was thinking about creating a pattern similar to this: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp#L192

In essence this is about creating a vectorization pattern for linalg.fill, not about converting to affine and using vectorization on that.
Note that the latter would also be a valid way but a bunch of things need to happen before it works and is not the shortest path.

[rriddle]

Is there a test case you can add?

[asaadaldien]

@rriddle This already landed (not sure why not showing that here!) the example in the summary I will added it as a test in a follow up diff.

mlir/lib/Transforms/Vectorize.cpp

if (it != state->replacementMap.end()) {
     return res;
   }
   // 2. TODO(ntv): broadcast needed.
   if (operand.getType().isa<VectorType>()) {
     LLVM_DEBUG(dbgs() << "-> non-vectorizable");
     return nullptr;
   }
   // 3. vectorize constant.
-  if (auto constant = dyn_cast<ConstantOp>(operand.getDefiningOp())) {
+  if (auto constant = dyn_cast_or_null<ConstantOp>(operand.getDefiningOp())) {
     return vectorizeConstant(
         op, constant,
         VectorType::get(state->strategy->vectorSizes, operand.getType()));
   }
   // 4. currently non-vectorizable.
   LLVM_DEBUG(dbgs() << "-> non-vectorizable: " << operand);
   return nullptr;
 }
 
 /// Encodes Operation-specific behavior for vectorization. In general we assume
 /// that all operands of an op must be vectorized but this is not always true.
static Operation *vectorizeOneOperation(Operation *opInst,
   assert(!isa<vector::TransferWriteOp>(opInst) &&
          "vector.transfer_write cannot be further vectorized");
 
   if (auto store = dyn_cast<AffineStoreOp>(opInst)) {
     OpBuilder b(opInst);
     auto memRef = store.getMemRef();
     auto value = store.getValueToStore();
     auto vectorValue = vectorizeOperand(value, opInst, state);
+    if (!vectorValue)
+      return nullptr;
 
     ValueRange mapOperands = store.getMapOperands();
     SmallVector<Value, 8> indices;
     indices.reserve(store.getMemRefType().getRank());
     if (store.getAffineMap() !=
         b.getMultiDimIdentityMap(store.getMemRefType().getRank())) {
       computeMemoryOpIndices(opInst, store.getAffineMap(), mapOperands,
                              indices);