[Mlir-commits] [mlir] [mlir][spirv] Fix spirv dialect to support Specialization constants as GlobalVar initializer (PR #75660)

Dimple Prajapati llvmlistbot at llvm.org
Fri Dec 15 13:48:14 PST 2023


https://github.com/drprajap created https://github.com/llvm/llvm-project/pull/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.

>From fd8c637f2b146ffce657307841f84a4123e351af Mon Sep 17 00:00:00 2001
From: Dimple Prajapati <dimpalben.r.prajapati at intel.com>
Date: Wed, 13 Dec 2023 22:33:23 +0000
Subject: [PATCH] [mlir][spirv] Fix spirv dialect to support Specialization
 constants in GlobalVar initializer

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 initialzed using specialized constants.
---
 mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp        |  5 ++--
 .../SPIRV/Deserialization/Deserializer.cpp    | 19 ++++++++++-----
 .../SPIRV/Serialization/SerializeOps.cpp      | 22 ++++++++++++-----
 mlir/test/Dialect/SPIRV/IR/structure-ops.mlir | 15 +++++++++++-
 mlir/test/Target/SPIRV/global-variable.mlir   | 24 +++++++++++++++++++
 5 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index 2a1d083308282a..66f1d6b2e12206 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
@@ -1163,9 +1163,10 @@ LogicalResult spirv::GlobalVariableOp::verify() {
     // constants and other variables is supported. They could be normal
     // constants in the module scope as well.
     if (!initOp ||
-        !isa<spirv::GlobalVariableOp, spirv::SpecConstantOp>(initOp)) {
+        !isa<spirv::GlobalVariableOp, spirv::SpecConstantOp, spirv::SpecConstantCompositeOp>(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 89e2e7ad52fa7d..ccea690a7c3ded 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
@@ -637,14 +637,21 @@ spirv::Deserializer::processGlobalVariable(ArrayRef<uint32_t> operands) {
 
   // Initializer.
   FlatSymbolRefAttr initializer = nullptr;
+  
   if (wordIndex < operands.size()) {
-    auto initializerOp = getGlobalVariable(operands[wordIndex]);
-    if (!initializerOp) {
-      return emitError(unknownLoc, "unknown <id> ")
-             << operands[wordIndex] << "used as initializer";
-    }
+    Operation *op = nullptr;
+
+    if((op = getGlobalVariable(operands[wordIndex])))
+      initializer = SymbolRefAttr::get((dyn_cast<spirv::GlobalVariableOp>(op)).getOperation());
+    else if ((op  = getSpecConstant(operands[wordIndex])))
+      initializer = SymbolRefAttr::get((dyn_cast<spirv::SpecConstantOp>(op)).getOperation());
+    else if((op = getSpecConstantComposite(operands[wordIndex])))
+      initializer = SymbolRefAttr::get((dyn_cast<spirv::SpecConstantCompositeOp>(op)).getOperation());
+    else
+      return emitError(unknownLoc,
+                        "Unknown op used as initializer");
+
     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 44538c38a41b83..bd1fc7a84fbd6a 100644
--- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
@@ -383,12 +383,22 @@ Serializer::processGlobalVariableOp(spirv::GlobalVariableOp varOp) {
   operands.push_back(static_cast<uint32_t>(varOp.storageClass()));
 
   // Encode initialization.
-  if (auto initializer = varOp.getInitializer()) {
-    auto initializerID = getVariableID(*initializer);
-    if (!initializerID) {
-      return emitError(varOp.getLoc(),
-                       "invalid usage of undefined variable as initializer");
-    }
+  if (auto initializerName = varOp.getInitializer()) {
+
+    uint32_t initializerID = 0;
+    auto init = varOp->getAttrOfType<FlatSymbolRefAttr>("initializer");
+    Operation *initOp = SymbolTable::lookupNearestSymbolFrom(varOp->getParentOp(), init.getAttr());
+
+    // Check if initializer is GlobalVariable or SpecConstant/SpecConstantComposite
+    if(isa<spirv::GlobalVariableOp>(initOp))
+      initializerID = getVariableID(*initializerName);
+    else
+      initializerID = getSpecConstID(*initializerName);
+
+    if (!initializerID)
+        return emitError(varOp.getLoc(),
+                      "invalid usage of undefined variable as initializer");
+   
     operands.push_back(initializerID);
     elidedAttrs.push_back("initializer");
   }
diff --git a/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir b/mlir/test/Dialect/SPIRV/IR/structure-ops.mlir
index 722e4434aeaf9f..77b605050e1442 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<f32, Private>
 
+
+// -----
+// 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<!spirv.array<3 x i8>, Private>
+  spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr<!spirv.array<3 x i8>, 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<f32, Private>
 }
 
diff --git a/mlir/test/Target/SPIRV/global-variable.mlir b/mlir/test/Target/SPIRV/global-variable.mlir
index 66d0782c205c7d..f22d2a9b3d14d9 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<v1.0, [Shader], []> {
 
 // -----
 
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+  // CHECK:         spirv.SpecConstant @sc = 1 : i8
+  // CHECK-NEXT:    spirv.GlobalVariable @var initializer(@sc) : !spirv.ptr<i8, Uniform>
+  spirv.SpecConstant @sc = 1 : i8
+
+  spirv.GlobalVariable @var initializer(@sc) : !spirv.ptr<i8, Uniform>
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+  // CHECK:         spirv.SpecConstantComposite @scc (@sc0, @sc1, @sc2) : !spirv.array<3 x i8>
+  // CHECK-NEXT:    spirv.GlobalVariable @var initializer(@scc) : !spirv.ptr<!spirv.array<3 x i8>, 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<!spirv.array<3 x i8>, Uniform>
+}
+
+// -----
+
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.GlobalVariable @globalInvocationID built_in("GlobalInvocationId") : !spirv.ptr<vector<3xi32>, Input>
   spirv.func @foo() "None" {



More information about the Mlir-commits mailing list