[mlir][linalg] Extend linalg vectorization to support non-identity input maps

[ThomasRaoux]

Summary

This handles cases where the linalg input needs to be transposed.

[asaadaldien]

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

363

if its isElementwise && isMinorIdentity we still need to broadcast. Can we push this down to merge the two path by inserting transposes if !isMinorIdentity ?

[ThomasRaoux]

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

363

if its isElementwise && isMinorIdentity we still need to broadcast. Can we push this down to merge the two path by inserting transposes if !isMinorIdentity ?

Correct, I was considering merging those two cases. That would simplify slightly the vectorization however that could generate less efficient code as right now the broadcasting is done lazily and if some operations don't need to be expanded on all the output dimensions it may pick a vector of smaller rank.
I'm not sure how likely this case is to happen so if people think it is not worth the extra complexity I can merge those two cases.

Note that this could also be optimized away by canonicalizations later on. I don't have a strong opinion, so I can change it based on what you, Nicolas and Aart think.

[asaadaldien]

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

273

Nit: Add a comment broadcasting into inner most dims.

363

If we have can one of operands at vectorized to a lower-rank that means all its consumers are at most at the same lower-rank. In other words we can split the linalg.generic into two independent ops operating at lower-rank & higher rank domains.....
IMO we can keep it that way, this is more general and the added complicity isn't much =)
Lets wait for @nicolasvasilache, @aartbik opinion.

[nicolasvasilache]

The vector.transfer_read / write operation have support for expressing the permutation and broadcast semantics as part of the permutation_map.
The transformation you are doing here looks like the kind of more general stuff needed to rewrite that permutation_map into vector.transpose + broadcast.

The tradeoff is that vector.transfer + permutation can lower better as loop + memory indexing reordering.
Introducing transpose + broadcast at the time of Linalg -> Vector is probably lowering too quickly.
OTOH we could also think of implementing some of this as canonicalizations of vector.transfer + transpose + broadcast.

This also has to be thought of in the context of vector.shape_cast and future scalable vectors.

So bottom line: I am a conflicted.
This looks more generally applicable and writing it in 2 steps as: 1) Linalg -> vector.transfer + permutation followed by 2) better support of permutation_map when lowering vector.transfer better satisfies my intuitions
OTOH it is also possible that the permutation_map in vector.transfer is more complicated to use; but this has to be weighed off against the alternative.

Can you put an extra time to chat in my cal tomorrow so we go a bit deeper with higher BW?
I'll put a blocker in the meantime because this sits at the intersection of a bunch of different things.

mlir/test/Dialect/Linalg/vectorization.mlir

346

args_in/out attributes are not needed anymore

[ThomasRaoux]

Use transfer_read maps instead of emitting directly broadcast/transpose

[ThomasRaoux]

In D98523#2625614, @nicolasvasilache wrote:

The vector.transfer_read / write operation have support for expressing the permutation and broadcast semantics as part of the permutation_map.
The transformation you are doing here looks like the kind of more general stuff needed to rewrite that permutation_map into vector.transpose + broadcast.

The tradeoff is that vector.transfer + permutation can lower better as loop + memory indexing reordering.
Introducing transpose + broadcast at the time of Linalg -> Vector is probably lowering too quickly.
OTOH we could also think of implementing some of this as canonicalizations of vector.transfer + transpose + broadcast.

This also has to be thought of in the context of vector.shape_cast and future scalable vectors.

So bottom line: I am a conflicted.
This looks more generally applicable and writing it in 2 steps as: 1) Linalg -> vector.transfer + permutation followed by 2) better support of permutation_map when lowering vector.transfer better satisfies my intuitions
OTOH it is also possible that the permutation_map in vector.transfer is more complicated to use; but this has to be weighed off against the alternative.

Can you put an extra time to chat in my cal tomorrow so we go a bit deeper with higher BW?
I'll put a blocker in the meantime because this sits at the intersection of a bunch of different things.

As discussed I switched to using affine map in vector transfer instead. I'll prepare another patch once we agree on this review to do the lowering of those transfer ops with affine maps.

[nicolasvasilache]

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

255

// -> ///

272

// -> ///

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

vectorizeOneOp(OpBuilder &builder, Operation *op,
   OperationState state(op->getLoc(), op->getName());
   state.addAttributes(op->getAttrs());
   state.addOperands(llvm::to_vector<4>(vectorizedOperands));
   state.addTypes(llvm::to_vector<4>(returnTypes));
   return VectorizationResult{VectorizationStatus::NewOp,
                              builder.createOperation(state)};
 }
 
+/// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
+static bool hasOnlyScalarElementwiseOp(Region &r) {
+  if (!llvm::hasSingleElement(r))
+    return false;
+  for (Operation &op : r.front()) {
+    if (!(isa<ConstantOp, linalg::YieldOp>(op) ||
+          OpTrait::hasElementwiseMappableTraits(&op)) ||
+        llvm::any_of(op.getResultTypes(),
+                     [](Type type) { return !type.isIntOrIndexOrFloat(); }))
+      return false;
+  }
+  return true;
+}
+
+// Return true if the op is an element-wise linalg op.
+static bool isElementwise(Operation *op) {
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+  if (!linalgOp)
+    return false;
+  if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
+    return false;
+  // TODO: relax the restrictions on indexing map.
+  for (unsigned i = 0, e = linalgOp.getNumOutputs(); i < e; i++) {
+    if (!linalgOp.getOutputIndexingMap(i).isIdentity())
+      return false;
+  }
+  if (linalgOp->getNumRegions() != 1)
+    return false;
+  return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
+}
+
+// Broadcast and transpose a linalgOp argument to a shape matching the linalg
+// output shape.
+static Value broadcastAndTranspose(OpBuilder &builder, Value vectorRead,
+                                   LinalgOp linalgOp, unsigned argIndex) {
+  Location loc = linalgOp.getLoc();
+  // Currently assume we don't support output permutations.
+  assert(linalgOp.getNumOutputs() > 0 &&
+         linalgOp.getOutputIndexingMap(0).isIdentity());
+  ArrayRef<int64_t> outputShape = linalgOp.getOutputShapedType(0).getShape();
+  AffineMap map = linalgOp.getIndexingMap(argIndex);
+  unsigned numMissingDim = map.getNumInputs() - map.getNumResults();
+  SmallVector<int64_t> permutation(map.getNumInputs());
+  llvm::SmallVector<bool, 8> existingDim(map.getNumInputs(), false);
+  for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
+    unsigned pos = map.getDimPosition(i);
+    existingDim[pos] = true;
+    permutation[pos] = numMissingDim + i;
+  }
+  // If there are missing dimensions first broadcast the vector.
+  if (numMissingDim > 0) {
+    // Insert the missing dimensions as high rank dimensions.
+    unsigned index = 0;
+    SmallVector<int64_t> broadcastShape;
+    for (unsigned i : llvm::seq<unsigned>(0, map.getNumInputs())) {
+      if (!existingDim[i]) {
+        broadcastShape.push_back(outputShape[i]);
+        permutation[i] = index++;
+      }
+    }
+    auto sourceType = vectorRead.getType().cast<VectorType>();
+    broadcastShape.append(sourceType.getShape().begin(),
+                          sourceType.getShape().end());
+    auto newVecType =
+        VectorType::get(broadcastShape, sourceType.getElementType());
+    vectorRead =
+        builder.create<vector::BroadcastOp>(loc, newVecType, vectorRead);
+  }
+  Value transposeOp =
+      builder.create<vector::TransposeOp>(loc, vectorRead, permutation);
+  return transposeOp;
+}
+
 /// Generic vectorization function that rewrites the body of a `linalgOp` into
 /// vector form. Generic vectorization proceeds as follows:
 ///   1. The region for the linalg op is created if necessary.
 ///   2. Values defined above the region are mapped to themselves and will be
 ///   broadcasted on a per-need basis by their consumers.
 ///   3. Each region argument is vectorized into a vector.transfer_read (or 0-d
 ///   load).
 ///   TODO: Reuse opportunities for RAR dependencies.
LogicalResult vectorizeAsLinalgGeneric(
 
   // 3. Turn all BBArgs into vector.transfer_read / load.
   SmallVector<AffineMap> indexings;
   for (auto bbarg : block->getArguments()) {
     Value vectorArg = linalgOp.getShapedOperand(bbarg.getArgNumber());
     Value vectorRead = buildVectorRead(builder, vectorArg);
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: new vectorized bbarg("
                       << bbarg.getArgNumber() << "): " << vectorRead);
+    // Apply broadcast and transpose if needed.
+    if (isElementwise(linalgOp) &&
+        !linalgOp.getIndexingMap(bbarg.getArgNumber()).isMinorIdentity())
+      vectorRead = broadcastAndTranspose(builder, vectorRead, linalgOp,
+                                         bbarg.getArgNumber());
     bvm.map(bbarg, vectorRead);
     bvm.map(vectorArg, vectorRead);
   }
 
   // 4. Register CustomVectorizationHook for yieldOp.
   CustomVectorizationHook vectorizeYield =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
if (result.status == VectorizationStatus::NewOp) {
                         << *result.newOp;);
       bvm.map(op.getResults(), result.newOp->getResults());
     }
   }
 
   return success();
 }
 
-/// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
-static bool hasOnlyScalarElementwiseOp(Region &r) {
-  if (!llvm::hasSingleElement(r))
-    return false;
-  for (Operation &op : r.front()) {
-    if (!(isa<ConstantOp, linalg::YieldOp>(op) ||
-          OpTrait::hasElementwiseMappableTraits(&op)) ||
-        llvm::any_of(op.getResultTypes(),
-                     [](Type type) { return !type.isIntOrIndexOrFloat(); }))
-      return false;
-  }
-  return true;
-}
-
-// Return true if the op is an element-wise linalg op.
-static bool isElementwise(Operation *op) {
-  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
-  if (!linalgOp)
-    return false;
-  if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
-    return false;
-  // TODO: relax the restrictions on indexing map.
-  for (unsigned i = 0, e = linalgOp.getNumOutputs(); i < e; i++) {
-    if (!linalgOp.getOutputIndexingMap(i).isIdentity())
-      return false;
-  }
-  // Currently bound the input indexing map to minor identity as other
-  // permutations might require adding transpose ops to convert the vector read
-  // to the right shape.
-  for (unsigned i = 0, e = linalgOp.getNumInputs(); i < e; i++) {
-    if (!linalgOp.getInputIndexingMap(i).isMinorIdentity())
-      return false;
-  }
-  if (linalgOp->getNumRegions() != 1)
-    return false;
-  return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
-}
-
 static LogicalResult vectorizeContraction(OpBuilder &builder, LinalgOp linalgOp,
                                           SmallVectorImpl<Value> &newResults) {
   assert(isaContractionOpInterface(linalgOp) &&
          "expected vectorizeContraction preconditions to be met");
   Location loc = linalgOp.getLoc();
   // Vectorize other ops as vector contraction.
   // TODO: interface.
   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "

mlir/test/Dialect/Linalg/vectorization.mlir

func @generic_vectorize_tensor(%arg0: tensor<4x256xf32>,
   return %r#0, %r#1, %r#2, %r#3, %r#4, %r#5, %r#6, %r#7, %r#8, %r#9:
     tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>,
     tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>,
     tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>
 }
 
 // -----
 
+// Test different input maps.
+#matmul_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [
+    affine_map<(d0, d1, d2, d3) -> (d1, d0)>,
+    affine_map<(d0, d1, d2, d3) -> (d3, d1)>,
+    affine_map<(d0, d1, d2, d3) -> (d3, d1, d0, d2)>,
+    affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+  ],
+  iterator_types = ["parallel", "parallel", "parallel", "parallel"]
+}
+// CHECK-LABEL: func @vectorization_transpose
+//       CHECK: vector.transfer_read {{.*}} : memref<14x7xf32>, vector<14x7xf32>
+//       CHECK: vector.broadcast {{.*}} : vector<14x7xf32> to vector<8x16x14x7xf32>
+//       CHECK: vector.transpose {{.*}}, [3, 2, 0, 1] : vector<8x16x14x7xf32> to vector<7x14x8x16xf32>
+//       CHECK: vector.transfer_read {{.*}} : memref<16x14xf32>, vector<16x14xf32>
+//       CHECK: vector.broadcast {{.*}} : vector<16x14xf32> to vector<7x8x16x14xf32>
+//       CHECK: vector.transpose {{.*}}, [0, 3, 1, 2] : vector<7x8x16x14xf32> to vector<7x14x8x16xf32>
+//       CHECK: vector.transfer_read {{.*}} : memref<16x14x7x8xf32>, vector<16x14x7x8xf32>
+//       CHECK: vector.transpose {{.*}}, [2, 1, 3, 0] : vector<16x14x7x8xf32> to vector<7x14x8x16xf32>
+//       CHECK: addf {{.*}} : vector<7x14x8x16xf32>
+//       CHECK: vector.transfer_write {{.*}} : vector<7x14x8x16xf32>, memref<7x14x8x16xf32>
+func @vectorization_transpose(%A: memref<14x7xf32>, %B: memref<16x14xf32>,
+                         %C: memref<16x14x7x8xf32>, %D: memref<7x14x8x16xf32>) {
+  linalg.generic #matmul_trait
+    ins(%A, %B, %C : memref<14x7xf32>, memref<16x14xf32>, memref<16x14x7x8xf32>)
+   outs(%D : memref<7x14x8x16xf32>) {
+    ^bb(%a: f32, %b: f32, %c: f32, %d: f32) :
+      %e = addf %a, %b: f32
+      %f = addf %e, %c: f32
+      linalg.yield %f : f32
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @matmul_tensors
 //  CHECK-SAME: (%[[ARG0:.*]]: tensor<8x4xf32>, %[[ARG1:.*]]: tensor<4x12xf32>,
 //  CHECK-SAME:  %[[ARG2:.*]]: tensor<8x12xf32>) -> tensor<8x12xf32>
 func @matmul_tensors(
   %arg0: tensor<8x4xf32>, %arg1: tensor<4x12xf32>, %arg2: tensor<8x12xf32>)
     -> tensor<8x12xf32> {
   //   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
   //   CHECK-DAG:   %[[VEC_C0:.*]] = constant dense<0.000000e+00> : vector<8x12xf32>