[mlir][spirv] Use larger range for target environment lookup function

Previously we only look at the directly passed-in op for a potential
spv.target_env attribute. This commit switches to use a larger range
and recursively check enclosing symbol tables.

Differential Revision: https://reviews.llvm.org/D75869

diff --git a/mlir/docs/Dialects/SPIR-V.md b/mlir/docs/Dialects/SPIR-V.md
index 4380c9f..0cb6281 100644 (file)
--- a/mlir/docs/Dialects/SPIR-V.md
+++ b/mlir/docs/Dialects/SPIR-V.md
@@ -792,9 +792,11 @@ spv.target_env = #spv.target_env<
 } { ... }
 ```
 
-Dialect conversion framework will utilize the information in `spv.target_env`
-to properly filter out patterns and ops not available in the target execution
-environment.
+Dialect conversion framework will utilize the information in `spv.target_env` to
+properly filter out patterns and ops not available in the target execution
+environment. When targeting SPIR-V, one needs to create a
+[`SPIRVConversionTarget`](#spirvconversiontarget) by providing such an
+attribute.
 
 ## Shader interface (ABI)
 
@@ -931,6 +933,10 @@ target satisfying a given [`spv.target_env`](#target-environment). It registers
 proper hooks to check the dynamic legality of SPIR-V ops. Users can further
 register other legality constraints into the returned `SPIRVConversionTarget`.
 
+`spirv::lookupTargetEnvOrDefault()` is a handy utility function to query an
+`spv.target_env` attached in the input IR or use the feault to construct a
+`SPIRVConversionTarget`.
+
 ### `SPIRVTypeConverter`
 
 The `mlir::SPIRVTypeConverter` derives from `mlir::TypeConverter` and provides

diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
index 1d3964a..63f7e26 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
+++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h
@@ -174,9 +174,9 @@ StringRef getTargetEnvAttrName();
 /// and no extra extensions.
 TargetEnvAttr getDefaultTargetEnv(MLIRContext *context);
 
-/// Queries the target environment from the given `op` or returns the default
-/// target environment (SPIR-V 1.0 with Shader capability and no extra
-/// extensions) if not provided.
+/// Queries the target environment recursively from enclosing symbol table ops
+/// containing the given `op` or returns the default target environment as
+/// returned by getDefaultTargetEnv() if not provided.
 TargetEnvAttr lookupTargetEnvOrDefault(Operation *op);
 
 } // namespace spirv

diff --git a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
index 3743cf4..1bd8729 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
@@ -11,6 +11,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
@@ -294,8 +295,18 @@ spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
-  if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
-          spirv::getTargetEnvAttrName()))
-    return attr;
+  Operation *symTable = op;
+  while (symTable) {
+    symTable = SymbolTable::getNearestSymbolTable(symTable);
+    if (!symTable)
+      break;
+
+    if (auto attr = symTable->getAttrOfType<spirv::TargetEnvAttr>(
+            spirv::getTargetEnvAttrName()))
+      return attr;
+
+    symTable = symTable->getParentOp();
+  }
+
   return getDefaultTargetEnv(op->getContext());
 }