[MLIR] Fix 0-dimensional case of conversion of vector ops to GPU

[bjacob]

This is a follow-up to #127844. That PR got vectors of arbitrary rank working, but I hadn't thought about the rank-0 case.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Benoit Jacob (bjacob)

Changes

This is a follow-up to #127844. That PR got vectors of arbitrary rank working, but I hadn't thought about the rank-0 case.

---

Full diff: https://github.com/llvm/llvm-project/pull/128075.diff

2 Files Affected:

• (modified) mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp (+2-1)
• (modified) mlir/test/Conversion/MathToROCDL/math-to-rocdl.mlir (+14)

diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
index c3b3a78abe7f7..8b6c62ca2e36d 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
@@ -624,7 +624,8 @@ LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
                                       const LLVMTypeConverter &converter) {
   TypeRange operandTypes(operands);
   if (llvm::any_of(operandTypes, llvm::IsaPred<VectorType>)) {
-    VectorType vectorType = cast<VectorType>(op->getResultTypes()[0]);
+    VectorType vectorType =
+        cast<VectorType>(converter.convertType(op->getResultTypes()[0]));
     rewriter.replaceOp(op, scalarizeVectorOpHelper(op, operands, vectorType,
                                                    rewriter, converter));
     return success();
diff --git a/mlir/test/Conversion/MathToROCDL/math-to-rocdl.mlir b/mlir/test/Conversion/MathToROCDL/math-to-rocdl.mlir
index 9448304f11dbd..313d7b086731e 100644
--- a/mlir/test/Conversion/MathToROCDL/math-to-rocdl.mlir
+++ b/mlir/test/Conversion/MathToROCDL/math-to-rocdl.mlir
@@ -516,6 +516,20 @@ module {
 
 // -----
 
+module @test_module {
+  // CHECK: llvm.func @__ocml_sin_f16(f16) -> f16
+  // CHECK-LABEL: func @math_sin_vector_0d
+  func.func @math_sin_vector_0d(%arg : vector<f16>) -> vector<f16> {
+    // CHECK: llvm.extractelement {{.*}} : vector<1xf16>
+    // CHECK: llvm.call @__ocml_sin_f16(%{{.*}}) : (f16) -> f16
+    // CHECK: llvm.insertelement {{.*}} : vector<1xf16>
+    %result = math.sin %arg : vector<f16>
+    func.return %result : vector<f16>
+  }
+}
+
+// -----
+
 module @test_module {
   // CHECK: llvm.func @__ocml_sin_f16(f16) -> f16
   // CHECK-LABEL: func @math_sin_vector_1d

[krzysz00]

lgtm