[mlir][Linalg] Add a Linalg DRR test to go from matmul to vectors

[nicolasvasilache]

Summary

This diff builds a simple fused pass consisting of 2 levels of tiling, memory promotion and
vectorization using linalg transformations written as composable pattern rewrites.

[nicolasvasilache]

Drop spurious change.

[merge_guards_bot]

Unit tests: pass. 62308 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: pass. 62308 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

39

would it help to have two passes of greedily applying patterns, each with their own set, or do we have interaction between the two sets?

[rriddle]

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

36

This sounds like a visitation ordering problem, i.e. you want to visit the source operand(AffineApplyOp) before the user(ViewOp). Greedy walks reverse order so you run into the problem here. Am I understanding this correctly?

[tetuante]

LGTM!

[mehdi_amini]

I don't understand why is this all a test for upstream, it seems rather like an "example" of something or some work in progress on a transformation.

[nicolasvasilache]

@mehdi_amini this is a custom test pass that tests all the patterns compose nicely all the way to vector.
Last I tried they actually didn't compose nicely and I need to debug.
How is this an example more than a test pass?

[mehdi_amini]

In D73656#1947477, @nicolasvasilache wrote:

@mehdi_amini this is a custom test pass that tests all the patterns compose nicely all the way to vector.
Last I tried they actually didn't compose nicely and I need to debug.
How is this an example more than a test pass?

You testing pass is testing patterns that are in the test directory as well. This looks all self-contained to me: so it isn't clear to me what is this testing that isn't already tested.

[nicolasvasilache]

it isn't clear to me what is this testing that isn't already tested.

It wants to test that these patterns compose and atm they don't, something is off that I have not been able to debug yet.
Once debugged and the pass can be standalone I sure want to guarantee these patterns will compose forever, hence a test.

[nicolasvasilache]

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

36

I still do not know, I'll pick up again next week.

39

the goal is to have a single pass.

[nicolasvasilache]

Make patterns compose end-to-end.

[nicolasvasilache]

Fix test.

[nicolasvasilache]

Rebase

[mehdi_amini]

mlir/test/Dialect/Linalg/matmul-to-vector.mlir

2

Did you solve why we need a canonicalize here?

[nicolasvasilache]

Update the test to use the fused pass.

[nicolasvasilache]

@mehdi_amini had forgotten to update the test, fixed.

Things now work as expected, my intuition is that a change interfered with this revision at the time I created it: 9dfcddfaae50a3138173311a15e7ab19d36c1860 was landed on Jan 29th which is the day I also created this revision.
I submitted 3cb1f35df2a560d5a8ce2047b87cba8e3c904170 recently to deal with the affine_min case and things now work properly.

[nicolasvasilache]

landed as 6fb6a4d7f972d8faacd6b2646fe15f2eea1e4915

mlir/test/Dialect/Linalg/matmul-to-vector.mlir

+// RUN: mlir-opt %s -linalg-matmul-to-vector | FileCheck %s
+
+func @matmul_perm(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+                  %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+                  %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
+  linalg.matmul(%A, %B, %C) {__internal_linalg_transform__ = "__with_perm__"} :
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>
+  return
+}
+
+// CHECK-LABEL:func @matmul_perm
+//      CHECK:        vector.contract
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: : vector<8x16xf32>, vector<16x12xf32> into vector<8x12xf32>

mlir/test/lib/DeclarativeTransforms/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS TestLinalgTransformPatterns.td)
 mlir_tablegen(TestLinalgTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestLinalgTransformPatternsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestVectorTransformPatterns.td)
 mlir_tablegen(TestVectorTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestVectorTransformPatternsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS TestLinalgMatmulToVectorPatterns.td)
+mlir_tablegen(TestLinalgMatmulToVectorPatterns.h.inc -gen-rewriters)
+add_public_tablegen_target(MLIRTestLinalgMatmulToVectorPatternsIncGen)

mlir/test/lib/DeclarativeTransforms/TestLinalgMatmulToVectorPatterns.td

+//===- TestLinalgMatmulToVectorPatterns.td - Test patterns -*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the pattern definition file for declarative Linalg transformations
+// tests.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS
+#define TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS
+
+include "mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td"
+include "mlir/Dialect/Vector/VectorTransformPatterns.td"
+
+//===----------------------------------------------------------------------===//
+// Linalg tiling and permutation patterns.
+//===----------------------------------------------------------------------===//
+def : Pat<(MatmulOp:$op $_, $_, $_),
+          (TileLinalgOp<[768, 264, 768], "L2__with_perm__", [1, 2, 0]>),
+         [(Constraint<HasLinalgTransformMarker<"__with_perm__">>)]>;
+def : Pat<(MatmulOp:$op $_, $_, $_),
+          (TileLinalgOp<[8, 12, 16], "L1__with_perm__", [1, 0, 2]>),
+         [(Constraint<HasLinalgTransformMarker<"L2__with_perm__">>)]>;
+def : Pat<(MatmulOp:$op $_, $_, $_),
+          (PromoteSubviewsLinalgOp),
+         [(Constraint<HasOperandsOfType<"SubViewOp">>),
+          (Constraint<HasLinalgTransformMarker<"L1__with_perm__">>)]>;
+
+//===----------------------------------------------------------------------===//
+// Linalg to vector contraction patterns.
+//===----------------------------------------------------------------------===//
+def : Pattern<(MatmulOp:$op $_, $_, $_),
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"L1__with_perm__">,
+                PreconditionVectorizeLinalgOp]>>)]>;
+
+#endif // TEST_LINALG_MATMUL_TO_VECTOR_PATTERNS

mlir/test/lib/Transforms/CMakeLists.txt

 add_llvm_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvertGPUKernelToCubin.cpp
   TestDominance.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
   TestInlining.cpp
+  TestLinalgMatmulToVector.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestOpaqueLoc.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestVectorToLoopsConversion.cpp
   TestVectorTransforms.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRStandardOpsIncGen
+  MLIRTestLinalgMatmulToVectorPatternsIncGen
   MLIRTestLinalgTransformPatternsIncGen
   MLIRTestVectorTransformPatternsIncGen
 )
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../DeclarativeTransforms)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../DeclarativeTransforms)

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

+//===- TestLinalgMatmulToVector.cpp - Test VectorTransfers lowering -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <type_traits>
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Transforms/LinalgTransforms.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Dialect/Vector/VectorTransforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+using namespace mlir::vector;
+
+namespace {
+#include "TestLinalgMatmulToVectorPatterns.h.inc"
+
+struct DeclarativeTransforms
+    : public PassWrapper<DeclarativeTransforms, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    auto *context = &getContext();
+    AffineApplyOp::getCanonicalizationPatterns(patterns, context);
+    AffineMinOp::getCanonicalizationPatterns(patterns, context);
+    AffineMaxOp::getCanonicalizationPatterns(patterns, context);
+    AllocOp::getCanonicalizationPatterns(patterns, context);
+    SubViewOp::getCanonicalizationPatterns(patterns, context);
+    ViewOp::getCanonicalizationPatterns(patterns, context);
+    populateWithGenerated(context, &patterns);
+    applyPatternsGreedily(getFunction(), patterns);
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+void registerTestLinalgMatmulToVectorPass() {
+  PassRegistration<DeclarativeTransforms> pass(
+      "linalg-matmul-to-vector",
+      "Test declarative transform patterns for matmul 3-D tiling + promotion"
+      " + vectorization");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerPassManagerTestPass();
 void registerPatternsTestPass();
 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAllReduceLoweringPass();
+void registerTestLinalgMatmulToVectorPass();
 void registerTestLoopPermutationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestDominancePass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLinalgTransforms();
void registerTestPasses() {
   registerPassManagerTestPass();
   registerPatternsTestPass();
   registerPrintOpAvailabilityPass();
   registerSideEffectTestPasses();
   registerSimpleParametricTilingPass();
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAllReduceLoweringPass();
+  registerTestLinalgMatmulToVectorPass();
   registerTestLoopPermutationPass();
   registerTestCallGraphPass();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
   registerTestDominancePass();
   registerTestFunc();