[llvm-branch-commits] [mlir] 1e17c40 - [MLIR] Introduce trait EquiMemRefAndEltTypeOperands

[Uday Bondhugula via llvm-branch-commits]

Author: Uday Bondhugula
Date: 2021-09-22T14:23:45+05:30
New Revision: 1e17c40c4ce8daf46bf5df28a04d1e9ea5690cb2

URL: https://github.com/llvm/llvm-project/commit/1e17c40c4ce8daf46bf5df28a04d1e9ea5690cb2
DIFF: https://github.com/llvm/llvm-project/commit/1e17c40c4ce8daf46bf5df28a04d1e9ea5690cb2.diff

LOG: [MLIR] Introduce trait EquiMemRefAndEltTypeOperands

This trait is used by LHLO ops where the operands and result may either
have operands of the same type or equivalent types where a 0-d memref of
an elemental type is considered equivalent to that elemental type.

TODO: add a test dialect test case here.

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/include/mlir/IR/OpDefinition.h
    mlir/lib/IR/Operation.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index ac99ac4d0897d..d82b734a329a8 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1889,6 +1889,9 @@ def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
 // Op has the same operand type.
 def SameTypeOperands : NativeOpTrait<"SameTypeOperands">;
+// Op has the same operand type or 0-d memref equivalent type.
+def EquiMemRefAndEltTypeOperands
+    : NativeOpTrait<"EquiMemRefAndEltTypeOperands">;
 // Op has same shape for all operands.
 def SameOperandsShape : NativeOpTrait<"SameOperandsShape">;
 // Op has same operand and result shape.

diff  --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index e7f644f456b22..932ef0542781e 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -248,6 +248,7 @@ LogicalResult verifyAtLeastNOperands(Operation *op, unsigned numOperands);
 LogicalResult verifyOperandsAreFloatLike(Operation *op);
 LogicalResult verifyOperandsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifySameTypeOperands(Operation *op);
+LogicalResult verifyEquiMemRefAndEltTypeOperands(Operation *op);
 LogicalResult verifyZeroRegion(Operation *op);
 LogicalResult verifyOneRegion(Operation *op);
 LogicalResult verifyNRegions(Operation *op, unsigned numRegions);
@@ -1143,6 +1144,18 @@ class SameTypeOperands : public TraitBase<ConcreteType, SameTypeOperands> {
   }
 };
 
+/// This class verifies that all operands of the specified op have the same
+/// type with the exception that 0-d memref type of an elemental type is treated
+/// the same as that elemental type.
+template <typename ConcreteType>
+class EquiMemRefAndEltTypeOperands
+    : public TraitBase<ConcreteType, EquiMemRefAndEltTypeOperands> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    return impl::verifyEquiMemRefAndEltTypeOperands(op);
+  }
+};
+
 /// This class provides the API for a sub-set of ops that are known to be
 /// constant-like. These are non-side effecting operations with one result and
 /// zero operands that can always be folded to a specific attribute value.

diff  --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 457c993b92ad3..657b612ccce16 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -792,6 +792,32 @@ LogicalResult OpTrait::impl::verifySameTypeOperands(Operation *op) {
   return success();
 }
 
+LogicalResult OpTrait::impl::verifyEquiMemRefAndEltTypeOperands(Operation *op) {
+  // Zero or one operand always have the "same" type.
+  unsigned nOperands = op->getNumOperands();
+  if (nOperands < 2)
+    return success();
+
+  StringRef msg =
+      "requires all operands to have the same or equivalent type where a 0-d "
+      "memref type is considered equivalent to its elemental type";
+
+  auto prevType = op->getOperand(0).getType();
+  for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) {
+    if (opType != prevType) {
+      if (auto opMemRefType = opType.dyn_cast<MemRefType>())
+        if (opMemRefType.getRank() != 0 ||
+            opMemRefType.getElementType() != prevType)
+          return op->emitOpError(msg);
+      if (auto prevMemRefType = prevType.dyn_cast<MemRefType>())
+        if (prevMemRefType.getRank() != 0 &&
+            prevMemRefType.getElementType() != opType)
+          return op->emitOpError(msg);
+    }
+  }
+  return success();
+}
+
 LogicalResult OpTrait::impl::verifyZeroRegion(Operation *op) {
   if (op->getNumRegions() != 0)
     return op->emitOpError() << "requires zero regions";