[mlir][spirv] Create builtin variable in nearest symbol table

[antiagainst]

Summary

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.

[merge_guards_bot]

Unit tests: pass. 62195 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

281

SymbolTable.cpp defines a static method for this(getNearestSymbolTable) that we could likely expose.

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

281

Good point! https://reviews.llvm.org/D73433

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

private:
   /// environment.
   bool isLegalOp(Operation *op);
 
   Version givenVersion;                            /// SPIR-V version to target
   llvm::SmallSet<Extension, 4> givenExtensions;    /// Allowed extensions
   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);
 
 /// Sets the InterfaceVarABIAttr and EntryPointABIAttr for a function and its
 /// arguments.
 LogicalResult setABIAttrs(FuncOp funcOp, EntryPointABIAttr entryPointInfo,
                           ArrayRef<InterfaceVarABIAttr> argABIInfo);
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVLOWERING_H

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

void mlir::populateBuiltinFuncToSPIRVPatterns(
     OwningRewritePatternList &patterns) {
   patterns.insert<FuncOpConversion>(context, typeConverter);
 }
 
 //===----------------------------------------------------------------------===//
 // 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))) {
       auto varBuiltIn = spirv::symbolizeBuiltIn(builtinAttr.getValue());
       if (varBuiltIn && varBuiltIn.getValue() == builtin) {
         return varOp;
       }
     }
   }
   return nullptr;
 }
 
 /// Gets name of global variable for a builtin.
 static std::string getBuiltinVarName(spirv::BuiltIn builtin) {
   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:
   case spirv::BuiltIn::WorkgroupSize:
   case spirv::BuiltIn::WorkgroupId:
   case spirv::BuiltIn::LocalInvocationId:
   case spirv::BuiltIn::GlobalInvocationId: {
     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;
   }
   default:
     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) {
-    op->emitError("expected operation to be within a SPIR-V module");
+  Operation *parent = op->getParentOp();
+  while (parent && !parent->hasTrait<OpTrait::SymbolTable>())
+    parent = parent->getParentOp();
+  if (!parent) {
+    op->emitError("expected operation to be within a module-like op");
     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);
 }
 
 //===----------------------------------------------------------------------===//
 // Set ABI attributes for lowering entry functions.
 //===----------------------------------------------------------------------===//