[mlir][Vector] Fix vectorization of generic ops with transposed outputs

[dcaballe]

Summary

This patch fixes a bug in the way we compute the vector type for vector
transfer writes when the value to store needs to be transposed.

[mravishankar]

Nice!

[nicolasvasilache]

mlir/lib/IR/AffineMap.cpp

570

This looks a bit weird and very specific from an API perspective, I wouldn't know when to use in other cases than your very specific use case.

How about something more reusable like (very rough sketch .. alternatives possible)

AffineMap AffineMap::getFilteredMultiDimIdentityMap(MLIRContext *ctx, int64_t numDims, llvm::function_ref<bool(AffineDimExpr)> keepDim) {
  auto m = getMultiDimIdentityMap(numDims, numSyms, ctx);
  // apply keepDim
  return m.dropResults(projectedDims);
}


// call 
auto m = AffineMap::getFilteredMultiDimIdentityMap(ctx, m.getNumDims, [](AffineDimExpr d){ return llvm::any_of(m.getResults(), [](AffineExpr e) { return e.isFunctionOf(d) ;})});

The double lambda nesting may be unappealing, sorry it's Sat night .. :p but you get the gist.
Feel free to introduce other helpers if they are generally reusable and not 1-off.

mlir/test/Dialect/Linalg/vectorization.mlir

1779

note: we should have an op able to simply target a particular op to vectorize, this would greatly increase applicability and avoid chasing parent ops and applying many patterns on the fly.
The existing transform.vectorize is dated and too monolithic but still used because "it works ok enough".
Bonus points if you quickly whip up a transform.vectorize_single_op.

I am open to renaming the existing op too but that will require more changes than we should spend effort on right now and is a good intro cleanup task.

[dcaballe]

Address feedback. Thanks!

[dcaballe]

Thanks for the quick review!

mlir/lib/IR/AffineMap.cpp

570

Thanks! Make sense!

mlir/test/Dialect/Linalg/vectorization.mlir

1779

I think @awarzynski have been looking a bit more into the vectorization ops in the transform dialect. Maybe he can help with this although I understand this is low priority.

[awarzynski]

mlir/test/Dialect/Linalg/vectorization.mlir

1779

I might have a few spare cycles for this. I assume that you are thinking of something similar to transform.structured.masked_vectorize:

transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
  transform.structured.masked_vectorize %0 vector_sizes [4] : !transform.any_op
}

mlir/include/mlir/IR/AffineMap.h

public:
   static AffineMap getMultiDimIdentityMap(unsigned numDims,
                                           MLIRContext *context);
 
   /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
   /// minor dimensions.
   static AffineMap getMinorIdentityMap(unsigned dims, unsigned results,
                                        MLIRContext *context);
 
+  /// Returns an identity affine map witn `numDims` input dimensions and
+  /// filtered results using `keepDimFilter`. If `keepDimFilter` returns true
+  /// for a dimension, the dimension is kept in the affine map results.
+  /// Otherwise, the dimension is dropped from the results.
+  ///
+  /// Examples:
+  ///   * getFilteredIdentityMap(4, [false, true, false, true])
+  ///       -> affine_map<(d0, d1, d2, d3) -> (d1, d3)>
+  ///   * getFilteredIdentityMap(3, [false, false, true])
+  ///       -> affine_map<(d0, d1, d2) -> (d2)>
+  static AffineMap
+  getFilteredIdentityMap(MLIRContext *ctx, unsigned numDims,
+                         llvm::function_ref<bool(AffineDimExpr)> keepDimFilter);
+
   /// Returns an AffineMap representing a permutation.
   /// The permutation is expressed as a non-empty vector of integers.
   /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with
   /// `permutation = [1,2,0]`. All values in `permutation` must be
   /// integers, in the range 0..`permutation.size()-1` without duplications
   /// (i.e. `[1,1,2]` is an invalid permutation).
   static AffineMap getPermutationMap(ArrayRef<unsigned> permutation,
                                      MLIRContext *context);

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

 // Note: this is a true builder that notifies the OpBuilder listener.
 // TODO: Consider moving as a static helper on the ReduceOp.
 static Value buildVectorWrite(RewriterBase &rewriter, Value value,
                               OpOperand *outputOperand,
                               VectorizationState &state) {
   Location loc = value.getLoc();
   auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
   AffineMap opOperandMap = linalgOp.getMatchingIndexingMap(outputOperand);
+
+  // Compute the vector type of the value to store. This type should be an
+  // identity or projection of the canonical vector type without any permutation
+  // applied, given that any permutation in a transfer write happens as part of
+  // the write itself.
+  AffineMap vectorTypeMap = AffineMap::getFilteredIdentityMap(
+      opOperandMap.getContext(), opOperandMap.getNumInputs(),
+      [&](AffineDimExpr dimExpr) -> bool {
+        return llvm::is_contained(opOperandMap.getResults(), dimExpr);
+      });
   auto vectorType = state.getCanonicalVecType(
-      getElementTypeOrSelf(outputOperand->get().getType()), opOperandMap);
+      getElementTypeOrSelf(outputOperand->get().getType()), vectorTypeMap);
 
   Operation *write;
   if (vectorType.getRank() > 0) {
     AffineMap writeMap = inversePermutation(reindexIndexingMap(opOperandMap));
     SmallVector<Value> indices(linalgOp.getRank(outputOperand),
                                rewriter.create<arith::ConstantIndexOp>(loc, 0));
     value = broadcastIfNeeded(rewriter, value, vectorType);
+    assert(value.getType() == vectorType && "Incorrect type");
     write = rewriter.create<vector::TransferWriteOp>(
         loc, value, outputOperand->get(), indices, writeMap);
   } else {
     // 0-d case is still special: do not invert the reindexing writeMap.
     if (!isa<VectorType>(value.getType()))
       value = rewriter.create<vector::BroadcastOp>(loc, vectorType, value);
-    assert(value.getType() == vectorType && "incorrect type");
+    assert(value.getType() == vectorType && "Incorrect type");
     write = rewriter.create<vector::TransferWriteOp>(
         loc, value, outputOperand->get(), ValueRange{});
   }
 
   write = state.maskOperation(rewriter, write, linalgOp, opOperandMap);
 
   // If masked, set in-bounds to true. Masking guarantees that the access will
   // be in-bounds.

mlir/lib/IR/AffineMap.cpp

 /// minor dimensions.
 AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
                                          MLIRContext *context) {
   assert(dims >= results && "Dimension mismatch");
   auto id = AffineMap::getMultiDimIdentityMap(dims, context);
   return AffineMap::get(dims, 0, id.getResults().take_back(results), context);
 }
 
+AffineMap AffineMap::getFilteredIdentityMap(
+    MLIRContext *ctx, unsigned numDims,
+    llvm::function_ref<bool(AffineDimExpr)> keepDimFilter) {
+  auto identityMap = getMultiDimIdentityMap(numDims, ctx);
+
+  // Apply filter to results.
+  llvm::SmallBitVector dropDimResults(numDims);
+  for (auto [idx, resultExpr] : llvm::enumerate(identityMap.getResults()))
+    dropDimResults[idx] = !keepDimFilter(resultExpr.cast<AffineDimExpr>());
+
+  return identityMap.dropResults(dropDimResults);
+}
+
 bool AffineMap::isMinorIdentity() const {
   return getNumDims() >= getNumResults() &&
          *this ==
              getMinorIdentityMap(getNumDims(), getNumResults(), getContext());
 }
 
 /// Returns true if this affine map is a minor identity up to broadcasted
 /// dimensions which are indicated by value 0 in the result.

 }
 
 bool AffineMap::isPermutation() const {
   if (getNumDims() != getNumResults())
     return false;
   return isProjectedPermutation();
 }
 
 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {
   SmallVector<AffineExpr, 4> exprs;
   exprs.reserve(resultPos.size());
   for (auto idx : resultPos)
     exprs.push_back(getResult(idx));
   return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
 }
 
 AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const {

mlir/test/Dialect/Linalg/vectorization.mlir

 //       CHECK:     vector.transfer_read {{.*}} : tensor<f32>, vector<f32>
 //       CHECK:     vector.extractelement
 //       CHECK:     vector.transfer_read {{.*}} : tensor<f32>, vector<f32>
 //       CHECK:     vector.extractelement
 //       CHECK:     arith.addf {{.*}} : f32
 //       CHECK:     vector.broadcast %{{.*}} : f32 to vector<f32>
 //       CHECK:     vector.transfer_write {{.*}} : vector<f32>, tensor<f32>
 
+// -----
+
+// Make sure we generate the right transfer writes for multi-output generic ops
+// with different permutation maps.
+
+func.func @multi_output_generic_different_perm_maps(%in0: tensor<4x1xf32>,
+                                                    %out0: tensor<4x1xf32>,
+                                                    %out1: tensor<1x4xf32>) -> (tensor<4x1xf32>, tensor<1x4xf32>) {
+  %13:2 = linalg.generic {indexing_maps = [ affine_map<(d0, d1) -> (d1, d0)>,
+                                            affine_map<(d0, d1) -> (d1, d0)>,
+                                            affine_map<(d0, d1) -> (d0, d1)> ],
+                          iterator_types = ["parallel", "parallel"]}
+                          ins(%in0 : tensor<4x1xf32>)
+                          outs(%out0, %out1 : tensor<4x1xf32>, tensor<1x4xf32>) {
+  ^bb0(%in: f32, %out: f32, %out_2: f32):
+    %16 = arith.addf %in, %in : f32
+    linalg.yield %16, %16 : f32, f32
+  } -> (tensor<4x1xf32>, tensor<1x4xf32>)
+  return %13#0, %13#1 : tensor<4x1xf32>, tensor<1x4xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %3 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %4 = get_closest_isolated_parent %3 : (!transform.any_op) -> !transform.any_op
+  %5 = transform.structured.vectorize %4 : (!transform.any_op) -> !transform.any_op
+}
+
+// CHECK-LABEL: func @multi_output_generic_different_perm_maps
+//       CHECK:     %[[VAL_5:.*]] = vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<4x1xf32>, vector<4x1xf32>
+//       CHECK:     %[[VAL_6:.*]] = arith.addf %[[VAL_5]], %[[VAL_5]] : vector<4x1xf32>
+//       CHECK:     %[[VAL_7:.*]] = vector.transpose %[[VAL_6]], [1, 0] : vector<4x1xf32> to vector<1x4xf32>
+//       CHECK:     %[[VAL_8:.*]] = vector.transpose %[[VAL_7]], [1, 0] : vector<1x4xf32> to vector<4x1xf32>
+//       CHECK:     vector.transfer_write %[[VAL_8]], %{{.*}} {in_bounds = [true, true]} : vector<4x1xf32>, tensor<4x1xf32>
+//       CHECK:     vector.transfer_write %[[VAL_7]], %{{.*}} {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>