[SPIRV] Implement bare bones lowering for G_PREFETCH

Project: llvm/llvm-project | Language: C++ | Concepts: api-design, backwards-compatibility | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/215505

The problem

The LLVM SPIR-V backend did not handle the G_PREFETCH generic prefetch instruction. As a result, code using llvm.prefetch was not correctly lowered to SPIR-V, causing failures or missing functionality when compiling code that uses prefetching on SPIR-V targets.

Why it's non-trivial

A naive fix might try to always lower G_PREFETCH to a SPIR-V prefetch instruction. However, the SPIR-V OpUntypedPrefetchKHR instruction requires the number of bytes to prefetch, which is not available from LLVM's llvm.prefetch. Additionally, OpenCL's prefetch expects the number of elements, not bytes, and cannot take untyped pointers. The lowering must handle these differences and avoid generating invalid SPIR-V.

The approach

The patch adds special handling for G_PREFETCH in the SPIR-V instruction selector and legalizer:

+  bool selectPrefetch(Register ResVReg, SPIRVTypeInst ResType,
+                       MachineInstr &I) const;

This adds a new method to handle prefetch selection in the instruction selector.

@@ -1451,6 +1453,8 @@ bool SPIRVInstructionSelector::spvSelect(Register ResVReg,
   case TargetOpcode::G_PREFETCH:
+    return selectPrefetch(ResVReg, ResType, I);

The main selection switch now delegates G_PREFETCH to the new handler.

+bool SPIRVInstructionSelector::selectPrefetch(Register ResVReg,
+                                              SPIRVTypeInst ResType,
+                                              MachineInstr &I) const {
+  // llvm.prefetch is missing info about how many bytes to prefetch that
+  // OpUntypedPrefetchKHR instruction requires, so just drop the instruction
+  // in case untyped pointers are enabled since OCL prefetch cannot take
+  // untyped pointers.
+  if (STI.canUseExtension(SPIRV::Extension::SPV_KHR_untyped_pointers))
+    return true;
+
+  if (STI.canUseExtInstSet(SPIRV::InstructionSet::OpenCL_std)) {
+    MachineIRBuilder MIRBuilder(I);
+    const unsigned PointerSize = GR.getPointerSize();
+    const SPIRVTypeInst PointerSizeType =
+        GR.getOrCreateSPIRVIntegerType(PointerSize, MIRBuilder);
+    bool ZeroAsNull = !STI.isShader();
+
+    // Assume that a single element needs prefetch.
+    Register ConstIntOne =
+        GR.getOrCreateConstInt(1, I, PointerSizeType, TII, ZeroAsNull);
+
+    Register AddrVal = I.getOperand(0).getReg();
+    Register ExtReg = MRI->createVirtualRegister(
+        GR.getRegClass(GR.getOpTypeVoid(MIRBuilder)));
+    return selectExtInst(ExtReg, GR.getOpTypeVoid(MIRBuilder), I, CL::prefetch,
+                         /*setMIFlags=*/false, /*useMSSrc=*/false,
+                         {AddrVal, ConstIntOne});
+  }
+  return true;
+}

If the target supports untyped pointers (SPV_KHR_untyped_pointers), the instruction is dropped (no code emitted). If OpenCL is available, it emits an OpenCL prefetch call, assuming a single element, since the number of bytes is not available. Otherwise, it does nothing.

@@ -473,6 +473,7 @@ SPIRVLegalizerInfo::SPIRVLegalizerInfo(const SPIRVSubtarget &ST) {
   getActionDefinitionsBuilder(G_ATOMIC_CMPXCHG).alwaysLegal();
+  getActionDefinitionsBuilder(G_PREFETCH).alwaysLegal();

Marks G_PREFETCH as always legal, so it is not rejected during legalization.

A new test is added to verify the lowering:

define spir_kernel void @foo(ptr addrspace(1) %p) {
  call void @llvm.prefetch.p1(ptr addrspace(1) %p, i32 0, i32 3, i32 1)
  ret void
}
; NO-EXT: OpExtInst %[[#VOID]] %[[#EXTSET]] prefetch %[[#P]] %[[#ONE]]
; EXT-NOT: OpExtInst{{.*}}prefetch

The test checks that, without the untyped pointer extension, an OpenCL prefetch is emitted; with the extension, nothing is emitted.

From the review thread

maarquitos14:
If I read the OpenCL spec correctly (see below quote), OCL needs the number of bytes (or a proxy) too, doesn't it?

void prefetch(const __global gentype *p, size_t num_gentypes)

Prefetch num_gentypes * sizeof(gentype) bytes into the global cache. The prefetch instruction is applied to a work-item in a work-group and does not affect the functional behavior of the kernel.

Fznamznon:
According to https://registry.khronos.org/SPIR-V/specs/unified1/OpenCL.ExtendedInstructionSet.100.html , prefetch SPIRV instruction takes amount of elements, not bytes. I'm a newbie in the area, so I discussed this with @MrSidims and we come to

Takeaways

• When lowering high-level IR to a target-specific IR, always check the requirements and limitations of the target instructions—don't assume argument semantics match.
• If a target feature or extension makes an operation ambiguous or unsupported (e.g., untyped pointers), it may be safer to drop the instruction than to emit potentially incorrect code.
• When information is missing (such as the number of bytes to prefetch), choose a conservative, safe default (like prefetching a single element) and document the limitation clearly.