[mlir] Fix SameOperandsAndResultType to check encoding.

[jpienaar]

Summary

Encoding was accidentally left out here even though it forms part of the type.
This is small tightening step and I'll look at follow on to tighten more.

[Mogball]

I'm getting confused as to what this trait actually means anymore. Why doesn't it just check that the types are pointer equal?

[rriddle]

I agree with Jeff here, and think this is getting further out of hand. I strongly feel like we should push https://reviews.llvm.org/D125051 forward.

[jpienaar]

In D140445#4009925, @rriddle wrote:

I agree with Jeff here, and think this is getting further out of hand. I strongly feel like we should push https://reviews.llvm.org/D125051 forward.

I'll look at that one after as its more invasive.

In D140445#4009643, @Mogball wrote:

I'm getting confused as to what this trait actually means anymore. Why doesn't it just check that the types are pointer equal?

This trait signifies equivalence and only flags where known non-matching (not where potentially not matching). E.g., tensor<?xf32> input to op where the other operands is tensor<2xf32> effectively means post type inference these should be equal or an error, or alternatively that we have effectively a runtime assert that they are equal. Which is also true for 2 tensor<?xf32> operands that pointer equality is not sufficient, it has to match executionally - the same static type doesn't suffice, we require the fully known type to be equal, which is something that the pointer check would be permissive off. But it would not fail verification with tensor<?xf32> != tensor<2xf32> . So it means one can express this type constraint and represents something different than the raw type being the same. Naming is difficult.

[jpienaar]

(and by that mean I agree it is confusing and I'm perfectly open to renaming)

mlir/lib/IR/Operation.cpp

 
 LogicalResult OpTrait::impl::verifySameOperandsAndResultType(Operation *op) {
   if (failed(verifyAtLeastNOperands(op, 1)) ||
       failed(verifyAtLeastNResults(op, 1)))
     return failure();
 
   auto type = op->getResult(0).getType();
   auto elementType = getElementTypeOrSelf(type);
+  Attribute encoding = nullptr;
+  if (auto rankedType = dyn_cast<RankedTensorType>(type))
+    encoding = rankedType.getEncoding();
   for (auto resultType : llvm::drop_begin(op->getResultTypes())) {
     if (getElementTypeOrSelf(resultType) != elementType ||
         failed(verifyCompatibleShape(resultType, type)))
       return op->emitOpError()
              << "requires the same type for all operands and results";
+    if (encoding)
+      if (auto rankedType = dyn_cast<RankedTensorType>(resultType);
+          encoding != rankedType.getEncoding())
+        return op->emitOpError()
+               << "requires the same encoding for all operands and results";
   }
   for (auto opType : op->getOperandTypes()) {
     if (getElementTypeOrSelf(opType) != elementType ||
         failed(verifyCompatibleShape(opType, type)))
       return op->emitOpError()
              << "requires the same type for all operands and results";
+    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::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)

mlir/test/IR/traits.mlir

 func.func @failedSameOperandAndResultType_operand_result_mismatch(%t10 : tensor<10xf32>, %t20 : tensor<20xf32>) {
   // expected-error@+1 {{requires the same type for all operands and results}}
   "test.same_operand_and_result_type"(%t10, %t20) : (tensor<10xf32>, tensor<20xf32>) -> tensor<10xf32>
   return
 }
 
 // -----
 
+func.func @failedSameOperandAndResultType_encoding_mismatch(%t10 : tensor<10xf32>, %t20 : tensor<10xf32>) {
+  // expected-error@+1 {{requires the same encoding for all operands and results}}
+  "test.same_operand_and_result_type"(%t10, %t20) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32, "enc">
+  return
+}
+
+// -----
+
 func.func @failedElementwiseMappable_different_rankedness(%arg0: tensor<?xf32>, %arg1: tensor<*xf32>) {
   // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type}}
   %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<*xf32>) -> tensor<*xf32>
   return
 }
 
 // -----