[mlir] Fix arith ops verifier for tensor with encoding

[ThomasRaoux]

Summary

The verifier for some arith ops were not considering that ranked
tensor types can have encodings.

[jpienaar]

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

132

This could just be a call to cloneWith with no shape and just element type too.

[ThomasRaoux]

Address comment

[ThomasRaoux]

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

132

ah yes much simpler indeed. Changed it.

[kuhar]

LGTM

mlir/test/Dialect/Arith/invalid.mlir

213

nit: Do we need to use the generic syntax here?

[ThomasRaoux]

mlir/test/Dialect/Arith/invalid.mlir

213

I believe that's the only way to explicitly specify both the input and output types? Do you see a different way to do this?

[kuhar]

mlir/test/Dialect/Arith/invalid.mlir

213

Ah, I missed the error line just above, nvm

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

 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Common helpers
 //===----------------------------------------------------------------------===//
 
 /// Return the type of the same shape (scalar, vector or tensor) containing i1.
 static Type getI1SameShape(Type type) {
   auto i1Type = IntegerType::get(type.getContext(), 1);
-  if (auto tensorType = llvm::dyn_cast<RankedTensorType>(type))
-    return RankedTensorType::get(tensorType.getShape(), i1Type);
+  if (auto shapedType = llvm::dyn_cast<ShapedType>(type))
+    return shapedType.cloneWith(std::nullopt, i1Type);
   if (llvm::isa<UnrankedTensorType>(type))
     return UnrankedTensorType::get(i1Type);
-  if (auto vectorType = llvm::dyn_cast<VectorType>(type))
-    return VectorType::get(vectorType.getShape(), i1Type,
-                           vectorType.getScalableDims());
   return i1Type;
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 void arith::ConstantOp::getAsmResultNames(

 /// Get allowed underlying types for vectors, tensors, and memrefs.
 template <typename... ElementTypes>
 static Type getTypeIfLikeOrMemRef(Type type) {
   return getUnderlyingType(type,
                            type_list<VectorType, TensorType, MemRefType>(),
                            type_list<ElementTypes...>());
 }
 
+/// Return false if both types are ranked tensor with mismatching encoding.
+static bool hasSameEncoding(Type typeA, Type typeB) {
+  auto rankedTensorA = dyn_cast<RankedTensorType>(typeA);
+  auto rankedTensorB = dyn_cast<RankedTensorType>(typeB);
+  if (!rankedTensorA || !rankedTensorB)
+    return true;
+  return rankedTensorA.getEncoding() == rankedTensorB.getEncoding();
+}
+
 static bool areValidCastInputsAndOutputs(TypeRange inputs, TypeRange outputs) {
-  return inputs.size() == 1 && outputs.size() == 1 &&
-         succeeded(verifyCompatibleShapes(inputs.front(), outputs.front()));
+  if (inputs.size() != 1 || outputs.size() != 1)
+    return false;
+  if (!hasSameEncoding(inputs.front(), outputs.front()))
+    return false;
+  return succeeded(verifyCompatibleShapes(inputs.front(), outputs.front()));
 }
 
 //===----------------------------------------------------------------------===//
 // Verifiers for integer and floating point extension/truncation ops
 //===----------------------------------------------------------------------===//
 
 // Extend ops can only extend to a wider type.
 template <typename ValType, typename Op>

mlir/test/Dialect/Arith/invalid.mlir

 ^bb0:
   %c = arith.constant dense<0> : vector<42 x i32>
   // expected-error@+1 {{op failed to verify that result type has i1 element type and same shape as operands}}
   %r = "arith.cmpi"(%c, %c) {predicate = 0} : (vector<42 x i32>, vector<42 x i32>) -> vector<41 x i1>
 }
 
 // -----
 
+func.func @func_with_ops() {
+^bb0:
+  %c = arith.constant dense<0> : tensor<42 x i32, "foo">
+  // expected-error@+1 {{op failed to verify that result type has i1 element type and same shape as operands}}
+  %r = "arith.cmpi"(%c, %c) {predicate = 0} : (tensor<42 x i32, "foo">, tensor<42 x i32, "foo">) -> tensor<42 x i1, "bar">
+}
+
+// -----
+
 func.func @invalid_cmp_shape(%idx : () -> ()) {
   // expected-error@+1 {{'lhs' must be signless-integer-like, but got '() -> ()'}}
   %cmp = arith.cmpi eq, %idx, %idx : () -> ()
 
 // -----
 
 func.func @invalid_cmp_attr(%idx : i32) {
   // expected-error@+1 {{expected string or keyword containing one of the following enum values}}

 func.func @fpext_vec_f32_to_i32(%arg0 : vector<2xf32>) {
   // expected-error@+1 {{op result #0 must be floating-point-like, but got 'vector<2xi32>'}}
   %0 = arith.extf %arg0 : vector<2xf32> to vector<2xi32>
   return
 }
 
 // -----
 
+func.func @fpext_vec_f32_to_i32(%arg0 : tensor<2xf32, "foo">) {
+  // expected-error@+1 {{op operand type 'tensor<2xf32, "foo">' and result type 'tensor<2xf64, "bar">' are cast incompatible}}
+  %0 = arith.extf %arg0 : tensor<2xf32, "foo"> to tensor<2xf64, "bar">
+  return
+}
+
+// -----
+
 func.func @fptrunc_f16_to_f32(%arg0 : f16) {
   // expected-error@+1 {{are cast incompatible}}
   %0 = arith.truncf %arg0 : f16 to f32
   return
 }
 
 // -----
 

 
 // -----
 
 func.func @disallow_zero_rank_tensor_with_unranked_tensor(%arg0 : tensor<i1>, %arg1 : tensor<2x?xi64>, %arg2 : tensor<2x?xi64>) -> tensor<2x?xi64> {
   // expected-error @+1 {{'arith.select' op failed to verify that condition is scalar or has matching shape}}
   %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<2x?xi64>
   return %0 : tensor<2x?xi64>
 }
+
+// -----
+
+func.func @select_tensor_encoding(
+  %arg0 : tensor<8xi1, "bar">, %arg1 : tensor<8xi32, "foo">, %arg2 : tensor<8xi32, "foo">) -> tensor<8xi32, "foo"> {
+  // expected-error @+1 {{'arith.select' op expected condition type to have the same shape as the result type}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<8xi1, "bar">, tensor<8xi32, "foo">
+  return %0 : tensor<8xi32, "foo">
+}

mlir/test/Dialect/Arith/ops.mlir

 }
 
 // CHECK-LABEL: test_extf_tensor
 func.func @test_extf_tensor(%arg0 : tensor<8x8xf32>) -> tensor<8x8xf64> {
   %0 = arith.extf %arg0 : tensor<8x8xf32> to tensor<8x8xf64>
   return %0 : tensor<8x8xf64>
 }
 
+// CHECK-LABEL: test_extf_tensor_encoding
+func.func @test_extf_tensor_encoding(%arg0 : tensor<8x8xf32, "foo">) -> tensor<8x8xf64, "foo"> {
+  %0 = arith.extf %arg0 : tensor<8x8xf32, "foo"> to tensor<8x8xf64, "foo">
+  return %0 : tensor<8x8xf64, "foo">
+}
+
 // CHECK-LABEL: test_extf_vector
 func.func @test_extf_vector(%arg0 : vector<8xf32>) -> vector<8xf64> {
   %0 = arith.extf %arg0 : vector<8xf32> to vector<8xf64>
   return %0 : vector<8xf64>
 }
 
 // CHECK-LABEL: test_extf_scalable_vector
 func.func @test_extf_scalable_vector(%arg0 : vector<[8]xf32>) -> vector<[8]xf64> {

 }
 
 // CHECK-LABEL: test_cmpi_tensor
 func.func @test_cmpi_tensor(%arg0 : tensor<8x8xi64>, %arg1 : tensor<8x8xi64>) -> tensor<8x8xi1> {
   %0 = arith.cmpi slt, %arg0, %arg1 : tensor<8x8xi64>
   return %0 : tensor<8x8xi1>
 }
 
+// CHECK-LABEL: test_cmpi_tensor_encoding
+func.func @test_cmpi_tensor_encoding(%arg0 : tensor<8x8xi64, "foo">, %arg1 : tensor<8x8xi64, "foo">) -> tensor<8x8xi1, "foo"> {
+  %0 = arith.cmpi slt, %arg0, %arg1 : tensor<8x8xi64, "foo">
+  return %0 : tensor<8x8xi1, "foo">
+}
+
 // CHECK-LABEL: test_cmpi_vector
 func.func @test_cmpi_vector(%arg0 : vector<8xi64>, %arg1 : vector<8xi64>) -> vector<8xi1> {
   %0 = arith.cmpi ult, %arg0, %arg1 : vector<8xi64>
   return %0 : vector<8xi1>
 }
 
 // CHECK-LABEL: test_cmpi_scalable_vector
 func.func @test_cmpi_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 : vector<[8]xi64>) -> vector<[8]xi1> {
// CHECK: {{.*}} = arith.addf %arg0, %arg1 : f32
   %6 = arith.addf %arg0, %arg1 fastmath<none> : f32
 // CHECK: {{.*}} = arith.addf %arg0, %arg1 fastmath<nnan,ninf> : f32
   %7 = arith.addf %arg0, %arg1 fastmath<nnan,ninf> : f32
 // CHECK: {{.*}} = arith.mulf %arg0, %arg1 fastmath<fast> : f32
   %8 = arith.mulf %arg0, %arg1 fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn> : f32
 
   return
 }
+
+// CHECK-LABEL: @select_tensor
+func.func @select_tensor(%arg0 : tensor<8xi1>, %arg1 : tensor<8xi32>, %arg2 : tensor<8xi32>) -> tensor<8xi32> {
+  // CHECK: = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xi1>, tensor<8xi32>
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<8xi1>, tensor<8xi32>
+  return %0 : tensor<8xi32>
+}
+
+// CHECK-LABEL: @select_tensor_encoding
+func.func @select_tensor_encoding(
+  %arg0 : tensor<8xi1, "foo">, %arg1 : tensor<8xi32, "foo">, %arg2 : tensor<8xi32, "foo">) -> tensor<8xi32, "foo"> {
+  // CHECK: = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xi1, "foo">, tensor<8xi32, "foo">
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<8xi1, "foo">, tensor<8xi32, "foo">
+  return %0 : tensor<8xi32, "foo">
+}