[mlir] Add IteratorType enum to StructuredOpsUtils.

[olegshyshkov]

Summary

Use the new enum in TilingIterface and verify that iterator_type attribute in
LinalgOp interface is compatible with the enum values. Later IteratorType enum
will be used in LinalgInterface to replace the current iterator_type attribute
array of string.

Existing enums in Linalg are moved into a separate td file and tablegen build
target. This is necessary, have one I32EnumAttr in a shared space that generated
enum class definition and EnumAttrs is dialect-specific location. Otherwise
there might be a conflict that I32EnumAttr generates enum definitions in
multiple places.

[olegshyshkov]

Added file header comments.

[olegshyshkov]

Remove LinalgBase.h. It wasn't used in this patch.

mlir/include/mlir/Dialect/CMakeLists.txt

-add_subdirectory(Affine)
 add_subdirectory(AMDGPU)
+add_subdirectory(AMX)
+add_subdirectory(Affine)
 add_subdirectory(Arithmetic)
-add_subdirectory(Async)
 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
-add_subdirectory(AMX)
+add_subdirectory(Async)
 add_subdirectory(Bufferization)
 add_subdirectory(Complex)
 add_subdirectory(ControlFlow)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(Func)
 add_subdirectory(GPU)
-add_subdirectory(Math)
-add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
-add_subdirectory(MemRef)
+add_subdirectory(Linalg)
 add_subdirectory(MLProgram)
+add_subdirectory(Math)
+add_subdirectory(MemRef)
 add_subdirectory(NVGPU)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
+add_subdirectory(SPIRV)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)
-add_subdirectory(SPIRV)
 add_subdirectory(Tensor)
 add_subdirectory(Tosa)
 add_subdirectory(Transform)
+add_subdirectory(Utils)
 add_subdirectory(Vector)
 add_subdirectory(X86Vector)

mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt

 add_custom_target(LinalgOdsGen
   DEPENDS
   MLIRLinalgNamedStructuredOpsYamlIncGen
 )
 add_dependencies(mlir-headers LinalgOdsGen)
 
 add_mlir_dialect(LinalgOps linalg)
 
-set(LLVM_TARGET_DEFINITIONS LinalgOps.td)
+set(LLVM_TARGET_DEFINITIONS LinalgEnums.td)
 mlir_tablegen(LinalgOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(LinalgOpsEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRLinalgOpsEnumsIncGen)
 add_dependencies(mlir-headers MLIRLinalgOpsEnumsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LinalgOps.td)
 mlir_tablegen(LinalgOpsAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(LinalgOpsAttrDefs.cpp.inc -gen-attrdef-defs)

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

 //===- LinalgBase.td - Linalg dialect base support ---------*- 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 definition file for base linear algebra support.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_BASE
 #define LINALG_BASE
 
+include "mlir/Dialect/Linalg/IR/LinalgEnums.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 
 def Linalg_Dialect : Dialect {
   let name = "linalg";
   let description = [{
     The `linalg` dialect groups together a set of types, operations and
     transformations that are useful to implement a structured abstraction on
let extraClassDeclaration = [{
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
 
   let emitAccessorPrefix = kEmitAccessorPrefix_Both;
 }
 
 // Define the function attribute enums matching the OpDSL functions.
-def UnaryFn : I32EnumAttr<"UnaryFn", "", [
-  I32EnumAttrCase<"exp", 0>,
-  I32EnumAttrCase<"log", 1>,
-  I32EnumAttrCase<"abs", 2>,
-  I32EnumAttrCase<"ceil", 3>,
-  I32EnumAttrCase<"floor", 4>,
-  I32EnumAttrCase<"negf", 5>
-]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::linalg";
-}
-def BinaryFn : I32EnumAttr<"BinaryFn", "", [
-  I32EnumAttrCase<"add", 0>,
-  I32EnumAttrCase<"sub", 1>,
-  I32EnumAttrCase<"mul", 2>,
-  I32EnumAttrCase<"max_signed", 3>,
-  I32EnumAttrCase<"min_signed", 4>,
-  I32EnumAttrCase<"max_unsigned", 5>,
-  I32EnumAttrCase<"min_unsigned", 6>
-]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::linalg";
-}
-def TypeFn : I32EnumAttr<"TypeFn", "", [
-  I32EnumAttrCase<"cast_signed", 0>,
-  I32EnumAttrCase<"cast_unsigned", 1>
-]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::linalg";
-}
-
 def UnaryFnAttr : EnumAttr<Linalg_Dialect, UnaryFn, "unary_fn"> {
   let assemblyFormat = "`<` $value `>`";
 }
 def BinaryFnAttr : EnumAttr<Linalg_Dialect, BinaryFn, "binary_fn"> {
   let assemblyFormat = "`<` $value `>`";
 }
 def TypeFnAttr : EnumAttr<Linalg_Dialect, TypeFn, "type_fn"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
 #endif // LINALG_BASE

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

+//===- LinalgBase.td - Linalg dialect base support ---------*- 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 definition file for enums used in linear algebra operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LINALG_ENUMS
+#define LINALG_ENUMS
+
+include "mlir/IR/EnumAttr.td"
+
+// Define the function attribute enums matching the OpDSL functions.
+def UnaryFn : I32EnumAttr<"UnaryFn", "", [
+  I32EnumAttrCase<"exp", 0>,
+  I32EnumAttrCase<"log", 1>,
+  I32EnumAttrCase<"abs", 2>,
+  I32EnumAttrCase<"ceil", 3>,
+  I32EnumAttrCase<"floor", 4>,
+  I32EnumAttrCase<"negf", 5>
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::linalg";
+}
+def BinaryFn : I32EnumAttr<"BinaryFn", "", [
+  I32EnumAttrCase<"add", 0>,
+  I32EnumAttrCase<"sub", 1>,
+  I32EnumAttrCase<"mul", 2>,
+  I32EnumAttrCase<"max_signed", 3>,
+  I32EnumAttrCase<"min_signed", 4>,
+  I32EnumAttrCase<"max_unsigned", 5>,
+  I32EnumAttrCase<"min_unsigned", 6>
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::linalg";
+}
+def TypeFn : I32EnumAttr<"TypeFn", "", [
+  I32EnumAttrCase<"cast_signed", 0>,
+  I32EnumAttrCase<"cast_unsigned", 1>
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::linalg";
+}
+
+#endif // LINALG_ENUMS

mlir/include/mlir/Dialect/Utils/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS StructuredOpsUtils.td)
+mlir_tablegen(DialectUtilsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(DialectUtilsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRDialectUtilsIncGen)
+add_dependencies(mlir-headers MLIRDialectUtilsIncGen)

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 #define MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H
 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 
+// Pull in all enum type definitions and utility function declarations.
+#include "mlir/Dialect/Utils/DialectUtilsEnums.h.inc"
+
 namespace mlir {
 
 class OpBuilder;
 
 /// Tests whether the given maps describe a row major matmul. The test is
 /// permutation-invariant. Note that this only checks the affine maps from an
 /// operation, so does not perform any checks on the math being performed within
 /// the reduction.

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.td

+//===- StructuredOpsUtils.td - structured ops enums --------*- 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 STRUCTURED_OPS_UTILS
+#define STRUCTURED_OPS_UTILS
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
+
+def IteratorType : I32EnumAttr<"IteratorType", "Iterator type", [
+  I32EnumAttrCase<"parallel", 0>,
+  I32EnumAttrCase<"reduction", 1>
+]> {
+    let genSpecializedAttr = 0;
+    let cppNamespace = "::mlir::utils";
+}
+
+#endif // STRUCTURED_OPS_UTILS

mlir/include/mlir/Interfaces/TilingInterface.h

 //===- TilingInterface.h - Interface for tiling operations ------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the definitions of the TilingInterface defined in
 // `TilingInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_TILINGINTERFACE_H_
 #define MLIR_INTERFACES_TILINGINTERFACE_H_
 
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
 /// Include the ODS generated interface header files.
 #include "mlir/Interfaces/TilingInterface.h.inc"
 
 #endif // MLIR_INTERFACES_TILINGINTERFACE_H_

mlir/include/mlir/Interfaces/TilingInterface.td

let methods = [
         /*retType=*/"SmallVector<Value>",
         /*methodName=*/"getDestinationOperands",
         /*args=*/(ins "OpBuilder &":$b),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return ValueRange{};"
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Returns a list of `StringRef`s that describe the number of
-          loops and the iterator types of the operation. The list is
-          expected to use
-          `getParallelIteratorTypeName()`/`getReductionIteratorTypeName()`
-          from MLIR Structured Op Utils.
+          Returns a list of iterator types that describe the number of loops.
         }],
-        /*retType=*/"SmallVector<StringRef>",
+        /*retType=*/"SmallVector<utils::IteratorType>",
         /*methodName=*/"getLoopIteratorTypes",
         /*args=*/(ins),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return {};"
       >,
       InterfaceMethod<
         /*desc=*/[{
           Returns a list of ranges that describe the loop bounds and

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

 
 LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   LinalgOp linalgOp = cast<LinalgOp>(op);
 
   // Check all iterator types are known.
   auto iteratorTypesRange =
       linalgOp.iterator_types().getAsValueRange<StringAttr>();
   for (StringRef iteratorType : iteratorTypesRange) {
-    if (!llvm::is_contained(getAllIteratorTypeNames(), iteratorType))
+    if (!llvm::is_contained(getAllIteratorTypeNames(), iteratorType) ||
+        !utils::symbolizeIteratorType(iteratorType).has_value())
       return op->emitOpError("unexpected iterator_type (")
              << iteratorType << ")";
   }
 
   // Before checking indexing maps, we need to make sure the attributes
   // referenced by it are valid.
   if (linalgOp.hasDynamicIndexingMaps())
     if (failed(linalgOp.verifyIndexingMapRequiredAttributes()))

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

struct LinalgOpTilingInterface
     : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
                                             LinalgOpTy> {
   /// Return the destination operands.
   SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
     return cast<LinalgOp>(op).getOutputOperands();
   }
 
   /// Return the loop iterator type.
-  SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
+  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
     return llvm::to_vector(
         llvm::map_range(concreteOp.iterator_types(), [](Attribute strAttr) {
-          return strAttr.cast<StringAttr>().getValue();
+          return utils::symbolizeIteratorType(
+                     strAttr.cast<StringAttr>().getValue())
+              .getValue();
         }));
   }
 
   /// Return the iteration domain range.
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     Location loc = op->getLoc();

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
 
     auto padOp = cast<PadOp>(op);
     SmallVector<OpFoldResult> mixedSizes = getAsOpFoldResult(reifiedShapes[0]);
     Value initTensor = b.create<linalg::InitTensorOp>(
         op->getLoc(), mixedSizes, padOp.getResultType().getElementType());
     return {initTensor};
   }
 
-  SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
+  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     auto padOp = cast<PadOp>(op);
-    SmallVector<StringRef> iteratorTypes(padOp.getResultType().getRank(),
-                                         getParallelIteratorTypeName());
+    SmallVector<utils::IteratorType> iteratorTypes(
+        padOp.getResultType().getRank(), utils::IteratorType::parallel);
     return iteratorTypes;
   }
 
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     ReifiedRankedShapedTypeDims reifiedShapes;
     ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
         dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
     (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);

mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp

 //===- StructuredOpsUtils.cpp - Utilities used by structured ops ----------===//
 //
 // 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 "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinAttributes.h"
 
+#include "mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc"
+
 using namespace mlir;
 
 bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
   if (indexingMaps.size() != 3)
     return false;
 
   auto map0 = indexingMaps[0].cast<AffineMapAttr>().getValue();
   auto map1 = indexingMaps[1].cast<AffineMapAttr>().getValue();

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     ]),
     includes = ["include"],
     deps = [
         ":IR",
         "//llvm:Support",
     ],
 )
 
+td_library(
+    name = "DialectUtilsTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Utils/StructuredOpsUtils.td",
+    ],
+    includes = ["include"],
+    deps = [":OpBaseTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "DialectUtilsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/Utils/DialectUtilsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Utils/StructuredOpsUtils.td",
+    deps = [":DialectUtilsTdFiles"],
+)
+
 cc_library(
     name = "DialectUtils",
     srcs = glob([
         "lib/Dialect/Utils/*.cpp",
         "lib/Dialect/Utils/*.h",
     ]),
     hdrs = glob([
         "include/mlir/Dialect/Utils/*.h",
     ]),
     includes = ["include"],
     deps = [
+        ":DialectUtilsIncGen",
         ":IR",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "AffineDialect",
deps = [
         "//llvm:Support",
     ],
 )
 
 td_library(
     name = "LinalgOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/Linalg/IR/LinalgBase.td",
+        "include/mlir/Dialect/Linalg/IR/LinalgEnums.td",
         "include/mlir/Dialect/Linalg/IR/LinalgOps.td",
     ],
     includes = ["include"],
     deps = [
         ":ControlFlowInterfacesTdFiles",
+        ":DialectUtilsTdFiles",
         ":InferTypeOpInterfaceTdFiles",
         ":LoopLikeInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
         ":TilingInterfaceTdFiles",
         ":ViewLikeInterfaceTdFiles",
     ],
 )
tbl_outs = [
         (
             [
                 "-gen-dialect-defs",
                 "-dialect=linalg",
             ],
             "include/mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc",
         ),
         (
-            ["-gen-enum-decls"],
-            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.h.inc",
-        ),
-        (
-            ["-gen-enum-defs"],
-            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc",
-        ),
-        (
             ["-gen-attrdef-decls"],
             "include/mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.h.inc",
         ),
         (
             ["-gen-attrdef-defs"],
             "include/mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgOps.td",
     deps = [":LinalgOpsTdFiles"],
 )
 
 gentbl_cc_library(
+    name = "LinalgEnumsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Linalg/IR/LinalgEnums.td",
+    deps = [":LinalgOpsTdFiles"],
+)
+
+gentbl_cc_library(
     name = "LinalgTransformOpsIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["-gen-op-decls"],
             "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h.inc",
         ),
         (
deps = [
         ":BufferizationDialect",
         ":ComplexDialect",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
+        ":LinalgEnumsIncGen",
         ":LinalgInterfacesIncGen",
         ":LinalgNamedStructuredOpsYamlIncGen",
         ":LinalgOpsIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MathDialect",
         ":MemRefDialect",
         ":Parser",
         ":SCFDialect",

 )
 
 cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
     includes = ["include"],
     deps = [
+        ":DialectUtils",
         ":IR",
         ":Support",
         ":TilingInterfaceIncGen",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )