| module { | |
| tt.func public @triton__0d1d2d3d4d5d6d7d8d9d10de11de(%arg0: !tt.ptr<bf16, 1> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<bf16, 1> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg4: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg5: !tt.ptr<bf16, 1> {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg7: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg8: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg9: !tt.ptr<f32, 1> {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32, tt.max_divisibility = 16 : i32}, %arg11: i32 {tt.divisibility = 16 : i32, tt.max_divisibility = 16 : i32}) attributes {noinline = false} { | |
| %cst = arith.constant dense<0.000000e+00> : tensor<64x8xbf16> | |
| %c8_i32 = arith.constant 8 : i32 | |
| %c128_i32 = arith.constant 128 : i32 | |
| %c0_i32 = arith.constant 0 : i32 | |
| %cst_0 = arith.constant dense<128> : tensor<64x1xi32> | |
| %cst_1 = arith.constant dense<32768> : tensor<64x1xi32> | |
| %cst_2 = arith.constant dense<256> : tensor<1x8xi32> | |
| %cst_3 = arith.constant dense<128> : tensor<1x8xi32> | |
| %cst_4 = arith.constant dense<0.000000e+00> : tensor<64x8xf32> | |
| %cst_5 = arith.constant dense<256> : tensor<64x1xi32> | |
| %c64_i32 = arith.constant 64 : i32 | |
| %0 = tt.get_program_id x : i32 | |
| %1 = arith.muli %0, %c64_i32 : i32 | |
| %2 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> | |
| %3 = tt.expand_dims %2 {axis = 1 : i32} : (tensor<64xi32>) -> tensor<64x1xi32> | |
| %4 = tt.splat %1 : (i32) -> tensor<64x1xi32> | |
| %5 = arith.addi %4, %3 : tensor<64x1xi32> | |
| %6 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> | |
| %7 = tt.expand_dims %6 {axis = 0 : i32} : (tensor<8xi32>) -> tensor<1x8xi32> | |
| %8 = arith.remsi %5, %cst_5 : tensor<64x1xi32> | |
| %9 = arith.divsi %5, %cst_5 : tensor<64x1xi32> | |
| %10 = tt.broadcast %8 : (tensor<64x1xi32>) -> tensor<64x8xi32> | |
| %11 = arith.muli %9, %cst_1 : tensor<64x1xi32> | |
| %12 = tt.broadcast %11 : (tensor<64x1xi32>) -> tensor<64x8xi32> | |
| %13 = tt.splat %arg0 : (!tt.ptr<bf16, 1>) -> tensor<64x8x!tt.ptr<bf16, 1>> | |
| %14 = tt.splat %arg1 : (!tt.ptr<f32, 1>) -> tensor<64x8x!tt.ptr<f32, 1>> | |
| %15 = tt.splat %arg2 : (!tt.ptr<bf16, 1>) -> tensor<64x8x!tt.ptr<bf16, 1>> | |
| %16 = arith.muli %9, %cst_0 : tensor<64x1xi32> | |
| %17 = tt.broadcast %16 : (tensor<64x1xi32>) -> tensor<64x8xi32> | |
| %18 = tt.splat %arg3 : (!tt.ptr<f32, 1>) -> tensor<64x8x!tt.ptr<f32, 1>> | |
| %19 = tt.splat %arg4 : (!tt.ptr<f32, 1>) -> tensor<64x8x!tt.ptr<f32, 1>> | |
| %20 = tt.splat %arg5 : (!tt.ptr<bf16, 1>) -> tensor<64x8x!tt.ptr<bf16, 1>> | |
| %21 = tt.splat %arg6 : (!tt.ptr<f32, 1>) -> tensor<64x8x!tt.ptr<f32, 1>> | |
| %22 = tt.splat %arg7 : (!tt.ptr<f32, 1>) -> tensor<64x8x!tt.ptr<f32, 1>> | |
| %23:2 = scf.for %arg12 = %c0_i32 to %c128_i32 step %c8_i32 iter_args(%arg13 = %cst_4, %arg14 = %cst_4) -> (tensor<64x8xf32>, tensor<64x8xf32>) : i32 { | |
| %32 = tt.splat %arg12 : (i32) -> tensor<1x8xi32> | |
| %33 = arith.addi %32, %7 : tensor<1x8xi32> | |
| %34 = arith.cmpi slt, %33, %cst_3 : tensor<1x8xi32> | |
| %35 = arith.muli %33, %cst_2 : tensor<1x8xi32> | |
| %36 = tt.broadcast %35 : (tensor<1x8xi32>) -> tensor<64x8xi32> | |
| %37 = arith.addi %10, %36 : tensor<64x8xi32> | |
| %38 = arith.addi %37, %12 : tensor<64x8xi32> | |
| %39 = tt.addptr %13, %38 : tensor<64x8x!tt.ptr<bf16, 1>>, tensor<64x8xi32> | |
| %40 = tt.broadcast %34 : (tensor<1x8xi1>) -> tensor<64x8xi1> | |
| %41 = tt.load %39, %40, %cst {cache = 1 : i32, evict = 2 : i32, isVolatile = false} : tensor<64x8xbf16> | |
| %42 = arith.extf %41 : tensor<64x8xbf16> to tensor<64x8xf32> | |
| %43 = tt.addptr %14, %38 : tensor<64x8x!tt.ptr<f32, 1>>, tensor<64x8xi32> | |
| %44 = tt.load %43, %40, %cst_4 {cache = 1 : i32, evict = 2 : i32, isVolatile = false} : tensor<64x8xf32> | |
| %45 = tt.addptr %15, %38 : tensor<64x8x!tt.ptr<bf16, 1>>, tensor<64x8xi32> | |
| %46 = tt.load %45, %40, %cst {cache = 1 : i32, evict = 2 : i32, isVolatile = false} : tensor<64x8xbf16> | |
| %47 = arith.extf %46 : tensor<64x8xbf16> to tensor<64x8xf32> | |
| %48 = tt.broadcast %33 : (tensor<1x8xi32>) -> tensor<64x8xi32> | |
| %49 = arith.addi %48, %17 : tensor<64x8xi32> | |
| %50 = tt.addptr %18, %49 : tensor<64x8x!tt.ptr<f32, 1>>, tensor<64x8xi32> | |
| %51 = tt.load %50, %40, %cst_4 {cache = 1 : i32, evict = 3 : i32, isVolatile = false} : tensor<64x8xf32> | |
| %52 = tt.addptr %19, %49 : tensor<64x8x!tt.ptr<f32, 1>>, tensor<64x8xi32> | |
| %53 = tt.load %52, %40, %cst_4 {cache = 1 : i32, evict = 3 : i32, isVolatile = false} : tensor<64x8xf32> | |
| %54 = tt.addptr %20, %38 : tensor<64x8x!tt.ptr<bf16, 1>>, tensor<64x8xi32> | |
| %55 = tt.load %54, %40, %cst {cache = 1 : i32, evict = 2 : i32, isVolatile = false} : tensor<64x8xbf16> | |
| %56 = arith.extf %55 : tensor<64x8xbf16> to tensor<64x8xf32> | |
| %57 = tt.addptr %21, %49 : tensor<64x8x!tt.ptr<f32, 1>>, tensor<64x8xi32> | |
| %58 = tt.load %57, %40, %cst_4 {cache = 1 : i32, evict = 3 : i32, isVolatile = false} : tensor<64x8xf32> | |
| %59 = tt.addptr %22, %49 : tensor<64x8x!tt.ptr<f32, 1>>, tensor<64x8xi32> | |
| %60 = tt.load %59, %40, %cst_4 {cache = 1 : i32, evict = 3 : i32, isVolatile = false} : tensor<64x8xf32> | |
| %61 = arith.addf %44, %47 : tensor<64x8xf32> | |
| %62 = arith.subf %61, %51 : tensor<64x8xf32> | |
| %63 = arith.mulf %62, %53 : tensor<64x8xf32> | |
| %64 = arith.mulf %42, %63 : tensor<64x8xf32> | |
| %65 = arith.addf %arg13, %64 : tensor<64x8xf32> | |
| %66 = arith.select %40, %65, %arg13 : tensor<64x8xi1>, tensor<64x8xf32> | |
| %67 = arith.subf %44, %58 : tensor<64x8xf32> | |
| %68 = arith.mulf %67, %60 : tensor<64x8xf32> | |
| %69 = arith.mulf %56, %68 : tensor<64x8xf32> | |
| %70 = arith.addf %arg14, %69 : tensor<64x8xf32> | |
| %71 = arith.select %40, %70, %arg14 : tensor<64x8xi1>, tensor<64x8xf32> | |
| scf.yield %66, %71 : tensor<64x8xf32>, tensor<64x8xf32> | |
| } | |
| %24 = "tt.reduce"(%23#0) <{axis = 1 : i32}> ({ | |
| ^bb0(%arg12: f32, %arg13: f32): | |
| %32 = arith.addf %arg12, %arg13 : f32 | |
| tt.reduce.return %32 : f32 | |
| }) : (tensor<64x8xf32>) -> tensor<64xf32> | |
| %25 = tt.expand_dims %24 {axis = 1 : i32} : (tensor<64xf32>) -> tensor<64x1xf32> | |
| %26 = tt.splat %arg8 : (!tt.ptr<f32, 1>) -> tensor<64x1x!tt.ptr<f32, 1>> | |
| %27 = tt.addptr %26, %5 : tensor<64x1x!tt.ptr<f32, 1>>, tensor<64x1xi32> | |
| tt.store %27, %25 {cache = 1 : i32, evict = 1 : i32} : tensor<64x1xf32> | |
| %28 = "tt.reduce"(%23#1) <{axis = 1 : i32}> ({ | |
| ^bb0(%arg12: f32, %arg13: f32): | |
| %32 = arith.addf %arg12, %arg13 : f32 | |
| tt.reduce.return %32 : f32 | |
| }) : (tensor<64x8xf32>) -> tensor<64xf32> | |
| %29 = tt.expand_dims %28 {axis = 1 : i32} : (tensor<64xf32>) -> tensor<64x1xf32> | |
| %30 = tt.splat %arg9 : (!tt.ptr<f32, 1>) -> tensor<64x1x!tt.ptr<f32, 1>> | |
| %31 = tt.addptr %30, %5 : tensor<64x1x!tt.ptr<f32, 1>>, tensor<64x1xi32> | |
| tt.store %31, %29 {cache = 1 : i32, evict = 1 : i32} : tensor<64x1xf32> | |
| tt.return | |
| } | |
| } | |