rocdl.mlir
11.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
// RUN: mlir-opt %s | FileCheck %s
func @rocdl_special_regs() -> !llvm.i32 {
// CHECK-LABEL: rocdl_special_regs
// CHECK: rocdl.workitem.id.x : !llvm.i32
%0 = rocdl.workitem.id.x : !llvm.i32
// CHECK: rocdl.workitem.id.y : !llvm.i32
%1 = rocdl.workitem.id.y : !llvm.i32
// CHECK: rocdl.workitem.id.z : !llvm.i32
%2 = rocdl.workitem.id.z : !llvm.i32
// CHECK: rocdl.workgroup.id.x : !llvm.i32
%3 = rocdl.workgroup.id.x : !llvm.i32
// CHECK: rocdl.workgroup.id.y : !llvm.i32
%4 = rocdl.workgroup.id.y : !llvm.i32
// CHECK: rocdl.workgroup.id.z : !llvm.i32
%5 = rocdl.workgroup.id.z : !llvm.i32
// CHECK: rocdl.workgroup.dim.x : !llvm.i32
%6 = rocdl.workgroup.dim.x : !llvm.i32
// CHECK: rocdl.workgroup.dim.y : !llvm.i32
%7 = rocdl.workgroup.dim.y : !llvm.i32
// CHECK: rocdl.workgroup.dim.z : !llvm.i32
%8 = rocdl.workgroup.dim.z : !llvm.i32
// CHECK: rocdl.grid.dim.x : !llvm.i32
%9 = rocdl.grid.dim.x : !llvm.i32
// CHECK: rocdl.grid.dim.y : !llvm.i32
%10 = rocdl.grid.dim.y : !llvm.i32
// CHECK: rocdl.grid.dim.z : !llvm.i32
%11 = rocdl.grid.dim.z : !llvm.i32
llvm.return %0 : !llvm.i32
}
func @rocdl.barrier() {
// CHECK: rocdl.barrier
rocdl.barrier
llvm.return
}
func @rocdl.xdlops(%arg0 : !llvm.float, %arg1 : !llvm.float,
%arg2 : !llvm<"<32 x float>">, %arg3 : !llvm.i32,
%arg4 : !llvm<"<16 x float>">, %arg5 : !llvm<"<4 x float>">,
%arg6 : !llvm<"<4 x half>">, %arg7 : !llvm<"<32 x i32>">,
%arg8 : !llvm<"<16 x i32>">, %arg9 : !llvm<"<4 x i32>">,
%arg10 : !llvm<"<2 x i16>">) -> !llvm<"<32 x float>"> {
// CHECK-LABEL: rocdl.xdlops
// CHECK: rocdl.mfma.f32.32x32x1f32 {{.*}} : (!llvm.float, !llvm.float, !llvm<"<32 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
%r0 = rocdl.mfma.f32.32x32x1f32 %arg0, %arg1, %arg2, %arg3, %arg3, %arg3 :
(!llvm.float, !llvm.float, !llvm<"<32 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
// CHECK: rocdl.mfma.f32.16x16x1f32 {{.*}} : (!llvm.float, !llvm.float, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r1 = rocdl.mfma.f32.16x16x1f32 %arg0, %arg1, %arg4, %arg3, %arg3, %arg3 :
(!llvm.float, !llvm.float, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.16x16x4f32 {{.*}} : (!llvm.float, !llvm.float, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r2 = rocdl.mfma.f32.16x16x4f32 %arg0, %arg1, %arg5, %arg3, %arg3, %arg3 :
(!llvm.float, !llvm.float, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
// CHECK: rocdl.mfma.f32.4x4x1f32 {{.*}} : (!llvm.float, !llvm.float, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r3 = rocdl.mfma.f32.4x4x1f32 %arg0, %arg1, %arg5, %arg3, %arg3, %arg3 :
(!llvm.float, !llvm.float, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
// CHECK: rocdl.mfma.f32.32x32x2f32 {{.*}} : (!llvm.float, !llvm.float, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r4= rocdl.mfma.f32.32x32x2f32 %arg0, %arg1, %arg4, %arg3, %arg3, %arg3 :
(!llvm.float, !llvm.float, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.32x32x4f16 {{.*}} : (!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<32 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
%r5 = rocdl.mfma.f32.32x32x4f16 %arg6, %arg6, %arg2, %arg3, %arg3, %arg3 :
(!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<32 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
// CHECK: rocdl.mfma.f32.16x16x4f16 {{.*}} : (!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r6 = rocdl.mfma.f32.16x16x4f16 %arg6, %arg6, %arg4, %arg3, %arg3, %arg3 :
(!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.4x4x4f16 {{.*}} : (!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r7 = rocdl.mfma.f32.4x4x4f16 %arg6, %arg6, %arg5, %arg3, %arg3, %arg3 :
(!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
// CHECK: rocdl.mfma.f32.32x32x8f16 {{.*}} : (!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r8 = rocdl.mfma.f32.32x32x8f16 %arg6, %arg6, %arg4, %arg3, %arg3, %arg3 :
(!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.16x16x16f16 {{.*}} : (!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r9 = rocdl.mfma.f32.16x16x16f16 %arg6, %arg6, %arg5, %arg3, %arg3, %arg3 :
(!llvm<"<4 x half>">, !llvm<"<4 x half>">, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
// CHECK: rocdl.mfma.i32.32x32x4i8 {{.*}} : (!llvm.i32, !llvm.i32, !llvm<"<32 x i32>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x i32>">
%r10 = rocdl.mfma.i32.32x32x4i8 %arg3, %arg3, %arg7, %arg3, %arg3, %arg3 :
(!llvm.i32, !llvm.i32, !llvm<"<32 x i32>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x i32>">
// CHECK: rocdl.mfma.i32.16x16x4i8 {{.*}} : (!llvm.i32, !llvm.i32, !llvm<"<16 x i32>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x i32>">
%r11 = rocdl.mfma.i32.16x16x4i8 %arg3, %arg3, %arg8, %arg3, %arg3, %arg3 :
(!llvm.i32, !llvm.i32, !llvm<"<16 x i32>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x i32>">
// CHECK: rocdl.mfma.i32.4x4x4i8 {{.*}} : (!llvm.i32, !llvm.i32, !llvm<"<4 x i32>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x i32>">
%r12 = rocdl.mfma.i32.4x4x4i8 %arg3, %arg3, %arg9, %arg3, %arg3, %arg3 :
(!llvm.i32, !llvm.i32, !llvm<"<4 x i32>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x i32>">
// CHECK: rocdl.mfma.i32.32x32x8i8 {{.*}} : (!llvm.i32, !llvm.i32, !llvm<"<16 x i32>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x i32>">
%r13 = rocdl.mfma.i32.32x32x8i8 %arg3, %arg3, %arg8, %arg3, %arg3, %arg3 :
(!llvm.i32, !llvm.i32, !llvm<"<16 x i32>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x i32>">
// CHECK: rocdl.mfma.i32.16x16x16i8 {{.*}} : (!llvm.i32, !llvm.i32, !llvm<"<4 x i32>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x i32>">
%r14 = rocdl.mfma.i32.16x16x16i8 %arg3, %arg3, %arg9, %arg3, %arg3, %arg3 :
(!llvm.i32, !llvm.i32, !llvm<"<4 x i32>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x i32>">
// CHECK: rocdl.mfma.f32.32x32x2bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<32 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
%r15 = rocdl.mfma.f32.32x32x2bf16 %arg10, %arg10, %arg2, %arg3, %arg3, %arg3 :
(!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<32 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<32 x float>">
// CHECK: rocdl.mfma.f32.16x16x2bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r16 = rocdl.mfma.f32.16x16x2bf16 %arg10, %arg10, %arg4, %arg3, %arg3, %arg3 :
(!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.4x4x2bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r17 = rocdl.mfma.f32.4x4x2bf16 %arg10, %arg10, %arg5, %arg3, %arg3, %arg3 :
(!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
// CHECK: rocdl.mfma.f32.32x32x4bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<16 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
%r18 = rocdl.mfma.f32.32x32x4bf16 %arg10, %arg10, %arg4, %arg3, %arg3, %arg3 :
(!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<16 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<16 x float>">
// CHECK: rocdl.mfma.f32.16x16x8bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
%r19 = rocdl.mfma.f32.16x16x8bf16 %arg10, %arg10, %arg5, %arg3, %arg3, %arg3 :
(!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">,
!llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
llvm.return %r0 : !llvm<"<32 x float>">
}
llvm.func @rocdl.mubuf(%rsrc : !llvm<"<4 x i32>">, %vindex : !llvm.i32,
%offset : !llvm.i32, %glc : !llvm.i1,
%slc : !llvm.i1, %vdata1 : !llvm<"<1 x float>">,
%vdata2 : !llvm<"<2 x float>">, %vdata4 : !llvm<"<4 x float>">) {
// CHECK-LABEL: rocdl.mubuf
// CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<1 x float>">
%r1 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
// CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<2 x float>">
%r2 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
// CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<4 x float>">
%r4 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
// CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<1 x float>">
rocdl.buffer.store %vdata1, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
// CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<2 x float>">
rocdl.buffer.store %vdata2, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
// CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<4 x float>">
rocdl.buffer.store %vdata4, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
llvm.return
}