Use a trait to avoid two-way dialect dependence

Author: cferry-AMD

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H
 
 #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
+#include "mlir/Dialect/Affine/IR/AffineTraits.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -13,6 +13,7 @@
 #ifndef AFFINE_OPS
 #define AFFINE_OPS
 
+include "mlir/Dialect/Affine/IR/AffineTraits.td"
 include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"

mlir/include/mlir/Dialect/Affine/IR/AffineTraits.h

@@ -0,0 +1,27 @@
+//===- AffineTraits.h - MLIR Affine Traits --------------------------------===//
+//
+// 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 defines traits brought in by the Affine dialect.
+//
+//===----------------------------------------------------------------------===//
+#ifndef AFFINE_TRAITS_H
+#define AFFINE_TRAITS_H
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir::OpTrait {
+
+template <typename ConcreteType>
+class AffineDim : public TraitBase<ConcreteType, AffineDim> {
+public:
+  static LogicalResult verifyTrait(Operation *op) { return success(); }
+};
+
+} // namespace mlir::OpTrait
+
+#endif // AFFINE_TRAITS_H

mlir/include/mlir/Dialect/Affine/IR/AffineTraits.td

@@ -0,0 +1,22 @@
+//===- AffineTraits.td - Affine dialect traits -------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines traits brought in by the MLIR Affine dialect.
+//
+//===----------------------------------------------------------------------===//
+#ifndef AFFINE_TRAITS
+#define AFFINE_TRAITS
+
+include "mlir/IR/OpBase.td"
+
+// Trait to declare that an op result is an affine dimension identifier.
+// Prevents the result from being seen as a symbol into AffineMaps
+// and IntegerSets.
+def AffineDim : NativeOpTrait<"AffineDim">;
+
+#endif // AFFINE_TRAITS

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_LINALG_IR_LINALG_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Affine/IR/AffineTraits.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"

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

@@ -13,6 +13,7 @@
 #ifndef LINALG_OPS
 #define LINALG_OPS
 
+include "mlir/Dialect/Affine/IR/AffineTraits.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
@@ -46,7 +47,7 @@ def Linalg_YieldOp : Linalg_Op<"yield", [Pure, ReturnLike, Terminator]>,
   let hasVerifier = 1;
 }
 
-def Linalg_IndexOp : Linalg_Op<"index", [Pure]>,
+def Linalg_IndexOp : Linalg_Op<"index", [Pure, AffineDim]>,
     Arguments<(ins ConfinedAttr<I64Attr, [IntMinValue<0>]>:$dim)>,
     Results<(outs Index:$result)> {
   let summary = "linalg index operation";

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -314,7 +314,7 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
   }
 
   // Remove me: linalg.index ops are valid affine dim identifiers
-  if (isa<linalg::IndexOp>(op))
+  if (op->hasTrait<OpTrait::AffineDim>())
     return true;
 
   // Affine apply operation is ok if all of its operands are ok.
@@ -445,7 +445,7 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
   }
 
   // Remove me: linalg.index ops are not valid affine symbols
-  if (isa<linalg::IndexOp>(defOp))
+  if (defOp->hasTrait<OpTrait::AffineDim>())
     return false;
 
   // Constant operation is ok.