[mlir][sparse] disable nonunit stride dense vectorization

[aartbik]

Summary

This is a temporary work-around to get our all-annotations-all-flags
stress testing effort run clean. In the long run, we want to provide
efficient implementations of strided loads and stores though

mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp

 }
 
 /// Generates end of a reduction.
 static void genReductionEnd(Merger &merger, CodeGen &codegen,
                             PatternRewriter &rewriter, linalg::GenericOp op) {
   Value red = codegen.redVal;
   if (!red)
     return;
+  assert(codegen.curVecLength == 1);
   codegen.redVal = merger.exp(codegen.redExp).val = Value(); // end chain
   unsigned lhs = op.getNumShapedOperands() - 1;
   if (red.getType().isa<VectorType>()) {
     // TODO: assumes + reductions for now
-    codegen.curVecLength = 1;
     Value ld = genTensorLoad(merger, codegen, rewriter, op, codegen.redExp);
     red = rewriter.create<vector::ReductionOp>(
         op.getLoc(), ld.getType(), rewriter.getStringAttr("add"), red, ld);
   }
   genTensorStore(merger, codegen, rewriter, op, lhs, red);
 }
 
 /// Recursively generates tensor expression.
static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isReduction,
   case linalg::SparseParallelizationStrategy::kDenseAnyLoop:
     return !isSparse && !isReduction && !isVector;
   case linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop:
     return !isReduction && !isVector;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
+/// Checks unit strides for dense tensors. The iteration graph may have ignored
+/// dense access patterns in order to avoid cycles (sparse access patterns are
+/// always placed innermost), but that means dense access has become strided.
+/// For now, we reject vectorization of such cases.
+/// TODO: implement strided load/stores on dense arrays
+static bool denseUnitStrides(Merger &merger, linalg::GenericOp op,
+                             unsigned idx) {
+  unsigned numTensors = op.getNumShapedOperands();
+  for (unsigned t = 0; t < numTensors; t++) {
+    if (!merger.isSparseTensor(t) && !linkedSparse(op, t)) {
+      auto map = op.getIndexingMap(t);
+      unsigned r = map.getNumResults();
+      if (r && map.getDimPosition(r - 1) != idx)
+        return false;
+    }
+  }
+  return true;
+}
+
 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen,
                          PatternRewriter &rewriter, linalg::GenericOp op,
                          bool isOuter, bool isInner, unsigned idx,
                          llvm::BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
   auto iteratorTypes = op.iterator_types().getValue();
   bool isReduction = linalg::isReductionIteratorType(iteratorTypes[idx]);
   bool isSparse = merger.isDim(fb, Dim::kSparse);
-  bool isVector = isVectorFor(codegen, isInner, isSparse);
+  bool isVector = isVectorFor(codegen, isInner, isSparse) &&
+                  denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
   // Prepare vector length.
   if (isVector)
     codegen.curVecLength = codegen.options.vectorLength;
 
   // Loop bounds and increment.
if (isWhile) {
         rewriter.create<scf::YieldOp>(loc, codegen.redVal);
         codegen.redVal = loop->getResult(0);
       }
     }
     rewriter.setInsertionPointAfter(loop);
   }
 
   // Wrap-up loop sequence.
+  codegen.curVecLength = 1;
   genReductionEnd(merger, codegen, rewriter, op);
   genInvariants(merger, codegen, rewriter, op, exp, ldx, /*hoist=*/false);
   codegen.loops[idx] = Value();
-  codegen.curVecLength = 1;
 }
 
 namespace {
 
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 public:
   GenericOpSparsifier(MLIRContext *context, linalg::SparsificationOptions o)