1# RUN: %PYTHON %s | FileCheck %s 2 3from mlir.ir import * 4from mlir.dialects import builtin 5from mlir.dialects import func 6from mlir.dialects import linalg 7 8from mlir.dialects.linalg.opdsl.lang import * 9 10T1 = TV.T1 11T2 = TV.T2 12 13 14@linalg_structured_op 15def matmul_mono( 16 A=TensorDef(T, S.M, S.K), 17 B=TensorDef(T, S.K, S.N), 18 C=TensorDef(T, S.M, S.N, output=True)): 19 domain(D.m, D.n, D.k) 20 C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n] 21 22 23@linalg_structured_op 24def matmul_poly( 25 A=TensorDef(T1, S.M, S.K), 26 B=TensorDef(T2, S.K, S.N), 27 C=TensorDef(U, S.M, S.N, output=True), 28 cast=TypeFnAttrDef(default=TypeFn.cast_signed)): 29 domain(D.m, D.n, D.k) 30 C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) 31 32 33with Context() as ctx, Location.unknown(): 34 module = Module.create() 35 f16 = F16Type.get() 36 f32 = F32Type.get() 37 f64 = F64Type.get() 38 i8 = IntegerType.get_signless(8) 39 i16 = IntegerType.get_signless(16) 40 i32 = IntegerType.get_signless(32) 41 with InsertionPoint(module.body): 42 43 # Multiplication indexing maps. We verify only the indexing maps of the 44 # first multiplication and then do additional tests on casting and body 45 # generation behavior. 46 # CHECK: #[[$MUL_MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> 47 # CHECK: #[[$MUL_MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> 48 # CHECK: #[[$MUL_MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> 49 50 # CHECK-LABEL: func @test_matmul_mono 51 # CHECK-SAME: %[[A:.+]]: tensor<4x16xf32> 52 # CHECK-SAME: %[[B:.+]]: tensor<16x8xf32> 53 # CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32> 54 # CHECK: linalg.generic 55 # CHECK-SAME: indexing_maps = [#[[$MUL_MAP_A]], #[[$MUL_MAP_B]], #[[$MUL_MAP_C]]] 56 # CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] 57 # CHECK-SAME: ins(%[[A]], %[[B]] 58 # CHECK-SAME: outs(%[[INITC]] 59 @func.FuncOp.from_py_func( 60 RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32)) 61 def test_matmul_mono(lhs, rhs): 62 init_result = linalg.InitTensorOp([4, 8], f32) 63 return matmul_mono(lhs, rhs, outs=[init_result.result]) 64 65 # CHECK-LABEL: @test_i8i8i32_matmul 66 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32) 67 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i32 68 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i8 to i32 69 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32 70 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32 71 # CHECK-NEXT: linalg.yield %[[ADD]] : i32 72 # CHECK-NEXT: -> tensor<4x8xi32> 73 @func.FuncOp.from_py_func( 74 RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8), 75 RankedTensorType.get((4, 8), i32)) 76 def test_i8i8i32_matmul(lhs, rhs, init_result): 77 return matmul_poly(lhs, rhs, outs=[init_result]) 78 79 # CHECK-LABEL: @test_i8i8i32_matmul_unsigned 80 # CHECK: = arith.extui 81 # CHECK: = arith.extui 82 @func.FuncOp.from_py_func( 83 RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8), 84 RankedTensorType.get((4, 8), i32)) 85 def test_i8i8i32_matmul_unsigned(lhs, rhs, init_result): 86 return matmul_poly( 87 lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned) 88 89 # CHECK-LABEL: @test_i8i16i32_matmul 90 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32) 91 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i32 92 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i16 to i32 93 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32 94 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32 95 # CHECK-NEXT: linalg.yield %[[ADD]] : i32 96 # CHECK-NEXT: -> tensor<4x8xi32> 97 @func.FuncOp.from_py_func( 98 RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i16), 99 RankedTensorType.get((4, 8), i32)) 100 def test_i8i16i32_matmul(lhs, rhs, init_result): 101 return matmul_poly(lhs, rhs, outs=[init_result]) 102 103 # CHECK-LABEL: @test_i32i32i16_matmul 104 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16) 105 # CHECK-NEXT: %[[A_CAST:.+]] = arith.trunci %[[A_ARG]] : i32 to i16 106 # CHECK-NEXT: %[[B_CAST:.+]] = arith.trunci %[[B_ARG]] : i32 to i16 107 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i16 108 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i16 109 # CHECK-NEXT: linalg.yield %[[ADD]] : i16 110 # CHECK-NEXT: -> tensor<4x8xi16> 111 @func.FuncOp.from_py_func( 112 RankedTensorType.get((4, 16), i32), RankedTensorType.get((16, 8), i32), 113 RankedTensorType.get((4, 8), i16)) 114 def test_i32i32i16_matmul(lhs, rhs, init_result): 115 return matmul_poly(lhs, rhs, outs=[init_result]) 116 117 # CHECK-LABEL: @test_i8i8f32_matmul 118 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32) 119 # CHECK-NEXT: %[[A_CAST:.+]] = arith.sitofp %[[A_ARG]] : i8 to f32 120 # CHECK-NEXT: %[[B_CAST:.+]] = arith.sitofp %[[B_ARG]] : i8 to f32 121 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32 122 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 123 # CHECK-NEXT: linalg.yield %[[ADD]] : f32 124 # CHECK-NEXT: -> tensor<4x8xf32> 125 @func.FuncOp.from_py_func( 126 RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8), 127 RankedTensorType.get((4, 8), f32)) 128 def test_i8i8f32_matmul(lhs, rhs, init_result): 129 return matmul_poly(lhs, rhs, outs=[init_result]) 130 131 # CHECK-LABEL: @test_i8i8f32_matmul_unsigned 132 # CHECK: = arith.uitofp 133 # CHECK: = arith.uitofp 134 @func.FuncOp.from_py_func( 135 RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8), 136 RankedTensorType.get((4, 8), f32)) 137 def test_i8i8f32_matmul_unsigned(lhs, rhs, init_result): 138 return matmul_poly( 139 lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned) 140 141 # CHECK-LABEL: @test_f16f16f32_matmul 142 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32) 143 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extf %[[A_ARG]] : f16 to f32 144 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extf %[[B_ARG]] : f16 to f32 145 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32 146 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 147 # CHECK-NEXT: linalg.yield %[[ADD]] : f32 148 # CHECK-NEXT: -> tensor<4x8xf32> 149 @func.FuncOp.from_py_func( 150 RankedTensorType.get((4, 16), f16), RankedTensorType.get((16, 8), f16), 151 RankedTensorType.get((4, 8), f32)) 152 def test_f16f16f32_matmul(lhs, rhs, init_result): 153 return matmul_poly(lhs, rhs, outs=[init_result]) 154 155 # CHECK-LABEL: @test_f64f64f32_matmul 156 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32) 157 # CHECK-NEXT: %[[A_CAST:.+]] = arith.truncf %[[A_ARG]] : f64 to f32 158 # CHECK-NEXT: %[[B_CAST:.+]] = arith.truncf %[[B_ARG]] : f64 to f32 159 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32 160 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 161 # CHECK-NEXT: linalg.yield %[[ADD]] : f32 162 # CHECK-NEXT: -> tensor<4x8xf32> 163 @func.FuncOp.from_py_func( 164 RankedTensorType.get((4, 16), f64), RankedTensorType.get((16, 8), f64), 165 RankedTensorType.get((4, 8), f32)) 166 def test_f64f64f32_matmul(lhs, rhs, init_result): 167 return matmul_poly(lhs, rhs, outs=[init_result]) 168 169 170print(module) 171