76 lines · plain
1// RUN: mlir-opt -convert-amdgpu-to-rocdl --canonicalize %s | FileCheck %s2 3// CHECK-LABEL: func @test_swizzle_i324// CHECK-SAME: (%[[ARG0:.*]]: i32)5func.func @test_swizzle_i32(%arg0 : i32) -> i32 {6// CHECK: %[[C:.*]] = llvm.mlir.constant(4161 : i32) : i327// CHECK: %[[RES:.*]] = rocdl.ds_swizzle %[[ARG0]], %[[C]] : (i32, i32) -> i328// CHECK: return %[[RES]] : i329 %0 = amdgpu.swizzle_bitmode %arg0 1 2 4 : i3210 return %0 : i3211}12 13// CHECK-LABEL: func @test_swizzle_f3214// CHECK-SAME: (%[[ARG0:.*]]: f32)15func.func @test_swizzle_f32(%arg0 : f32) -> f32 {16// CHECK: %[[C:.*]] = llvm.mlir.constant(4161 : i32) : i3217// CHECK: %[[CAST:.*]] = llvm.bitcast %[[ARG0]] : f32 to i3218// CHECK: %[[RES:.*]] = rocdl.ds_swizzle %[[CAST]], %[[C]] : (i32, i32) -> i3219// CHECK: %[[RES_CAST:.*]] = llvm.bitcast %[[RES]] : i32 to f3220// CHECK: return %[[RES_CAST]] : f3221 %0 = amdgpu.swizzle_bitmode %arg0 1 2 4 : f3222 return %0 : f3223}24 25// CHECK-LABEL: func @test_swizzle_f1626// CHECK-SAME: (%[[ARG0:.*]]: f16)27func.func @test_swizzle_f16(%arg0 : f16) -> f16 {28// CHECK: %[[C:.*]] = llvm.mlir.constant(4161 : i32) : i3229// CHECK: %[[CAST:.*]] = llvm.bitcast %[[ARG0]] : f16 to i1630// CHECK: %[[ZEXT:.*]] = llvm.zext %[[CAST]] : i16 to i3231// CHECK: %[[RES:.*]] = rocdl.ds_swizzle %[[ZEXT]], %[[C]] : (i32, i32) -> i3232// CHECK: %[[TRUNC:.*]] = llvm.trunc %[[RES]] : i32 to i1633// CHECK: %[[RES_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i16 to f1634// CHECK: return %[[RES_CAST]] : f1635 %0 = amdgpu.swizzle_bitmode %arg0 1 2 4 : f1636 return %0 : f1637}38 39// CHECK-LABEL: func @test_swizzle_2xi3240// CHECK-SAME: (%[[ARG0:.*]]: vector<2xi32>)41func.func @test_swizzle_2xi32(%arg0 : vector<2xi32>) -> vector<2xi32> {42// CHECK-DAG: %[[V1:.*]] = llvm.mlir.poison : vector<2xi32>43// CHECK-DAG: %[[C:.*]] = llvm.mlir.constant(4161 : i32) : i3244// CHECK-DAG: %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i3245// CHECK-DAG: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i3246// CHECK: %[[E0:.*]] = llvm.extractelement %[[ARG0]][%[[C0]] : i32] : vector<2xi32>47// CHECK: %[[E1:.*]] = llvm.extractelement %[[ARG0]][%[[C1]] : i32] : vector<2xi32>48// CHECK: %[[S1:.*]] = rocdl.ds_swizzle %[[E0]], %[[C]] : (i32, i32) -> i3249// CHECK: %[[S2:.*]] = rocdl.ds_swizzle %[[E1]], %[[C]] : (i32, i32) -> i3250// CHECK: %[[V2:.*]] = llvm.insertelement %[[S1]], %[[V1]][%[[C0]] : i32] : vector<2xi32>51// CHECK: %[[V3:.*]] = llvm.insertelement %[[S2]], %[[V2]][%[[C1]] : i32] : vector<2xi32>52// CHECK: return %[[V3]] : vector<2xi32>53 %0 = amdgpu.swizzle_bitmode %arg0 1 2 4 : vector<2xi32>54 return %0 : vector<2xi32>55}56 57// CHECK-LABEL: func @test_swizzle_4xf1658// CHECK-SAME: (%[[ARG0:.*]]: vector<4xf16>)59func.func @test_swizzle_4xf16(%arg0 : vector<4xf16>) -> vector<4xf16> {60// CHECK-DAG: %[[V1:.*]] = llvm.mlir.poison : vector<2xi32>61// CHECK-DAG: %[[C:.*]] = llvm.mlir.constant(4161 : i32) : i3262// CHECK-DAG: %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i3263// CHECK-DAG: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i3264// CHECK: %[[CAST1:.*]] = llvm.bitcast %[[ARG0]] : vector<4xf16> to vector<2xi32>65// CHECK: %[[E0:.*]] = llvm.extractelement %[[CAST1]][%[[C0]] : i32] : vector<2xi32>66// CHECK: %[[E1:.*]] = llvm.extractelement %[[CAST1]][%[[C1]] : i32] : vector<2xi32>67// CHECK: %[[S1:.*]] = rocdl.ds_swizzle %[[E0]], %[[C]] : (i32, i32) -> i3268// CHECK: %[[S2:.*]] = rocdl.ds_swizzle %[[E1]], %[[C]] : (i32, i32) -> i3269// CHECK: %[[V2:.*]] = llvm.insertelement %[[S1]], %[[V1]][%[[C0]] : i32] : vector<2xi32>70// CHECK: %[[V3:.*]] = llvm.insertelement %[[S2]], %[[V2]][%[[C1]] : i32] : vector<2xi32>71// CHECK: %[[CAST2:.*]] = llvm.bitcast %[[V3]] : vector<2xi32> to vector<4xf16>72// CHECK: return %[[CAST2]] : vector<4xf16>73 %0 = amdgpu.swizzle_bitmode %arg0 1 2 4 : vector<4xf16>74 return %0 : vector<4xf16>75}76