From 66c378d66e93b67c71aaf99974577f2bd72f336e Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst@google.com>
Date: Wed, 11 Mar 2020 16:02:56 -0400
Subject: [PATCH] [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
---
 mlir/docs/Dialects/SPIR-V.md                   | 12 +++++++++---
 mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h |  6 +++---
 mlir/lib/Dialect/SPIRV/TargetAndABI.cpp        | 17 ++++++++++++++---
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/mlir/docs/Dialects/SPIR-V.md b/mlir/docs/Dialects/SPIR-V.md
index 4380c9f006d1..0cb628177750 100644
--- 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 1d3964a67fe5..63f7e267ecff 100644
--- 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 3743cf44348c..1bd8729347ac 100644
--- 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());
 }