Implement SPV_INTEL_predicated_io extension (#3370)

This extension adds predicated load and store instructions. Predicated
load performs load from memory if predicate is true; otherwise, it uses
default_value as a result. Predicated store performs store of value to
memory if predicate is true; otherwise, it does nothing.

Spec: intel/llvm#20158

Signed-off-by: Zhang, Yixing <[email protected]>

Author: YixingZhang007

include/LLVMSPIRVExtensions.inc

@@ -82,3 +82,4 @@ EXT(SPV_INTEL_ternary_bitwise_function)
 EXT(SPV_INTEL_int4)
 EXT(SPV_INTEL_function_variants)
 EXT(SPV_INTEL_shader_atomic_bfloat16)
+EXT(SPV_INTEL_predicated_io)

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -4480,5 +4480,22 @@ class SPIRVTernaryBitwiseFunctionINTELInst : public SPIRVInstTemplateBase {
 _SPIRV_OP(BitwiseFunction, true, 7)
 #undef _SPIRV_OP
 
+class SPIRVPredicatedIOINTELInst : public SPIRVInstTemplateBase {
+public:
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_predicated_io;
+  }
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(internal::CapabilityPredicatedIOINTEL);
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVPredicatedIOINTELInst,                        \
+                            internal::Op##x##INTEL, __VA_ARGS__>               \
+      SPIRV##x##INTEL;
+_SPIRV_OP(PredicatedLoad, true, 6, true)
+_SPIRV_OP(PredicatedStore, false, 4, true)
+#undef _SPIRV_OP
 } // namespace SPIRV
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -691,6 +691,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityFunctionVariantsINTEL, "FunctionVariantsINTEL");
   add(CapabilitySpecConditionalINTEL, "SpecConditionalINTEL");
   add(internal::CapabilityBFloat16ArithmeticINTEL, "BFloat16ArithmeticINTEL");
+  add(internal::CapabilityPredicatedIOINTEL, "PredicatedIOINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/SPIRVOpCodeEnumInternal.h

@@ -41,3 +41,7 @@ _SPIRV_OP_INTERNAL(ConvertHandleToSamplerINTEL,
                    internal::ConvertHandleToSamplerINTEL)
 _SPIRV_OP_INTERNAL(ConvertHandleToSampledImageINTEL,
                    internal::ConvertHandleToSampledImageINTEL)
+_SPIRV_OP_INTERNAL(PredicatedLoadINTEL,
+                   internal::OpPredicatedLoadINTEL)
+_SPIRV_OP_INTERNAL(PredicatedStoreINTEL,
+                   internal::OpPredicatedStoreINTEL)

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -76,6 +76,8 @@ enum InternalOp {
   IOpCooperativeMatrixConstructCheckedINTEL = 6195,
   IOpCooperativeMatrixLoadOffsetINTEL = 6239,
   IOpCooperativeMatrixStoreOffsetINTEL = 6240,
+  IOpPredicatedLoadINTEL = 6258,
+  IOpPredicatedStoreINTEL = 6259,
   IOpJointMatrixWorkItemLengthINTEL = 6410,
   IOpTypeTaskSequenceINTEL = 6199,
   IOpMaskedGatherINTEL = 6428,
@@ -110,6 +112,7 @@ enum InternalCapability {
   ICapabilityCooperativeMatrixOffsetInstructionsINTEL = 6238,
   ICapabilityAtomicBFloat16AddINTEL = 6255,
   ICapabilityAtomicBFloat16MinMaxINTEL = 6256,
+  ICapabilityPredicatedIOINTEL = 6257,
   ICapabilityCooperativeMatrixPrefetchINTEL = 6411,
   ICapabilityMaskedGatherScatterINTEL = 6427,
   ICapabilityJointMatrixWIInstructionsINTEL = 6435,
@@ -208,6 +211,11 @@ _SPIRV_OP(Op, ConvertHandleToSampledImageINTEL)
 
 _SPIRV_OP(Capability, AtomicBFloat16AddINTEL)
 _SPIRV_OP(Capability, AtomicBFloat16MinMaxINTEL)
+
+_SPIRV_OP(Capability, PredicatedIOINTEL)
+_SPIRV_OP(Op, PredicatedLoadINTEL)
+_SPIRV_OP(Op, PredicatedStoreINTEL)
+
 #undef _SPIRV_OP
 
 constexpr SourceLanguage SourceLanguagePython =

test/extensions/INTEL/SPV_INTEL_predicated_io/predicated_io_generic.ll

@@ -0,0 +1,75 @@
+; Generated with:
+; source.cl:
+; int __spirv_PredicatedLoadINTEL(const __global int* pointer, bool predicate, int default_value);
+; int __spirv_PredicatedLoadINTEL(const __global int* pointer, bool predicate, int default_value, int memory_operands);
+; void __spirv_PredicatedStoreINTEL(const __global int* pointer, int object, bool predicate);
+; void __spirv_PredicatedStoreINTEL(const __global int* pointer, int object, bool predicate, int memory_operands);
+;
+; void foo(const __global int* load_pointer, __global int* store_pointer, int default_value, int store_object, bool predicate) {
+;   const int memory_ops = 0;
+;    int result1 = __spirv_PredicatedLoadINTEL(load_pointer, predicate, default_value);
+;    int result2 = __spirv_PredicatedLoadINTEL(load_pointer, predicate, default_value, memory_ops);
+;    __spirv_PredicatedStoreINTEL(store_pointer, store_object, predicate);
+;    __spirv_PredicatedStoreINTEL(store_pointer, store_object, predicate, memory_ops);
+; }
+; clang -cc1 -cl-std=clc++2021 -triple spir64-unknown-unknown -emit-llvm -finclude-default-header source.cl -o tmp.ll
+
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv --spirv-ext=+SPV_INTEL_predicated_io
+; RUN: llvm-spirv %t.spv -o %t.spt --to-text
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv %t.spv -o %t.rev.bc -r --spirv-target-env=SPV-IR
+; RUN: llvm-dis %t.rev.bc -o %t.rev.ll
+; RUN: FileCheck < %t.rev.ll %s --check-prefix=CHECK-LLVM
+
+; RUN: not llvm-spirv %t.bc 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+; CHECK-ERROR: RequiresExtension: Feature requires the following SPIR-V extension:
+; CHECK-ERROR-NEXT: SPV_INTEL_predicated_io
+
+; CHECK-SPIRV: Capability PredicatedIOINTEL
+; CHECK-SPIRV: Extension "SPV_INTEL_predicated_io"
+; CHECK-SPIRV-DAG: TypeInt [[#Int32Ty:]] 32 0
+; CHECK-SPIRV-DAG: Constant [[#Int32Ty]] [[#Const0:]] 0
+; CHECK-SPIRV-DAG: TypeVoid [[#VoidTy:]]
+; CHECK-SPIRV-DAG: TypePointer [[#IntPtrTy:]] 5 [[#Int32Ty]]
+; CHECK-SPIRV-DAG: TypeBool [[#BoolTy:]]
+; CHECK-SPIRV: FunctionParameter [[#IntPtrTy]] [[#LoadPtr:]]
+; CHECK-SPIRV: FunctionParameter [[#IntPtrTy]] [[#StorePtr:]]
+; CHECK-SPIRV: FunctionParameter [[#Int32Ty]] [[#DefaultVal:]]
+; CHECK-SPIRV: FunctionParameter [[#Int32Ty]] [[#StoreObj:]]
+; CHECK-SPIRV: FunctionParameter [[#BoolTy]] [[#Predicate:]]
+; CHECK-SPIRV: PredicatedLoadINTEL [[#Int32Ty]] [[#Result1:]] [[#LoadPtr]] [[#Predicate]] [[#DefaultVal]]
+; CHECK-SPIRV: PredicatedLoadINTEL [[#Int32Ty]] [[#Result2:]] [[#LoadPtr]] [[#Predicate]] [[#DefaultVal]] [[#Const0]]
+; CHECK-SPIRV: PredicatedStoreINTEL [[#StorePtr]] [[#StoreObj]] [[#Predicate]]
+; CHECK-SPIRV: PredicatedStoreINTEL [[#StorePtr]] [[#StoreObj]] [[#Predicate]] [[#Const0]]
+
+; CHECK-LLVM: call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1ibi(ptr addrspace(1) %{{.*}}, i1 %{{.*}}, i32 %{{.*}})
+; CHECK-LLVM: call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1ibii(ptr addrspace(1) %{{.*}}, i1 %{{.*}}, i32 %{{.*}}, i32 0)
+; CHECK-LLVM: call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1iib(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, i1 %{{.*}})
+; CHECK-LLVM: call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1iibi(ptr addrspace(1) %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i32 0)
+
+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"
+target triple = "spir64-unknown-unknown"
+
+define spir_func void @foo(ptr addrspace(1) %load_pointer, ptr addrspace(1) %store_pointer, i32  %default_value, i32 %store_object, i1 zeroext %predicate) {
+entry:
+  %1 = call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibi(ptr addrspace(1) %load_pointer, i1 %predicate, i32 %default_value)
+  %2 = call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibii(ptr addrspace(1) %load_pointer, i1 %predicate, i32 %default_value, i32 0)
+  call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiib(ptr addrspace(1) %store_pointer, i32 %store_object, i1 %predicate)
+  call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiibi(ptr addrspace(1) %store_pointer, i32 %store_object, i1 %predicate, i32 0)
+  ret void
+}
+
+declare spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibi(ptr addrspace(1), i1, i32)
+declare spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibii(ptr addrspace(1), i1, i32, i32)
+declare spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiib(ptr addrspace(1), i32, i1)
+declare spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiibi(ptr addrspace(1), i32, i1, i32)
+
+!opencl.spir.version = !{!0}
+!spirv.Source = !{!1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}
+!2 = !{!"clang version 17.0.0"}