[Backport to 14] Implement SPV_INTEL_predicated_io extension (#3370) (#3378)

YixingZhang007 · web-flow · commit 139247938147 · 2025-10-09T20:22:01.000+02:00
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 <yixing.zhang@intel.com>
diff --git a/include/LLVMSPIRVExtensions.inc b/include/LLVMSPIRVExtensions.inc
@@ -74,3 +74,4 @@ EXT(SPV_INTEL_subgroup_matrix_multiply_accumulate)
 EXT(SPV_KHR_bfloat16)
 EXT(SPV_INTEL_bfloat16_arithmetic)
 EXT(SPV_INTEL_shader_atomic_bfloat16)
+EXT(SPV_INTEL_predicated_io)
diff --git a/lib/SPIRV/libSPIRV/SPIRVInstruction.h b/lib/SPIRV/libSPIRV/SPIRVInstruction.h
@@ -4050,5 +4050,22 @@ class SPIRVSubgroupMatrixMultiplyAccumulateINTELInst
 _SPIRV_OP(SubgroupMatrixMultiplyAccumulate, true, 7, true, 4)
 #undef _SPIRV_OP
 
+class SPIRVPredicatedIOINTELInst : public SPIRVInstTemplateBase {
+public:
+  llvm::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
diff --git a/lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h b/lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h
@@ -650,6 +650,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
       "CooperativeMatrixCheckedInstructionsINTEL");
   add(internal::CapabilityBindlessImagesINTEL, "BindlessImagesINTEL");
   add(internal::CapabilityBFloat16ArithmeticINTEL, "BFloat16ArithmeticINTEL");
+  add(internal::CapabilityPredicatedIOINTEL, "PredicatedIOINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 
diff --git a/lib/SPIRV/libSPIRV/SPIRVOpCodeEnumInternal.h b/lib/SPIRV/libSPIRV/SPIRVOpCodeEnumInternal.h
@@ -35,3 +35,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)
diff --git a/lib/SPIRV/libSPIRV/spirv_internal.hpp b/lib/SPIRV/libSPIRV/spirv_internal.hpp
@@ -72,6 +72,8 @@ enum InternalOp {
   IOpCooperativeMatrixLoadCheckedINTEL = 6193,
   IOpCooperativeMatrixStoreCheckedINTEL = 6194,
   IOpCooperativeMatrixConstructCheckedINTEL = 6195,
+  IOpPredicatedLoadINTEL = 6258,
+  IOpPredicatedStoreINTEL = 6259,
   IOpJointMatrixWorkItemLengthINTEL = 6410,
   IOpComplexFMulINTEL = 6415,
   IOpComplexFDivINTEL = 6416,
@@ -117,6 +119,7 @@ enum InternalCapability {
   ICapabilityBFloat16ArithmeticINTEL = 6226,
   ICapabilityAtomicBFloat16AddINTEL = 6255,
   ICapabilityAtomicBFloat16MinMaxINTEL = 6256,
+  ICapabilityPredicatedIOINTEL = 6257,
   ICapabilityCooperativeMatrixPrefetchINTEL = 6411,
   ICapabilityComplexFloatMulDivINTEL = 6414,
   ICapabilityTensorFloat32RoundingINTEL = 6425,
@@ -222,6 +225,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 =
diff --git a/test/extensions/INTEL/SPV_INTEL_predicated_io/predicated_io_generic.ll b/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(i32 addrspace(1)* %{{.*}}, i1 %{{.*}}, i32 %{{.*}})
+; CHECK-LLVM: call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1ibii(i32 addrspace(1)* %{{.*}}, i1 %{{.*}}, i32 %{{.*}}, i32 0)
+; CHECK-LLVM: call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1iib(i32 addrspace(1)* %{{.*}}, i32 %{{.*}}, i1 %{{.*}})
+; CHECK-LLVM: call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1iibi(i32 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(i32 addrspace(1)* %load_pointer, i32 addrspace(1)* %store_pointer, i32  %default_value, i32 %store_object, i1 zeroext %predicate) {
+entry:
+  %0 = call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibi(i32 addrspace(1)* %load_pointer, i1 %predicate, i32 %default_value)
+  %1 = call spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibii(i32 addrspace(1)* %load_pointer, i1 %predicate, i32 %default_value, i32 0)
+  call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiib(i32 addrspace(1)* %store_pointer, i32 %store_object, i1 %predicate)
+  call spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiibi(i32 addrspace(1)* %store_pointer, i32 %store_object, i1 %predicate, i32 0)
+  ret void
+}
+
+declare spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibi(i32 addrspace(1)*, i1, i32)
+declare spir_func i32 @_Z27__spirv_PredicatedLoadINTELPU3AS1Kibii(i32 addrspace(1)*, i1, i32, i32)
+declare spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiib(i32 addrspace(1)*, i32, i1)
+declare spir_func void @_Z28__spirv_PredicatedStoreINTELPU3AS1Kiibi(i32 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"}

Original file line number	Diff line number	Diff line change
`@@ -650,6 +650,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {`
`650`	`650`	`"CooperativeMatrixCheckedInstructionsINTEL");`
`651`	`651`	`add(internal::CapabilityBindlessImagesINTEL, "BindlessImagesINTEL");`
`652`	`652`	`add(internal::CapabilityBFloat16ArithmeticINTEL, "BFloat16ArithmeticINTEL");`
	`653`	`+ add(internal::CapabilityPredicatedIOINTEL, "PredicatedIOINTEL");`
`653`	`654`	`}`
`654`	`655`	`SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)`
`655`	`656`