[mlir][Vector] Support 0-D vectors in `VectorPrintOpConversion`

[nicolasvasilache]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

60

rank <= 1 plz

1015

I think the API needs to change here.
emitRanks(...VectorType) -> emitRanks(...Type)

and then inside the function:

VectorType vectorType = type.dyn_cast<VectorType>() ... 
if (!vectorType) {
  assert(rank == 0 && "scalar case expects rank == 0");
 ...
}

1056

The ternary conditions already look too tricky here I find
I'd early-exit on:

if (rank <= 1) {
      auto reducedType = vector.getElementType();
      auto llvmType = typeConverter->convertType(reducedType);
      Value nestedVal = extractOne(rewriter, *getTypeConverter(), loc, value,
                                   llvmType, rank, d);
     return emitRanks(rewriter, op, nestedVal, reducedType, printer, /*lowerRank=*/0,
}

and then you have an opportunity to make the ternaries below simpler.

[michalt]

Address review comments

[nicolasvasilache]

Thanks!

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

static Value insertOne(ConversionPatternRewriter &rewriter,
   return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2,
                                               rewriter.getI64ArrayAttr(pos));
 }
 
 // Helper that picks the proper sequence for extracting.
 static Value extractOne(ConversionPatternRewriter &rewriter,
                         LLVMTypeConverter &typeConverter, Location loc,
                         Value val, Type llvmType, int64_t rank, int64_t pos) {
-  assert(rank > 0 && "0-D vector corner case should have been handled already");
-  if (rank == 1) {
+  if (rank <= 1) {
     auto idxType = rewriter.getIndexType();
     auto constant = rewriter.create<LLVM::ConstantOp>(
         loc, typeConverter.convertType(idxType),
         rewriter.getIntegerAttr(idxType, pos));
     return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val,
                                                    constant);
   }
   return rewriter.create<LLVM::ExtractValueOp>(loc, llvmType, val,
if (eltType.isF32()) {
         }
       }
     } else {
       return failure();
     }
 
     // Unroll vector into elementary print calls.
     int64_t rank = vectorType ? vectorType.getRank() : 0;
-    emitRanks(rewriter, printOp, adaptor.source(), vectorType, printer, rank,
+    Type type = vectorType ? vectorType : eltType;
+    emitRanks(rewriter, printOp, adaptor.source(), type, printer, rank,
               conversion);
     emitCall(rewriter, printOp->getLoc(),
              LLVM::lookupOrCreatePrintNewlineFn(
                  printOp->getParentOfType<ModuleOp>()));
     rewriter.eraseOp(printOp);
     return success();
   }
 
 private:
   enum class PrintConversion {
     // clang-format off
     None,
     ZeroExt64,
     SignExt64
     // clang-format on
   };
 
   void emitRanks(ConversionPatternRewriter &rewriter, Operation *op,
-                 Value value, VectorType vectorType, Operation *printer,
-                 int64_t rank, PrintConversion conversion) const {
+                 Value value, Type type, Operation *printer, int64_t rank,
+                 PrintConversion conversion) const {
+    VectorType vectorType = type.dyn_cast<VectorType>();
     Location loc = op->getLoc();
-    if (rank == 0) {
+    if (!vectorType) {
+      assert(rank == 0 && "The scalar case expects rank == 0");
       switch (conversion) {
       case PrintConversion::ZeroExt64:
         value = rewriter.create<arith::ExtUIOp>(
             loc, value, IntegerType::get(rewriter.getContext(), 64));
         break;
       case PrintConversion::SignExt64:
         value = rewriter.create<arith::ExtSIOp>(
             loc, value, IntegerType::get(rewriter.getContext(), 64));
         break;
       case PrintConversion::None:
         break;
       }
       emitCall(rewriter, loc, printer, value);
       return;
     }
 
     emitCall(rewriter, loc,
              LLVM::lookupOrCreatePrintOpenFn(op->getParentOfType<ModuleOp>()));
     Operation *printComma =
         LLVM::lookupOrCreatePrintCommaFn(op->getParentOfType<ModuleOp>());
+
+    if (rank <= 1) {
+      auto reducedType = vectorType.getElementType();
+      auto llvmType = typeConverter->convertType(reducedType);
+      int64_t dim = rank == 0 ? 1 : vectorType.getDimSize(0);
+      for (int64_t d = 0; d < dim; ++d) {
+        Value nestedVal = extractOne(rewriter, *getTypeConverter(), loc, value,
+                                     llvmType, /*rank=*/0, /*pos=*/d);
+        emitRanks(rewriter, op, nestedVal, reducedType, printer, /*rank=*/0,
+                  conversion);
+        if (d != dim - 1)
+          emitCall(rewriter, loc, printComma);
+      }
+      emitCall(
+          rewriter, loc,
+          LLVM::lookupOrCreatePrintCloseFn(op->getParentOfType<ModuleOp>()));
+      return;
+    }
+
     int64_t dim = vectorType.getDimSize(0);
     for (int64_t d = 0; d < dim; ++d) {
-      auto reducedType =
-          rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr;
-      auto llvmType = typeConverter->convertType(
-          rank > 1 ? reducedType : vectorType.getElementType());
+      auto reducedType = reducedVectorTypeFront(vectorType);
+      auto llvmType = typeConverter->convertType(reducedType);
       Value nestedVal = extractOne(rewriter, *getTypeConverter(), loc, value,
                                    llvmType, rank, d);
       emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1,
                 conversion);
       if (d != dim - 1)
         emitCall(rewriter, loc, printComma);
     }
     emitCall(rewriter, loc,

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 }
 // CHECK-LABEL: @vector_print_scalar_f64(
 // CHECK-SAME: %[[A:.*]]: f64)
 //       CHECK:    llvm.call @printF64(%[[A]]) : (f64) -> ()
 //       CHECK:    llvm.call @printNewline() : () -> ()
 
 // -----
 
+func @vector_print_vector_0d(%arg0: vector<f32>) {
+  vector.print %arg0 : vector<f32>
+  return
+}
+// CHECK-LABEL: @vector_print_vector_0d(
+// CHECK-SAME: %[[A:.*]]: vector<f32>)
+//       CHECK: %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<f32> to vector<1xf32>
+//       CHECK: llvm.call @printOpen() : () -> ()
+//       CHECK: %[[T1:.*]] = llvm.mlir.constant(0 : index) : i64
+//       CHECK: %[[T2:.*]] = llvm.extractelement %[[T0]][%[[T1]] : i64] : vector<1xf32>
+//       CHECK: llvm.call @printF32(%[[T2]]) : (f32) -> ()
+//       CHECK: llvm.call @printClose() : () -> ()
+//       CHECK: llvm.call @printNewline() : () -> ()
+//       CHECK: return
+
+// -----
+
 func @vector_print_vector(%arg0: vector<2x2xf32>) {
   vector.print %arg0 : vector<2x2xf32>
   return
 }
 // CHECK-LABEL: @vector_print_vector(
 // CHECK-SAME: %[[A:.*]]: vector<2x2xf32>)
 //       CHECK:    %[[VAL_1:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<2x2xf32> to !llvm.array<2 x vector<2xf32>>
 //       CHECK:    llvm.call @printOpen() : () -> ()

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

func @extract_element_0d(%a: vector<f32>) {
   return
 }
 
 func @insert_element_0d(%a: f32, %b: vector<f32>) -> (vector<f32>) {
   %1 = vector.insertelement %a, %b[] : vector<f32>
   return %1: vector<f32>
 }
 
+func @print_vector_0d(%a: vector<f32>) {
+  // CHECK: ( 42 )
+  vector.print %a: vector<f32>
+  return
+}
+
 func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
+
+  %3 = arith.constant dense<42.0> : vector<f32>
+  call  @print_vector_0d(%3) : (vector<f32>) -> ()
+
   return
 }