Skip to content

Commit 6959c2f

Browse files
matrix72cmatrix72c
committed
[examples] Add conv2d gpu tranform example.
1 parent c3e9d6a commit 6959c2f

File tree

3 files changed

+357
-0
lines changed

3 files changed

+357
-0
lines changed

‎examples/BuddyGPU/conv2d.mlir

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
!input_tensor_t = tensor<1x128x66x66xf32>
2+
!weight_tensor_t = tensor<256x128x3x3xf32>
3+
!output_tensor_t = tensor<1x256x64x64xf32>
4+
5+
func.func @conv_2d_nchw_fchw(%in: !input_tensor_t, %wei: !weight_tensor_t,
6+
%out: !output_tensor_t) -> !output_tensor_t {
7+
%res = linalg.conv_2d_nchw_fchw
8+
{dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
9+
ins(%in, %wei: !input_tensor_t, !weight_tensor_t)
10+
outs(%out: !output_tensor_t) -> !output_tensor_t
11+
return %res : !output_tensor_t
12+
}

‎examples/BuddyGPU/makefile

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,12 @@ buddy-gpu-matmul-lower:
1111
-transform-interpreter="entry-point=codegen" \
1212
-o log.mlir
1313

14+
buddy-gpu-conv2d-lower:
15+
@${BUDDY_OPT} conv2d.mlir \
16+
-transform-preload-library="transform-library-paths=transform-conv2d.mlir" \
17+
-transform-interpreter="entry-point=codegen" \
18+
-o log.mlir
19+
1420
buddy-gpu-matmul:
1521
@${BUDDY_OPT} matmul.mlir -transform-preload-library="transform-library-paths=transform.mlir" -transform-interpreter="entry-point=codegen" | \
1622
${BUDDY_OPT} --pass-pipeline='builtin.module(func.func(nvgpu-optimize-shared-memory))' | \

‎examples/BuddyGPU/transform-conv2d.mlir

Lines changed: 339 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,339 @@
1+
module attributes {transform.with_named_sequence} {
2+
3+
transform.named_sequence @__transform_main(
4+
%arg0: !transform.any_op,
5+
%conv: !transform.op<"linalg.conv_2d_nchw_fchw">) {
6+
7+
transform.debug.emit_remark_at %conv, "Input conv" : !transform.op<"linalg.conv_2d_nchw_fchw">
8+
9+
%conv2, %loops2:5 = transform.structured.tile_using_for %conv
10+
// N, F, OH, OW, C, KH, KW
11+
tile_sizes [1, 64, 1, 32, 16, 0, 0] // 16 canais , 32 colunas, 64 filtros, 2o, 1 tile de uma linha
12+
interchange = [0, 4, 3, 2, 1] // 4 = F, 3: OH, 2:OW, 1:C
13+
: (!transform.op<"linalg.conv_2d_nchw_fchw">)
14+
-> (!transform.op<"linalg.conv_2d_nchw_fchw">, !transform.any_op,
15+
!transform.any_op, !transform.any_op, !transform.any_op,
16+
!transform.any_op)
17+
18+
transform.debug.emit_remark_at %conv2, "conv2" : !transform.op<"linalg.conv_2d_nchw_fchw">
19+
20+
%conv3, %loops3:2 = transform.structured.tile_using_for %conv2
21+
// N, F, OH, OW, C, KH, KW
22+
tile_sizes [0, 8, 0, 16, 0, 0, 0]
23+
interchange = [1, 0]
24+
: (!transform.op<"linalg.conv_2d_nchw_fchw">)
25+
-> (!transform.op<"linalg.conv_2d_nchw_fchw">, !transform.any_op,
26+
!transform.any_op)
27+
28+
transform.debug.emit_remark_at %conv3, "conv3" : !transform.op<"linalg.conv_2d_nchw_fchw">
29+
30+
%conv4, %matmul = transform.structured.convert_conv2d_to_img2col %conv3
31+
: (!transform.op<"linalg.conv_2d_nchw_fchw">) -> (!transform.any_op, !transform.any_op)
32+
33+
transform.debug.emit_remark_at %conv4, "img2col" : !transform.any_op
34+
35+
transform.apply_patterns to %conv4 {
36+
transform.apply_patterns.canonicalization
37+
transform.apply_patterns.linalg.tiling_canonicalization
38+
} : !transform.any_op
39+
40+
// Perform tiling for the grid.
41+
// For the matrix multiplication of 5376x2048 and 2048x5376, the compilation
42+
// strategy sets the tile size for grid-based partitioning to 128x256.
43+
// This means that each [128, 2048] @ [2048, 256] matmul tile is computed within a GPU block,
44+
// while multiple such blocks are computed in parallel across the grid.
45+
// `tile_sizes` specify the dimensions of the tiled matmul result.
46+
// `%tiled_op` is the tiled matmul operation within the `scf.forall` loop.
47+
// `%forall_op` is the `scf.forall` loop that maintains tile information.
48+
%tiled_op, %forall_op = transform.structured.tile_using_forall %conv4
49+
tile_sizes [128, 256] (mapping = [#gpu.block<y>, #gpu.block<x>])
50+
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
51+
52+
// Perform canonicalization.
53+
%1 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
54+
transform.apply_patterns to %1 {
55+
transform.apply_patterns.linalg.tiling_canonicalization
56+
transform.apply_patterns.scf.for_loop_canonicalization
57+
transform.apply_patterns.canonicalization
58+
} : !transform.any_op
59+
transform.apply_cse to %1 : !transform.any_op
60+
%all_loops = transform.structured.match interface{LoopLikeInterface}
61+
in %arg0
62+
: (!transform.any_op) -> !transform.any_op
63+
transform.apply_licm to %all_loops : !transform.any_op
64+
transform.apply_patterns to %1 {
65+
transform.apply_patterns.linalg.tiling_canonicalization
66+
} : !transform.any_op
67+
68+
// Further tile the tiled matmul
69+
// Tile the third dimension in matmul.
70+
// [128, 2048] @ [2048, 256] matmul is further tiled into [128, 16] @ [16, 256] matmul.
71+
%tiled_linalg_op, %loops = transform.structured.tile_using_for %tiled_op [0, 0, 16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
72+
73+
// Create pad op and prepare for mapping to GPU.
74+
// Nothing has changed in the operation.
75+
%padded, %pad, %copy = transform.structured.pad %tiled_linalg_op {copy_back_op = "none", pack_paddings = [1, 1, 1], pad_to_multiple_of = [1, 1, 1], padding_dimensions = [0, 1, 2], padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
76+
77+
// Rewrite tensor.pad into linalg.copy.
78+
%3 = transform.get_producer_of_operand %padded[0] : (!transform.any_op) -> !transform.any_op
79+
%4 = transform.get_producer_of_operand %padded[1] : (!transform.any_op) -> !transform.any_op
80+
%5 = transform.get_producer_of_operand %padded[2] : (!transform.any_op) -> !transform.any_op
81+
%6 = transform.structured.rewrite_in_destination_passing_style %3 : (!transform.any_op) -> !transform.any_op
82+
%7 = transform.structured.rewrite_in_destination_passing_style %4 : (!transform.any_op) -> !transform.any_op
83+
%8 = transform.structured.rewrite_in_destination_passing_style %5 : (!transform.any_op) -> !transform.any_op
84+
85+
// Tile the linalg.copy op and map it to GPU thread level,
86+
// such that the tiled matrix are copied to GPU shared memory.
87+
// num_threads is different from tile_sizes used above,
88+
// as it specifies the number of tile instead of the size of the tile.
89+
// The first transform tile the [128, 16] into [4, 4],
90+
// and the second transform tile the [16, 256] into [2, 16].
91+
%tiled_op_0, %forall_op_1 = transform.structured.tile_using_forall %6 num_threads [32, 4](mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
92+
%tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %7 num_threads [8, 16](mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
93+
94+
// Tile the linalg.matmul op and map it to GPU warp level.
95+
%tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %padded num_threads [2, 2](mapping = [#gpu.warp<y>, #gpu.warp<x>]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
96+
// Tile the linalg.fill op and map it to GPU warp level.
97+
%tiled_op_6, %forall_op_7 = transform.structured.tile_using_forall %fused_op num_threads [2, 2](mapping = [#gpu.warp<y>, #gpu.warp<x>]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
98+
99+
// Perform canonicalization.
100+
%9 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
101+
transform.apply_patterns to %9 {
102+
transform.apply_patterns.linalg.tiling_canonicalization
103+
transform.apply_patterns.scf.for_loop_canonicalization
104+
transform.apply_patterns.canonicalization
105+
} : !transform.any_op
106+
transform.apply_cse to %9 : !transform.any_op
107+
%all_loops_2 = transform.structured.match interface{LoopLikeInterface}
108+
in %9
109+
: (!transform.any_op) -> !transform.any_op
110+
transform.apply_licm to %all_loops_2 : !transform.any_op
111+
transform.apply_patterns to %9 {
112+
transform.apply_patterns.linalg.tiling_canonicalization
113+
transform.apply_patterns.vector.lower_masked_transfers
114+
} : !transform.any_op
115+
116+
// Perform vectorization.
117+
// Vectorize the linalg.copy, linalg.fill, and linalg.matmul operations.
118+
%10 = transform.structured.vectorize_children_and_apply_patterns %9 : (!transform.any_op) -> !transform.any_op
119+
120+
// Perform canonicalization.
121+
transform.apply_patterns to %10 {
122+
transform.apply_patterns.linalg.tiling_canonicalization
123+
transform.apply_patterns.scf.for_loop_canonicalization
124+
transform.apply_patterns.canonicalization
125+
} : !transform.any_op
126+
transform.apply_cse to %10 : !transform.any_op
127+
%all_loops_3 = transform.structured.match interface{LoopLikeInterface}
128+
in %10
129+
: (!transform.any_op) -> !transform.any_op
130+
transform.apply_licm to %all_loops_3 : !transform.any_op
131+
transform.apply_patterns to %10 {
132+
transform.apply_patterns.linalg.tiling_canonicalization
133+
transform.apply_patterns.vector.lower_masked_transfers
134+
} : !transform.any_op
135+
136+
// Match bufferization.alloc_tensors inside the forall op
137+
%scf_forall = transform.structured.match ops{["scf.forall"]} attributes{mapping = [#gpu.block<y>, #gpu.block<x>]} in %arg0 : (!transform.any_op) -> !transform.any_op
138+
%alloc_tensor_ops = transform.structured.match ops{["bufferization.alloc_tensor"]} in %scf_forall : (!transform.any_op) -> !transform.any_op
139+
140+
// Bufferize the alloc_tensor ops to memref.alloc ops.
141+
// The memory_space attribute for GPU Dialect 0 means global memory, 3 means workgroup memory address, 5 means private memory address.
142+
// According to https://discourse.llvm.org/t/rfc-memref-memory-shape-as-attribute/2229
143+
%buffer, %new_ops = transform.structured.bufferize_to_allocation %alloc_tensor_ops {memory_space = 3 } : !transform.any_op
144+
145+
// Eliminate empty tensors and erase unnecessary inputs.
146+
transform.structured.eliminate_empty_tensors %arg0 : !transform.any_op
147+
%func_eras = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
148+
transform.apply_patterns to %func_eras {
149+
transform.apply_patterns.linalg.erase_unnecessary_inputs
150+
} : !transform.any_op
151+
152+
// Bufferize the remaining operations in one time.
153+
%11 = transform.bufferization.one_shot_bufferize %arg0 { bufferize_function_boundaries = true, function_boundary_type_conversion = 1 : i32} : (!transform.any_op) -> !transform.any_op
154+
155+
// Erase dead alloc and stores.
156+
%12 = transform.structured.match ops{["func.func"]} in %11 : (!transform.any_op) -> !transform.any_op
157+
transform.memref.erase_dead_alloc_and_stores %12 : (!transform.any_op) -> ()
158+
159+
// Generate GPU launch.
160+
%13 = transform.structured.match ops{["func.func"]} in %11 : (!transform.any_op) -> !transform.any_op
161+
%gpu_launch = transform.gpu.map_forall_to_blocks %13 { generate_gpu_launch } : (!transform.any_op) -> !transform.any_op
162+
163+
// Rewrite bufferized scf.forall ops to distributed gpu.thread_id attribute.
164+
%mapped = transform.gpu.map_nested_forall_to_threads %gpu_launch block_dims = [64, 2, 1] warp_size = 32 : (!transform.any_op) -> !transform.any_op
165+
166+
%15 = transform.structured.match ops{["func.func"]} in %11 : (!transform.any_op) -> !transform.any_op
167+
168+
// Removes unnecessary GPU barriers from the function.
169+
// %15 = transform.buddy.eliminate_gpu_barriers %14 : (!transform.any_op) -> !transform.any_op
170+
171+
// Perform canonicalization.
172+
transform.apply_patterns to %15 {
173+
transform.apply_patterns.linalg.tiling_canonicalization
174+
transform.apply_patterns.scf.for_loop_canonicalization
175+
transform.apply_patterns.canonicalization
176+
} : !transform.any_op
177+
transform.apply_cse to %15 : !transform.any_op
178+
%all_loops_4 = transform.structured.match interface{LoopLikeInterface}
179+
in %15
180+
: (!transform.any_op) -> !transform.any_op
181+
transform.apply_licm to %all_loops_4 : !transform.any_op
182+
transform.apply_patterns to %15 {
183+
transform.apply_patterns.linalg.tiling_canonicalization
184+
transform.apply_patterns.vector.lower_masked_transfers
185+
} : !transform.any_op
186+
187+
// Identify static memory allocations within the given region,
188+
// and move them to a higher level (hoisting).
189+
transform.buddy.hoist_static_alloc %15 : (!transform.any_op) -> ()
190+
191+
// Collects patterns for folding memref aliasing ops (memref.subview) into consumer load/store ops (affine.load, memref.load, nvgpu.ldmatrix, vector.load, vector.transfer_read, affine.store, memref.store, etc.) and other ops (e.g., memref.subview).
192+
transform.apply_patterns to %15 {
193+
transform.apply_patterns.memref.fold_memref_alias_ops
194+
} : !transform.any_op
195+
// Collects patterns for extracting address computations from operations with memory accesses such that these memory accesses use only a base pointer.
196+
transform.apply_patterns to %15 {
197+
transform.apply_patterns.memref.extract_address_computations
198+
} : !transform.any_op
199+
// Perform canonicalization.
200+
transform.apply_patterns to %15 {
201+
transform.apply_patterns.linalg.tiling_canonicalization
202+
transform.apply_patterns.scf.for_loop_canonicalization
203+
transform.apply_patterns.canonicalization
204+
} : !transform.any_op
205+
transform.apply_cse to %15 : !transform.any_op
206+
%all_loops_5 = transform.structured.match interface{LoopLikeInterface}
207+
in %15
208+
: (!transform.any_op) -> !transform.any_op
209+
transform.apply_licm to %all_loops_5 : !transform.any_op
210+
transform.apply_patterns to %15 {
211+
transform.apply_patterns.linalg.tiling_canonicalization
212+
transform.apply_patterns.vector.lower_masked_transfers
213+
} : !transform.any_op
214+
215+
// Adds patterns that unroll vectors to a native tile size for GPUs with mma operations
216+
transform.apply_patterns to %15 {
217+
transform.apply_patterns.buddy.unroll_vectors_gpu_mma_sync
218+
} : !transform.any_op
219+
220+
// Insert a gpu.barrier after a given scf.for loop
221+
%16 = transform.structured.match ops{["scf.for"]} in %15 : (!transform.any_op) -> !transform.op<"scf.for">
222+
// transform.buddy.synchronize_loop %16 : (!transform.op<"scf.for">) -> ()
223+
224+
225+
transform.apply_patterns to %15 {
226+
transform.apply_patterns.memref.fold_memref_alias_ops
227+
} : !transform.any_op
228+
transform.apply_cse to %15 : !transform.any_op
229+
230+
// Hoist vector.transfer_read / vector.transfer_write pairs out of immediately enclosing scf::ForOp iteratively
231+
// Warning: Deprecated
232+
%17 = transform.structured.hoist_redundant_vector_transfers %15 : (!transform.any_op) -> !transform.any_op
233+
234+
// Perform canonicalization.
235+
transform.apply_patterns to %17 {
236+
transform.apply_patterns.linalg.tiling_canonicalization
237+
transform.apply_patterns.scf.for_loop_canonicalization
238+
transform.apply_patterns.canonicalization
239+
} : !transform.any_op
240+
transform.apply_cse to %17 : !transform.any_op
241+
%all_loops_6 = transform.structured.match interface{LoopLikeInterface}
242+
in %17
243+
: (!transform.any_op) -> !transform.any_op
244+
transform.apply_licm to %all_loops_6 : !transform.any_op
245+
transform.apply_patterns to %17 {
246+
transform.apply_patterns.linalg.tiling_canonicalization
247+
transform.apply_patterns.vector.lower_masked_transfers
248+
} : !transform.any_op
249+
250+
// This converts slices of operations containing vector.contract op into
251+
// mma operations, targetting warp level tensorcore operations.
252+
transform.buddy.vector.vector_to_mma_conversion %17 {use_mma_sync} : (!transform.any_op) -> ()
253+
254+
// %18 = transform.buddy.eliminate_gpu_barriers %17 : (!transform.any_op) -> !transform.any_op
255+
256+
// Perform canonicalization.
257+
transform.apply_patterns to %17 {
258+
transform.apply_patterns.linalg.tiling_canonicalization
259+
transform.apply_patterns.scf.for_loop_canonicalization
260+
transform.apply_patterns.canonicalization
261+
} : !transform.any_op
262+
transform.apply_cse to %17 : !transform.any_op
263+
%all_loops_7 = transform.structured.match interface{LoopLikeInterface}
264+
in %17
265+
: (!transform.any_op) -> !transform.any_op
266+
transform.apply_licm to %all_loops_7 : !transform.any_op
267+
transform.apply_patterns to %17 {
268+
transform.apply_patterns.linalg.tiling_canonicalization
269+
transform.apply_patterns.vector.lower_masked_transfers
270+
} : !transform.any_op
271+
272+
%19 = transform.structured.match ops{["gpu.launch"]} in %17 : (!transform.any_op) -> !transform.any_op
273+
%fwfa = transform.structured.match ops{["memref.alloc"]} in %19 : (!transform.any_op) -> !transform.op<"memref.alloc">
274+
275+
// Do multi-buffering/array expansion to remove dependencies on the temporary allocation between consecutive loop iterations.
276+
transform.memref.multibuffer %fwfa {factor = 3 : i64, skip_analysis} : (!transform.op<"memref.alloc">) -> !transform.any_op
277+
278+
transform.apply_patterns to %17 {
279+
transform.apply_patterns.vector.transfer_to_scf full_unroll = true
280+
} : !transform.any_op
281+
transform.apply_patterns to %17 {
282+
transform.apply_patterns.linalg.tiling_canonicalization
283+
transform.apply_patterns.scf.for_loop_canonicalization
284+
transform.apply_patterns.canonicalization
285+
} : !transform.any_op
286+
transform.apply_cse to %17 : !transform.any_op
287+
%all_loops_8 = transform.structured.match interface{LoopLikeInterface}
288+
in %17
289+
: (!transform.any_op) -> !transform.any_op
290+
transform.apply_licm to %all_loops_8 : !transform.any_op
291+
transform.apply_patterns to %17 {
292+
transform.apply_patterns.linalg.tiling_canonicalization
293+
transform.apply_patterns.vector.lower_masked_transfers
294+
} : !transform.any_op
295+
296+
// Convert sync copies to shared memory to async.
297+
// transform.buddy.create_async_groups %17 {use_mma_sync} : (!transform.any_op) -> ()
298+
transform.apply_patterns to %17 {
299+
transform.apply_patterns.linalg.tiling_canonicalization
300+
transform.apply_patterns.scf.for_loop_canonicalization
301+
transform.apply_patterns.canonicalization
302+
transform.apply_patterns.memref.fold_memref_alias_ops
303+
} : !transform.any_op
304+
%all_loops_9 = transform.structured.match interface{LoopLikeInterface}
305+
in %17
306+
: (!transform.any_op) -> !transform.any_op
307+
transform.apply_licm to %all_loops_9 : !transform.any_op
308+
transform.apply_cse to %17 : !transform.any_op
309+
310+
311+
%20 = transform.structured.match ops{["nvgpu.mma.sync"]} in %17 : (!transform.any_op) -> !transform.any_op
312+
%21 = transform.get_parent_op %20 {deduplicate, op_name = "scf.for"} : (!transform.any_op) -> !transform.any_op
313+
// This applies software pipelining to a given scf.for loop.
314+
// The pipelining strategy will look for a copy to shared memory and pipeline it to overlap it with the rest of the loop.
315+
// %22 = transform.buddy.pipeline_shared_memory_copies %21 {depth = 3 : i64, use_mma_sync, peel_epilogue} : (!transform.any_op) -> !transform.any_op
316+
317+
// Perform canonicalization.
318+
transform.apply_patterns to %17 {
319+
transform.apply_patterns.vector.lower_masks
320+
} : !transform.any_op
321+
transform.apply_patterns to %17 {
322+
transform.apply_patterns.vector.materialize_masks
323+
} : !transform.any_op
324+
transform.apply_patterns to %17 {
325+
transform.apply_patterns.linalg.tiling_canonicalization
326+
transform.apply_patterns.scf.for_loop_canonicalization
327+
transform.apply_patterns.canonicalization
328+
transform.apply_patterns.memref.fold_memref_alias_ops
329+
} : !transform.any_op
330+
331+
%all_loops_10 = transform.structured.match interface{LoopLikeInterface}
332+
in %17
333+
: (!transform.any_op) -> !transform.any_op
334+
transform.apply_licm to %all_loops_10 : !transform.any_op
335+
transform.apply_cse to %17 : !transform.any_op
336+
337+
transform.yield
338+
}
339+
} // module

0 commit comments

Comments
 (0)