-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathoriginal_ttgir.txt
125 lines (125 loc) · 10.6 KB
/
original_ttgir.txt
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
#blocked = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1]}>
#blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [0, 1]}>
#blocked2 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
#blocked3 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [0, 1]}>
#blocked4 = #triton_gpu.blocked<{sizePerThread = [4, 4], threadsPerWarp = [2, 16], warpsPerCTA = [4, 1], order = [1, 0]}>
module attributes {"triton_gpu.num-warps" = 4 : i32} {
tt.func public @matmul_kernel_0d1d2d3d4d5d6d7c8d9c10d11c(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} {
%c31_i32 = arith.constant 31 : i32
%c63_i32 = arith.constant 63 : i32
%c127_i32 = arith.constant 127 : i32
%c0_i32 = arith.constant 0 : i32
%c1_i32 = arith.constant 1 : i32
%cst = arith.constant dense<32> : tensor<128x32xi32, #blocked>
%cst_0 = arith.constant dense<0.000000e+00> : tensor<128x64xf32, #blocked>
%cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32, #blocked1>
%cst_2 = arith.constant dense<0.000000e+00> : tensor<128x32xf32, #blocked>
%c32_i32 = arith.constant 32 : i32
%c64_i32 = arith.constant 64 : i32
%c128_i32 = arith.constant 128 : i32
%c8_i32 = arith.constant 8 : i32
%0 = tt.get_program_id {axis = 0 : i32} : i32
%1 = arith.addi %arg3, %c127_i32 : i32
%2 = arith.divsi %1, %c128_i32 : i32
%3 = arith.addi %arg4, %c63_i32 : i32
%4 = arith.divsi %3, %c64_i32 : i32
%5 = arith.muli %4, %c8_i32 : i32
%6 = arith.divsi %0, %5 : i32
%7 = arith.muli %6, %c8_i32 : i32
%8 = arith.subi %2, %7 : i32
%9 = "triton_gpu.cmpi"(%8, %c8_i32) {predicate = 2 : i64} : (i32, i32) -> i1
%10 = arith.select %9, %8, %c8_i32 : i32
%11 = arith.remsi %0, %10 : i32
%12 = arith.addi %7, %11 : i32
%13 = arith.remsi %0, %5 : i32
%14 = arith.divsi %13, %10 : i32
%15 = arith.muli %12, %c128_i32 : i32
%16 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #blocked2>
%17 = tt.splat %15 : (i32) -> tensor<128xi32, #blocked2>
%18 = arith.addi %17, %16 : tensor<128xi32, #blocked2>
%19 = tt.splat %arg3 : (i32) -> tensor<128xi32, #blocked2>
%20 = arith.remsi %18, %19 : tensor<128xi32, #blocked2>
%21 = arith.muli %14, %c64_i32 : i32
%22 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #blocked2>
%23 = tt.splat %21 : (i32) -> tensor<64xi32, #blocked2>
%24 = arith.addi %23, %22 : tensor<64xi32, #blocked2>
%25 = tt.splat %arg4 : (i32) -> tensor<64xi32, #blocked2>
%26 = arith.remsi %24, %25 : tensor<64xi32, #blocked2>
%27 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #blocked2>
%28 = triton_gpu.convert_layout %20 : (tensor<128xi32, #blocked2>) -> tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>
%29 = tt.expand_dims %28 {axis = 1 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>) -> tensor<128x1xi32, #blocked>
%30 = tt.splat %arg6 : (i32) -> tensor<128x1xi32, #blocked>
%31 = arith.muli %29, %30 : tensor<128x1xi32, #blocked>
%32 = triton_gpu.convert_layout %27 : (tensor<32xi32, #blocked2>) -> tensor<32xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
%33 = tt.expand_dims %32 {axis = 0 : i32} : (tensor<32xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>) -> tensor<1x32xi32, #blocked3>
%34 = tt.broadcast %31 : (tensor<128x1xi32, #blocked>) -> tensor<128x32xi32, #blocked>
%35 = tt.broadcast %33 : (tensor<1x32xi32, #blocked3>) -> tensor<128x32xi32, #blocked3>
%36 = triton_gpu.convert_layout %35 : (tensor<128x32xi32, #blocked3>) -> tensor<128x32xi32, #blocked>
%37 = arith.addi %34, %36 : tensor<128x32xi32, #blocked>
%38 = tt.splat %arg0 : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #blocked>
%39 = tt.addptr %38, %37 : tensor<128x32x!tt.ptr<f16>, #blocked>, tensor<128x32xi32, #blocked>
%40 = triton_gpu.convert_layout %27 : (tensor<32xi32, #blocked2>) -> tensor<32xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>
%41 = tt.expand_dims %40 {axis = 1 : i32} : (tensor<32xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>) -> tensor<32x1xi32, #blocked>
%42 = triton_gpu.convert_layout %41 : (tensor<32x1xi32, #blocked>) -> tensor<32x1xi32, #blocked1>
%43 = tt.splat %arg7 : (i32) -> tensor<32x1xi32, #blocked1>
%44 = arith.muli %42, %43 : tensor<32x1xi32, #blocked1>
%45 = triton_gpu.convert_layout %26 : (tensor<64xi32, #blocked2>) -> tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
%46 = tt.expand_dims %45 {axis = 0 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>) -> tensor<1x64xi32, #blocked3>
%47 = tt.broadcast %44 : (tensor<32x1xi32, #blocked1>) -> tensor<32x64xi32, #blocked1>
%48 = tt.broadcast %46 : (tensor<1x64xi32, #blocked3>) -> tensor<32x64xi32, #blocked3>
%49 = triton_gpu.convert_layout %48 : (tensor<32x64xi32, #blocked3>) -> tensor<32x64xi32, #blocked1>
%50 = arith.addi %47, %49 : tensor<32x64xi32, #blocked1>
%51 = tt.splat %arg1 : (!tt.ptr<f16>) -> tensor<32x64x!tt.ptr<f16>, #blocked1>
%52 = tt.addptr %51, %50 : tensor<32x64x!tt.ptr<f16>, #blocked1>, tensor<32x64xi32, #blocked1>
%53 = arith.addi %arg5, %c31_i32 : i32
%54 = arith.divsi %53, %c32_i32 : i32
%55 = arith.truncf %cst_2 : tensor<128x32xf32, #blocked> to tensor<128x32xf16, #blocked>
%56 = arith.truncf %cst_1 : tensor<32x64xf32, #blocked1> to tensor<32x64xf16, #blocked1>
%57 = arith.muli %arg7, %c32_i32 : i32
%58 = tt.splat %57 : (i32) -> tensor<32x64xi32, #blocked1>
%59:3 = scf.for %arg9 = %c0_i32 to %54 step %c1_i32 iter_args(%arg10 = %cst_0, %arg11 = %39, %arg12 = %52) -> (tensor<128x64xf32, #blocked>, tensor<128x32x!tt.ptr<f16>, #blocked>, tensor<32x64x!tt.ptr<f16>, #blocked1>) : i32 {
%81 = arith.muli %arg9, %c32_i32 : i32
%82 = arith.subi %arg5, %81 : i32
%83 = tt.splat %82 : (i32) -> tensor<1x32xi32, #blocked3>
%84 = "triton_gpu.cmpi"(%33, %83) {predicate = 2 : i64} : (tensor<1x32xi32, #blocked3>, tensor<1x32xi32, #blocked3>) -> tensor<1x32xi1, #blocked3>
%85 = tt.broadcast %84 : (tensor<1x32xi1, #blocked3>) -> tensor<128x32xi1, #blocked3>
%86 = triton_gpu.convert_layout %85 : (tensor<128x32xi1, #blocked3>) -> tensor<128x32xi1, #blocked>
%87 = tt.load %arg11, %86, %55 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x32xf16, #blocked>
%88 = tt.splat %82 : (i32) -> tensor<32x1xi32, #blocked1>
%89 = "triton_gpu.cmpi"(%42, %88) {predicate = 2 : i64} : (tensor<32x1xi32, #blocked1>, tensor<32x1xi32, #blocked1>) -> tensor<32x1xi1, #blocked1>
%90 = tt.broadcast %89 : (tensor<32x1xi1, #blocked1>) -> tensor<32x64xi1, #blocked1>
%91 = tt.load %arg12, %90, %56 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<32x64xf16, #blocked1>
%92 = triton_gpu.convert_layout %87 : (tensor<128x32xf16, #blocked>) -> tensor<128x32xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #blocked4}>>
%93 = triton_gpu.convert_layout %91 : (tensor<32x64xf16, #blocked1>) -> tensor<32x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #blocked4}>>
%94 = triton_gpu.convert_layout %arg10 : (tensor<128x64xf32, #blocked>) -> tensor<128x64xf32, #blocked4>
%95 = tt.dot %92, %93, %94 {allowTF32 = true} : tensor<128x32xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #blocked4}>> * tensor<32x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #blocked4}>> -> tensor<128x64xf32, #blocked4>
%96 = triton_gpu.convert_layout %95 : (tensor<128x64xf32, #blocked4>) -> tensor<128x64xf32, #blocked>
%97 = tt.addptr %arg11, %cst : tensor<128x32x!tt.ptr<f16>, #blocked>, tensor<128x32xi32, #blocked>
%98 = tt.addptr %arg12, %58 : tensor<32x64x!tt.ptr<f16>, #blocked1>, tensor<32x64xi32, #blocked1>
scf.yield %96, %97, %98 : tensor<128x64xf32, #blocked>, tensor<128x32x!tt.ptr<f16>, #blocked>, tensor<32x64x!tt.ptr<f16>, #blocked1>
}
%60 = arith.truncf %59#0 : tensor<128x64xf32, #blocked> to tensor<128x64xf16, #blocked>
%61 = triton_gpu.convert_layout %18 : (tensor<128xi32, #blocked2>) -> tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>
%62 = tt.expand_dims %61 {axis = 1 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked}>>) -> tensor<128x1xi32, #blocked>
%63 = tt.splat %arg8 : (i32) -> tensor<128x1xi32, #blocked>
%64 = arith.muli %63, %62 : tensor<128x1xi32, #blocked>
%65 = tt.splat %arg2 : (!tt.ptr<f16>) -> tensor<128x1x!tt.ptr<f16>, #blocked>
%66 = tt.addptr %65, %64 : tensor<128x1x!tt.ptr<f16>, #blocked>, tensor<128x1xi32, #blocked>
%67 = triton_gpu.convert_layout %24 : (tensor<64xi32, #blocked2>) -> tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
%68 = tt.expand_dims %67 {axis = 0 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>) -> tensor<1x64xi32, #blocked3>
%69 = tt.broadcast %66 : (tensor<128x1x!tt.ptr<f16>, #blocked>) -> tensor<128x64x!tt.ptr<f16>, #blocked>
%70 = tt.broadcast %68 : (tensor<1x64xi32, #blocked3>) -> tensor<128x64xi32, #blocked3>
%71 = triton_gpu.convert_layout %70 : (tensor<128x64xi32, #blocked3>) -> tensor<128x64xi32, #blocked>
%72 = tt.addptr %69, %71 : tensor<128x64x!tt.ptr<f16>, #blocked>, tensor<128x64xi32, #blocked>
%73 = tt.splat %arg3 : (i32) -> tensor<128x1xi32, #blocked>
%74 = "triton_gpu.cmpi"(%62, %73) {predicate = 2 : i64} : (tensor<128x1xi32, #blocked>, tensor<128x1xi32, #blocked>) -> tensor<128x1xi1, #blocked>
%75 = tt.splat %arg4 : (i32) -> tensor<1x64xi32, #blocked3>
%76 = "triton_gpu.cmpi"(%68, %75) {predicate = 2 : i64} : (tensor<1x64xi32, #blocked3>, tensor<1x64xi32, #blocked3>) -> tensor<1x64xi1, #blocked3>
%77 = tt.broadcast %74 : (tensor<128x1xi1, #blocked>) -> tensor<128x64xi1, #blocked>
%78 = tt.broadcast %76 : (tensor<1x64xi1, #blocked3>) -> tensor<128x64xi1, #blocked3>
%79 = triton_gpu.convert_layout %78 : (tensor<128x64xi1, #blocked3>) -> tensor<128x64xi1, #blocked>
%80 = arith.andi %77, %79 : tensor<128x64xi1, #blocked>
tt.store %72, %60, %80 {cache = 1 : i32, evict = 1 : i32} : tensor<128x64xf16, #blocked>
tt.return
}
}