[mlir][spirv] Implement gpu::TargetAttrInterface (#69949)

This commit implements gpu::TargetAttrInterface for SPIR-V target
attribute. The plan is to use this to enable GPU compilation pipeline
for OpenCL kernels later.

The changes do not impact Vulkan shaders using milr-vulkan-runner.
New GPU Dialect transform pass spirv-attach-target is implemented for
attaching attribute from CLI.

gpu-module-to-binary pass now works with GPU module that has SPIR-V
module with OpenCL kernel functions inside.

Author: silee2

mlir/include/mlir/Dialect/GPU/Transforms/Passes.h

@@ -15,6 +15,7 @@
 
 #include "Utils.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Pass/Pass.h"
 #include <optional>
 

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

@@ -188,4 +188,48 @@ def GpuROCDLAttachTarget: Pass<"rocdl-attach-target", ""> {
   ];
 }
 
+def GpuSPIRVAttachTarget: Pass<"spirv-attach-target", ""> {
+  let summary = "Attaches an SPIR-V target attribute to a GPU Module.";
+  let description = [{
+    This pass searches for all GPU Modules in the immediate regions and attaches
+    an SPIR-V target if the module matches the name specified by the `module` argument.
+
+    Example:
+    ```
+    // Given the following file: in1.mlir:
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @spirv_module_1 {...}
+    // With
+    // mlir-opt --spirv-attach-target="module=spirv.* ver=v1.0 caps=Kernel" in1.mlir
+    // it will generate,
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @spirv_module_1 [#spirv.target<#spirv.vce<v1.0, [Kernel], []>, #spirv.resource_limits<>>] {...}
+    ```
+  }];
+  let options = [
+    Option<"moduleMatcher", "module", "std::string",
+           /*default=*/ [{""}],
+           "Regex used to identify the modules to attach the target to.">,
+    Option<"spirvVersion", "ver", "std::string",
+           /*default=*/ "\"v1.0\"",
+           "SPIR-V Version.">,
+    ListOption<"spirvCapabilities", "caps", "std::string",
+           "List of supported SPIR-V Capabilities">,
+    ListOption<"spirvExtensions", "exts", "std::string",
+           "List of supported SPIR-V Extensions">,
+    Option<"clientApi", "client_api", "std::string",
+           /*default=*/ "\"Unknown\"",
+           "Client API">,
+    Option<"deviceVendor", "vendor", "std::string",
+           /*default=*/ "\"Unknown\"",
+           "Device Vendor">,
+    Option<"deviceType", "device_type", "std::string",
+           /*default=*/ "\"Unknown\"",
+           "Device Type">,
+    Option<"deviceId", "device_id", "uint32_t",
+           /*default=*/ "mlir::spirv::TargetEnvAttr::kUnknownDeviceID",
+           "Device ID">,
+  ];
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.h

@@ -17,6 +17,12 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Support/LLVM.h"
 
+namespace mlir {
+namespace spirv {
+class VerCapExtAttr;
+}
+} // namespace mlir
+
 // Pull in TableGen'erated SPIR-V attribute definitions for target and ABI.
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h.inc"

mlir/include/mlir/InitAllDialects.h

@@ -91,6 +91,7 @@
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Target/LLVM/NVVM/Target.h"
 #include "mlir/Target/LLVM/ROCDL/Target.h"
+#include "mlir/Target/SPIRV/Target.h"
 
 namespace mlir {
 
@@ -175,6 +176,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   vector::registerSubsetOpInterfaceExternalModels(registry);
   NVVM::registerNVVMTargetInterfaceExternalModels(registry);
   ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
+  spirv::registerSPIRVTargetInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.

mlir/include/mlir/Target/SPIRV/Target.h

@@ -0,0 +1,30 @@
+//===- Target.h - MLIR SPIR-V target registration ---------------*- C++ -*-===//
+//
+// 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 provides registration calls for attaching the SPIR-V target interface.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_SPIRV_TARGET_H
+#define MLIR_TARGET_SPIRV_TARGET_H
+
+namespace mlir {
+class DialectRegistry;
+class MLIRContext;
+namespace spirv {
+/// Registers the `TargetAttrInterface` for the `#spirv.target_env` attribute in
+/// the given registry.
+void registerSPIRVTargetInterfaceExternalModels(DialectRegistry &registry);
+
+/// Registers the `TargetAttrInterface` for the `#spirv.target_env` attribute in
+/// the registry associated with the given context.
+void registerSPIRVTargetInterfaceExternalModels(MLIRContext &context);
+} // namespace spirv
+} // namespace mlir
+
+#endif // MLIR_TARGET_SPIRV_TARGET_H

mlir/lib/Dialect/GPU/CMakeLists.txt

@@ -60,6 +60,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
   Transforms/ShuffleRewriter.cpp
+  Transforms/SPIRVAttachTarget.cpp
   Transforms/ROCDLAttachTarget.cpp
 
   ADDITIONAL_HEADER_DIRS
@@ -95,6 +96,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   MLIRPass
   MLIRSCFDialect
   MLIRSideEffectInterfaces
+  MLIRSPIRVTarget
   MLIRSupport
   MLIRROCDLTarget
   MLIRTransformUtils

mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp

@@ -18,6 +18,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
@@ -53,6 +54,7 @@ void GpuModuleToBinaryPass::getDependentDialects(
 #if MLIR_ROCM_CONVERSIONS_ENABLED == 1
   registry.insert<ROCDL::ROCDLDialect>();
 #endif
+  registry.insert<spirv::SPIRVDialect>();
 }
 
 void GpuModuleToBinaryPass::runOnOperation() {

mlir/lib/Dialect/GPU/Transforms/SPIRVAttachTarget.cpp

@@ -0,0 +1,95 @@
+//===- SPIRVAttachTarget.cpp - Attach an SPIR-V target --------------------===//
+//
+// 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 the `GPUSPIRVAttachTarget` pass, attaching
+// `#spirv.target_env` attributes to GPU modules.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Target/SPIRV/Target.h"
+#include "llvm/Support/Regex.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_GPUSPIRVATTACHTARGET
+#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::spirv;
+
+namespace {
+struct SPIRVAttachTarget
+    : public impl::GpuSPIRVAttachTargetBase<SPIRVAttachTarget> {
+  using Base::Base;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<spirv::SPIRVDialect>();
+  }
+};
+} // namespace
+
+void SPIRVAttachTarget::runOnOperation() {
+  OpBuilder builder(&getContext());
+  auto versionSymbol = symbolizeVersion(spirvVersion);
+  if (!versionSymbol)
+    return signalPassFailure();
+  auto apiSymbol = symbolizeClientAPI(clientApi);
+  if (!apiSymbol)
+    return signalPassFailure();
+  auto vendorSymbol = symbolizeVendor(deviceVendor);
+  if (!vendorSymbol)
+    return signalPassFailure();
+  auto deviceTypeSymbol = symbolizeDeviceType(deviceType);
+  if (!deviceTypeSymbol)
+    return signalPassFailure();
+
+  Version version = versionSymbol.value();
+  SmallVector<Capability, 4> capabilities;
+  SmallVector<Extension, 8> extensions;
+  for (auto cap : spirvCapabilities) {
+    auto capSymbol = symbolizeCapability(cap);
+    if (capSymbol)
+      capabilities.push_back(capSymbol.value());
+  }
+  ArrayRef<Capability> caps(capabilities);
+  for (auto ext : spirvExtensions) {
+    auto extSymbol = symbolizeExtension(ext);
+    if (extSymbol)
+      extensions.push_back(extSymbol.value());
+  }
+  ArrayRef<Extension> exts(extensions);
+  VerCapExtAttr vce = VerCapExtAttr::get(version, caps, exts, &getContext());
+  auto target = TargetEnvAttr::get(vce, getDefaultResourceLimits(&getContext()),
+                                   apiSymbol.value(), vendorSymbol.value(),
+                                   deviceTypeSymbol.value(), deviceId);
+  llvm::Regex matcher(moduleMatcher);
+  getOperation()->walk([&](gpu::GPUModuleOp gpuModule) {
+    // Check if the name of the module matches.
+    if (!moduleMatcher.empty() && !matcher.match(gpuModule.getName()))
+      return;
+    // Create the target array.
+    SmallVector<Attribute> targets;
+    if (std::optional<ArrayAttr> attrs = gpuModule.getTargets())
+      targets.append(attrs->getValue().begin(), attrs->getValue().end());
+    targets.push_back(target);
+    // Remove any duplicate targets.
+    targets.erase(std::unique(targets.begin(), targets.end()), targets.end());
+    // Update the target attribute array.
+    gpuModule.setTargetsAttr(builder.getArrayAttr(targets));
+  });
+}

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

@@ -14,6 +14,7 @@
 
 #include "SPIRVParsingUtils.h"
 
+#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
@@ -133,6 +134,7 @@ void SPIRVDialect::initialize() {
 
   // Allow unknown operations because SPIR-V is extensible.
   allowUnknownOperations();
+  declarePromisedInterface<TargetEnvAttr, gpu::TargetAttrInterface>();
 }
 
 std::string SPIRVDialect::getAttributeName(Decoration decoration) {

mlir/lib/Target/SPIRV/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(Serialization)
 set(LLVM_OPTIONAL_SOURCES
   SPIRVBinaryUtils.cpp
   TranslateRegistration.cpp
+  Target.cpp
   )
 
 add_mlir_translation_library(MLIRSPIRVBinaryUtils
@@ -26,3 +27,15 @@ add_mlir_translation_library(MLIRSPIRVTranslateRegistration
   MLIRSupport
   MLIRTranslateLib
   )
+
+add_mlir_dialect_library(MLIRSPIRVTarget
+  Target.cpp
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRSPIRVDialect
+  MLIRSPIRVSerialization
+  MLIRSPIRVDeserialization
+  MLIRSupport
+  MLIRTranslateLib
+  )

mlir/lib/Target/SPIRV/Target.cpp

@@ -0,0 +1,114 @@
+//===- Target.cpp - MLIR SPIR-V target compilation --------------*- C++ -*-===//
+//
+// 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 files defines SPIR-V target related functions including registration
+// calls for the `#spirv.target_env` compilation attribute.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/SPIRV/Target.h"
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Export.h"
+#include "mlir/Target/SPIRV/Serialization.h"
+
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
+#include "llvm/Support/Program.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include <cstdlib>
+#include <cstring>
+
+using namespace mlir;
+using namespace mlir::spirv;
+
+namespace {
+// SPIR-V implementation of the gpu:TargetAttrInterface.
+class SPIRVTargetAttrImpl
+    : public gpu::TargetAttrInterface::FallbackModel<SPIRVTargetAttrImpl> {
+public:
+  std::optional<SmallVector<char, 0>>
+  serializeToObject(Attribute attribute, Operation *module,
+                    const gpu::TargetOptions &options) const;
+
+  Attribute createObject(Attribute attribute,
+                         const SmallVector<char, 0> &object,
+                         const gpu::TargetOptions &options) const;
+};
+} // namespace
+
+// Register the SPIR-V dialect, the SPIR-V translation & the target interface.
+void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, spirv::SPIRVDialect *dialect) {
+    spirv::TargetEnvAttr::attachInterface<SPIRVTargetAttrImpl>(*ctx);
+  });
+}
+
+void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
+    MLIRContext &context) {
+  DialectRegistry registry;
+  registerSPIRVTargetInterfaceExternalModels(registry);
+  context.appendDialectRegistry(registry);
+}
+
+// Reuse from existing serializer
+std::optional<SmallVector<char, 0>> SPIRVTargetAttrImpl::serializeToObject(
+    Attribute attribute, Operation *module,
+    const gpu::TargetOptions &options) const {
+  if (!module)
+    return std::nullopt;
+  auto gpuMod = dyn_cast<gpu::GPUModuleOp>(module);
+  if (!gpuMod) {
+    module->emitError("expected to be a gpu.module op");
+    return std::nullopt;
+  }
+  auto spvMods = gpuMod.getOps<spirv::ModuleOp>();
+  if (spvMods.empty())
+    return std::nullopt;
+
+  auto spvMod = *spvMods.begin();
+  llvm::SmallVector<uint32_t, 0> spvBinary;
+
+  spvBinary.clear();
+  // Serialize the spirv.module op to SPIR-V blob.
+  if (mlir::failed(spirv::serialize(spvMod, spvBinary))) {
+    spvMod.emitError() << "failed to serialize SPIR-V module";
+    return std::nullopt;
+  }
+
+  SmallVector<char, 0> spvData(spvBinary.size() * sizeof(uint32_t), 0);
+  std::memcpy(spvData.data(), spvBinary.data(), spvData.size());
+
+  spvMod.erase();
+  return spvData;
+}
+
+// Prepare Attribute for gpu.binary with serialized kernel object
+Attribute
+SPIRVTargetAttrImpl::createObject(Attribute attribute,
+                                  const SmallVector<char, 0> &object,
+                                  const gpu::TargetOptions &options) const {
+  gpu::CompilationTarget format = options.getCompilationTarget();
+  DictionaryAttr objectProps;
+  Builder builder(attribute.getContext());
+  return builder.getAttr<gpu::ObjectAttr>(
+      attribute, format,
+      builder.getStringAttr(StringRef(object.data(), object.size())),
+      objectProps);
+}