From 3dbb6678a5729df24529d580ec1a36cb313e68fa Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Fri, 12 Jun 2020 06:50:06 -0700
Subject: [PATCH] [mlir] Mark CastOp class's shape constraint

These ops have the same operands and result shapes.

Differential Revision: https://reviews.llvm.org/D81664
---
 mlir/include/mlir/Dialect/StandardOps/IR/Ops.td | 3 ++-
 mlir/test/Dialect/Standard/invalid.mlir         | 4 ++--
 mlir/test/IR/invalid-ops.mlir                   | 4 ++--
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 35bc0cd4..167205e 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -42,7 +42,8 @@ class Std_Op<string mnemonic, list<OpTrait> traits = []> :
 // Base class for standard cast operations. Requires single operand and result,
 // but does not constrain them to specific types.
 class CastOp<string mnemonic, list<OpTrait> traits = []> :
-    Std_Op<mnemonic, !listconcat(traits, [NoSideEffect])> {
+    Std_Op<mnemonic,
+           !listconcat(traits, [NoSideEffect, SameOperandsAndResultShape])> {
 
   let results = (outs AnyType);
 
diff --git a/mlir/test/Dialect/Standard/invalid.mlir b/mlir/test/Dialect/Standard/invalid.mlir
index 471ffeb..f2b71f6 100644
--- a/mlir/test/Dialect/Standard/invalid.mlir
+++ b/mlir/test/Dialect/Standard/invalid.mlir
@@ -2,7 +2,7 @@
 
 // CHECK-LABEL: test_index_cast_shape_error
 func @test_index_cast_shape_error(%arg0 : tensor<index>) -> tensor<2xi64> {
-  // expected-error @+1 {{operand type 'tensor<index>' and result type 'tensor<2xi64>' are cast incompatible}}
+  // expected-error @+1 {{requires the same shape for all operands and results}}
   %0 = index_cast %arg0 : tensor<index> to tensor<2xi64>
   return %0 : tensor<2xi64>
 }
@@ -11,7 +11,7 @@ func @test_index_cast_shape_error(%arg0 : tensor<index>) -> tensor<2xi64> {
 
 // CHECK-LABEL: test_index_cast_tensor_error
 func @test_index_cast_tensor_error(%arg0 : tensor<index>) -> i64 {
-  // expected-error @+1 {{operand type 'tensor<index>' and result type 'i64' are cast incompatible}}
+  // expected-error @+1 {{requires the same shape for all operands and results}}
   %0 = index_cast %arg0 : tensor<index> to i64
   return %0 : i64
 }
diff --git a/mlir/test/IR/invalid-ops.mlir b/mlir/test/IR/invalid-ops.mlir
index c8908a0..3668c25 100644
--- a/mlir/test/IR/invalid-ops.mlir
+++ b/mlir/test/IR/invalid-ops.mlir
@@ -685,7 +685,7 @@ func @fpext_f32_to_i32(%arg0 : f32) {
 // -----
 
 func @fpext_vec(%arg0 : vector<2xf16>) {
-  // expected-error@+1 {{are cast incompatible}}
+  // expected-error@+1 {{requires the same shape for all operands and results}}
   %0 = fpext %arg0 : vector<2xf16> to vector<3xf32>
   return
 }
@@ -757,7 +757,7 @@ func @fptrunc_f32_to_i32(%arg0 : f32) {
 // -----
 
 func @fptrunc_vec(%arg0 : vector<2xf16>) {
-  // expected-error@+1 {{are cast incompatible}}
+  // expected-error@+1 {{requires the same shape for all operands and results}}
   %0 = fptrunc %arg0 : vector<2xf16> to vector<3xf32>
   return
 }
-- 
2.7.4