[Mlir-commits] [mlir] [mlir][SPIRV] Add SPIRVToLLVM conversion for selection with yielding values (PR #210600)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Vito Secona (secona)

<details>
<summary>Changes</summary>

The current SPIRVToLLVM conversion for SelectionOp does not handle merge blocks with yielding values. This change implements that by adding arguments to the continue block in the SelectionPattern.

Closes #<!-- -->204714

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


2 Files Affected:

- (modified) mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp (+8) 
- (modified) mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir (+56) 


``````````diff
diff --git a/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp b/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
index c5f9305a5c226..2982061c957e0 100644
--- a/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
+++ b/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
@@ -1470,6 +1470,14 @@ class SelectionPattern : public SPIRVToLLVMConversion<spirv::SelectionOp> {
     auto position = rewriter.getInsertionPoint();
     auto *continueBlock = rewriter.splitBlock(currentBlock, position);
 
+    // Add arguments to the continue block for selections that yield values.
+    for (auto ty : op.getResultTypes()) {
+      Type dstTy = getTypeConverter()->convertType(ty);
+      if (!dstTy)
+        return rewriter.notifyMatchFailure(op, "failed to convert type");
+      continueBlock->addArgument(dstTy, loc);
+    }
+
     // Extract conditional branch information from the header block. By SPIR-V
     // dialect spec, it should contain `spirv.BranchConditional` or
     // `spirv.Switch` op. Note that `spirv.Switch op` is not supported at the
diff --git a/mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir b/mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir
index ec47dd31cc637..389a21681bafb 100644
--- a/mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir
+++ b/mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir
@@ -214,6 +214,62 @@ spirv.module Logical GLSL450 {
     %one = spirv.Constant 1 : i32
     spirv.ReturnValue %one : i32
   }
+
+  spirv.func @selection_with_yielding_value(%cond: i1) -> i32 "None" {
+    // CHECK: llvm.cond_br %{{.*}}, ^bb1, ^bb2
+    %0 = spirv.mlir.selection -> i32 {
+      spirv.BranchConditional %cond, ^true, ^false
+    // CHECK: ^bb1:
+    ^true:
+      // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i32
+      %cst1 = spirv.Constant 1 : i32
+      // CHECK: llvm.br ^bb3(%[[C1]] : i32)
+      spirv.Branch ^merge(%cst1 : i32)
+    // CHECK: ^bb2:
+    ^false:
+      // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32
+      %cst2 = spirv.Constant 2 : i32
+      // CHECK: llvm.br ^bb3(%[[C2]] : i32)
+      spirv.Branch ^merge(%cst2 : i32)
+    // CHECK: ^bb3(%[[ARG:.*]]: i32):
+    ^merge(%1: i32):
+      // CHECK: llvm.br ^bb4(%[[ARG]] : i32)
+      spirv.mlir.merge %1 : i32
+    }
+    // CHECK: ^bb4({{.*}}):
+    %one = spirv.Constant 1 : i32
+    spirv.ReturnValue %one : i32
+  }
+
+  spirv.func @selection_with_multiple_yielding_values(%cond: i1) -> i32 "None" {
+    // CHECK: llvm.cond_br %{{.*}}, ^bb1, ^bb2
+    %0:2 = spirv.mlir.selection -> i32, i32 {
+      spirv.BranchConditional %cond, ^true, ^false
+    // CHECK: ^bb1:
+    ^true:
+      // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i32
+      %cst1 = spirv.Constant 1 : i32
+      // CHECK: %[[C3:.*]] = llvm.mlir.constant(3 : i32) : i32
+      %cst3 = spirv.Constant 3 : i32
+      // CHECK: llvm.br ^bb3(%[[C1]], %[[C3]] : i32, i32)
+      spirv.Branch ^merge(%cst1, %cst3 : i32, i32)
+    // CHECK: ^bb2:
+    ^false:
+      // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32
+      %cst2 = spirv.Constant 2 : i32
+      // CHECK: %[[C4:.*]] = llvm.mlir.constant(4 : i32) : i32
+      %cst4 = spirv.Constant 4 : i32
+      // CHECK: llvm.br ^bb3(%[[C2]], %[[C4]] : i32, i32)
+      spirv.Branch ^merge(%cst2, %cst4 : i32, i32)
+    // CHECK: ^bb3(%[[ARG:.*]]: i32, %[[ARG1:.*]]: i32):
+    ^merge(%1: i32, %2: i32):
+      // CHECK: llvm.br ^bb4(%[[ARG]], %[[ARG1]] : i32, i32)
+      spirv.mlir.merge %1, %2 : i32, i32
+    }
+    // CHECK: ^bb4({{.*}}):
+    %one = spirv.Constant 1 : i32
+    spirv.ReturnValue %one : i32
+  }
 }
 
 // -----

``````````

</details>


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