[mlir][Linalg] NFC - Make markers use Identifier instead of StringRef

[nicolasvasilache]

Summary

This removes string ownership worries by putting everything into the context and allows more constructing identifiers programmatically.

[ftynse]

Have you considered constructing identifiers in the constructor of Linalg dialect and just using them? That's what we do for llvm types.

[nicolasvasilache]

Not really but the use case I have for this is constructing strings dynamically.

[nicolasvasilache]

Appease MSVC.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 //===- Transforms.h - Linalg transformations as patterns --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 #define DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/IR/Identifier.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 namespace mlir {
 namespace linalg {
 
 struct LinalgTilingOptions;
 

 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
 struct LinalgTransforms {
   static const StringLiteral kLinalgTransformMarker;
 };
 
 /// Helper class to control common attribute matching and setting behavior.
 struct LinalgMarker {
-  LinalgMarker(ArrayRef<StringRef> matchDisjunction = {},
-               Optional<StringRef> replacement = None);
-  LinalgMarker(ArrayRef<StringRef> matchDisjunction, StringRef replacement);
+  explicit LinalgMarker(ArrayRef<Identifier> matchDisjunction = {},
+                        Optional<Identifier> replacement = None);
+  LinalgMarker(LinalgMarker &&) = default;
+  LinalgMarker(const LinalgMarker &) = default;
   LogicalResult checkAndNotify(PatternRewriter &rewriter, Operation *op) const;
   void replaceLinalgMarker(PatternRewriter &rewriter, Operation *op) const;
 
 private:
-  SmallVector<StringRef, 4> matchDisjunction;
-  Optional<StringRef> replacement;
+  SmallVector<Identifier, 4> matchDisjunction;
+  Optional<Identifier> replacement;
 };
 
 ///
 /// Linalg tiling patterns.
 ///
 /// Apply the `tileLinalgOp` transformation as a pattern.
 /// `marker` controls LinalgTransformMarker matching and update when specified.
 /// See `tileLinalgOp` for more details.

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

static void insert(OwningRewritePatternList &patterns,
                      const LinalgTilingOptions &options, MLIRContext *ctx) {}
 };
 
 template <typename OpTy, typename... OpTypes>
 class RewritePatternList<OpTy, OpTypes...> {
 public:
   static void insert(OwningRewritePatternList &patterns,
                      const LinalgTilingOptions &options, MLIRContext *ctx) {
-    patterns.insert<LinalgTilingPattern<OpTy>>(ctx, options,
-                                               LinalgMarker({}, "tiled"));
+    patterns.insert<LinalgTilingPattern<OpTy>>(
+        ctx, options, LinalgMarker({}, Identifier::get("tiled", ctx)));
     RewritePatternList<OpTypes...>::insert(patterns, options, ctx);
   }
 };
 } // namespace
 
 OwningRewritePatternList
 mlir::linalg::getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx) {
   OwningRewritePatternList patterns;

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

 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
 const StringLiteral mlir::linalg::LinalgTransforms::kLinalgTransformMarker =
     "__internal_linalg_transform__";
 
-mlir::linalg::LinalgMarker::LinalgMarker(ArrayRef<StringRef> matchDisjunction,
-                                         Optional<StringRef> replacement)
+mlir::linalg::LinalgMarker::LinalgMarker(ArrayRef<Identifier> matchDisjunction,
+                                         Optional<Identifier> replacement)
     : matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
       replacement(replacement) {}
 
-mlir::linalg::LinalgMarker::LinalgMarker(ArrayRef<StringRef> matchDisjunction,
-                                         StringRef replacement)
-    : LinalgMarker(matchDisjunction, Optional<StringRef>{replacement}) {}
-
 LogicalResult
 mlir::linalg::LinalgMarker::checkAndNotify(PatternRewriter &rewriter,
                                            Operation *op) const {
   auto attr = op->template getAttrOfType<StringAttr>(
       LinalgTransforms::kLinalgTransformMarker);
 
   if (!attr) {
     // 1. Has no marker case and matchDisjunction is empty.
     if (matchDisjunction.empty())
       return success();
 
-    // 2. Has no marker and matchDisjuntion matches the no-moarker case.
-    for (auto marker : matchDisjunction)
-      if (marker.empty())
-        return success();
-
-    // 3. Has no marker but was expecting a marker.
+    // 2. Has no marker but was expecting a marker.
     return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
       diag << " does not have any marker from list: ";
       interleaveComma(matchDisjunction, diag);
     });
   }
 
   // 4. Match explicit marker.
   for (auto marker : matchDisjunction)

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-patterns | FileCheck %s
 
 // CHECK-DAG: #[[STRIDED_1D:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1  + s0)>
 // Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
 // CHECK-DAG: #[[STRIDED_2D_u_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // Map corresponding to a 2D memory access where the stride along all dims are unknown.
 // CHECK-DAG: #[[STRIDED_2D:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 // CHECK-DAG: #[[mk:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
 // CHECK-DAG: #[[kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
 // CHECK-DAG: #[[mn:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK-DAG: #[[nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
 // CHECK-DAG: #[[km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
 
 func @dot(%x: memref<?xf32, offset: ?, strides: [1]>,
           %y: memref<?xf32, offset: ?, strides: [1]>,
           %v: memref<f32>) {
-  linalg.dot(%x, %y, %v) : memref<?xf32, offset: ?, strides: [1]>,
-                           memref<?xf32, offset: ?, strides: [1]>,
-                           memref<f32>
+  linalg.dot(%x, %y, %v) { __internal_linalg_transform__ = "MEM" } :
+             memref<?xf32, offset: ?, strides: [1]>,
   return
 }
 // CHECK-LABEL: func @dot
 // CHECK-DAG:     %[[c0:.*]] = constant 0 : index
 // CHECK-DAG:     %[[c1:.*]] = constant 1 : index
 // CHECK-DAG:     %[[c8000:.*]] = constant 8000 : index
 // CHECK:         scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] {
 // CHECK:             scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c1]] {
 // CHECK:               load
 // CHECK:               load
 // CHECK:               mulf
 // CHECK:               load
 // CHECK:               addf
 // CHECK:               store
 
 func @matvec(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %x: memref<?xf32, offset: ?, strides: [1]>,
              %y: memref<?xf32, offset: ?, strides: [1]>) {
-  linalg.matvec(%A, %x, %y) : memref<?x?xf32, offset: ?, strides: [?, 1]>,
-                              memref<?xf32, offset: ?, strides: [1]>,
-                              memref<?xf32, offset: ?, strides: [1]>
+  linalg.matvec(%A, %x, %y) :
+                memref<?x?xf32, offset: ?, strides: [?, 1]>,
   return
 }
 // CHECK-LABEL: func @matvec
 // CHECK-DAG:     %[[c0:.*]] = constant 0 : index
 // CHECK-DAG:     %[[c5:.*]] = constant 5 : index
 // CHECK-DAG:     %[[c6:.*]] = constant 6 : index
 // CHECK:         scf.parallel {{.*}} step (%[[c5]])
 // CHECK:           scf.for {{.*}} step %[[c6]]
 // CHECK:             linalg.matvec({{.*}}, {{.*}}, {{.*}}) : memref<?x?xf32, #[[STRIDED_2D]]>, memref<?xf32, #[[STRIDED_1D]]>, memref<?xf32, #[[STRIDED_1D]]>
 
 func @matmul(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
-  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: ?, strides: [?, 1]>,
-                              memref<?x?xf32, offset: ?, strides: [?, 1]>,
-                              memref<?x?xf32, offset: ?, strides: [?, 1]>
+  linalg.matmul(%A, %B, %C) { __internal_linalg_transform__ = "MEM" } :
+                memref<?x?xf32, offset: ?, strides: [?, 1]>,
   return
 }
 // CHECK-LABEL: func @matmul
 // CHECK-DAG:     %[[c0:.*]] = constant 0 : index
 // CHECK-DAG:     %[[c2:.*]] = constant 2 : index
 // CHECK-DAG:     %[[c3:.*]] = constant 3 : index
 // CHECK-DAG:     %[[c4:.*]] = constant 4 : index
 // CHECK-DAG:     %[[c20:.*]] = constant 20 : index

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//
   // Linalg tiling patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx, LinalgTilingOptions().setTileSizes({2000, 3000, 4000}),
-      LinalgMarker({"MEM", {}}, "L3"));
+      LinalgMarker(Identifier::get("MEM", ctx), Identifier::get("L3", ctx)));
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx, LinalgTilingOptions().setTileSizes({200, 300, 400}),
-      LinalgMarker({"L3"}, "L2"));
+      LinalgMarker(Identifier::get("L3", ctx), Identifier::get("L2", ctx)));
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
-      LinalgMarker({"L2"}, "L1"));
+      LinalgMarker(Identifier::get("L2", ctx), Identifier::get("L1", ctx)));
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx, LinalgTilingOptions().setTileSizes({2, 3, 4}),
-      LinalgMarker({"L1"}, "REG"));
+      LinalgMarker(Identifier::get("L1", ctx), Identifier::get("REG", ctx)));
 
   patterns.insert<LinalgTilingPattern<MatvecOp>>(
       ctx,
       LinalgTilingOptions().setTileSizes({5, 6}).setLoopType(
           LinalgTilingLoopType::ParallelLoops),
-      LinalgMarker({}, "L1"));
+      LinalgMarker({}, Identifier::get("L1", ctx)));
 
   patterns.insert<LinalgTilingPattern<DotOp>>(
       ctx, LinalgTilingOptions().setTileSizes(8000),
-      LinalgMarker({"MEM", "L3", "L2", {}}, "REG"));
+      LinalgMarker(ArrayRef<Identifier>{Identifier::get("MEM", ctx),
+                                        Identifier::get("L3", ctx),
+                                        Identifier::get("L2", ctx)},
+                   Identifier::get("REG", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg tiling and permutation patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx,
       LinalgTilingOptions()
           .setTileSizes({2000, 3000, 4000})
           .setInterchange({1, 2, 0}),
-      LinalgMarker({"__with_perm__"}, "L2__with_perm__"));
+      LinalgMarker(Identifier::get("__with_perm__", ctx),
+                   Identifier::get("L2__with_perm__", ctx)));
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx,
       LinalgTilingOptions()
           .setTileSizes({200, 300, 400})
           .setInterchange({1, 0, 2}),
-      LinalgMarker({"L2__with_perm__"}, "L1__with_perm__"));
+      LinalgMarker(Identifier::get("L2__with_perm__", ctx),
+                   Identifier::get("L1__with_perm__", ctx)));
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
-      LinalgMarker({"L1__with_perm__"}, "REG__with_perm__"));
+      LinalgMarker(Identifier::get("L1__with_perm__", ctx),
+                   Identifier::get("REG__with_perm__", ctx)));
 
   patterns.insert<LinalgTilingPattern<MatvecOp>>(
       ctx, LinalgTilingOptions().setTileSizes({5, 6}).setInterchange({1, 0}),
-      LinalgMarker({"__with_perm__"}, "L1__with_perm__"));
+      LinalgMarker(Identifier::get("__with_perm__", ctx),
+                   Identifier::get("L1__with_perm__", ctx)));
 
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
       ctx,
       LinalgTilingOptions()
           .setTileSizes({16, 8, 4})
           .setInterchange({1, 2, 0})
           .setLoopType(LinalgTilingLoopType::ParallelLoops),
-      LinalgMarker({"par__with_perm__"}, "after_par__with_perm__"));
+      LinalgMarker(Identifier::get("par__with_perm__", ctx),
+                   Identifier::get("after_par__with_perm__", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg to loops patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgLoweringPattern<DotOp>>(
       ctx,
-      /*loweringType=*/LinalgLoweringType::Loops, LinalgMarker({"REG"}));
+      /*loweringType=*/LinalgLoweringType::Loops,
+      LinalgMarker(Identifier::get("REG", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgVectorizationPattern<MatmulOp>,
                   LinalgVectorizationPattern<FillOp>,
                   LinalgVectorizationPattern<GenericOp>>(
-      ctx, LinalgMarker({"VECTORIZE"}));
+      ctx, LinalgMarker(Identifier::get("VECTORIZE", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg generic permutation patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgInterchangePattern<GenericOp>>(
       ctx,
       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
-      LinalgMarker({}, "PERMUTED"));
+      LinalgMarker({}, Identifier::get("PERMUTED", ctx)));
   patterns.insert<LinalgInterchangePattern<IndexedGenericOp>>(
       ctx,
       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
-      LinalgMarker({}, "PERMUTED"));
+      LinalgMarker({}, Identifier::get("PERMUTED", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg subview operands promotion.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgPromotionPattern<MatmulOp>>(
       ctx, LinalgPromotionOptions().useFullTileBuffersByDefault(),
-      LinalgMarker({"_promote_views_"}, "_views_promoted_"));
+      LinalgMarker(Identifier::get("_promote_views_", ctx),
+                   Identifier::get("_views_promoted_", ctx)));
   patterns.insert<LinalgPromotionPattern<MatmulOp>>(
       ctx,
       LinalgPromotionOptions()
           .setOperandsToPromote({0})
           .useFullTileBuffersByDefault(),
-      LinalgMarker({"_promote_first_view_"}, "_first_view_promoted_"));
+      LinalgMarker(Identifier::get("_promote_first_view_", ctx),
+                   Identifier::get("_first_view_promoted_", ctx)));
   patterns.insert<LinalgPromotionPattern<FillOp>>(
       ctx,
       LinalgPromotionOptions()
           .setOperandsToPromote({0})
           .setUseFullTileBuffers({true})
           .setAlignment(32),
-      LinalgMarker({"_promote_views_aligned_"}, "_views_aligned_promoted_"));
+      LinalgMarker(Identifier::get("_promote_views_aligned_", ctx),
+                   Identifier::get("_views_aligned_promoted_", ctx)));
 
   applyPatternsAndFoldGreedily(funcOp, patterns);
 
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }
 
 static void fillL1TilingAndMatmulToVectorPatterns(
     FuncOp funcOp, StringRef startMarker,
     SmallVectorImpl<OwningRewritePatternList> &patternsVector) {
-  MLIRContext *context = funcOp.getContext();
+  MLIRContext *ctx = funcOp.getContext();
   patternsVector.emplace_back(LinalgTilingPattern<MatmulOp>(
-      context,
+      ctx,
       LinalgTilingOptions().setTileSizes({8, 12, 16}).setInterchange({1, 0, 2}),
-      LinalgMarker({startMarker}, "L1")));
+      LinalgMarker(Identifier::get(startMarker, ctx),
+                   Identifier::get("L1", ctx))));
 
   patternsVector.emplace_back(LinalgPromotionPattern<MatmulOp>(
-      context, LinalgPromotionOptions().useFullTileBuffersByDefault(),
-      LinalgMarker({"L1"}, "VEC")));
+      ctx, LinalgPromotionOptions().useFullTileBuffersByDefault(),
+      LinalgMarker(Identifier::get("L1", ctx), Identifier::get("VEC", ctx))));
 
-  patternsVector.emplace_back(
-      LinalgVectorizationPattern<MatmulOp>(context, LinalgMarker({"VEC"})));
+  patternsVector.emplace_back(LinalgVectorizationPattern<MatmulOp>(
+      ctx, LinalgMarker(Identifier::get("VEC", ctx))));
   patternsVector.back()
       .insert<LinalgVectorizationPattern<FillOp>,
-              LinalgVectorizationPattern<CopyOp>>(context);
+              LinalgVectorizationPattern<CopyOp>>(ctx);
 }
 
 //===----------------------------------------------------------------------===//
 // Test promotion callbacks
 //===----------------------------------------------------------------------===//
 
 // Allocation call back
 static Optional<Value> allocCallBackFn(OpBuilder &b, SubViewOp subView,
static LogicalResult copyCallBackFn(OpBuilder &b, Value src, Value dst,
   if (!isOutput)
     b.create<FillOp>(
         src.getLoc(), dst,
         b.create<ConstantOp>(src.getLoc(), FloatAttr::get(floatType, 42.0)));
   b.create<CopyOp>(src.getLoc(), src, dst);
   return success();
 }
 
-void fillPromotionCallBackPatterns(MLIRContext *context,
+void fillPromotionCallBackPatterns(MLIRContext *ctx,
                                    OwningRewritePatternList &patterns) {
   patterns.insert<LinalgTilingPattern<MatmulOp>>(
-      context, LinalgTilingOptions().setTileSizes({16, 16, 16}),
-      LinalgMarker({"START"}, "PROMOTE"));
+      ctx, LinalgTilingOptions().setTileSizes({16, 16, 16}),
+      LinalgMarker(Identifier::get("START", ctx),
+                   Identifier::get("PROMOTE", ctx)));
   patterns.insert<LinalgPromotionPattern<MatmulOp>>(
-      context,
+      ctx,
       LinalgPromotionOptions()
           .setOperandsToPromote({0, 2})
           .setUseFullTileBuffers({false, false})
           .setAllocationDeallocationFns(allocCallBackFn, deallocCallBackFn)
           .setCopyInOutFns(
               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                 copyCallBackFn(b, src, dst, false);
                 return success();
               },
               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                 copyCallBackFn(b, src, dst, true);
                 return success();
               }),
-      LinalgMarker({"PROMOTE"}));
+      LinalgMarker(Identifier::get("PROMOTE", ctx)));
 }
 
 static void
 applyMatmulToVectorPatterns(FuncOp funcOp,
                             bool testMatmulToVectorPatterns1dTiling,
                             bool testMatmulToVectorPatterns2dTiling) {
   MLIRContext *ctx = funcOp.getContext();
   SmallVector<OwningRewritePatternList, 4> stage1Patterns;
   if (testMatmulToVectorPatterns1dTiling) {
-    fillL1TilingAndMatmulToVectorPatterns(funcOp, "START", stage1Patterns);
+    fillL1TilingAndMatmulToVectorPatterns(funcOp, Identifier::get("START", ctx),
+                                          stage1Patterns);
   } else if (testMatmulToVectorPatterns2dTiling) {
-    stage1Patterns.emplace_back(
-        LinalgTilingPattern<MatmulOp>(ctx,
-                                      LinalgTilingOptions()
-                                          .setTileSizes({768, 264, 768})
-                                          .setInterchange({1, 2, 0}),
-                                      LinalgMarker({"START"}, "L2")));
-    fillL1TilingAndMatmulToVectorPatterns(funcOp, "L2", stage1Patterns);
+    stage1Patterns.emplace_back(LinalgTilingPattern<MatmulOp>(
+        ctx,
+        LinalgMarker(Identifier::get("START", ctx),
+                     Identifier::get("L2", ctx))));
+    fillL1TilingAndMatmulToVectorPatterns(funcOp, Identifier::get("L2", ctx),
+                                          stage1Patterns);
   }
   OwningRewritePatternList stage2Patterns =
       getLinalgTilingCanonicalizationPatterns(ctx);
   applyStagedPatterns(funcOp, stage1Patterns, stage2Patterns);
 }
 
 static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   OwningRewritePatternList forwardPattern;