[NFC][MLIR][Linalg] Refactor linalg.matmul tablegen ODS and related C++ code. #116377

shahidact · 2024-11-15T12:04:02Z

This commit refactors part of the code in preparation for the migration of other matmul variants from OpDSL to ODS.
Moves getDefaultIndexingmaps() helper into the MatmulOp class.

llvmbot · 2024-11-16T06:31:42Z

@llvm/pr-subscribers-mlir-linalg

Author: Md Asghar Ahmad Shahid (shahidact)

Changes

This commit refactors part of the code in preparation for the migration of other matmul variants from OpDSL to ODS.
Moves getDefaultIndexingmaps() helper into the MatmulOp class.

Full diff: https://github.com/llvm/llvm-project/pull/116377.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td (+15-4)
(modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+4-25)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index e578f4b956ef5e..aed28c72ea5bde 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -622,7 +622,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
             CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
       [{
         buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
-          attributes, MatmulOp::getRegionBuilder());
+          attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
@@ -630,7 +631,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
             CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
       [{
         buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, MatmulOp::getRegionBuilder());
+          inputs, outputs, attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
@@ -648,7 +650,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
       [{
         $_state.addAttribute("cast", cast);
         buildStructuredOp($_builder, $_state, resultTensorTypes, inputs, outputs,
-          attributes, MatmulOp::getRegionBuilder());
+          attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>
 
     ];
@@ -664,7 +667,15 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
                                 Block &block, ArrayRef<NamedAttribute> attrs);
 
       /// Returns a list of AffineMap with the typical matmul indexing charactristic.
-      SmallVector<AffineMap> getDefaultIndexingMaps();
+      static SmallVector<AffineMap> getDefaultIndexingMaps(MLIRContext *context){
+        AffineExpr d0, d1, d2;
+        SmallVector<AffineMap, 3> indexingMaps;
+        bindDims(context, d0, d1, d2);
+        indexingMaps.push_back(AffineMap::get(3, 0, {d0, d2}, context));
+        indexingMaps.push_back(AffineMap::get(3, 0, {d2, d1}, context));
+        indexingMaps.push_back(AffineMap::get(3, 0, {d0, d1}, context));
+        return indexingMaps;
+      }
 
       /// Returns true if the given broadcast map \p bcastMap is valid for this op.
       bool isValidLhsRhsBroadcastMap(AffineMap bcastMap);
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index c909d13e4314b4..259e060933b436 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -155,23 +155,10 @@ static void fillStructuredOpRegion(OpBuilder &opBuilder, Region &region,
   // iterator_types is an auto-generated method.
 }
 
-/// Helper to create a typical indexing map for MatmulOp. Returns a list of
-/// AffineMap.
-static SmallVector<AffineMap, 3>
-getDefaultIndexingMapsForMatmul(MLIRContext *context) {
-  AffineExpr d0, d1, d2;
-  SmallVector<AffineMap, 3> indexingMaps;
-  bindDims(context, d0, d1, d2);
-  indexingMaps.push_back(AffineMap::get(3, 0, {d0, d2}, context));
-  indexingMaps.push_back(AffineMap::get(3, 0, {d2, d1}, context));
-  indexingMaps.push_back(AffineMap::get(3, 0, {d0, d1}, context));
-  return indexingMaps;
-}
-
 /// Wrapper to return the typical indexing map array attribute for MatmulOp.
 static SmallVector<Attribute> getDefaultIndexingMapAttr(MLIRContext *context) {
   return llvm::map_to_vector(
-      getDefaultIndexingMapsForMatmul(context),
+      MatmulOp::getDefaultIndexingMaps(context),
       [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); });
 }
 
@@ -204,9 +191,6 @@ static void buildStructuredOp(
       indexingMapsAttrVal.push_back(AffineMapAttr::get(map));
     }
     state.addAttribute("indexing_maps", b.getArrayAttr(indexingMapsAttrVal));
-  } else {
-    indexingMapsAttrVal = getDefaultIndexingMapAttr(b.getContext());
-    state.addAttribute("indexing_maps", b.getArrayAttr(indexingMapsAttrVal));
   }
 
   state.addAttributes(attributes);
@@ -3481,7 +3465,7 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
                                                   unsigned opIndex) {
   SmallVector<AffineMap, 3> opIndexingMaps = matmulOp.getIndexingMapsArray();
   SmallVector<AffineMap, 3> defaultIndexingMaps =
-      matmulOp.getDefaultIndexingMaps();
+      matmulOp.getDefaultIndexingMaps(matmulOp->getContext());
 
   auto opIndexingMap = opIndexingMaps[opIndex];
   auto defaultIndexingMap = defaultIndexingMaps[opIndex];
@@ -3523,7 +3507,8 @@ bool MatmulOp::hasDynamicIndexingMaps() { return true; }
 /// Check if the op has broadcast and/or transpose semantic. Returns true if the
 /// user defined indexing maps are not equal to default map.
 bool MatmulOp::hasUserDefinedMaps() {
-  SmallVector<AffineMap, 3> defaultMaps = getDefaultIndexingMaps();
+  SmallVector<AffineMap, 3> defaultMaps =
+      MatmulOp::getDefaultIndexingMaps(this->getContext());
   SmallVector<AffineMap, 3> explicitMaps = getIndexingMapsArray();
   return defaultMaps != explicitMaps;
 }
@@ -3557,12 +3542,6 @@ void MatmulOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
   helper.yieldOutputs(yields);
 }
 
-/// Returns a list of AffineMap with the typical matmul indexing charactristic.
-SmallVector<AffineMap> MatmulOp::getDefaultIndexingMaps() {
-  MLIRContext *context = this->getContext();
-  return getDefaultIndexingMapsForMatmul(context);
-}
-
 /// Returns true if the given broadcast map \p bcastMap is valid for this op.
 bool MatmulOp::isValidLhsRhsBroadcastMap(AffineMap bcastMap) {
   assert(bcastMap.getNumResults() == 1 && "Expected single result dim expr.");

llvmbot · 2024-11-16T06:31:42Z

@llvm/pr-subscribers-mlir

Author: Md Asghar Ahmad Shahid (shahidact)

Changes

This commit refactors part of the code in preparation for the migration of other matmul variants from OpDSL to ODS.
Moves getDefaultIndexingmaps() helper into the MatmulOp class.

Full diff: https://github.com/llvm/llvm-project/pull/116377.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td (+15-4)
(modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+4-25)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index e578f4b956ef5e..aed28c72ea5bde 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -622,7 +622,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
             CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
       [{
         buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
-          attributes, MatmulOp::getRegionBuilder());
+          attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
@@ -630,7 +631,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
             CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
       [{
         buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, MatmulOp::getRegionBuilder());
+          inputs, outputs, attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
@@ -648,7 +650,8 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
       [{
         $_state.addAttribute("cast", cast);
         buildStructuredOp($_builder, $_state, resultTensorTypes, inputs, outputs,
-          attributes, MatmulOp::getRegionBuilder());
+          attributes, MatmulOp::getRegionBuilder(),
+          MatmulOp::getDefaultIndexingMaps($_builder.getContext()));
       }]>
 
     ];
@@ -664,7 +667,15 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
                                 Block &block, ArrayRef<NamedAttribute> attrs);
 
       /// Returns a list of AffineMap with the typical matmul indexing charactristic.
-      SmallVector<AffineMap> getDefaultIndexingMaps();
+      static SmallVector<AffineMap> getDefaultIndexingMaps(MLIRContext *context){
+        AffineExpr d0, d1, d2;
+        SmallVector<AffineMap, 3> indexingMaps;
+        bindDims(context, d0, d1, d2);
+        indexingMaps.push_back(AffineMap::get(3, 0, {d0, d2}, context));
+        indexingMaps.push_back(AffineMap::get(3, 0, {d2, d1}, context));
+        indexingMaps.push_back(AffineMap::get(3, 0, {d0, d1}, context));
+        return indexingMaps;
+      }
 
       /// Returns true if the given broadcast map \p bcastMap is valid for this op.
       bool isValidLhsRhsBroadcastMap(AffineMap bcastMap);
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index c909d13e4314b4..259e060933b436 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -155,23 +155,10 @@ static void fillStructuredOpRegion(OpBuilder &opBuilder, Region &region,
   // iterator_types is an auto-generated method.
 }
 
-/// Helper to create a typical indexing map for MatmulOp. Returns a list of
-/// AffineMap.
-static SmallVector<AffineMap, 3>
-getDefaultIndexingMapsForMatmul(MLIRContext *context) {
-  AffineExpr d0, d1, d2;
-  SmallVector<AffineMap, 3> indexingMaps;
-  bindDims(context, d0, d1, d2);
-  indexingMaps.push_back(AffineMap::get(3, 0, {d0, d2}, context));
-  indexingMaps.push_back(AffineMap::get(3, 0, {d2, d1}, context));
-  indexingMaps.push_back(AffineMap::get(3, 0, {d0, d1}, context));
-  return indexingMaps;
-}
-
 /// Wrapper to return the typical indexing map array attribute for MatmulOp.
 static SmallVector<Attribute> getDefaultIndexingMapAttr(MLIRContext *context) {
   return llvm::map_to_vector(
-      getDefaultIndexingMapsForMatmul(context),
+      MatmulOp::getDefaultIndexingMaps(context),
       [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); });
 }
 
@@ -204,9 +191,6 @@ static void buildStructuredOp(
       indexingMapsAttrVal.push_back(AffineMapAttr::get(map));
     }
     state.addAttribute("indexing_maps", b.getArrayAttr(indexingMapsAttrVal));
-  } else {
-    indexingMapsAttrVal = getDefaultIndexingMapAttr(b.getContext());
-    state.addAttribute("indexing_maps", b.getArrayAttr(indexingMapsAttrVal));
   }
 
   state.addAttributes(attributes);
@@ -3481,7 +3465,7 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
                                                   unsigned opIndex) {
   SmallVector<AffineMap, 3> opIndexingMaps = matmulOp.getIndexingMapsArray();
   SmallVector<AffineMap, 3> defaultIndexingMaps =
-      matmulOp.getDefaultIndexingMaps();
+      matmulOp.getDefaultIndexingMaps(matmulOp->getContext());
 
   auto opIndexingMap = opIndexingMaps[opIndex];
   auto defaultIndexingMap = defaultIndexingMaps[opIndex];
@@ -3523,7 +3507,8 @@ bool MatmulOp::hasDynamicIndexingMaps() { return true; }
 /// Check if the op has broadcast and/or transpose semantic. Returns true if the
 /// user defined indexing maps are not equal to default map.
 bool MatmulOp::hasUserDefinedMaps() {
-  SmallVector<AffineMap, 3> defaultMaps = getDefaultIndexingMaps();
+  SmallVector<AffineMap, 3> defaultMaps =
+      MatmulOp::getDefaultIndexingMaps(this->getContext());
   SmallVector<AffineMap, 3> explicitMaps = getIndexingMapsArray();
   return defaultMaps != explicitMaps;
 }
@@ -3557,12 +3542,6 @@ void MatmulOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
   helper.yieldOutputs(yields);
 }
 
-/// Returns a list of AffineMap with the typical matmul indexing charactristic.
-SmallVector<AffineMap> MatmulOp::getDefaultIndexingMaps() {
-  MLIRContext *context = this->getContext();
-  return getDefaultIndexingMapsForMatmul(context);
-}
-
 /// Returns true if the given broadcast map \p bcastMap is valid for this op.
 bool MatmulOp::isValidLhsRhsBroadcastMap(AffineMap bcastMap) {
   assert(bcastMap.getNumResults() == 1 && "Expected single result dim expr.");

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

MaheshRavishankar

Thanks!

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

…ated CPP file.

uweigand · 2024-11-20T12:44:49Z

Looks like this introduced a build break on s390x (https://lab.llvm.org/buildbot/#/builders/117/builds/3919):

In static member function 'static llvm::SmallVector<mlir::AffineMap> mlir::linalg::MatmulOp::getDefaultIndexingMaps(mlir::MLIRContext*)':
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:3468:10: error: could not convert 'indexingMaps' from 'SmallVector<[...],3>' to 'SmallVector<[...],6>'
 3468 |   return indexingMaps;
      |          ^~~~~~~~~~~~
      |          |
      |          SmallVector<[...],3>

shahidact · 2024-11-20T13:02:05Z

Looks like this introduced a build break on s390x (https://lab.llvm.org/buildbot/#/builders/117/builds/3919):

In static member function 'static llvm::SmallVector<mlir::AffineMap> mlir::linalg::MatmulOp::getDefaultIndexingMaps(mlir::MLIRContext*)':
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:3468:10: error: could not convert 'indexingMaps' from 'SmallVector<[...],3>' to 'SmallVector<[...],6>'
 3468 |   return indexingMaps;
      |          ^~~~~~~~~~~~
      |          |
      |          SmallVector<[...],3>

Here a PR to fix this build issue.
#116958

sylvestre · 2024-11-20T13:08:05Z

i opened this #116966 issue

shahidact · 2024-11-20T13:47:17Z

The fix has been merged.
#116958

shahidact requested review from dcaballe, nicolasvasilache and rengolin as code owners November 15, 2024 12:04

llvmbot added mlir:linalg mlir labels Nov 16, 2024

Resolve conflict.

0b1cbee

shahidact force-pushed the nfc-matmul branch from 55b2bf4 to 0b1cbee Compare November 18, 2024 04:19

rengolin requested review from MaheshRavishankar, banach-space and Groverkss November 18, 2024 08:51

Groverkss reviewed Nov 18, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Show resolved Hide resolved

Make the optional argument of buildMatmulOp() helper as non-optional.

881675e

MaheshRavishankar approved these changes Nov 18, 2024

View reviewed changes

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td Outdated Show resolved Hide resolved

shahidact added 2 commits November 18, 2024 20:58

Move out the definition of getDefaultIndexingMaps() function into rel…

1758aa2

…ated CPP file.

Merge branch 'main' into nfc-matmul

fe22f45

rengolin merged commit 288f05f into llvm:main Nov 20, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[NFC][MLIR][Linalg] Refactor linalg.matmul tablegen ODS and related C++ code. #116377

[NFC][MLIR][Linalg] Refactor linalg.matmul tablegen ODS and related C++ code. #116377

Uh oh!

shahidact commented Nov 15, 2024

Uh oh!

llvmbot commented Nov 16, 2024

Uh oh!

llvmbot commented Nov 16, 2024

Uh oh!

Uh oh!

Uh oh!

MaheshRavishankar left a comment

Uh oh!

Uh oh!

Uh oh!

uweigand commented Nov 20, 2024

Uh oh!

shahidact commented Nov 20, 2024

Uh oh!

sylvestre commented Nov 20, 2024

Uh oh!

shahidact commented Nov 20, 2024

Uh oh!

Uh oh!

[NFC][MLIR][Linalg] Refactor linalg.matmul tablegen ODS and related C++ code. #116377

[NFC][MLIR][Linalg] Refactor linalg.matmul tablegen ODS and related C++ code. #116377

Uh oh!

Conversation

shahidact commented Nov 15, 2024

Uh oh!

llvmbot commented Nov 16, 2024

Uh oh!

llvmbot commented Nov 16, 2024

Uh oh!

Uh oh!

Uh oh!

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

uweigand commented Nov 20, 2024

Uh oh!

shahidact commented Nov 20, 2024

Uh oh!

sylvestre commented Nov 20, 2024

Uh oh!

shahidact commented Nov 20, 2024

Uh oh!

Uh oh!