[Mlir-commits] [mlir] [mlir][SPIR-V] Unroll multi-rank vectors in VectorToSPIRV and GPUToSPIRV (PR #214153)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

Ops like vector.transpose on multi-rank vectors were left unconverted since these passes never invoked the existing unrolling patterns

Only unroll when a multi-rank vector is actually present, since unrolling also runs whole module canonicalization as a side effect

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


5 Files Affected:

- (modified) mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h (+4) 
- (modified) mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp (+9) 
- (modified) mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp (+9) 
- (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+15) 
- (modified) mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir (+14) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
index 03ae54a8ae30a..8a0fc3bce71a1 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
@@ -207,6 +207,10 @@ SmallVector<int64_t> getNativeVectorShapeImpl(vector::TransposeOp op);
 // For general ops.
 std::optional<SmallVector<int64_t>> getNativeVectorShape(Operation *op);
 
+// True if `op` has a vector type of rank > 1, which needs unrolling before
+// SPIR-V conversion patterns (rank-1 only) can apply.
+bool hasMultiRankVectorType(Operation *op);
+
 // Unroll vectors in function signatures to native size.
 LogicalResult unrollVectorsInSignatures(Operation *op);
 
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
index 1b49e9d6305f1..2589e2d5e9848 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -137,6 +137,15 @@ void GPUToSPIRVPass::runOnOperation() {
       });
     }
 
+    // Skip unrolling when nothing needs it: the unroll functions also run
+    // whole-module canonicalization as a side effect.
+    if (spirv::hasMultiRankVectorType(gpuModule)) {
+      if (failed(spirv::unrollVectorsInSignatures(gpuModule)))
+        return signalPassFailure();
+      if (failed(spirv::unrollVectorsInFuncBodies(gpuModule)))
+        return signalPassFailure();
+    }
+
     std::unique_ptr<ConversionTarget> target =
         SPIRVConversionTarget::get(targetAttr);
 
diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp
index b3ef23085c186..7d767d438b054 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp
@@ -37,6 +37,15 @@ void ConvertVectorToSPIRVPass::runOnOperation() {
   MLIRContext *context = &getContext();
   Operation *op = getOperation();
 
+  // Skip unrolling when nothing needs it: the unroll functions also run
+  // whole-module canonicalization as a side effect.
+  if (spirv::hasMultiRankVectorType(op)) {
+    if (failed(spirv::unrollVectorsInSignatures(op)))
+      return signalPassFailure();
+    if (failed(spirv::unrollVectorsInFuncBodies(op)))
+      return signalPassFailure();
+  }
+
   auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index ef4d79c827bb2..005d76e2ea161 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1493,6 +1493,21 @@ mlir::spirv::getNativeVectorShape(Operation *op) {
       .Default(std::nullopt);
 }
 
+bool mlir::spirv::hasMultiRankVectorType(Operation *op) {
+  auto isMultiRankVector = [](Type type) {
+    auto vectorType = dyn_cast<VectorType>(type);
+    return vectorType && vectorType.getRank() > 1;
+  };
+  return op
+      ->walk([&](Operation *nestedOp) {
+        if (llvm::any_of(nestedOp->getOperandTypes(), isMultiRankVector) ||
+            llvm::any_of(nestedOp->getResultTypes(), isMultiRankVector))
+          return WalkResult::interrupt();
+        return WalkResult::advance();
+      })
+      .wasInterrupted();
+}
+
 LogicalResult mlir::spirv::unrollVectorsInSignatures(Operation *op) {
   MLIRContext *context = op->getContext();
   RewritePatternSet patterns(context);
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index f904dd9d35c37..9908d8e2d6b68 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -1039,6 +1039,20 @@ func.func @shape_cast_size1_vector(%arg0 : vector<f32>) -> vector<1xf32> {
 
 // -----
 
+// Transpose is unrolled to element-wise extracts/inserts, which fold away
+// entirely for this static case.
+// CHECK-LABEL: @transpose
+//  CHECK-SAME: (%[[A:.+]]: f32, %[[B:.+]]: f32, %[[C:.+]]: f32, %[[D:.+]]: f32)
+//       CHECK:   return %[[A]], %[[C]], %[[B]], %[[D]]
+func.func @transpose(%a: f32, %b: f32, %c: f32, %d: f32) -> (f32, f32, f32, f32) {
+  %v = vector.from_elements %a, %b, %c, %d : vector<2x2xf32>
+  %t = vector.transpose %v, [1, 0] : vector<2x2xf32> to vector<2x2xf32>
+  %r:4 = vector.to_elements %t : vector<2x2xf32>
+  return %r#0, %r#1, %r#2, %r#3 : f32, f32, f32, f32
+}
+
+// -----
+
 // CHECK-LABEL: @step()
 //       CHECK:   %[[CST0:.*]] = spirv.Constant 0 : i32
 //       CHECK:   %[[CST1:.*]] = spirv.Constant 1 : i32

``````````

</details>


https://github.com/llvm/llvm-project/pull/214153