From 5e54319b7be3e8aa035836098e0a9defc0a41c3a Mon Sep 17 00:00:00 2001 From: Dimple Prajapati Date: Fri, 5 Jan 2024 16:27:30 -0800 Subject: [PATCH] [mlir][spirv] Support spec constants as GlobalVar initializer (#75660) Changes include: - spirv serialization and deserialization needs handling in cases when GlobalVariableOp initializer is defined using spirv SpecConstant or SpecConstantComposite op, currently even though it allows SpecConst, it only looked up in for GlobalVariable Map to find initializer symbol reference, change is fixing this and extending the support to SpecConstantComposite as an initializer. - Adds tests to make sure GlobalVariable can be initialized using specialized constants. --------- Co-authored-by: Lei Zhang --- mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp | 7 +++--- .../SPIRV/Deserialization/Deserializer.cpp | 16 +++++++++---- .../SPIRV/Serialization/SerializeOps.cpp | 23 +++++++++++++----- mlir/test/Dialect/SPIRV/IR/structure-ops.mlir | 15 +++++++++++- mlir/test/Target/SPIRV/global-variable.mlir | 24 +++++++++++++++++++ 5 files changed, 71 insertions(+), 14 deletions(-) diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp index 2a1d08330828..5343a12132a9 100644 --- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp +++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp @@ -1162,10 +1162,11 @@ LogicalResult spirv::GlobalVariableOp::verify() { // TODO: Currently only variable initialization with specialization // constants and other variables is supported. They could be normal // constants in the module scope as well. - if (!initOp || - !isa(initOp)) { + if (!initOp || !isa(initOp)) { return emitOpError("initializer must be result of a " - "spirv.SpecConstant or spirv.GlobalVariable op"); + "spirv.SpecConstant or spirv.GlobalVariable or " + "spirv.SpecConstantCompositeOp op"); } } diff --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp index 89e2e7ad52fa..00645d2c4551 100644 --- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp +++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp @@ -637,14 +637,22 @@ spirv::Deserializer::processGlobalVariable(ArrayRef operands) { // Initializer. FlatSymbolRefAttr initializer = nullptr; + if (wordIndex < operands.size()) { - auto initializerOp = getGlobalVariable(operands[wordIndex]); - if (!initializerOp) { + Operation *op = nullptr; + + if (auto initOp = getGlobalVariable(operands[wordIndex])) + op = initOp; + else if (auto initOp = getSpecConstant(operands[wordIndex])) + op = initOp; + else if (auto initOp = getSpecConstantComposite(operands[wordIndex])) + op = initOp; + else return emitError(unknownLoc, "unknown ") << operands[wordIndex] << "used as initializer"; - } + + initializer = SymbolRefAttr::get(op); wordIndex++; - initializer = SymbolRefAttr::get(initializerOp.getOperation()); } if (wordIndex != operands.size()) { return emitError(unknownLoc, diff --git a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp index 44538c38a41b..7bfcca5b4dcd 100644 --- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp +++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp @@ -383,20 +383,31 @@ Serializer::processGlobalVariableOp(spirv::GlobalVariableOp varOp) { operands.push_back(static_cast(varOp.storageClass())); // Encode initialization. - if (auto initializer = varOp.getInitializer()) { - auto initializerID = getVariableID(*initializer); - if (!initializerID) { + StringRef initAttrName = varOp.getInitializerAttrName().getValue(); + if (std::optional initSymbolName = varOp.getInitializer()) { + uint32_t initializerID = 0; + auto initRef = varOp->getAttrOfType(initAttrName); + Operation *initOp = SymbolTable::lookupNearestSymbolFrom( + varOp->getParentOp(), initRef.getAttr()); + + // Check if initializer is GlobalVariable or SpecConstant* cases. + if (isa(initOp)) + initializerID = getVariableID(*initSymbolName); + else + initializerID = getSpecConstID(*initSymbolName); + + if (!initializerID) return emitError(varOp.getLoc(), "invalid usage of undefined variable as initializer"); - } + operands.push_back(initializerID); - elidedAttrs.push_back("initializer"); + elidedAttrs.push_back(initAttrName); } if (failed(emitDebugLine(typesGlobalValues, varOp.getLoc()))) return failure(); encodeInstructionInto(typesGlobalValues, spirv::Opcode::OpVariable, operands); - elidedAttrs.push_back("initializer"); + elidedAttrs.push_back(initAttrName); // Encode decorations. for (auto attr : varOp->getAttrs()) { diff --git a/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir b/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir index 722e4434aeaf..77b605050e14 100644 --- a/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir +++ b/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir @@ -349,6 +349,19 @@ spirv.SpecConstant @sc = 4.0 : f32 // CHECK: spirv.GlobalVariable @var initializer(@sc) spirv.GlobalVariable @var initializer(@sc) : !spirv.ptr + +// ----- +// Allow SpecConstantComposite as initializer + spirv.module Logical GLSL450 { + spirv.SpecConstant @sc1 = 1 : i8 + spirv.SpecConstant @sc2 = 2 : i8 + spirv.SpecConstant @sc3 = 3 : i8 + spirv.SpecConstantComposite @scc (@sc1, @sc2, @sc3) : !spirv.array<3 x i8> + + // CHECK: spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr, Private> + spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr, Private> +} + // ----- spirv.module Logical GLSL450 { @@ -410,7 +423,7 @@ spirv.module Logical GLSL450 { // ----- spirv.module Logical GLSL450 { - // expected-error @+1 {{op initializer must be result of a spirv.SpecConstant or spirv.GlobalVariable op}} + // expected-error @+1 {{op initializer must be result of a spirv.SpecConstant or spirv.GlobalVariable or spirv.SpecConstantCompositeOp op}} spirv.GlobalVariable @var0 initializer(@var1) : !spirv.ptr } diff --git a/mlir/test/Target/SPIRV/global-variable.mlir b/mlir/test/Target/SPIRV/global-variable.mlir index 66d0782c205c..f22d2a9b3d14 100644 --- a/mlir/test/Target/SPIRV/global-variable.mlir +++ b/mlir/test/Target/SPIRV/global-variable.mlir @@ -23,6 +23,30 @@ spirv.module Logical GLSL450 requires #spirv.vce { // ----- +spirv.module Logical GLSL450 requires #spirv.vce { + // CHECK: spirv.SpecConstant @sc = 1 : i8 + // CHECK-NEXT: spirv.GlobalVariable @var initializer(@sc) : !spirv.ptr + spirv.SpecConstant @sc = 1 : i8 + + spirv.GlobalVariable @var initializer(@sc) : !spirv.ptr +} + +// ----- + +spirv.module Logical GLSL450 requires #spirv.vce { + // CHECK: spirv.SpecConstantComposite @scc (@sc0, @sc1, @sc2) : !spirv.array<3 x i8> + // CHECK-NEXT: spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr, Uniform> + spirv.SpecConstant @sc0 = 1 : i8 + spirv.SpecConstant @sc1 = 2 : i8 + spirv.SpecConstant @sc2 = 3 : i8 + + spirv.SpecConstantComposite @scc (@sc0, @sc1, @sc2) : !spirv.array<3 x i8> + + spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr, Uniform> +} + +// ----- + spirv.module Logical GLSL450 requires #spirv.vce { spirv.GlobalVariable @globalInvocationID built_in("GlobalInvocationId") : !spirv.ptr, Input> spirv.func @foo() "None" {