[SYCL] Fix ESIMD split detection in module properties computation (in…

…tel#15527)

We need to know if a module is an ESIMD split or not split at all to set
some image properties.

Right now we compute if a module is ESIMD or not by looking at the
functions in the module and checking they are all ESIMD, except for
allowed exceptions.

A new exception was found related to using `assert` in user code where
there is a scalar SYCL function that correctly remains in the ESIMD
split. The end result is we don't set the ESIMD property because we
don't think this is the ESIMD split and do the wrong thing at runtime.
Instead of just adding this new exception to the list, I reworked what I
consider to be flaky logic (that I wrote originally, whoops) to figure
out the splits.

Just save metadata in the module of what kind of split it is before we
try to compute module properties.

We already do something similar for the spec constant default split, and
I moved that to a centralized place as part of this change.

The reason we don't just pass in the `ModuleDesc` object as an argument
is because we want to untie properties creation from the
`sycl-post-link` tool so that it can be called in other places (which we
will do for thinLTO).

---------

Signed-off-by: Sarnie, Nick <nick.sarnie@intel.com>

llvm/include/llvm/SYCLLowerIR/ModuleSplitter.h

@@ -37,6 +37,8 @@ class OptionCategory;
 
 namespace module_split {
 
+constexpr char SYCL_ESIMD_SPLIT_MD_NAME[] = "sycl-esimd-split-status";
+
 extern cl::OptionCategory &getModuleSplitCategory();
 
 enum IRSplitMode {
@@ -221,6 +223,8 @@ class ModuleDesc {
     return *Reqs;
   }
 
+  void saveSplitInformationAsMetadata();
+
 #ifndef NDEBUG
   void verifyESIMDProperty() const;
   void dump() const;

llvm/lib/SYCLLowerIR/ComputeModuleRuntimeInfo.cpp

@@ -28,6 +28,21 @@ constexpr int DebugModuleProps = 0;
 #endif
 
 namespace llvm::sycl {
+namespace {
+module_split::SyclEsimdSplitStatus
+getSYCLESIMDSplitStatusFromMetadata(const Module &M) {
+  auto *SplitMD = M.getNamedMetadata(module_split::SYCL_ESIMD_SPLIT_MD_NAME);
+  assert(SplitMD && "Unexpected metadata");
+  auto *MDOp = SplitMD->getOperand(0);
+  assert(MDOp && "Unexpected metadata operand");
+  const auto &MDConst = MDOp->getOperand(0);
+  auto *MDVal = mdconst::dyn_extract_or_null<ConstantInt>(MDConst);
+  uint8_t Val = MDVal->getZExtValue();
+  assert(Val < 3 && "Unexpected value for split metadata");
+  auto AsEnum = static_cast<module_split::SyclEsimdSplitStatus>(Val);
+  return AsEnum;
+}
+} // namespace
 
 bool isModuleUsingAsan(const Module &M) {
   for (const auto &F : M) {
@@ -305,16 +320,11 @@ PropSetRegTy computeModuleProperties(const Module &M,
                   GV.getName());
     }
   }
-  bool SeenESIMDFunction = false;
-  bool SeenSYCLFunction = false;
-  for (const auto &F : M) {
-    if (llvm::module_split::isESIMDFunction(F))
-      SeenESIMDFunction = true;
-    else if (utils::isSYCLExternalFunction(&F) &&
-             !F.getName().starts_with("__itt"))
-      SeenSYCLFunction = true;
-  }
-  if (SeenESIMDFunction && !SeenSYCLFunction)
+
+  module_split::SyclEsimdSplitStatus SplitType =
+      getSYCLESIMDSplitStatusFromMetadata(M);
+
+  if (SplitType == module_split::SyclEsimdSplitStatus::ESIMD_ONLY)
     PropSet.add(PropSetRegTy::SYCL_MISC_PROP, "isEsimdImage", true);
   {
     StringRef RegAllocModeAttr = "sycl-register-alloc-mode";
@@ -359,7 +369,7 @@ PropSetRegTy computeModuleProperties(const Module &M,
   // 'if' below essentially preserves the behavior (presumably mistakenly)
   // implemented in intel/llvm#8763: ignore 'optLevel' property for images which
   // were produced my merge after ESIMD split
-  if (!SeenESIMDFunction || !SeenSYCLFunction) {
+  if (SplitType != module_split::SyclEsimdSplitStatus::SYCL_AND_ESIMD) {
     // Handle sycl-optlevel property
     int OptLevel = -1;
     for (const Function *F : EntryPoints) {

llvm/lib/SYCLLowerIR/ModuleSplitter.cpp

@@ -26,6 +26,7 @@
 #include "llvm/SYCLLowerIR/DeviceGlobals.h"
 #include "llvm/SYCLLowerIR/LowerInvokeSimd.h"
 #include "llvm/SYCLLowerIR/SYCLUtils.h"
+#include "llvm/SYCLLowerIR/SpecConstants.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
@@ -798,6 +799,23 @@ void ModuleDesc::dump() const {
 }
 #endif // NDEBUG
 
+void ModuleDesc::saveSplitInformationAsMetadata() {
+  // Add metadata to the module so we can identify what kind of SYCL/ESIMD split
+  // later.
+  auto *SplitMD = M->getOrInsertNamedMetadata(SYCL_ESIMD_SPLIT_MD_NAME);
+  auto *SplitMDOp = MDNode::get(
+      M->getContext(), ConstantAsMetadata::get(ConstantInt::get(
+                           Type::getInt8Ty(M->getContext()),
+                           static_cast<uint8_t>(EntryPoints.Props.HasESIMD))));
+  SplitMD->addOperand(SplitMDOp);
+
+  // Add metadata to the module so we can identify it as the default value spec
+  // constants split later.
+  if (isSpecConstantDefault())
+    M->getOrInsertNamedMetadata(
+        SpecConstantsPass::SPEC_CONST_DEFAULT_VAL_MODULE_MD_STRING);
+}
+
 void EntryPointGroup::saveNames(std::vector<std::string> &Dest) const {
   Dest.reserve(Dest.size() + Functions.size());
   std::transform(Functions.begin(), Functions.end(),
@@ -1291,6 +1309,7 @@ static Expected<SplitModule> saveModuleDesc(ModuleDesc &MD, std::string Prefix,
                                             bool OutputAssembly) {
   SplitModule SM;
   Prefix += OutputAssembly ? ".ll" : ".bc";
+  MD.saveSplitInformationAsMetadata();
   Error E = saveModuleIRInFile(MD.getModule(), Prefix, OutputAssembly);
   if (E)
     return E;

llvm/test/tools/sycl-post-link/sycl-esimd/assert.ll

@@ -0,0 +1,43 @@
+; RUN: sycl-post-link -properties -split-esimd  -S < %s -o %t.table
+; RUN: FileCheck %s -input-file=%t_esimd_0.prop
+
+; Verify we mark a image with an ESIMD kernel with the isEsimdImage property
+
+; CHECK: isEsimdImage=1|1
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
+target triple = "spir64-unknown-unknown"
+
+%"struct.sycl::_V1::detail::AssertHappened" = type { i32, [257 x i8], [257 x i8], [129 x i8], i32, i64, i64, i64, i64, i64, i64 }
+%"class.sycl::_V1::id" = type { %"class.sycl::_V1::detail::array" }
+%"class.sycl::_V1::detail::array" = type { [1 x i64] }
+
+@.str = private unnamed_addr addrspace(1) constant [10 x i8] c"Id != 400\00", align 1
+@.str.1 = private unnamed_addr addrspace(1) constant [8 x i8] c"foo.cpp\00", align 1
+@__PRETTY_FUNCTION__ = private unnamed_addr addrspace(1) constant [56 x i8] c"auto main()::(anonymous class)::operator()(id<1>) const\00", align 1
+@SPIR_AssertHappenedMem = linkonce_odr dso_local addrspace(1) global %"struct.sycl::_V1::detail::AssertHappened" zeroinitializer, align 8
+
+declare void @llvm.assume(i1 noundef) #2
+
+define weak_odr dso_local spir_kernel void @esimd_kernel() local_unnamed_addr #0 !sycl_explicit_simd !0 {
+entry:
+  tail call spir_func void @__assert_fail(ptr addrspace(4) noundef addrspacecast (ptr addrspace(1) @.str to ptr addrspace(4)), ptr addrspace(4) noundef addrspacecast (ptr addrspace(1) @.str.1 to ptr addrspace(4)), i32 noundef 13, ptr addrspace(4) noundef addrspacecast (ptr addrspace(1) @__PRETTY_FUNCTION__ to ptr addrspace(4))) #12
+  ret void
+}
+
+define weak dso_local spir_func void @__assert_fail(ptr addrspace(4) noundef %expr, ptr addrspace(4) noundef %file, i32 noundef %line, ptr addrspace(4) noundef %func) #1 {
+entry:
+  tail call spir_func void @__devicelib_assert_fail(ptr addrspace(4) noundef %expr, ptr addrspace(4) noundef %file, i32 noundef %line, ptr addrspace(4) noundef %func) #1
+  ret void
+}
+
+define weak dso_local spir_func void @__devicelib_assert_fail(ptr addrspace(4) noundef %expr, ptr addrspace(4) noundef %file, i32 noundef %line, ptr addrspace(4) noundef %func) #2 {
+entry:                                 
+  ret void
+}
+
+attributes #0 = { convergent mustprogress norecurse nounwind "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "sycl-module-id"="foo.cpp" "sycl-optlevel"="2" "uniform-work-group-size"="true" }
+attributes #1 = { convergent mustprogress norecurse nounwind "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "sycl-module-id"="bar.cpp" "sycl-optlevel"="2" }
+attributes #2 = { convergent nounwind }
+
+!0 = !{}

llvm/test/tools/sycl-post-link/sycl-esimd/basic-esimd-lower.ll

@@ -55,7 +55,7 @@ attributes #0 = { "sycl-module-id"="a.cpp" }
 ; CHECK-NO-LOWERING: }
 
 ; With -O0, we only lower ESIMD code, but no other optimizations
-; CHECK-O0: define dso_local spir_kernel void @ESIMD_kernel() #{{[0-9]}} !sycl_explicit_simd !3 !intel_reqd_sub_group_size !4 {
+; CHECK-O0: define dso_local spir_kernel void @ESIMD_kernel() #{{[0-9]}} !sycl_explicit_simd !{{[0-9]}} !intel_reqd_sub_group_size !{{[0-9]}} {
 ; CHECK-O0: entry:
 ; CHECK-O0:   %0 = load <3 x i64>, {{.*}} addrspacecast {{.*}} @__spirv_BuiltInGlobalInvocationId
 ; CHECK-O0:   %1 = extractelement <3 x i64> %0, i64 0

llvm/tools/sycl-post-link/sycl-post-link.cpp

@@ -419,6 +419,7 @@ void saveModule(std::vector<std::unique_ptr<util::SimpleTable>> &OutTables,
                 module_split::ModuleDesc &MD, int I, StringRef IRFilename) {
   IrPropSymFilenameTriple BaseTriple;
   StringRef Suffix = getModuleSuffix(MD);
+  MD.saveSplitInformationAsMetadata();
   if (!IRFilename.empty()) {
     // don't save IR, just record the filename
     BaseTriple.Ir = IRFilename.str();
@@ -509,10 +510,6 @@ processSpecConstantsWithDefaultValues(const module_split::ModuleDesc &MD) {
   assert(NewModuleDesc->Props.SpecConstsMet &&
          "This property should be true since the presence of SpecConsts "
          "has been checked before the run of the pass");
-  // Add metadata to the module so we can identify it as the default value split
-  // later.
-  NewModuleDesc->getModule().getOrInsertNamedMetadata(
-      SpecConstantsPass::SPEC_CONST_DEFAULT_VAL_MODULE_MD_STRING);
   NewModuleDesc->rebuildEntryPoints();
   return NewModuleDesc;
 }
@@ -791,7 +788,7 @@ processInputModule(std::unique_ptr<Module> M) {
   // to keep the optimizer from wrongfully removing them. llvm.compiler.used
   // symbols are usually removed at backend lowering, but this is handled here
   // for SPIR-V since SYCL compilation uses llvm-spirv, not the SPIR-V backend.
-  if (auto Triple = M->getTargetTriple().find("spir") != std::string::npos)
+  if (M->getTargetTriple().find("spir") != std::string::npos)
     Modified |= removeDeviceGlobalFromCompilerUsed(*M.get());
 
   // Instrument each image scope device globals if the module has been