[Mlir-commits] [mlir] df47f3e - [mlir][spirv] Add lowering for gpu shuffle idx

[Thomas Raoux]

Author: Thomas Raoux
Date: 2022-11-28T22:17:19Z
New Revision: df47f3ea0dcb5c483db67d4a620cf6cd8fe3badb

URL: https://github.com/llvm/llvm-project/commit/df47f3ea0dcb5c483db67d4a620cf6cd8fe3badb
DIFF: https://github.com/llvm/llvm-project/commit/df47f3ea0dcb5c483db67d4a620cf6cd8fe3badb.diff

LOG: [mlir][spirv] Add lowering for gpu shuffle idx

Differential Revision: https://reviews.llvm.org/D138863

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
    mlir/test/Conversion/GPUToSPIRV/shuffle.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index b6c64f3df7a3c..e78d7e87b6ef8 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -414,6 +414,10 @@ LogicalResult GPUShuffleConversion::matchAndRewrite(
     result = rewriter.create<spirv::GroupNonUniformShuffleXorOp>(
         loc, scope, adaptor.getValue(), adaptor.getOffset());
     break;
+  case gpu::ShuffleMode::IDX:
+    result = rewriter.create<spirv::GroupNonUniformShuffleOp>(
+        loc, scope, adaptor.getValue(), adaptor.getOffset());
+    break;
   default:
     return rewriter.notifyMatchFailure(shuffleOp, "unimplemented shuffle mode");
   }

diff  --git a/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir b/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
index 718539626c29d..21858844673de 100644
--- a/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/shuffle.mlir
@@ -46,3 +46,29 @@ gpu.module @kernels {
 }
 
 }
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, GroupNonUniformShuffle], []>, #spirv.resource_limits<subgroup_size = 16>>
+} {
+
+gpu.module @kernels {
+  // CHECK-LABEL:  spirv.func @shuffle_idx()
+  gpu.func @shuffle_idx() kernel
+    attributes {spirv.entry_point_abi = #spirv.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:.+]] = spirv.Constant 8 : i32
+    // CHECK: %[[VAL:.+]] = spirv.Constant 4.200000e+01 : f32
+    // CHECK: %{{.+}} = spirv.Constant true
+    // CHECK: %{{.+}} = spirv.GroupNonUniformShuffle <Subgroup> %[[VAL]], %[[MASK]] : f32, i32
+    %result, %valid = gpu.shuffle idx %val, %mask, %width : f32
+    gpu.return
+  }
+}
+
+}