[SPIR-V] Enable structurizer for kernel environment

[MrSidims]

No description provided.

[llvmbot]

@llvm/pr-subscribers-backend-spir-v

Author: Dmitry Sidorov (MrSidims)

Changes

---

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

2 Files Affected:

• (modified) llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp (+2)
• (added) llvm/test/CodeGen/SPIRV/structurizer/kernel-loop.ll (+80)

diff --git a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp b/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
index 7dd0b95cd9763..eba99f6dacb59 100644
--- a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
@@ -181,7 +181,9 @@ void SPIRVPassConfig::addISelPrepare() {
     // If an address space cast is not removed while targeting Vulkan, lowering
     // will fail during MIR lowering.
     addPass(createInferAddressSpacesPass());
+  }
 
+  if (TM.getSubtargetImpl()->isShader() || TM.getSubtargetImpl()->isKernel()) {
     // 1.  Simplify loop for subsequent transformations. After this steps, loops
     // have the following properties:
     //  - loops have a single entry edge (pre-header to loop header).
diff --git a/llvm/test/CodeGen/SPIRV/structurizer/kernel-loop.ll b/llvm/test/CodeGen/SPIRV/structurizer/kernel-loop.ll
new file mode 100644
index 0000000000000..56dc916ac7bd3
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/structurizer/kernel-loop.ll
@@ -0,0 +1,80 @@
+; RUN: llc -mtriple=spirv64-unknown-opencl -O0 %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-opencl %s -o - -filetype=obj | spirv-val %}
+
+; This test verifies that the structurizer pass runs for OpenCL kernels,
+; generating structured control flow with OpLoopMerge instructions and
+; translating loop metadata to appropriate LoopControl operands.
+
+; CHECK: OpEntryPoint Kernel %[[#kernel_unroll:]] "test_kernel_unroll"
+; CHECK: OpEntryPoint Kernel %[[#kernel_dontunroll:]] "test_kernel_dontunroll"
+
+; Verify unroll metadata is translated to Unroll LoopControl
+; CHECK: %[[#kernel_unroll]] = OpFunction
+; CHECK: OpLabel
+; CHECK: OpLoopMerge %[[#]] %[[#]] Unroll
+; CHECK: OpFunctionEnd
+
+; Verify dont_unroll metadata is translated to DontUnroll LoopControl
+; CHECK: %[[#kernel_dontunroll]] = OpFunction
+; CHECK: OpLabel
+; CHECK: OpLoopMerge %[[#]] %[[#]] DontUnroll
+; CHECK: OpFunctionEnd
+
+define spir_kernel void @test_kernel_unroll(ptr addrspace(1) %out) {
+entry:
+  %i = alloca i32, align 4
+  store i32 0, ptr %i, align 4
+  br label %for.cond
+
+for.cond:
+  %0 = load i32, ptr %i, align 4
+  %cmp = icmp slt i32 %0, 10
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:
+  %1 = load i32, ptr %i, align 4
+  %arrayidx = getelementptr inbounds i32, ptr addrspace(1) %out, i64 0
+  store i32 %1, ptr addrspace(1) %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %2 = load i32, ptr %i, align 4
+  %inc = add nsw i32 %2, 1
+  store i32 %inc, ptr %i, align 4
+  br label %for.cond, !llvm.loop !0
+
+for.end:
+  ret void
+}
+
+define spir_kernel void @test_kernel_dontunroll(ptr addrspace(1) %out) {
+entry:
+  %i = alloca i32, align 4
+  store i32 0, ptr %i, align 4
+  br label %for.cond
+
+for.cond:
+  %0 = load i32, ptr %i, align 4
+  %cmp = icmp slt i32 %0, 10
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:
+  %1 = load i32, ptr %i, align 4
+  %arrayidx = getelementptr inbounds i32, ptr addrspace(1) %out, i64 0
+  store i32 %1, ptr addrspace(1) %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %2 = load i32, ptr %i, align 4
+  %inc = add nsw i32 %2, 1
+  store i32 %inc, ptr %i, align 4
+  br label %for.cond, !llvm.loop !2
+
+for.end:
+  ret void
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.full"}
+!2 = distinct !{!2, !3}
+!3 = !{!"llvm.loop.unroll.disable"}

[Keenuts]

Curious, why is this needed now? Is OpenCL going to requires SCFG?

[Keenuts]

Is there another valid target? Should the condition be removed instead?

[MrSidims]

Indeed, will remove

[MrSidims]

Or actually, given there is the following PR, let me summon @AlexVlx to answer if it can break something on their end (TargetTriple.getVendor() == Triple::AMD falls under IsKernel, but still).

[AlexVlx]

Hmm, this does indeed cause issues, I think. We don't / cannot run passes over SPIR-V, asides from those that are absolutely necessary for successful lowering / obtaining valid SPIR-V (see #154765 for a bit more on this). We use -disable-llvm-optzns, but it'd not work here. Would it be at all possible to guard against enabling these for AMDGCNSPIRV? I.e. have the current check be replaced for a check on the vendor?

[MrSidims]

Let me move it to draft actually. May be just implementing https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/INTEL/SPV_INTEL_unstructured_loop_controls.asciidoc is the right direction. I haven't though about it for some reasons before Nathan's comment.

[AlexVlx]

Cheers, apologies for the inconvenience.

[MrSidims]

Is OpenCL going to requires SCFG?

Not really, yet (AFAIK) it's the only way to preserve unroll and other loop metadata (well, this is also 'unstructured CF extension' and future of it is a bit uncertain).