[mlir][vector] add proper verification to vector.print operation

[aartbik]

Summary

Rationale:
Only proper vectors and scalars of floating-point or integral types
are actually lowered to calls into the light-weight output library.

[ThomasRaoux]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

5194–5206

I hate to be that guy but that could be done at the tablegen level by making the input type a VectorOf<[...?

[aartbik]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

5194–5206

Is there an easy way to say

VectorOf or ScalarOf

?

[ThomasRaoux]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

5194–5206

Looks like it doesn't exist and it would need to be added like it is on LLVM dialect:
https://github.com/llvm/llvm-project-release-prs/blob/0604154e006e88e9e7f82d8ee5fd076bda206613/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td#L213

[nicolasvasilache]

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

2443

This should be AnyVectorOf Or Scalar type of thing ?

[nicolasvasilache]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

5194–5206

ah sorry for parroting .. +1 to adding to Tablegen .. here's a potential way

Try Or<[AnyVector.predicate, AnySignlessIntegerOrIndex.predicate, AnyFloat.predicate]> ?

We should have a proper:

def AnyScalar : Or<AnySignlessIntegerOrIndex.predicate, AnyFloat.predicate, AnyComplex.predicate>;

5194–5206

Note that AnyVectorOfAnyRank is the proper thing to use until all ops support 0-D vectors..

[aartbik]

C++ verification -> declarative verification

[aartbik]

Replaced C++ verification with declarative verification. Note that the C++ version was perhaps slightly more precise but we can refine the declarative part as we get better traits.
This at least moves my code away from the now obsoleted C++ way of checking stuff.

PTAL

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

let assemblyFormat = [{
     $vector `,` $transp attr-dict `:` type($vector) `to` type($result)
   }];
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def Vector_PrintOp :
   Vector_Op<"print", []>, Arguments<(ins AnyType:$source)> {
   let summary = "print operation (for testing and debugging)";
   let description = [{
     Prints the source vector (or scalar) to stdout in human readable
     format (for testing and debugging). No return value.
 
     Example:
 
     ```mlir
let description = [{
     ```
   }];
   let extraClassDeclaration = [{
     Type getPrintType() {
       return getSource().getType();
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source)";
+  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Ops used for supporting progressive lowering and conversion type changes.
 // The Ops are typically not used directly by higher level dialects, but are
 // used by intra-dialect rewriting rules to bring vector operations closer
 // to the hardware ISA.
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

results
           context);
 }
 
 void vector::TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(getTransp(), results);
 }
 
 //===----------------------------------------------------------------------===//
+// PrintOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult vector::PrintOp::verify() {
+  Type tp = getPrintType();
+
+  // Find element type of a vector type. Assume direct scalar type otherwise.
+  VectorType vtp = tp.dyn_cast<VectorType>();
+  Type eltType = vtp ? vtp.getElementType() : tp;
+
+  // Now verify if element type is supported
+  if (eltType.isF32())
+    return success();
+  if (eltType.isF64())
+    return success();
+  if (eltType.isIndex())
+    return success();
+  if (auto intTy = eltType.dyn_cast<IntegerType>())
+    return success();
+
+  // Anything else not supported.
+  return emitError("operand must be a vector or scalar with floating-point or "
+                   "integral type");
+}
+
+//===----------------------------------------------------------------------===//
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantMaskOp::verify() {
   auto resultType = getResult().getType().cast<VectorType>();
   // Check the corner case of 0-D vectors first.
   if (resultType.getRank() == 0) {
     if (getMaskDimSizes().size() != 1)

mlir/test/Dialect/Vector/invalid.mlir

 
 func.func @print_no_result(%arg0 : f32) -> i32 {
   // expected-error@+1 {{cannot name an operation with no results}}
   %0 = vector.print %arg0 : f32
 }
 
 // -----
 
+func.func private @print_needs_buffer(%arg0: tensor<8xf32>) {
+  // expected-error@+1 {{operand must be a vector or scalar with floating-point or integral type}}
+  vector.print %arg0 : tensor<8xf32>
+  return
+}
+
+// -----
+
 func.func @reshape_bad_input_shape(%arg0 : vector<3x2x4xf32>) {
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   %c6 = arith.constant 6 : index
   %c9 = arith.constant 9 : index
   // expected-error@+1 {{invalid input shape for vector type}}
   %1 = vector.reshape %arg0, [%c3, %c6, %c3], [%c2, %c9], [4]
     : vector<3x2x4xf32> to vector<2x3x4xf32>

mlir/test/Dialect/Vector/ops.mlir

 func.func @constant_vector_mask() {
   // CHECK: vector.constant_mask [3, 2] : vector<4x3xi1>
   %0 = vector.constant_mask [3, 2] : vector<4x3xi1>
   // CHECK: vector.constant_mask [0] : vector<[4]xi1>
   %1 = vector.constant_mask [0] : vector<[4]xi1>
   return
 }
 
-// CHECK-LABEL: @vector_print
-func.func @vector_print(%arg0: vector<8x4xf32>) {
+// CHECK-LABEL: @vector_print_on_vector
+func.func @vector_print_on_vector(%arg0: vector<8x4xf32>) {
   // CHECK: vector.print %{{.*}} : vector<8x4xf32>
   vector.print %arg0 : vector<8x4xf32>
   return
 }
 
+// CHECK-LABEL: @vector_print_on_scalar
+func.func @vector_print_on_scalar(%arg0: i64) {
+  // CHECK: vector.print %{{.*}} : i64
+  vector.print %arg0 : i64
+  return
+}
+
 // CHECK-LABEL: @reshape
 func.func @reshape(%arg0 : vector<3x2x4xf32>) -> (vector<2x3x4xf32>) {
   // CHECK:      %[[C2:.*]] = arith.constant 2 : index
   %c2 = arith.constant 2 : index
   // CHECK:      %[[C3:.*]] = arith.constant 3 : index
   %c3 = arith.constant 3 : index
   // CHECK:      %[[C6:.*]] = arith.constant 6 : index
   %c6 = arith.constant 6 : index