@@ -60,44 +60,34 @@ def test_kernel_from_09_tutorial(device, tmp_path: pathlib.Path):
60
60
ir = """
61
61
#blocked = #ttg.blocked<{sizePerThread = [4, 4], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}>
62
62
#blocked1 = #ttg.blocked<{sizePerThread = [8, 1], threadsPerWarp = [8, 4], warpsPerCTA = [1, 4], order = [0, 1]}>
63
- #blocked2 = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
64
63
#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}>
65
64
#shared1 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}>
66
65
#smem = #ttg.shared_memory
67
66
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "xpu", "ttg.threads-per-warp" = 32 : i32, ttig.min_sg_size = 8 : i32, ttig.support_bf16_conversion, ttig.support_dpas, ttig.target_arch = "spir64"} {
68
- tt.func public @matmul_kernel(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} {
67
+ tt.func public @matmul_kernel(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) {
69
68
%cst = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked>
70
- %c63_i32 = arith.constant 63 : i32
71
- %c127_i32 = arith.constant 127 : i32
72
- %c1_i32 = arith.constant 1 : i32
73
- %c0_i32 = arith.constant 0 : i32
74
- %c64_i32 = arith.constant 64 : i32
75
69
%cst_0 = arith.constant dense<0.000000e+00> : tensor<64x128xf32, #blocked1>
76
- %c8_i32 = arith.constant 8 : i32
77
- %c128_i32 = arith.constant 128 : i32
78
70
%18 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #ttg.slice<{dim = 0, parent = #blocked1}>>
79
- %42 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<128x64x!tt.ptr<f32>, #blocked2 >
71
+ %42 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<128x64x!tt.ptr<f32>, #blocked1 >
80
72
%44 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>>
81
73
%45 = tt.expand_dims %44 {axis = 1 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<64x1xi32, #blocked1>
82
74
%46 = tt.expand_dims %18 {axis = 0 : i32} : tensor<128xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x128xi32, #blocked1>
83
75
%50 = tt.broadcast %46 : tensor<1x128xi32, #blocked1> -> tensor<64x128xi32, #blocked1>
84
76
%52 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<64x128x!tt.ptr<f32>, #blocked1>
85
77
%53 = tt.addptr %52, %50 : tensor<64x128x!tt.ptr<f32>, #blocked1>, tensor<64x128xi32, #blocked1>
86
78
87
- %81 = arith.subi %arg5, %c64_i32 : i32
88
- %85 = tt.load %42: tensor<128x64x!tt.ptr<f32>, #blocked2>
89
- %86 = tt.splat %81 : i32 -> tensor<64x1xi32, #blocked1>
79
+ %85 = tt.load %42: tensor<128x64x!tt.ptr<f32>, #blocked1>
80
+ %86 = tt.splat %arg5 : i32 -> tensor<64x1xi32, #blocked1>
90
81
%87 = arith.cmpi slt, %45, %86 : tensor<64x1xi32, #blocked1>
91
82
%88 = tt.broadcast %87 : tensor<64x1xi1, #blocked1> -> tensor<64x128xi1, #blocked1>
92
83
%89 = tt.load %53, %88, %cst_0 : tensor<64x128x!tt.ptr<f32>, #blocked1>
93
- %91 = ttg.local_alloc %85 : (tensor<128x64xf32, #blocked2 >) -> !ttg.memdesc<128x64xf32, #shared, #smem>
84
+ %91 = ttg.local_alloc %85 : (tensor<128x64xf32, #blocked1 >) -> !ttg.memdesc<128x64xf32, #shared, #smem>
94
85
%92 = ttg.local_load %91 : !ttg.memdesc<128x64xf32, #shared, #smem> -> tensor<128x64xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>
95
86
%94 = ttg.local_alloc %89 : (tensor<64x128xf32, #blocked1>) -> !ttg.memdesc<64x128xf32, #shared1, #smem>
96
- %cst_test = arith.constant dense<1.11111116> : tensor<128x64xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>
97
87
%cst_test2 = arith.constant dense<1.11111116> : tensor<64x128xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>
98
88
%96 = tt.dot %92, %cst_test2, %cst, inputPrecision = tf32 : tensor<128x64xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<64x128xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<128x128xf32, #blocked>
99
89
100
- %78 = ttg.convert_layout %96 : tensor<128x128xf32, #blocked> -> tensor<128x128xf32, #blocked2 >
90
+ %78 = ttg.convert_layout %96 : tensor<128x128xf32, #blocked> -> tensor<128x128xf32, #blocked1 >
101
91
tt.return
102
92
}
103
93
}
0 commit comments