[mlir][spirv] Add support for converting gpu.shuffle xor

[kuhar]

LGTM

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

public:
   };
 
   explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                               Options options = {});
 
   /// Gets the SPIR-V correspondence for the standard index type.
   Type getIndexType() const;
 
+  const spirv::TargetEnv &getTargetEnv() const { return targetEnv; }
+
   /// Returns the options controlling the SPIR-V type converter.
-  const Options &getOptions() const;
+  const Options &getOptions() const { return options; }
 
 private:
   spirv::TargetEnv targetEnv;
   Options options;
 
   MLIRContext *getContext() const;
 };
 

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

 //===- GPUToSPIRV.cpp - GPU to SPIR-V Patterns ----------------------------===//
 //
 // 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 patterns to convert GPU dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.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/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 static constexpr const char kSPIRVModule[] = "__spv__";
 
 namespace {
 /// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation

 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(gpu::BarrierOp barrierOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Pattern to convert a gpu.shuffle op into a spv.GroupNonUniformShuffle op.
+class GPUShuffleConversion final : public OpConversionPattern<gpu::ShuffleOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::ShuffleOp shuffleOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Builtins.
 //===----------------------------------------------------------------------===//
 
 template <typename SourceOp, spirv::BuiltIn builtin>
 LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
auto memorySemantics = spirv::MemorySemanticsAttr::get(
       context, spirv::MemorySemantics::WorkgroupMemory |
                    spirv::MemorySemantics::AcquireRelease);
   rewriter.replaceOpWithNewOp<spirv::ControlBarrierOp>(barrierOp, scope, scope,
                                                        memorySemantics);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Shuffle
+//===----------------------------------------------------------------------===//
+
+LogicalResult GPUShuffleConversion::matchAndRewrite(
+    gpu::ShuffleOp shuffleOp, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  // Require the shuffle width to be the same as the target's subgroup size,
+  // given that for SPIR-V non-uniform subgroup ops, we cannot select
+  // participating invocations.
+  auto targetEnv = getTypeConverter<SPIRVTypeConverter>()->getTargetEnv();
+  unsigned subgroupSize =
+      targetEnv.getAttr().getResourceLimits().getSubgroupSize();
+  IntegerAttr widthAttr;
+  if (!matchPattern(shuffleOp.width(), m_Constant(&widthAttr)) ||
+      widthAttr.getValue().getZExtValue() != subgroupSize)
+    return rewriter.notifyMatchFailure(
+        shuffleOp, "shuffle width and target subgroup size mismatch");
+
+  Location loc = shuffleOp.getLoc();
+  Value trueVal = spirv::ConstantOp::getOne(rewriter.getI1Type(),
+                                            shuffleOp.getLoc(), rewriter);
+  auto scope = rewriter.getAttr<spirv::ScopeAttr>(spirv::Scope::Subgroup);
+  Value result;
+
+  switch (shuffleOp.mode()) {
+  case gpu::ShuffleMode::XOR:
+    result = rewriter.create<spirv::GroupNonUniformShuffleXorOp>(
+        loc, scope, adaptor.value(), adaptor.offset());
+    break;
+  default:
+    return rewriter.notifyMatchFailure(shuffleOp, "unimplemented shuffle mode");
+  }
+
+  rewriter.replaceOp(shuffleOp, {result, trueVal});
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // GPU To SPIRV Patterns.
 //===----------------------------------------------------------------------===//
 
 void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                       RewritePatternSet &patterns) {
   patterns.add<
       GPUBarrierConversion, GPUFuncOpConversion, GPUModuleConversion,
-      GPUModuleEndConversion, GPUReturnOpConversion,
+      GPUModuleEndConversion, GPUReturnOpConversion, GPUShuffleConversion,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>,
       LaunchConfigConversion<gpu::GlobalIdOp,
                              spirv::BuiltIn::GlobalInvocationId>,
       SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
                                       spirv::BuiltIn::SubgroupId>,
       SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
                                       spirv::BuiltIn::NumSubgroups>,
       SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
                                       spirv::BuiltIn::SubgroupSize>,
       WorkGroupSizeConversion>(typeConverter, patterns.getContext());
 }

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 Type SPIRVTypeConverter::getIndexType() const {
   return IntegerType::get(getContext(), options.use64bitIndex ? 64 : 32);
 }
 
-const SPIRVTypeConverter::Options &SPIRVTypeConverter::getOptions() const {
-  return options;
-}
-
 MLIRContext *SPIRVTypeConverter::getContext() const {
   return targetEnv.getAttr().getContext();
 }
 
 // TODO: This is a utility function that should probably be exposed by the
 // SPIR-V dialect. Keeping it local till the use case arises.
 static Optional<int64_t>
 getTypeNumBytes(const SPIRVTypeConverter::Options &options, Type type) {

mlir/test/Conversion/GPUToSPIRV/shuffle.mlir

+// RUN: mlir-opt -split-input-file -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
+
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader, GroupNonUniformShuffle], []>, #spv.resource_limits<subgroup_size = 16>>
+} {
+
+gpu.module @kernels {
+  // CHECK-LABEL:  spv.func @shuffle_xor()
+  gpu.func @shuffle_xor() kernel
+    attributes {spv.entry_point_abi = #spv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+    %mask = arith.constant 8 : i32
+    %width = arith.constant 16 : i32
+    %val = arith.constant 42.0 : f32
+
+    // CHECK: %[[MASK:.+]] = spv.Constant 8 : i32
+    // CHECK: %[[VAL:.+]] = spv.Constant 4.200000e+01 : f32
+    // CHECK: %{{.+}} = spv.Constant true
+    // CHECK: %{{.+}} = spv.GroupNonUniformShuffleXor <Subgroup> %[[VAL]], %[[MASK]] : f32, i32
+    %result, %valid = gpu.shuffle xor %val, %mask, %width : f32
+    gpu.return
+  }
+}
+
+}
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader, GroupNonUniformShuffle], []>, #spv.resource_limits<subgroup_size = 32>>
+} {
+
+gpu.module @kernels {
+  gpu.func @shuffle_xor() kernel
+    attributes {spv.entry_point_abi = #spv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>} {
+    %mask = arith.constant 8 : i32
+    %width = arith.constant 16 : i32
+    %val = arith.constant 42.0 : f32
+
+    // Cannot convert due to shuffle width and target subgroup size mismatch
+    // expected-error @+1 {{failed to legalize operation 'gpu.shuffle'}}
+    %result, %valid = gpu.shuffle xor %val, %mask, %width : f32
+    gpu.return
+  }
+}
+
+}