55 lines
5.4 KiB
Plaintext
55 lines
5.4 KiB
Plaintext
#blocked0 = #triton_gpu.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
|
|
#blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [2, 2], order = [0, 1]}>
|
|
#mma0 = #triton_gpu.mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [4, 1]}>
|
|
#mma1 = #triton_gpu.mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [2, 2]}>
|
|
#shared = #triton_gpu.shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}>
|
|
module attributes {"triton_gpu.num-warps" = 4 : i32} {
|
|
func public @kernel_0d1d2c3d4d5c6d7d8c9d10d11c(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}) {
|
|
%cst = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma0>
|
|
%cst_0 = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma1>
|
|
%0 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #triton_gpu.slice<{dim = 1, parent = #blocked0}>>
|
|
%1 = tt.expand_dims %0 {axis = 1 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 1, parent = #blocked0}>>) -> tensor<64x1xi32, #blocked0>
|
|
%2 = tt.splat %arg1 : (i32) -> tensor<64x1xi32, #blocked0>
|
|
%3 = tt.splat %arg0 : (!tt.ptr<f16>) -> tensor<64x1x!tt.ptr<f16>, #blocked0>
|
|
%4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked0}>>
|
|
%5 = tt.expand_dims %4 {axis = 0 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked0}>>) -> tensor<1x64xi32, #blocked0>
|
|
%6 = tt.broadcast %5 : (tensor<1x64xi32, #blocked0>) -> tensor<64x64xi32, #blocked0>
|
|
%7 = tt.splat %arg3 : (i32) -> tensor<64x1xi32, #blocked0>
|
|
%8 = tt.splat %arg2 : (!tt.ptr<f16>) -> tensor<64x1x!tt.ptr<f16>, #blocked0>
|
|
%9 = tt.splat %arg5 : (i32) -> tensor<64x1xi32, #blocked0>
|
|
%10 = tt.splat %arg4 : (!tt.ptr<f16>) -> tensor<64x1x!tt.ptr<f16>, #blocked0>
|
|
%11 = tt.splat %arg7 : (i32) -> tensor<64x1xi32, #blocked0>
|
|
%12 = tt.splat %arg6 : (!tt.ptr<f16>) -> tensor<64x1x!tt.ptr<f16>, #blocked0>
|
|
%13 = arith.muli %1, %2 : tensor<64x1xi32, #blocked0>
|
|
%14 = tt.addptr %3, %13 : tensor<64x1x!tt.ptr<f16>, #blocked0>, tensor<64x1xi32, #blocked0>
|
|
%15 = tt.broadcast %14 : (tensor<64x1x!tt.ptr<f16>, #blocked0>) -> tensor<64x64x!tt.ptr<f16>, #blocked0>
|
|
%16 = tt.addptr %15, %6 : tensor<64x64x!tt.ptr<f16>, #blocked0>, tensor<64x64xi32, #blocked0>
|
|
%17 = tt.load %16 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64x64xf16, #blocked0>
|
|
%18 = arith.muli %1, %7 : tensor<64x1xi32, #blocked0>
|
|
%19 = tt.addptr %8, %18 : tensor<64x1x!tt.ptr<f16>, #blocked0>, tensor<64x1xi32, #blocked0>
|
|
%20 = tt.broadcast %19 : (tensor<64x1x!tt.ptr<f16>, #blocked0>) -> tensor<64x64x!tt.ptr<f16>, #blocked0>
|
|
%21 = tt.addptr %20, %6 : tensor<64x64x!tt.ptr<f16>, #blocked0>, tensor<64x64xi32, #blocked0>
|
|
%22 = tt.load %21 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64x64xf16, #blocked0>
|
|
%23 = triton_gpu.convert_layout %17 : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>>
|
|
%24 = triton_gpu.convert_layout %22 : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>>
|
|
%25 = tt.dot %23, %24, %cst {allowTF32 = false} : tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>> * tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>> -> tensor<64x64xf32, #mma0>
|
|
%27 = arith.muli %1, %9 : tensor<64x1xi32, #blocked0>
|
|
%28 = tt.addptr %10, %27 : tensor<64x1x!tt.ptr<f16>, #blocked0>, tensor<64x1xi32, #blocked0>
|
|
%29 = tt.broadcast %28 : (tensor<64x1x!tt.ptr<f16>, #blocked0>) -> tensor<64x64x!tt.ptr<f16>, #blocked0>
|
|
%30 = tt.addptr %29, %6 : tensor<64x64x!tt.ptr<f16>, #blocked0>, tensor<64x64xi32, #blocked0>
|
|
%31 = tt.load %30 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64x64xf16, #blocked0>
|
|
%32 = arith.truncf %25 : tensor<64x64xf32, #mma0> to tensor<64x64xf16, #mma0>
|
|
%133 = triton_gpu.convert_layout %32 : (tensor<64x64xf16, #mma0>) -> tensor<64x64xf16, #shared>
|
|
%33 = triton_gpu.convert_layout %133 : (tensor<64x64xf16, #shared>) -> tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>>
|
|
%34 = triton_gpu.convert_layout %31 : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>>
|
|
%35 = tt.dot %33, %34, %cst_0 {allowTF32 = true} : tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>> * tensor<64x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>> -> tensor<64x64xf32, #mma1>
|
|
%36 = triton_gpu.convert_layout %35 : (tensor<64x64xf32, #mma1>) -> tensor<64x64xf32, #blocked0>
|
|
%37 = arith.muli %1, %11 : tensor<64x1xi32, #blocked0>
|
|
%38 = tt.addptr %12, %37 : tensor<64x1x!tt.ptr<f16>, #blocked0>, tensor<64x1xi32, #blocked0>
|
|
%39 = tt.broadcast %38 : (tensor<64x1x!tt.ptr<f16>, #blocked0>) -> tensor<64x64x!tt.ptr<f16>, #blocked0>
|
|
%40 = tt.addptr %39, %6 : tensor<64x64x!tt.ptr<f16>, #blocked0>, tensor<64x64xi32, #blocked0>
|
|
%41 = arith.truncf %36 : tensor<64x64xf32, #blocked0> to tensor<64x64xf16, #blocked0>
|
|
tt.store %40, %41 : tensor<64x64xf16, #blocked0>
|
|
return
|
|
}
|
|
} |