[llvm] [SPIR-V] Don't change switch condition type in CodeGen opts (PR #94959)

[Michal Paszkowski via llvm-commits]

https://github.com/michalpaszkowski created https://github.com/llvm/llvm-project/pull/94959

This change makes sure the preferred switch condition int type size remains the same throughout CodeGen optimizations.

The change fixes running several OpenCL applications with -O2 or higher opt levels, and fixes Basic/stream/stream_max_stmt_exceed.cpp DPC++ E2E test with -O2.

>From be4b6dca7bbbef189e7fcc52fe021b3729c79d81 Mon Sep 17 00:00:00 2001
From: Michal Paszkowski <michal at paszkowski.org>
Date: Mon, 10 Jun 2024 03:51:18 -0700
Subject: [PATCH] [SPIR-V] Don't change switch condition type in CodeGen opts

This change makes sure the preferresd switch condition int type size
remains the same throughout CodeGen optimizations.

The change fixes running several OpenCL applications with -O2 or higher
opt levels, and fixes Basic/stream/stream_max_stmt_exceed.cpp DPC++ E2E
test with -O2.
---
 llvm/lib/Target/SPIRV/SPIRVISelLowering.h      |  5 +++++
 .../optimizations/switch-condition-type.ll     | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+)
 create mode 100644 llvm/test/CodeGen/SPIRV/optimizations/switch-condition-type.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVISelLowering.h b/llvm/lib/Target/SPIRV/SPIRVISelLowering.h
index 6fc200abf4627..77356b7512a73 100644
--- a/llvm/lib/Target/SPIRV/SPIRVISelLowering.h
+++ b/llvm/lib/Target/SPIRV/SPIRVISelLowering.h
@@ -68,6 +68,11 @@ class SPIRVTargetLowering : public TargetLowering {
   // extra instructions required to preserve validity of SPIR-V code imposed by
   // the standard.
   void finalizeLowering(MachineFunction &MF) const override;
+
+  MVT getPreferredSwitchConditionType(LLVMContext &Context,
+                                      EVT ConditionVT) const override {
+    return ConditionVT.getSimpleVT();
+  }
 };
 } // namespace llvm
 
diff --git a/llvm/test/CodeGen/SPIRV/optimizations/switch-condition-type.ll b/llvm/test/CodeGen/SPIRV/optimizations/switch-condition-type.ll
new file mode 100644
index 0000000000000..054520d2021b9
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/optimizations/switch-condition-type.ll
@@ -0,0 +1,18 @@
+; RUN: llc -O2 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O2 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: %[[#INT16:]] = OpTypeInt 16 0
+; CHECK: %[[#PARAM:]] = OpFunctionParameter %[[#INT16]]
+; CHECK: OpSwitch %[[#PARAM]] %[[#]] 1 %[[#]] 2 %[[#]]
+
+define i32 @test_switch(i16 %cond) {
+entry:
+  switch i16 %cond, label %default [ i16 1, label %case_one
+                                     i16 2, label %case_two ]
+case_one:
+  ret i32 1
+case_two:
+  ret i32 2
+default:
+  ret i32 3
+}