[mlir][vector] Share enums with the transform dialect

[qcolombet]

Summary

Refactor the definition of the enums that are used in the lower_vectors operation of the transformation dialect.
This avoid duplicating the definition of all the configurations that this operation can trigger.

NFC

[qcolombet]

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

54

We have to explicitly list all the fields here to have the proper parsing of the enums.

This may have an impact on the supported syntax. I.e., I don't know if attr-dict allows to put the attributes in whatever order, but if it does, then we lose this ability here.

mlir/include/mlir/Dialect/Vector/Transforms/VectorTransformsEnums.h

2 ↗

(On Diff #489606)

This file is a work around the fact that tablegen'ed enums header are not guarded.
I don't know how we want to move forward here.

[nicolasvasilache]

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

54

attr-dict allows any order and not enforcing a particular order indeed seems better.

Lookup oilist in the ODS doc, it does what you want here: https://mlir.llvm.org/docs/DefiningDialects/Operations/#directives

mlir/include/mlir/Dialect/Vector/Transforms/VectorTransformsEnums.h

2 ↗

(On Diff #489606)

Just add the include to VectorTransforms.h and avoid a special include file for this?

[qcolombet]

• Make the change really NFC thanks to oilist (thanks @nicolasvasilache for the pointer)
• Include directly VectorTransforms.h to avoid creating a VectorTransformsEnums header

[qcolombet]

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

54

Great!

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.h

 //===- VectorTransformOps.h - Vector transform ops --------------*- 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 MLIR_DIALECT_VECTOR_TRANSFORMOPS_VECTORTRANSFORMOPS_H
 #define MLIR_DIALECT_VECTOR_TRANSFORMOPS_VECTORTRANSFORMOPS_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 namespace vector {
 class VectorOp;
 } // namespace vector
 } // namespace mlir
 

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

 //===- VectorTransformOps.td - Vector transform ops --------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef VECTOR_TRANSFORM_OPS
 #define VECTOR_TRANSFORM_OPS
 
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def LowerVectorsOp : Op<Transform_Dialect, "vector.lower_vectors",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let description = [{
     Indicates that the vector operations nested under the isolated from above op
     `target` should be lowered to finer-grained vector primitives.
 
     At this time, the transform is all or nothing.
 
     This is usally a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   // TODO: evolve this to proper enums.
   let arguments = (ins PDL_Operation:$target,
-     DefaultValuedAttr<StrAttr, "\"outerproduct\"">:$contraction_lowering,
-     DefaultValuedAttr<StrAttr, "\"innerparallel\"">:$multireduction_lowering,
-     DefaultValuedAttr<StrAttr, "\"linalg-copy\"">:$split_transfers,
-     DefaultValuedAttr<StrAttr, "\"eltwise\"">:$transpose_lowering,
+     DefaultValuedAttr<VectorContractLoweringAttr,
+       "vector::VectorContractLowering::OuterProduct">:$contraction_lowering,
+     DefaultValuedAttr<VectorMultiReductionLoweringAttr,
+       "vector::VectorMultiReductionLowering::InnerParallel">:
+         $multireduction_lowering,
+     DefaultValuedAttr<VectorTransferSplitAttr,
+       "vector::VectorTransferSplit::LinalgCopy">:$split_transfers,
+     DefaultValuedAttr<VectorTransposeLoweringAttr,
+       "vector::VectorTransposeLowering::EltWise">:$transpose_lowering,
      DefaultValuedAttr<BoolAttr, "false">:$transpose_avx2_lowering,
      DefaultValuedAttr<BoolAttr, "true">:$unroll_vector_transfers
   );
   let results = (outs PDL_Operation:$results);
 
-  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat = [{
+    $target
+    oilist (
+      `contraction_lowering` `=` $contraction_lowering
+      | `multireduction_lowering` `=` $multireduction_lowering
+      | `split_transfers` `=` $split_transfers
+      | `transpose_lowering` `=` $transpose_lowering
+    )
+    attr-dict
+  }];
 }
 
 #endif // VECTOR_TRANSFORM_OPS

mlir/include/mlir/Dialect/Vector/Transforms/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS VectorTransformsBase.td)
+mlir_tablegen(VectorTransformsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(VectorTransformsEnums.cpp.inc -gen-enum-defs)
+
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Vector)
 add_public_tablegen_target(MLIRVectorTransformsIncGen)
 
 add_mlir_doc(Passes VectorPasses ./ -gen-pass-doc)

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

 //===- VectorRewritePatterns.h - Vector rewrite 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 MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORREWRITEPATTERNS_H
 #define MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORREWRITEPATTERNS_H
 
 #include <utility>
 #include <optional>
 
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransformsEnums.h.inc"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 class RewritePatternSet;
 
 namespace vector {
 
 //===----------------------------------------------------------------------===//
 // Vector transformation options exposed as auxiliary structs.
 //===----------------------------------------------------------------------===//
-/// Enum to control the lowering of `vector.transpose` operations.
-enum class VectorTransposeLowering {
-  /// Lower transpose into element-wise extract and inserts.
-  EltWise = 0,
-  /// Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix
-  /// intrinsics.
-  Flat = 1,
-  /// Lower 2-D transpose to `vector.shuffle`.
-  Shuffle = 2,
-};
-/// Enum to control the lowering of `vector.multi_reduction` operations.
-enum class VectorMultiReductionLowering {
-  /// Lower multi_reduction into outer-reduction and inner-parallel ops.
-  InnerParallel = 0,
-  /// Lower multi_reduction into outer-parallel and inner-reduction ops.
-  InnerReduction = 1,
-};
-/// Enum to control the lowering of `vector.contract` operations.
-enum class VectorContractLowering {
-  /// Progressively lower to finer grained `vector.contract` and dot-products.
-  Dot = 0,
-  /// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
-  Matmul = 1,
-  /// Lower to `vector.outerproduct`.
-  OuterProduct = 2,
-  /// Lower contract with all reduction dimensions unrolled to 1 to a vector
-  /// elementwise operations.
-  ParallelArith = 3,
-};
-/// Enum to control the splitting of `vector.transfer` operations into
-/// in-bounds and out-of-bounds variants.
-enum class VectorTransferSplit {
-  /// Do not split vector transfer operations.
-  None = 0,
-  /// Split using in-bounds + out-of-bounds vector.transfer operations.
-  VectorTransfer = 1,
-  /// Split using an in-bounds vector.transfer + linalg.fill + linalg.copy
-  /// operations.
-  LinalgCopy = 2,
-  /// Do not split vector transfer operation but instead mark it as "in-bounds".
-  ForceInBounds = 3
-};
 /// Structure to control the behavior of vector transform patterns.
 struct VectorTransformsOptions {
   /// Option to control the lowering of vector.contract.
   VectorContractLowering vectorContractLowering = VectorContractLowering::Dot;
   VectorTransformsOptions &
   setVectorTransformsOptions(VectorContractLowering opt) {
     vectorContractLowering = opt;
     return *this;

mlir/include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td

+//===- VectorTransformBase.td - Vector transform ops --------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef VECTOR_TRANSFORMS_BASE
+#define VECTOR_TRANSFORMS_BASE
+
+include "mlir/IR/EnumAttr.td"
+
+// Lower transpose into element-wise extract and inserts.
+def VectorTransposeLowering_Elementwise:
+  I32EnumAttrCase<"EltWise",  0, "eltwise">;
+// Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix
+// intrinsics.
+def VectorTransposeLowering_FlatTranspose:
+  I32EnumAttrCase<"Flat",  1, "flat_transpose">;
+// Lower 2-D transpose to `vector.shuffle`.
+def VectorTransposeLowering_Shuffle:
+  I32EnumAttrCase<"Shuffle",  2, "shuffle">;
+def VectorTransposeLoweringAttr : I32EnumAttr<
+    "VectorTransposeLowering",
+    "control the lowering of `vector.transpose` operations.",
+    [VectorTransposeLowering_Elementwise, VectorTransposeLowering_FlatTranspose,
+     VectorTransposeLowering_Shuffle]> {
+  let cppNamespace = "::mlir::vector";
+}
+
+// Lower multi_reduction into outer-reduction and inner-parallel ops.
+def VectorMultiReductionLowering_InnerParallel:
+  I32EnumAttrCase<"InnerParallel", 0, "innerparallel">;
+// Lower multi_reduction into outer-parallel and inner-reduction ops.
+def VectorMultiReductionLowering_InnerReduction:
+  I32EnumAttrCase<"InnerReduction", 1, "innerreduction">;
+def VectorMultiReductionLoweringAttr: I32EnumAttr<
+    "VectorMultiReductionLowering",
+    "control the lowering of `vector.multi_reduction`.",
+  [VectorMultiReductionLowering_InnerParallel,
+   VectorMultiReductionLowering_InnerReduction]> {
+  let cppNamespace = "::mlir::vector";
+}
+
+// Progressively lower to finer grained `vector.contract` and dot-products.
+def VectorContractLowering_Dot: I32EnumAttrCase<"Dot", 0, "dot">;
+// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
+def VectorContractLowering_Matmul:
+  I32EnumAttrCase<"Matmul", 1, "matmulintrinsics">;
+// Lower to `vector.outerproduct`.
+def VectorContractLowering_OuterProduct:
+  I32EnumAttrCase<"OuterProduct", 2, "outerproduct">;
+// Lower contract with all reduction dimensions unrolled to 1 to a vector
+// elementwise operations.
+def VectorContractLowering_ParallelArith:
+  I32EnumAttrCase<"ParallelArith", 3, "parallelarith">;
+def VectorContractLoweringAttr: I32EnumAttr<
+    "VectorContractLowering",
+    "control the lowering of `vector.contract` operations.",
+  [VectorContractLowering_Dot, VectorContractLowering_Matmul,
+   VectorContractLowering_OuterProduct, VectorContractLowering_ParallelArith]> {
+  let cppNamespace = "::mlir::vector";
+}
+
+// Do not split vector transfer operations.
+def VectorTransferSplit_None: I32EnumAttrCase<"None", 0, "none">;
+// Split using in-bounds + out-of-bounds vector.transfer operations.
+def VectorTransferSplit_VectorTransfer:
+  I32EnumAttrCase<"VectorTransfer", 1, "vector-transfer">;
+// Split using an in-bounds vector.transfer + linalg.fill + linalg.copy
+// operations.
+def VectorTransferSplit_LinalgCopy:
+  I32EnumAttrCase<"LinalgCopy", 2, "linalg-copy">;
+// Do not split vector transfer operation but instead mark it as "in-bounds".
+def VectorTransferSplit_ForceInBounds:
+  I32EnumAttrCase<"ForceInBounds", 3, "force-in-bounds">;
+def VectorTransferSplitAttr: I32EnumAttr<
+    "VectorTransferSplit",
+    "control the splitting of `vector.transfer` operations into in-bounds"
+    " and out-of-bounds variants.",
+  [VectorTransferSplit_None, VectorTransferSplit_VectorTransfer,
+   VectorTransferSplit_LinalgCopy, VectorTransferSplit_ForceInBounds]> {
+  let cppNamespace = "::mlir::vector";
+}
+#endif

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
                                        "applies only to isolated-from-above "
                                        "targets because it needs to apply "
                                        "patterns greedily");
     }
 
     MLIRContext *ctx = getContext();
     RewritePatternSet patterns(ctx);
     vector::VectorTransposeLowering vectorTransposeLowering =
-        llvm::StringSwitch<vector::VectorTransposeLowering>(
-            getTransposeLowering())
-            .Case("eltwise", vector::VectorTransposeLowering::EltWise)
-            .Case("flat_transpose", vector::VectorTransposeLowering::Flat)
-            .Case("shuffle", vector::VectorTransposeLowering::Shuffle)
-            .Default(vector::VectorTransposeLowering::EltWise);
+        getTransposeLowering();
     vector::VectorMultiReductionLowering vectorMultiReductionLowering =
-        llvm::StringSwitch<vector::VectorMultiReductionLowering>(
-            getMultireductionLowering())
-            .Case("innerreduction",
-                  vector::VectorMultiReductionLowering::InnerReduction)
-            .Default(vector::VectorMultiReductionLowering::InnerParallel);
+        getMultireductionLowering();
     vector::VectorContractLowering vectorContractLowering =
-        llvm::StringSwitch<vector::VectorContractLowering>(
-            getContractionLowering())
-            .Case("matrixintrinsics", vector::VectorContractLowering::Matmul)
-            .Case("dot", vector::VectorContractLowering::Dot)
-            .Case("outerproduct", vector::VectorContractLowering::OuterProduct)
-            .Default(vector::VectorContractLowering::OuterProduct);
-    vector::VectorTransferSplit vectorTransferSplit =
-        llvm::StringSwitch<vector::VectorTransferSplit>(getSplitTransfers())
-            .Case("none", vector::VectorTransferSplit::None)
-            .Case("linalg-copy", vector::VectorTransferSplit::LinalgCopy)
-            .Case("vector-transfers",
-                  vector::VectorTransferSplit::VectorTransfer)
-            .Default(vector::VectorTransferSplit::None);
+        getContractionLowering();
+    vector::VectorTransferSplit vectorTransferSplit = getSplitTransfers();
 
     vector::VectorTransformsOptions vectorTransformOptions;
     vectorTransformOptions.setVectorTransformsOptions(vectorContractLowering)
         .setVectorMultiReductionLowering(vectorMultiReductionLowering)
         .setVectorTransposeLowering(vectorTransposeLowering)
         .setVectorTransferSplit(vectorTransferSplit);
 
     VectorTransferToSCFOptions vectorTransferToSCFOptions =

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

patterns
       .add<VectorLoadToMemrefLoadLowering, VectorStoreToMemrefStoreLowering>(
           patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateVectorScanLoweringPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<ScanToArithOps>(patterns.getContext(), benefit);
 }
+
+//===----------------------------------------------------------------------===//
+// TableGen'd enum attribute definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/Transforms/VectorTransformsEnums.cpp.inc"

mlir/test/Dialect/LLVM/transform-e2e.mlir

 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op
   %1, %loops:3 = transform.structured.tile %0 [2, 2, 2]
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
   transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module_op
     {bufferize_function_boundaries = true}
   %func = transform.structured.match ops{["func.func"]} in %module_op
-  transform.vector.lower_vectors %func { multireduction_lowering = "innerreduce"}
+  transform.vector.lower_vectors %func multireduction_lowering = "innerreduction"
 }

mlir/test/Dialect/Vector/transform-vector.mlir

 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op
   %1, %loops:3 = transform.structured.tile %0 [8, 4, 2]
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
   transform.bufferization.one_shot_bufferize %module_op
 
   %func = transform.structured.match ops{["func.func"]} in %module_op
-  transform.vector.lower_vectors %func { multireduction_lowering = "innerreduce"}
+  transform.vector.lower_vectors %func multireduction_lowering = "innerreduction"
 }