[mlir] Fix asan issue in Vectorization.cpp of Linalg.

[dcaballe]

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

531

I'm curious... What is asan complaining about here?

[pifon2a]

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

531

For linalg.map there is no block argument that corresponds to the init/output and the old code was still trying to call getArgument() with the Idx which is out-of-bounds.

[dcaballe]

Thanks for clarifying, Alex!

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

vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
   SmallVector<int64_t> commonVectorShape = linalgOp.computeStaticLoopSizes();
 
   // 3. Turn all BBArgs into vector.transfer_read / load.
   Location loc = linalgOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
-  for (OpOperand &opOperand : linalgOp->getOpOperands()) {
-    BlockArgument bbarg = block->getArgument(opOperand.getOperandNumber());
-    if (linalgOp.isScalar(&opOperand)) {
-      bvm.map(bbarg, opOperand.get());
+  for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
+    BlockArgument bbarg = linalgOp.getMatchingBlockArgument(opOperand);
+    if (linalgOp.isScalar(opOperand)) {
+      bvm.map(bbarg, opOperand->get());
       continue;
     }
     VectorType readType;
     AffineMap map;
     // TODO: can we keep this simplification?
     // if (linalgOp.getShape(&opOperand).empty()) {
     //   readType = VectorType::get({}, bbarg.getType());
     // } else {
-    if (opOperand.getOperandNumber() < linalgOp.getNumInputs()) {
+    if (opOperand->getOperandNumber() < linalgOp.getNumInputs()) {
       map = inverseAndBroadcastProjectedPermutation(
-          linalgOp.getMatchingIndexingMap(&opOperand));
+          linalgOp.getMatchingIndexingMap(opOperand));
       readType = VectorType::get(commonVectorShape,
-                                 getElementTypeOrSelf(opOperand.get()));
+                                 getElementTypeOrSelf(opOperand->get()));
     } else {
       map = inversePermutation(
-          reindexIndexingMap(linalgOp.getMatchingIndexingMap(&opOperand)));
-      readType = VectorType::get(map.compose(linalgOp.getShape(&opOperand)),
-                                 getElementTypeOrSelf(opOperand.get()));
+          reindexIndexingMap(linalgOp.getMatchingIndexingMap(opOperand)));
+      readType = VectorType::get(map.compose(linalgOp.getShape(opOperand)),
+                                 getElementTypeOrSelf(opOperand->get()));
     }
     // }
 
-    auto shape = linalgOp.getShape(&opOperand);
+    auto shape = linalgOp.getShape(opOperand);
     SmallVector<Value> indices(shape.size(), zero);
     Value readValue = b.create<vector::TransferReadOp>(
-        loc, readType, opOperand.get(), indices, map);
+        loc, readType, opOperand->get(), indices, map);
     // Not all ops support 0-d vectors, extract the scalar for now.
     // TODO: remove this.
     if (readValue.getType().cast<VectorType>().getRank() == 0)
       readValue = b.create<vector::ExtractElementOp>(loc, readValue);
 
     LDBG("new vectorized bbarg(" << bbarg.getArgNumber() << "): " << readValue);
     bvm.map(bbarg, readValue);
-    bvm.map(opOperand.get(), readValue);
+    bvm.map(opOperand->get(), readValue);
   }
 
   SmallVector<CustomVectorizationHook> hooks;
   // 4a. Register CustomVectorizationHook for yieldOp.
   CustomVectorizationHook vectorizeYield =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeLinalgYield(b, op, bvm, linalgOp, newResults);