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

[antiagainst]

Summary

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.

Depends On D75868

[denis13]

mlir/docs/Dialects/SPIR-V.md

936

Do you mean handy here?

[mravishankar]

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

298–310

Nit: make this a while loop?

[antiagainst]

Address comments

[antiagainst]

mlir/docs/Dialects/SPIR-V.md

936

Oops, yes, good catch!

mlir/docs/Dialects/SPIR-V.md

spv.target_env = #spv.target_env<
     #spv.vce<v1.3, [Shader, GroupNonUniform], [SPV_KHR_8bit_storage]>,
     {
       max_compute_workgroup_invocations = 128 : i32,
       max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
     }>
 } { ... }
 ```
 
-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)
 
 SPIR-V itself is just expressing computation happening on GPU device. SPIR-V
 programs themselves are not enough for running workloads on GPU; a companion
 host application is needed to manage the resources referenced by SPIR-V programs
 and dispatch the workload. For the Vulkan execution environment, the host
 application will be written using Vulkan API. Unlike CUDA, the SPIR-V program

 ### `SPIRVConversionTarget`
 
 The `mlir::spirv::SPIRVConversionTarget` class derives from the
 `mlir::ConversionTarget` class and serves as a utility to define a conversion
 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
 type conversion for standard types to SPIR-V types:
 
 *   [Standard Integer][MlirIntegerType] -> Standard Integer
 *   [Standard Float][MlirFloatType] -> Standard Float
 *   [Vector Type][MlirVectorType] -> Vector Type

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

 
 /// Returns the attribute name for specifying SPIR-V target environment.
 StringRef getTargetEnvAttrName();
 
 /// Returns the default target environment: SPIR-V 1.0 with Shader capability
 /// 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
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_TARGETANDABI_H

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

 //===- TargetAndABI.cpp - SPIR-V target and ABI utilities -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // DictionaryDict derived attributes
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
   auto triple = spirv::VerCapExtAttr::get(spirv::Version::V_1_0,
                                           {spirv::Capability::Shader},
                                           ArrayRef<Extension>(), context);
   return spirv::TargetEnvAttr::get(triple,
                                    spirv::getDefaultResourceLimits(context));
 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
-  if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
+  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());
 }