[mlir][spirv] Let SPIRVConversionTarget consider type availability

[antiagainst]

Summary

Previously we only consider the version/extension/capability requirement
on the op itself. This commit updates SPIRVConversionTarget to also
take into consideration the values' types when deciding op legality.

Depends On D75875

[antiagainst]

Updates

[mravishankar]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

364

Could we add a method that sets some targetEnv by default (the one that are used in the tests here for example).

[mravishankar]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

364

Just to add to this. Then we can remove the "lookupTargetEnvOrDefault" method right?

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

364

I'd like to keep the responsibility separate. We already have three helper functions for getting the target env: getDefaultTargetEnv, lookupTargetEnv, and lookupTargetEnvOrDefault. This class should be focusing on conversion target. If one want to just use the default target env, one can use getDefaultTargetEnv to create a SPIRVConversionTarget.

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

-//===- SPIRVLowering.cpp - Standard to SPIR-V dialect conversion--===//
+//===- SPIRVLowering.cpp - SPIR-V lowering 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities used to lower to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/SPIRVLowering.h"
 #include "mlir/Dialect/SPIRV/LayoutUtils.h"
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #include <functional>
 
 #define DEBUG_TYPE "mlir-spirv-lowering"
 
 using namespace mlir;
 
mlir::spirv::setABIAttrs(spirv::FuncOp funcOp,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V ConversionTarget
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<spirv::SPIRVConversionTarget>
 spirv::SPIRVConversionTarget::get(spirv::TargetEnvAttr targetEnv,
                                   MLIRContext *context) {
   std::unique_ptr<SPIRVConversionTarget> target(
       // std::make_unique does not work here because the constructor is private.
       new SPIRVConversionTarget(targetEnv, context));
   SPIRVConversionTarget *targetPtr = target.get();
   target->addDynamicallyLegalDialect<SPIRVDialect>(
       Optional<ConversionTarget::DynamicLegalityCallbackFn>(
           // We need to capture the raw pointer here because it is stable:
for (spirv::Capability cap : targetEnv.getCapabilities()) {
     givenCapabilities.insert(cap);
 
     // Add capabilities implied by the current capability.
     for (spirv::Capability c : spirv::getRecursiveImpliedCapabilities(cap))
       givenCapabilities.insert(c);
   }
 }
 
+/// Checks that `candidates` extension requirements are possible to be satisfied
+/// with the given `allowedExtensions`.
+///
+///  `candidates` is a vector of vector for extension requirements following
+/// ((Extension::A OR Extension::B) AND (Extension::C OR Extension::D))
+/// convention.
+static LogicalResult checkExtensionRequirements(
+    Operation *op, const llvm::SmallSet<spirv::Extension, 4> &allowedExtensions,
+    const spirv::SPIRVType::ExtensionArrayRefVector &candidates) {
+  for (const auto &ors : candidates) {
+    auto chosen = llvm::find_if(ors, [&](spirv::Extension ext) {
+      return allowedExtensions.count(ext);
+    });
+
+    if (chosen == ors.end()) {
+      SmallVector<StringRef, 4> extStrings;
+      for (spirv::Extension ext : ors)
+        extStrings.push_back(spirv::stringifyExtension(ext));
+
+      LLVM_DEBUG(llvm::dbgs() << op->getName()
+                              << "illegal: requires at least one extension in ["
+                              << llvm::join(extStrings, ", ")
+                              << "] but none allowed in target environment\n");
+      return failure();
+    }
+  }
+  return success();
+}
+
+/// Checks that `candidates`capability requirements are possible to be satisfied
+/// with the given `allowedCapabilities`.
+///
+///  `candidates` is a vector of vector for capability requirements following
+/// ((Capability::A OR Capability::B) AND (Capability::C OR Capability::D))
+/// convention.
+static LogicalResult checkCapabilityRequirements(
+    Operation *op,
+    const llvm::SmallSet<spirv::Capability, 8> &allowedCapabilities,
+    const spirv::SPIRVType::CapabilityArrayRefVector &candidates) {
+  for (const auto &ors : candidates) {
+    auto chosen = llvm::find_if(ors, [&](spirv::Capability cap) {
+      return allowedCapabilities.count(cap);
+    });
+
+    if (chosen == ors.end()) {
+      SmallVector<StringRef, 4> capStrings;
+      for (spirv::Capability cap : ors)
+        capStrings.push_back(spirv::stringifyCapability(cap));
+
+      LLVM_DEBUG(llvm::dbgs()
+                 << op->getName()
+                 << "illegal: requires at least one capability in ["
+                 << llvm::join(capStrings, ", ")
+                 << "] but none allowed in target environment\n");
+      return failure();
+    }
+  }
+  return success();
+}
+
 bool spirv::SPIRVConversionTarget::isLegalOp(Operation *op) {
   // Make sure this op is available at the given version. Ops not implementing
   // QueryMinVersionInterface/QueryMaxVersionInterface are available to all
   // SPIR-V versions.
   if (auto minVersion = dyn_cast<spirv::QueryMinVersionInterface>(op))
     if (minVersion.getMinVersion() > givenVersion) {
       LLVM_DEBUG(llvm::dbgs()
                  << op->getName() << " illegal: requiring min version "
                  << spirv::stringifyVersion(minVersion.getMinVersion())
                  << "\n");
       return false;
     }
   if (auto maxVersion = dyn_cast<spirv::QueryMaxVersionInterface>(op))
     if (maxVersion.getMaxVersion() < givenVersion) {
       LLVM_DEBUG(llvm::dbgs()
                  << op->getName() << " illegal: requiring max version "
                  << spirv::stringifyVersion(maxVersion.getMaxVersion())
                  << "\n");
       return false;
     }
 
-  // Make sure this op's required extensions are allowed to use. For each op,
-  // we return a vector of vector for its extension requirements following
-  // ((Extension::A OR Extension::B) AND (Extension::C OR Extension::D))
-  // convention. Ops not implementing QueryExtensionInterface do not require
-  // extensions to be available.
-  if (auto extensions = dyn_cast<spirv::QueryExtensionInterface>(op)) {
-    auto exts = extensions.getExtensions();
-    for (const auto &ors : exts)
-      if (llvm::all_of(ors, [this](spirv::Extension ext) {
-            return this->givenExtensions.count(ext) == 0;
-          })) {
-        LLVM_DEBUG(llvm::dbgs() << op->getName()
-                                << " illegal: missing required extension\n");
-        return false;
-      }
-  }
+  // Make sure this op's required extensions are allowed to use. Ops not
+  // implementing QueryExtensionInterface do not require extensions to be
+  // available.
+  if (auto extensions = dyn_cast<spirv::QueryExtensionInterface>(op))
+    if (failed(checkExtensionRequirements(op, this->givenExtensions,
+                                          extensions.getExtensions())))
 
-  // Make sure this op's required extensions are allowed to use. For each op,
-  // we return a vector of vector for its capability requirements following
-  // ((Capability::A OR Extension::B) AND (Capability::C OR Capability::D))
-  // convention. Ops not implementing QueryExtensionInterface do not require
-  // extensions to be available.
-  if (auto capabilities = dyn_cast<spirv::QueryCapabilityInterface>(op)) {
-    auto caps = capabilities.getCapabilities();
-    for (const auto &ors : caps)
-      if (llvm::all_of(ors, [this](spirv::Capability cap) {
-            return this->givenCapabilities.count(cap) == 0;
-          })) {
-        LLVM_DEBUG(llvm::dbgs() << op->getName()
-                                << " illegal: missing required capability\n");
+  // Make sure this op's required extensions are allowed to use. Ops not
+  // implementing QueryCapabilityInterface do not require capabilities to be
+  // available.
+  if (auto capabilities = dyn_cast<spirv::QueryCapabilityInterface>(op))
+    if (failed(checkCapabilityRequirements(op, this->givenCapabilities,
+                                           capabilities.getCapabilities())))
+      return false;
+
+    // Special treatment for global variables, whose type requirements are
+    // conveyed by type attributes.
+    SmallVector<Type, 1> attrTypes;
+    if (auto globalVar = dyn_cast<spirv::GlobalVariableOp>(op))
+      attrTypes.push_back(globalVar.type());
+
+    // Make sure the op's operands/results use types that are allowed by the
+    // target environment.
+    SmallVector<ArrayRef<spirv::Extension>, 4> typeExtensions;
+    SmallVector<ArrayRef<spirv::Capability>, 8> typeCapabilities;
+    for (Type valueType : llvm::concat<const Type>(
+             op->getOperandTypes(), op->getResultTypes(), attrTypes)) {
+      typeExtensions.clear();
+      valueType.cast<spirv::SPIRVType>().getExtensions(typeExtensions);
+      if (failed(checkExtensionRequirements(op, this->givenExtensions,
+                                            typeExtensions)))
+        return false;
+
+      typeCapabilities.clear();
+      valueType.cast<spirv::SPIRVType>().getCapabilities(typeCapabilities);
+      if (failed(checkCapabilityRequirements(op, this->givenCapabilities,
+                                             typeCapabilities)))
         return false;
-      }
     }
 
   return true;
 }

mlir/test/Conversion/GPUToSPIRV/if.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
-module attributes {gpu.container_module} {
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
   func @main(%arg0 : memref<10xf32>, %arg1 : i1) {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = "kernel_simple_selection", kernel_module = @kernels} : (index, index, index, index, index, index, memref<10xf32>, i1) -> ()
     return
   }
 
   gpu.module @kernels {
     // CHECK-LABEL: @kernel_simple_selection

mlir/test/Conversion/GPUToSPIRV/load-store.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
-module attributes {gpu.container_module} {
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
   func @load_store(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>) {
     %c0 = constant 0 : index
     %c12 = constant 12 : index
     %0 = subi %c12, %c0 : index
     %c1 = constant 1 : index
     %c0_0 = constant 0 : index
     %c4 = constant 4 : index
     %1 = subi %c4, %c0_0 : index

mlir/test/Conversion/GPUToSPIRV/loop.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
-module attributes {gpu.container_module} {
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
   func @loop(%arg0 : memref<10xf32>, %arg1 : memref<10xf32>) {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = "loop_kernel", kernel_module = @kernels} : (index, index, index, index, index, index, memref<10xf32>, memref<10xf32>) -> ()
     return
   }
 
   gpu.module @kernels {
     gpu.func @loop_kernel(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>)

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

 // RUN: mlir-opt -convert-std-to-spirv %s -o - | FileCheck %s
 
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader, Int64, Float64], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+
 //===----------------------------------------------------------------------===//
 // std binary arithmetic ops
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @add_sub
 func @add_sub(%arg0 : i32, %arg1 : i32) {
   // CHECK: spv.IAdd
   %0 = addi %arg0, %arg1 : i32
func @load_store_zero_rank_int(%arg0: memref<i32>, %arg1: memref<i32>) {
   //      CHECK: [[ZERO2:%.*]] = spv.constant 0 : i32
   //      CHECK: spv.AccessChain [[ARG1]][
   // CHECK-SAME: [[ZERO2]], [[ZERO2]]
   // CHECK-SAME: ] :
   //      CHECK: spv.Store "StorageBuffer" %{{.*}} : i32
   store %0, %arg1[] : memref<i32>
   return
 }
+
+} // end module

mlir/test/Conversion/StandardToSPIRV/subview-to-spirv.mlir

 // RUN: mlir-opt -legalize-std-for-spirv -convert-std-to-spirv %s -o - | FileCheck %s
 
 // TODO: For these examples running these passes separately produces
 // the desired output. Adding all of patterns within a single pass does
 // not seem to work.
 
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+
 //===----------------------------------------------------------------------===//
 // std.subview
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @fold_static_stride_subview_with_load
 // CHECK-SAME: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<384 x f32 [4]> [0]>, StorageBuffer>, [[ARG1:%.*]]: i32, [[ARG2:%.*]]: i32, [[ARG3:%.*]]: i32, [[ARG4:%.*]]: i32
 func @fold_static_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) {
   // CHECK: [[C2:%.*]] = spv.constant 2
func @fold_static_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {
   // CHECK: [[T10:%.*]] = spv.IAdd [[T7]], [[T9]]
   // CHECK: [[C0:%.*]] = spv.constant 0
   // CHECK: [[T12:%.*]] = spv.AccessChain [[ARG0]]{{\[}}[[C0]], [[T10]]
   // CHECK: spv.Store "StorageBuffer" [[T12]], [[ARG5]] : f32
   %0 = subview %arg0[%arg1, %arg2][][] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
   store %arg5, %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [64, 3]>
   return
 }
+
+} // end module