[PATCH] D73416: [mlir][spirv] Create builtin variable in nearest symbol table

[Lei Zhang via Phabricator via llvm-commits]

antiagainst created this revision.
antiagainst added reviewers: benvanik, denis13.
Herald added subscribers: llvm-commits, liufengdb, lucyrfox, mgester, arpith-jacob, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: mravishankar.
Herald added a reviewer: nicolasvasilache.
Herald added a project: LLVM.

This commit changes the logic of `getBuiltinVariableValue` to get
or create the builtin variable in the nearest symbol table. This
will allow us to use this function in other partial conversion
cases where we haven't created the spv.module yet.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D73416

Files:
  mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h
  mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp


Index: mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
===================================================================
--- mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
+++ mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
@@ -174,11 +174,11 @@
 // Builtin Variables
 //===----------------------------------------------------------------------===//
 
-/// Look through all global variables in `moduleOp` and check if there is a
-/// spv.globalVariable that has the same `builtin` attribute.
-static spirv::GlobalVariableOp getBuiltinVariable(spirv::ModuleOp &moduleOp,
+static spirv::GlobalVariableOp getBuiltinVariable(Block &body,
                                                   spirv::BuiltIn builtin) {
-  for (auto varOp : moduleOp.getBlock().getOps<spirv::GlobalVariableOp>()) {
+  // Look through all global variables in the given `body` block and check if
+  // there is a spv.globalVariable that has the same `builtin` attribute.
+  for (auto varOp : body.getOps<spirv::GlobalVariableOp>()) {
     if (auto builtinAttr = varOp.getAttrOfType<StringAttr>(
             spirv::SPIRVDialect::getAttributeName(
                 spirv::Decoration::BuiltIn))) {
@@ -196,16 +196,16 @@
   return std::string("__builtin_var_") + stringifyBuiltIn(builtin).str() + "__";
 }
 
-/// Gets or inserts a global variable for a builtin within a module.
+/// Gets or inserts a global variable for a builtin within `body` block.
 static spirv::GlobalVariableOp
-getOrInsertBuiltinVariable(spirv::ModuleOp &moduleOp, Location loc,
-                           spirv::BuiltIn builtin, OpBuilder &builder) {
-  if (auto varOp = getBuiltinVariable(moduleOp, builtin)) {
+getOrInsertBuiltinVariable(Block &body, Location loc, spirv::BuiltIn builtin,
+                           OpBuilder &builder) {
+  if (auto varOp = getBuiltinVariable(body, builtin))
     return varOp;
-  }
-  auto ip = builder.saveInsertionPoint();
-  builder.setInsertionPointToStart(&moduleOp.getBlock());
-  auto name = getBuiltinVarName(builtin);
+
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(&body);
+
   spirv::GlobalVariableOp newVarOp;
   switch (builtin) {
   case spirv::BuiltIn::NumWorkgroups:
@@ -216,6 +216,7 @@
     auto ptrType = spirv::PointerType::get(
         VectorType::get({3}, builder.getIntegerType(32)),
         spirv::StorageClass::Input);
+    std::string name = getBuiltinVarName(builtin);
     newVarOp =
         builder.create<spirv::GlobalVariableOp>(loc, ptrType, name, builtin);
     break;
@@ -224,22 +225,22 @@
     emitError(loc, "unimplemented builtin variable generation for ")
         << stringifyBuiltIn(builtin);
   }
-  builder.restoreInsertionPoint(ip);
   return newVarOp;
 }
 
-/// Gets the global variable associated with a builtin and add
-/// it if it doesn't exist.
 Value mlir::spirv::getBuiltinVariableValue(Operation *op,
                                            spirv::BuiltIn builtin,
                                            OpBuilder &builder) {
-  auto moduleOp = op->getParentOfType<spirv::ModuleOp>();
-  if (!moduleOp) {
+  Operation *parent = op->getParentOp();
+  while (parent && !parent->hasTrait<OpTrait::SymbolTable>())
+    parent = parent->getParentOp();
+
+  if (!parent) {
     op->emitError("expected operation to be within a SPIR-V module");
     return nullptr;
   }
-  spirv::GlobalVariableOp varOp =
-      getOrInsertBuiltinVariable(moduleOp, op->getLoc(), builtin, builder);
+  spirv::GlobalVariableOp varOp = getOrInsertBuiltinVariable(
+      *parent->getRegion(0).begin(), op->getLoc(), builtin, builder);
   Value ptr = builder.create<spirv::AddressOfOp>(op->getLoc(), varOp);
   return builder.create<spirv::LoadOp>(op->getLoc(), ptr,
                                        /*memory_access =*/nullptr,
Index: mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h
===================================================================
--- mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h
+++ mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h
@@ -75,8 +75,10 @@
   llvm::SmallSet<Capability, 8> givenCapabilities; /// Allowed capabilities
 };
 
-/// Returns a value that represents a builtin variable value within the SPIR-V
-/// module.
+/// Returns the value for the given `builtin` variable. This function gets or
+/// inserts the global variable associated for the builtin within the nearest
+/// enclosing op that has a symbol table. Returns null Value if such an
+/// enclosing op cannot be found.
 Value getBuiltinVariableValue(Operation *op, BuiltIn builtin,
                               OpBuilder &builder);
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D73416.240379.patch
Type: text/x-patch
Size: 4577 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200125/7038a6c5/attachment.bin>