disable index type for bitcast

Author: jacquesguan

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

@@ -1392,11 +1392,10 @@ def Arith_IndexCastUIOp
 // BitcastOp
 //===----------------------------------------------------------------------===//
 
-// Bitcast can convert between memrefs of signless integers, indices, and
-// floats too.
+// Bitcast can convert between memrefs of signless integers and floats.
 def BitcastTypeConstraint : TypeConstraint<Or<[
         SignlessIntegerOrFloatLike.predicate,
-        MemRefOf<[AnySignlessInteger, Index, AnyFloat]>.predicate]>,
+        MemRefOf<[AnySignlessInteger, AnyFloat]>.predicate]>,
     "signless-integer-or-float-like or memref of signless-integer or float">;
 
 def Arith_BitcastOp : Arith_CastOp<"bitcast", BitcastTypeConstraint,

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -1740,25 +1740,12 @@ bool arith::BitcastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (!areValidCastInputsAndOutputs(inputs, outputs))
     return false;
 
-  auto srcType =
-      getTypeIfLikeOrMemRef<IntegerType, IndexType, FloatType>(inputs.front());
-  auto dstType =
-      getTypeIfLikeOrMemRef<IntegerType, IndexType, FloatType>(outputs.front());
+  auto srcType = getTypeIfLikeOrMemRef<IntegerType, FloatType>(inputs.front());
+  auto dstType = getTypeIfLikeOrMemRef<IntegerType, FloatType>(outputs.front());
   if (!srcType || !dstType)
     return false;
 
-  unsigned srcWidth, dstWidth;
-  if (auto indexTy = dyn_cast<IndexType>(srcType))
-    srcWidth = IndexType::kInternalStorageBitWidth;
-  else
-    srcWidth = srcType.getIntOrFloatBitWidth();
-
-  if (auto indexTy = dyn_cast<IndexType>(dstType))
-    dstWidth = IndexType::kInternalStorageBitWidth;
-  else
-    dstWidth = dstType.getIntOrFloatBitWidth();
-
-  return srcWidth == dstWidth;
+  return srcType.getIntOrFloatBitWidth() == dstType.getIntOrFloatBitWidth();
 }
 
 OpFoldResult arith::BitcastOp::fold(FoldAdaptor adaptor) {

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

@@ -219,6 +219,10 @@ bool BitcastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (!aT || !bT)
     return false;
 
+  if (isa<IndexType>(aT.getElementType()) ||
+      isa<IndexType>(bT.getElementType()))
+    return false;
+
   if (aT.getElementTypeBitWidth() != bT.getElementTypeBitWidth())
     return false;
 

mlir/test/Dialect/Arith/invalid.mlir

@@ -853,3 +853,19 @@ func.func @select_tensor_encoding(
   %0 = arith.select %arg0, %arg1, %arg2 : tensor<8xi1, "bar">, tensor<8xi32, "foo">
   return %0 : tensor<8xi32, "foo">
 }
+
+// -----
+
+func.func @bitcast_index_0(%arg0 : i64) -> index {
+  // expected-error @+1 {{'arith.bitcast' op operand type 'i64' and result type 'index' are cast incompatible}}
+  %0 = arith.bitcast %arg0 : i64 to index
+  return %0 : index
+}
+
+// -----
+
+func.func @bitcast_index_1(%arg0 : index) -> i64 {
+  // expected-error @+1 {{'arith.bitcast' op operand type 'index' and result type 'i64' are cast incompatible}}
+  %0 = arith.bitcast %arg0 : index to i64
+  return %0 : i64
+}

mlir/test/Dialect/Arith/ops.mlir

@@ -954,12 +954,6 @@ func.func @test_bitcast_scalable_vector1(%arg0 : vector<[8]xf32>) -> vector<[8]x
   return %0 : vector<[8]xi32>
 }
 
-// CHECK-LABEL: test_bitcast_index
-func.func @test_bitcast_index(%arg0 : i64) -> index {
-  %0 = arith.bitcast %arg0 : i64 to index
-  return %0 : index
-}
-
 // CHECK-LABEL: test_cmpi
 func.func @test_cmpi(%arg0 : i64, %arg1 : i64) -> i1 {
   %0 = arith.cmpi ne, %arg0, %arg1 : i64

mlir/test/Dialect/Tensor/invalid.mlir

@@ -807,3 +807,19 @@ func.func @unpack_static_inner_tile_size_and_dynamic_output_shape(
   %0 = tensor.unpack %input inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %output : tensor<?x?x?x4xf32> -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
+
+// -----
+
+func.func @bitcast_index_0(%arg0 : tensor<?xi64>) -> tensor<?xindex> {
+  // expected-error @+1 {{'tensor.bitcast' op operand type 'tensor<?xi64>' and result type 'tensor<?xindex>' are cast incompatible}}
+  %0 = tensor.bitcast %arg0 : tensor<?xi64> to tensor<?xindex>
+  return %0 : tensor<?xindex>
+}
+
+// -----
+
+func.func @bitcast_index_1(%arg0 : tensor<?xindex>) -> tensor<?xi64> {
+  // expected-error @+1 {{'tensor.bitcast' op operand type 'tensor<?xindex>' and result type 'tensor<?xi64>' are cast incompatible}}
+  %0 = tensor.bitcast %arg0 : tensor<?xindex> to tensor<?xi64>
+  return %0 : tensor<?xi64>
+}