[mlir] Add trait SameOperandsAndResultRank

[Tai78641]

Summary

This adds a native op trait SameOperandsAndResultRank
and associated verifier that checks that an operator's
operands and result types have same ranks if their ranks
are known.

Signed-off-by: Tai Ly <[e-mail]>
Change-Id: I2d536f77be10f3710d0c8d84c907ff492a984fda

[Tai78641]

please see RFC https://discourse.llvm.org/t/rfc-adding-allranksmatchifknown-trait-to-tosa-broadcastable-operators/71500 for use case of this new trait

[rsuderman]

mlir/include/mlir/IR/OpDefinition.h

344

Try to match the naming convention above. Should be verifySameOperandsAndResultRank

[Tai78641]

renamed verify function to verifySameOperandsAndResultRank

[rsuderman]

mlir/include/mlir/IR/OpDefinition.h

1124

You should also have the class name match. I.e. SameOperandsAndResultRank as above

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

373

Ditto.

mlir/test/lib/Dialect/Test/TestOps.td

697

Verify that the name convention maps here.

[Tai78641]

renamed trait to SameOperandsAndResultRank

[jmorse]

Apologies, copy-pasted your Phab URL into my commit there without updating, sorry for the noise -- a case of mistaken identity ._.

mlir/include/mlir/IR/OpDefinition.h

 LogicalResult verifyOneResult(Operation *op);
 LogicalResult verifyNResults(Operation *op, unsigned numOperands);
 LogicalResult verifyAtLeastNResults(Operation *op, unsigned numOperands);
 LogicalResult verifySameOperandsShape(Operation *op);
 LogicalResult verifySameOperandsAndResultShape(Operation *op);
 LogicalResult verifySameOperandsElementType(Operation *op);
 LogicalResult verifySameOperandsAndResultElementType(Operation *op);
 LogicalResult verifySameOperandsAndResultType(Operation *op);
+LogicalResult verifyAllRanksMatchIfKnown(Operation *op);
 LogicalResult verifyResultsAreBoolLike(Operation *op);
 LogicalResult verifyResultsAreFloatLike(Operation *op);
 LogicalResult verifyResultsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifyIsTerminator(Operation *op);
 LogicalResult verifyZeroSuccessors(Operation *op);
 LogicalResult verifyOneSuccessor(Operation *op);
 LogicalResult verifyNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyAtLeastNSuccessors(Operation *op, unsigned numSuccessors);

 class SameOperandsAndResultType
     : public TraitBase<ConcreteType, SameOperandsAndResultType> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySameOperandsAndResultType(op);
   }
 };
 
+/// This class verifies that op has same ranks for all
+/// operands and results types, if known.
+template <typename ConcreteType>
+class AllRanksMatchIfKnown
+    : public TraitBase<ConcreteType, AllRanksMatchIfKnown> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    return impl::verifyAllRanksMatchIfKnown(op);
+  }
+};
+
 /// This class verifies that any results of the specified op have a boolean
 /// type, a vector thereof, or a tensor thereof.
 template <typename ConcreteType>
 class ResultsAreBoolLike : public TraitBase<ConcreteType, ResultsAreBoolLike> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyResultsAreBoolLike(op);
   }

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

let methods = [
     >
   ];
 }
 
 // Op has the same operand and result type.
 // TODO: Change from hard coded to utilizing type inference trait.
 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 
+// Op has the same ranks for all operands and results types, if known.
+def AllRanksMatchIfKnown : NativeOpTrait<"AllRanksMatchIfKnown">;
+
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/lib/IR/Operation.cpp

if (encoding)
       if (auto rankedType = dyn_cast<RankedTensorType>(opType);
           encoding != rankedType.getEncoding())
         return op->emitOpError()
                << "requires the same encoding for all operands and results";
   }
   return success();
 }
 
+LogicalResult OpTrait::impl::verifyAllRanksMatchIfKnown(Operation *op) {
+  if (failed(verifyAtLeastNOperands(op, 1)))
+    return failure();
+
+  // delegate function that returns true if type is a shaped type with known
+  // rank
+  auto hasRank = [](const Type type) {
+    if (auto shaped_type = dyn_cast<ShapedType>(type))
+      return shaped_type.hasRank();
+
+    return false;
+  };
+
+  auto rankedOperandTypes =
+      llvm::make_filter_range(op->getOperandTypes(), hasRank);
+  auto rankedResultTypes =
+      llvm::make_filter_range(op->getResultTypes(), hasRank);
+
+  // If all operands and results are unranked, then no further verification.
+  if (rankedOperandTypes.empty() && rankedResultTypes.empty())
+    return success();
+
+  // delegate function that returns rank of shaped type with known rank
+  auto getRank = [](const Type type) {
+    return type.cast<ShapedType>().getRank();
+  };
+
+  auto rank = !rankedOperandTypes.empty() ? getRank(*rankedOperandTypes.begin())
+                                          : getRank(*rankedResultTypes.begin());
+
+  for (const auto type : rankedOperandTypes) {
+    if (rank != getRank(type)) {
+      return op->emitOpError("operands don't have matching ranks");
+    }
+  }
+
+  for (const auto type : rankedResultTypes) {
+    if (rank != getRank(type)) {
+      return op->emitOpError("result type has different rank than operands");
+    }
+  }
+
+  return success();
+}
+
 LogicalResult OpTrait::impl::verifyIsTerminator(Operation *op) {
   Block *block = op->getBlock();
   // Verify that the operation is at the end of the respective parent block.
   if (!block || &block->back() != op)
     return op->emitOpError("must be the last operation in the parent block");
   return success();
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

 
 def OperandZeroAndResultHaveSameRank :
     TEST_Op<"operand0_and_result_have_same_rank",
             [AllRanksMatch<["x", "res"]>]> {
   let arguments = (ins AnyShaped:$x, AnyShaped:$y);
   let results = (outs AnyShaped:$res);
 }
 
+def OperandsAndResultsHaveSameRankIfKnown :
+    TEST_Op<"operands_and_result_have_same_rank_if_known", [AllRanksMatchIfKnown]> {
+  let arguments = (ins AnyShaped:$x, AnyShaped:$y);
+  let results = (outs AnyShaped:$res);
+}
+
 def OperandZeroAndResultHaveSameShape :
     TEST_Op<"operand0_and_result_have_same_shape",
             [AllShapesMatch<["x", "res"]>]> {
   let arguments = (ins AnyShaped:$x, AnyShaped:$y);
   let results = (outs AnyShaped:$res);
 }
 
 def OperandZeroAndResultHaveSameElementCount :

mlir/test/mlir-tblgen/types.mlir

 func.func @same_rank_failure(%arg0: tensor<1x2xi32>, %arg1: tensor<1x2xf32>) {
   // expected-error@+1 {{all of {x, res} have same rank}}
   "test.operand0_and_result_have_same_rank"(%arg0, %arg1) : (tensor<1x2xi32>, tensor<1x2xf32>) -> (tensor<3xi32>)
   return
 }
 
 // -----
 
+// CHECK-LABEL: same_rank_if_known_success
+func.func @same_rank_if_known_success(%t1xi : tensor<1xi32>, %t2xf : tensor<2xf32>, %m3xi : memref<3xi32>, %t1x2xf : tensor<1x2xf32>, %tuxi : tensor<*xi32>) {
+  %0 = "test.operands_and_result_have_same_rank_if_known"(%t1xi, %t2xf) : (tensor<1xi32>, tensor<2xf32>) -> (tensor<3xf64>)
+  %1 = "test.operands_and_result_have_same_rank_if_known"(%t1xi, %m3xi) : (tensor<1xi32>, memref<3xi32>) -> (tensor<3xi64>)
+  %3 = "test.operands_and_result_have_same_rank_if_known"(%tuxi, %t2xf) : (tensor<*xi32>, tensor<2xf32>) -> (tensor<2xf32>)
+  %4 = "test.operands_and_result_have_same_rank_if_known"(%t1x2xf, %tuxi) : (tensor<1x2xf32>, tensor<*xi32>) -> (tensor<1x2xf64>)
+  return
+}
+
+// -----
+
+func.func @same_rank_if_known_failure(%arg0: tensor<1xi32>, %arg1: tensor<1x2xf32>) {
+  // expected-error@+1 {{operands don't have matching ranks}}
+  %0 = "test.operands_and_result_have_same_rank_if_known"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xf32>) -> (tensor<*xf32>)
+  return
+}
+
+// -----
+
+func.func @same_rank_if_known_failure(%arg0: tensor<1xi32>, %arg1: tensor<1x2xf32>) {
+  // expected-error@+1 {{result type has different rank than operands}}
+  %0 = "test.operands_and_result_have_same_rank_if_known"(%arg1, %arg1) : (tensor<1x2xf32>, tensor<1x2xf32>) -> (tensor<1x2x3xf32>)
+  return
+}
+
+// -----
+
 // CHECK-LABEL: same_shape_success
 func.func @same_shape_success(%t2x3: tensor<2x3xi32>, %m2x3: memref<2x3xf32>, %v2x3 : vector<2x3xi32>, %t4x5 : tensor<4x5xi32>) {
   "test.operand0_and_result_have_same_shape"(%t2x3, %t4x5) : (tensor<2x3xi32>, tensor<4x5xi32>) -> (tensor<2x3xf32>)
   "test.operand0_and_result_have_same_shape"(%t2x3, %t4x5) : (tensor<2x3xi32>, tensor<4x5xi32>) -> (memref<2x3xf32>)
   "test.operand0_and_result_have_same_shape"(%t2x3, %t4x5) : (tensor<2x3xi32>, tensor<4x5xi32>) -> (vector<2x3xf32>)
   return
 }