[Mlir-commits] [mlir] 3d56f16 - [mlir][StandardOps] Updated IndexCastOp to support tensor<index> cast

[River Riddle]

Author: Rob Suderman
Date: 2020-06-10T17:19:08-07:00
New Revision: 3d56f166bd3c00e1b05931d57d385ef2b41c092d

URL: https://github.com/llvm/llvm-project/commit/3d56f166bd3c00e1b05931d57d385ef2b41c092d
DIFF: https://github.com/llvm/llvm-project/commit/3d56f166bd3c00e1b05931d57d385ef2b41c092d.diff

LOG: [mlir][StandardOps] Updated IndexCastOp to support tensor<index> cast

Summary:
We now support index casting for tensor<index> to tensor<int>. This
better supports compatibility with the Shape dialect.

Differential Revision: https://reviews.llvm.org/D81611

Added: 
    mlir/test/Dialect/Standard/invalid.mlir
    mlir/test/Dialect/Standard/ops.mlir

Modified: 
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 0177dfca5460..47532c5bc5f3 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -1765,6 +1765,15 @@ bool FPTruncOp::areCastCompatible(Type a, Type b) {
 
 // Index cast is applicable from index to integer and backwards.
 bool IndexCastOp::areCastCompatible(Type a, Type b) {
+  if (a.isa<ShapedType>() && b.isa<ShapedType>()) {
+    auto aShaped = a.cast<ShapedType>();
+    auto bShaped = b.cast<ShapedType>();
+
+    return (aShaped.getShape() == bShaped.getShape()) &&
+           areCastCompatible(aShaped.getElementType(),
+                             bShaped.getElementType());
+  }
+
   return (a.isIndex() && b.isSignlessInteger()) ||
          (a.isSignlessInteger() && b.isIndex());
 }

diff  --git a/mlir/test/Dialect/Standard/invalid.mlir b/mlir/test/Dialect/Standard/invalid.mlir
new file mode 100644
index 000000000000..471ffeb14538
--- /dev/null
+++ b/mlir/test/Dialect/Standard/invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt -split-input-file %s -verify-diagnostics
+
+// 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}}
+  %0 = index_cast %arg0 : tensor<index> to tensor<2xi64>
+  return %0 : 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}}
+  %0 = index_cast %arg0 : tensor<index> to i64
+  return %0 : i64
+}

diff  --git a/mlir/test/Dialect/Standard/ops.mlir b/mlir/test/Dialect/Standard/ops.mlir
new file mode 100644
index 000000000000..3b098eb960ac
--- /dev/null
+++ b/mlir/test/Dialect/Standard/ops.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-opt -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: test_index_cast
+func @test_index_cast(%arg0 : index) -> i64 {
+  %0 = index_cast %arg0 : index to i64
+  return %0 : i64
+}
+
+// CHECK-LABEL: test_index_cast_tensor
+func @test_index_cast_tensor(%arg0 : tensor<index>) -> tensor<i64> {
+  %0 = index_cast %arg0 : tensor<index> to tensor<i64>
+  return %0 : tensor<i64>
+}
+
+// CHECK-LABEL: test_index_cast_tensor_reverse
+func @test_index_cast_tensor_reverse(%arg0 : tensor<i64>) -> tensor<index> {
+  %0 = index_cast %arg0 : tensor<i64> to tensor<index>
+  return %0 : tensor<index>
+}
+