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

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

LogicalResult GPUShuffleConversion::matchAndRewrite(
   auto scope = rewriter.getAttr<spirv::ScopeAttr>(spirv::Scope::Subgroup);
   Value result;
 
   switch (shuffleOp.getMode()) {
   case gpu::ShuffleMode::XOR:
     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");
   }
 
   rewriter.replaceOp(shuffleOp, {result, trueVal});
   return success();
 }
 

mlir/test/Conversion/GPUToSPIRV/shuffle.mlir

gpu.func @shuffle_xor() kernel
     // 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
   }
 }
 
 }
+
+// -----
+
+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
+  }
+}
+
+}