Updated IndexCastOp to support tensor<index> cast

[rsuderman]

Summary

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

[rriddle]

Can you add test cases?

[rsuderman]

Updated to include IR tests for index_cast.

[rriddle]

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

1782

Can you add a test for this?

mlir/test/Dialect/Standard/ops.mlir

5

This looks more like a test for the core parser than for the Standard dialect. Not sure that we should include this here.

[rriddle]

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

1772

Is this intended to support mismatch static/dynamic dimensions?

[rsuderman]

Removed unneeded folder.

[rsuderman]

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

1782

Turns out its unneeded as index_cast does not allow noop casts.

mlir/test/Dialect/Standard/ops.mlir

5

My goal is to match the invalid tests and have correctly parsed tests pass through (both with and without tensors). Is there a better way to include examples of valid index_casts?

[rriddle]

Can you revise the commit description if you aren't updating the folder?

mlir/test/Dialect/Standard/ops.mlir

5

I meant the last RUN line. Parsing the generic input seems like it is testing the core parser and nothing related to the standard dialect that isn't tested by the other two tests.

[rsuderman]

Removed unneeded parts othe Standard/ops.mlir checks.

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

1772

I wasn't planning to. It the input/output shapes should match and it would be strange to have a incorrectly shape determined side of the case.

[rsuderman]

Comment updated.

[jpienaar]

Thanks!

mlir/test/Dialect/Standard/ops.mlir

2

What does split input do here given fed back into mlir-opt? (and also no file splits in here)

[rsuderman]

jpienaar@ comment for removing pipe into mlir-opt

[rsuderman]

mlir/test/Dialect/Standard/ops.mlir

2

Extra from copy-paste of another test.

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

 }
 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 // 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());
 }
 
 OpFoldResult IndexCastOp::fold(ArrayRef<Attribute> cstOperands) {
+  // Cast op does not change the type.
+  if (getOperand().getType() == getType())
+    return getOperand();
+
   // Fold IndexCast(IndexCast(x)) -> x
   auto cast = getOperand().getDefiningOp<IndexCastOp>();
   if (cast && cast.getOperand().getType() == getType())
     return cast.getOperand();
 
   // Fold IndexCast(constant) -> constant
   // A little hack because we go through int.  Otherwise, the size
   // of the constant might need to change.

mlir/test/Dialect/Standard/invalid.mlir

+// 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
+}

mlir/test/Dialect/Standard/ops.mlir

+// RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s
+// Verify the printed output can be parsed.
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// Verify the generic form can be parsed.
+// RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | 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>
+}
+