[mlir] Correct verifyCompatibleShapes

[tpopp]

Summary

verifyCompatibleShapes is not transitive. Create an n-ary version and
update SameOperandShapes and SameOperandAndResultShapes traits to use
it.

[tpopp]

Add a missing newline and remove an unneeded include.

[frgossen]

A test case that would have passed before and fails now would be nice:

%2 = some_same_shape_op %0, %1 : tensor<?xf32>, tensor<4xf32> -> tensor<7xf32>

mlir/lib/IR/TypeUtilities.cpp

170

nit: extra line?

174

Should the same comment exists twice? I think it would be enough in the header.

182

any_of?

199–201

If you extract this, it'd be easier to read, I feel.

[tpopp]

Respond to feedback.

[tpopp]

mlir/lib/IR/TypeUtilities.cpp

174

I'm just following the other examples in the file.

182

No. Any mixture of ShapedTypes and non-ShapedTypes is not allowed in the other functions.

[tpopp]

Add test case.

[frgossen]

Thanks!

mlir/include/mlir/IR/TypeUtilities.h

 /// compatible if at least one is dynamic or both are equal. The element type
 /// does not matter.
 LogicalResult verifyCompatibleShape(Type type1, Type type2);
 
 /// Returns success if the given two arrays have the same number of elements and
 /// each pair wise entries have compatible shape.
 LogicalResult verifyCompatibleShapes(TypeRange types1, TypeRange types2);
 
+/// Returns success if all given types have compatible shapes. That is, they are
+/// all scalars (not shaped), or they are all shaped types and any ranked shapes
+/// have compatible dimensions. The element type does not matter.
+LogicalResult verifyCompatibleShapes(TypeRange types);
+
+/// Dimensions are compatible if all non-dynamic dims are equal.
+LogicalResult verifyCompatibleDims(ArrayRef<int64_t> dims);
 //===----------------------------------------------------------------------===//
 // Utility Iterators
 //===----------------------------------------------------------------------===//
 
 // An iterator for the element types of an op's operands of shaped types.
 class OperandElementTypeIterator final
     : public llvm::mapped_iterator<Operation::operand_iterator,
                                    Type (*)(Value)> {

mlir/lib/IR/Operation.cpp

return op->emitOpError()
            << "expected " << numOperands << " or more results";
   return success();
 }
 
 LogicalResult OpTrait::impl::verifySameOperandsShape(Operation *op) {
   if (failed(verifyAtLeastNOperands(op, 1)))
     return failure();
 
-  auto type = op->getOperand(0).getType();
-  for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) {
-    if (failed(verifyCompatibleShape(opType, type)))
-      return op->emitOpError() << "requires the same shape for all operands";
-  }
+  if (failed(verifyCompatibleShapes(op->getOperandTypes())))
+
   return success();
 }
 
 LogicalResult OpTrait::impl::verifySameOperandsAndResultShape(Operation *op) {
   if (failed(verifyAtLeastNOperands(op, 1)) ||
       failed(verifyAtLeastNResults(op, 1)))
     return failure();
 
-  auto type = op->getOperand(0).getType();
-  for (auto resultType : op->getResultTypes()) {
-    if (failed(verifyCompatibleShape(resultType, type)))
-      return op->emitOpError()
-             << "requires the same shape for all operands and results";
-  }
-  for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) {
-    if (failed(verifyCompatibleShape(opType, type)))
-      return op->emitOpError()
-             << "requires the same shape for all operands and results";
-  }
+  SmallVector<Type, 8> types(op->getOperandTypes());
+  types.append(llvm::to_vector<4>(op->getResultTypes()));
+
+  if (failed(verifyCompatibleShapes(types)))
+
   return success();
 }
 
 LogicalResult OpTrait::impl::verifySameOperandsElementType(Operation *op) {
   if (failed(verifyAtLeastNOperands(op, 1)))
     return failure();
   auto elementType = getElementTypeOrSelf(op->getOperand(0));
 

mlir/lib/IR/TypeUtilities.cpp

 //===- TypeUtilities.cpp - Helper function for type queries ---------------===//
 //
 // 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 generic type utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/TypeUtilities.h"
+
+#include <numeric>
+
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 
 using namespace mlir;
 
 Type mlir::getElementTypeOrSelf(Type type) {
LogicalResult mlir::verifyCompatibleShapes(TypeRange types1, TypeRange types2) {
   if (types1.size() != types2.size())
     return failure();
   for (auto it : llvm::zip_first(types1, types2))
     if (failed(verifyCompatibleShape(std::get<0>(it), std::get<1>(it))))
       return failure();
   return success();
 }
 
+LogicalResult mlir::verifyCompatibleDims(ArrayRef<int64_t> dims) {
+  if (dims.empty())
+    return success();
+  auto staticDim = std::accumulate(
+      dims.begin(), dims.end(), dims.front(), [](auto fold, auto dim) {
+        return ShapedType::isDynamic(dim) ? fold : dim;
+      });
+  return success(llvm::all_of(dims, [&](auto dim) {
+    return ShapedType::isDynamic(dim) || dim == staticDim;
+  }));
+}
+
+/// Returns success if all given types have compatible shapes. That is, they are
+/// all scalars (not shaped), or they are all shaped types and any ranked shapes
+/// have compatible dimensions. Dimensions are compatible if all non-dynamic
+/// dims are equal. The element type does not matter.
+LogicalResult mlir::verifyCompatibleShapes(TypeRange types) {
+  auto shapedTypes = llvm::to_vector<8>(llvm::map_range(
+      types, [](auto type) { return type.template dyn_cast<ShapedType>(); }));
+  // Return failure if some, but not all are not shaped. Return early if none
+  // are shaped also.
+  if (llvm::none_of(shapedTypes, [](auto t) { return t; }))
+    return success();
+  if (!llvm::all_of(shapedTypes, [](auto t) { return t; }))
+    return failure();
+
+  // Remove all unranked shapes
+  auto shapes = llvm::to_vector<8>(llvm::make_filter_range(
+      shapedTypes, [](auto shapedType) { return shapedType.hasRank(); }));
+  if (shapes.empty())
+    return success();
+
+  // All ranks should be equal
+  auto firstRank = shapes.front().getRank();
+  if (llvm::any_of(shapes,
+                   [&](auto shape) { return firstRank != shape.getRank(); }))
+    return failure();
+
+  for (unsigned i = 0; i < firstRank; ++i) {
+    // Retrieve all ranked dimensions
+    auto dims = llvm::to_vector<8>(llvm::map_range(
+        llvm::make_filter_range(
+            shapes, [&](auto shape) { return shape.getRank() >= i; }),
+        [&](auto shape) { return shape.getDimSize(i); }));
+    if (verifyCompatibleDims(dims).failed())
+      return failure();
+  }
+
+  return success();
+}
+
 OperandElementTypeIterator::OperandElementTypeIterator(
     Operation::operand_iterator it)
     : llvm::mapped_iterator<Operation::operand_iterator, Type (*)(Value)>(
           it, &unwrap) {}
 
 Type OperandElementTypeIterator::unwrap(Value value) {
   return value.getType().cast<ShapedType>().getElementType();
 }
 
 ResultElementTypeIterator::ResultElementTypeIterator(
     Operation::result_iterator it)
     : llvm::mapped_iterator<Operation::result_iterator, Type (*)(Value)>(
           it, &unwrap) {}
 
 Type ResultElementTypeIterator::unwrap(Value value) {
   return value.getType().cast<ShapedType>().getElementType();
 }

mlir/test/IR/traits.mlir

 
 func @failedSameOperandAndResultShape_operand_result_mismatch(%t10x10 : tensor<10x10xf32>, %t1: tensor<1xf32>) {
   // expected-error@+1 {{requires the same shape for all operands and results}}
   "test.same_operand_and_result_shape"(%t1, %t10x10) : (tensor<1xf32>, tensor<10x10xf32>) -> tensor<10x10xf32>
 }
 
 // -----
 
+func @failedSameOperandAndResultShape_operand_result_mismatch(%t10 : tensor<10xf32>, %t1: tensor<?xf32>) {
+  // expected-error@+1 {{requires the same shape for all operands and results}}
+  "test.same_operand_and_result_shape"(%t1, %t10) : (tensor<?xf32>, tensor<10xf32>) -> tensor<3xf32>
+}
+
+// -----
+
 func @failedSameOperandAndResultShape_no_operands() {
   // expected-error@+1 {{expected 1 or more operands}}
   "test.same_operand_and_result_shape"() : () -> (tensor<1xf32>)
 }
 
 // -----
 
 func @failedSameOperandAndResultShape_no_operands(%t1: tensor<1xf32>) {

 
 // -----
 
 // Test that operation with the SymbolTable Trait define a new symbol scope.
 "test.symbol_scope"() ({
   func private @foo()
   "test.finish" () : () -> ()
 }) : () -> ()
-func private @foo() 
+func private @foo()
 
 // -----
 
 // Test that operation with the SymbolTable Trait fails with  too many blocks.
 // expected-error@+1 {{Operations with a 'SymbolTable' must have exactly one block}}
 "test.symbol_scope"() ({
   ^entry:
     "test.finish" () : () -> ()