Introduce linalg.vecmat

[burmako]

Summary

This patch adds a new named structured op to accompany linalg.matmul and
linalg.matvec. We needed it for our codegen, so I figured it would be useful
to add it to Linalg.

[mravishankar]

Not specific to the patch itself, but I am really concerned about the ballooning surface area of named ops in Linalg. I mentioned it before with the autogeneration from TC language approach to begin with, that one of the great things about Linalg is that many of the ops from dialects like MHLO (which have more justification for wide surface area of ops) can be boiled down to a "few" ops in Linalg. Adding more such operations here seems to be going away from that really useful feature. For example in this patch itself

1. The ops to loops method now needs to be updated.
2. The vectorization patterns now needs to be updated.

This is only going to get worse, and writing transformations is going to be really painful. For example in IREE, we have been working with linalg.conv. Now there are 7 versions of convolution in Linalg. We can use interfaces to solve a lot of this problem, but looks like the transformations above did not use interfaces. Either there is a fundamental reason for that, or its not implemented. If its the latter, that should probably addressed sooner rather than later.

Specifically for this patch, the op definition seems to just be linalg.matvec with operands swapped. Not sure I understand why this necessitates a separate operation.

[nicolasvasilache]

@burmako ah my apologies, I thought you wanted a linalg op that would operate on memref<vector>.

Vecmat indeed looks just like matvec with permuted operands, as @mravishankar points out.
It seems your uses of vecmat could easily be turned into matvec ?

[burmako]

Fixed the definition of linalg.vecmat - the original definition was wrong.

[burmako]

@mravishankar @nicolasvasilache Sorry for the confusion! My original patch messed up the definition of vecmat, making it look like it was just matvec with the parameters swapped. Please take a look at the new version of the patch.

I understand the concerns with increasing the surface area of named ops in Linalg, and I understand that this alone might be ground for rejecting this patch. Please let me know what you think about the new formulation though.

[mravishankar]

@burmako This makes sense. I suspected this is what you were planning for, but as I said, it is not specific to this patch. I understand why you need this.
FWIW

vecmat(y, A) = transpose(matvec(transpose(A), transpose(y)))

So we can always fold these away into a smaller op surface area, but this is effectively what the MHLO -> Linalg conversion does.

As I said earlier, not related to this patch, but it would be good to have some guidance/documentation about when a new named op is needed.

Accepting from my side. Please wait for @nicolasvasilache .

[bondhugula]

I chatted with Nicolas about this earlier today, and he was okay with this idea.

Please drop these from the commit summary - you could include it in a separate comment below.

[nicolasvasilache]

I am not concerned. by having many ops as long as the underlying mechanisms are shared.
Of course, I agree that more core work is needed to get new op definitions to be just the 4 lines in .tc.
This is just a bit of tablegen + macro away so I am not worried that a future CL will make things more automatic.

[burmako]

Rebase to account for recent changes

mlir/include/mlir/Dialect/Linalg/EDSC/Intrinsics.h

 namespace edsc {
 namespace intrinsics {
 
 using linalg_copy = OperationBuilder<linalg::CopyOp>;
 using linalg_dot = OperationBuilder<linalg::DotOp>;
 using linalg_fill = OperationBuilder<linalg::FillOp>;
 using linalg_matmul = OperationBuilder<linalg::MatmulOp>;
 using linalg_matvec = OperationBuilder<linalg::MatvecOp>;
+using linalg_vecmat = OperationBuilder<linalg::VecmatOp>;
 using linalg_range = ValueBuilder<linalg::RangeOp>;
 using linalg_reshape = ValueBuilder<linalg::ReshapeOp>;
 using linalg_slice = ValueBuilder<linalg::SliceOp>;
 using linalg_yield = OperationBuilder<linalg::YieldOp>;
 
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_EDSC_INTRINSICS_H_

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

 ods_def<MatmulOp>:
 def matmul(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
   C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
 }
 
 ods_def<MatvecOp>:
 def matvec(A: f32(M, N), y: f32(N)) -> (x: f32(M)) {
   x(m) = std_addf<n>(std_mulf(A(m, n), y(n)));
 }
 
+ods_def<VecmatOp>:
+def vecmat(y: f32(M), A: f32(M, N)) -> (x: f32(N)) {
+  x(n) = std_addf<m>(std_mulf(y(m), A(m, n)));
+}
+
 ods_def<DotOp>:
 def dot(A: f32(M), B: f32(M)) -> (C: f32()) {
   C() = std_addf<m>(std_mulf(A(m), B(m)));
 }
 
 ods_def<BatchMatmulOp>:
 def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(b, k, n)));
def conv_3d_ndhwc(I: f32(N, D, H, W, C), K: f32(F, KD, KH, KW, C)) -> (O: f32(N, D, H, W, F)) {
   O(n, d, h, w, f) = std_addf<kd, kh, kw>(std_mulf(
     I(n, d + kd, h + kh, w + kw, c), K(f, kd, kh, kw, c)));
 }
 
 ods_def<ConvNCDHWOp>:
 def conv_3d_ncdhw(I: f32(N, C, D, H, W), K: f32(F, C, KD, KH, KW)) -> (O: f32(N, F, D, H, W)) {
   O(n, f, d, h, w) = std_addf<kd, kh, kw>(std_mulf(
     I(n, c, d + kd, h + kh, w + kw), K(f, c, kd, kh, kw)));
-}
+}
 

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

patterns.insert<
       LinalgOpConversion<FillOp>,
       LinalgOpConversion<GenericOp>,
       LinalgOpConversion<IndexedGenericOp>>(ctx);
   // TODO: collect all auto-generated named ops with a tblgen directive.
   patterns.insert<
       LinalgOpConversion<DotOp>,
       LinalgOpConversion<BatchMatmulOp>,
       LinalgOpConversion<MatvecOp>,
+      LinalgOpConversion<VecmatOp>,
       LinalgOpConversion<MatmulOp>,
       LinalgOpConversion<ConvWOp>,
       LinalgOpConversion<ConvNWCOp>,
       LinalgOpConversion<ConvNCWOp>,
       LinalgOpConversion<ConvHWOp>,
       LinalgOpConversion<ConvNHWCOp>,
       LinalgOpConversion<ConvNCHWOp>,
       LinalgOpConversion<ConvDHWOp>,

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 #include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.cpp.inc"
 
 // TODO: Determine whether we can generate the folders and verifiers.
 CANONICALIZERS_AND_FOLDERS(BatchMatmulOp)
 CANONICALIZERS_AND_FOLDERS(DotOp)
 CANONICALIZERS_AND_FOLDERS(MatmulOp)
 CANONICALIZERS_AND_FOLDERS(MatvecOp)
+CANONICALIZERS_AND_FOLDERS(VecmatOp)
 CANONICALIZERS_AND_FOLDERS(ConvWOp)
 CANONICALIZERS_AND_FOLDERS(ConvNWCOp)
 CANONICALIZERS_AND_FOLDERS(ConvNCWOp)
 CANONICALIZERS_AND_FOLDERS(ConvHWOp)
 CANONICALIZERS_AND_FOLDERS(ConvNHWCOp)
 CANONICALIZERS_AND_FOLDERS(ConvNCHWOp)
 CANONICALIZERS_AND_FOLDERS(ConvDHWOp)
 CANONICALIZERS_AND_FOLDERS(ConvNDHWCOp)
 CANONICALIZERS_AND_FOLDERS(ConvNCDHWOp)

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

static Optional<LinalgLoops> linalgOpToLoopsImplSwitch(Operation *op,
 
   // TODO: Cases below are generic and need a LinalgStructuredOpInterface.
   if (isa<GenericOp>(op))
     return linalgOpToLoopsImpl<LoopTy, GenericOp>(op, builder);
   if (isa<MatmulOp>(op))
     return linalgOpToLoopsImpl<LoopTy, MatmulOp>(op, builder);
   if (isa<MatvecOp>(op))
     return linalgOpToLoopsImpl<LoopTy, MatvecOp>(op, builder);
+  if (isa<VecmatOp>(op))
+    return linalgOpToLoopsImpl<LoopTy, VecmatOp>(op, builder);
   if (isa<DotOp>(op))
     return linalgOpToLoopsImpl<LoopTy, DotOp>(op, builder);
   if (isa<BatchMatmulOp>(op))
     return linalgOpToLoopsImpl<LoopTy, BatchMatmulOp>(op, builder);
   if (isa<ConvWOp>(op))
     return linalgOpToLoopsImpl<LoopTy, ConvWOp>(op, builder);
   if (isa<ConvNWCOp>(op))
     return linalgOpToLoopsImpl<LoopTy, ConvNWCOp>(op, builder);

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

return pattern1.match(&r.front().back()) ||
          pattern6.match(&r.front().back()) ||
          pattern7.match(&r.front().back()) || pattern8.match(&r.front().back());
 }
 
 // TODO: Should be Tablegen'd from a single source that generates the op itself.
 static LogicalResult isContraction(Operation *op) {
   // TODO: interface for named ops.
   if (isa<linalg::BatchMatmulOp, linalg::MatmulOp, linalg::MatvecOp,
-          linalg::DotOp>(op))
+          linalg::VecmatOp, linalg::DotOp>(op))
     return success();
 
   auto genericOp = dyn_cast<linalg::GenericOp>(op);
   if (!genericOp)
     return failure();
 
   auto mapRange = genericOp.indexing_maps().getAsValueRange<AffineMapAttr>();
   return success(

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   applyPatternsAndFoldGreedily(funcOp, forwardPattern);
 }
 
 static void applyContractionToVectorPatterns(FuncOp funcOp) {
   OwningRewritePatternList patterns;
   patterns.insert<LinalgVectorizationPattern<BatchMatmulOp>,
                   LinalgVectorizationPattern<MatmulOp>,
                   LinalgVectorizationPattern<MatvecOp>,
+                  LinalgVectorizationPattern<VecmatOp>,
                   LinalgVectorizationPattern<DotOp>,
                   LinalgVectorizationPattern<GenericOp>>(funcOp.getContext());
   applyPatternsAndFoldGreedily(funcOp, patterns);
 }
 
 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   OwningRewritePatternList foldPattern;
   foldPattern.insert<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());