189 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import builtin5from mlir.dialects import func6from mlir.dialects import linalg7from mlir.dialects import tensor8 9from mlir.dialects.linalg.opdsl.lang import *10 11T1 = TV.T112T2 = TV.T213 14 15@linalg_structured_op16def matmul_mono(17 A=TensorDef(T, S.M, S.K),18 B=TensorDef(T, S.K, S.N),19 C=TensorDef(T, S.M, S.N, output=True),20):21 domain(D.m, D.n, D.k)22 C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]23 24 25@linalg_structured_op26def matmul_poly(27 A=TensorDef(T1, S.M, S.K),28 B=TensorDef(T2, S.K, S.N),29 C=TensorDef(U, S.M, S.N, output=True),30 cast=TypeFnAttrDef(default=TypeFn.cast_signed),31):32 domain(D.m, D.n, D.k)33 C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])34 35 36with Context() as ctx, Location.unknown():37 module = Module.create()38 f16 = F16Type.get()39 f32 = F32Type.get()40 f64 = F64Type.get()41 i8 = IntegerType.get_signless(8)42 i16 = IntegerType.get_signless(16)43 i32 = IntegerType.get_signless(32)44 with InsertionPoint(module.body):45 46 # Multiplication indexing maps. We verify only the indexing maps of the47 # first multiplication and then do additional tests on casting and body48 # generation behavior.49 # CHECK: #[[$MUL_MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>50 # CHECK: #[[$MUL_MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>51 # CHECK: #[[$MUL_MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>52 53 # CHECK-LABEL: func @test_matmul_mono54 # CHECK-SAME: %[[A:.+]]: tensor<4x16xf32>55 # CHECK-SAME: %[[B:.+]]: tensor<16x8xf32>56 # CHECK: %[[INITC:.+]] = tensor.empty() : tensor<4x8xf32>57 # CHECK: linalg.generic58 # CHECK-SAME: indexing_maps = [#[[$MUL_MAP_A]], #[[$MUL_MAP_B]], #[[$MUL_MAP_C]]]59 # CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]60 # CHECK-SAME: ins(%[[A]], %[[B]]61 # CHECK-SAME: outs(%[[INITC]]62 @func.FuncOp.from_py_func(63 RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32)64 )65 def test_matmul_mono(lhs, rhs):66 init_result = tensor.empty([4, 8], f32)67 return matmul_mono(lhs, rhs, outs=[init_result])68 69 # CHECK-LABEL: @test_i8i8i32_matmul70 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)71 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i3272 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i8 to i3273 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i3274 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i3275 # CHECK-NEXT: linalg.yield %[[ADD]] : i3276 # CHECK-NEXT: -> tensor<4x8xi32>77 @func.FuncOp.from_py_func(78 RankedTensorType.get((4, 16), i8),79 RankedTensorType.get((16, 8), i8),80 RankedTensorType.get((4, 8), i32),81 )82 def test_i8i8i32_matmul(lhs, rhs, init_result):83 return matmul_poly(lhs, rhs, outs=[init_result])84 85 # CHECK-LABEL: @test_i8i8i32_matmul_unsigned86 # CHECK: = arith.extui87 # CHECK: = arith.extui88 @func.FuncOp.from_py_func(89 RankedTensorType.get((4, 16), i8),90 RankedTensorType.get((16, 8), i8),91 RankedTensorType.get((4, 8), i32),92 )93 def test_i8i8i32_matmul_unsigned(lhs, rhs, init_result):94 return matmul_poly(lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned)95 96 # CHECK-LABEL: @test_i8i16i32_matmul97 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)98 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i3299 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i16 to i32100 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32101 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32102 # CHECK-NEXT: linalg.yield %[[ADD]] : i32103 # CHECK-NEXT: -> tensor<4x8xi32>104 @func.FuncOp.from_py_func(105 RankedTensorType.get((4, 16), i8),106 RankedTensorType.get((16, 8), i16),107 RankedTensorType.get((4, 8), i32),108 )109 def test_i8i16i32_matmul(lhs, rhs, init_result):110 return matmul_poly(lhs, rhs, outs=[init_result])111 112 # CHECK-LABEL: @test_i32i32i16_matmul113 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16)114 # CHECK-NEXT: %[[A_CAST:.+]] = arith.trunci %[[A_ARG]] : i32 to i16115 # CHECK-NEXT: %[[B_CAST:.+]] = arith.trunci %[[B_ARG]] : i32 to i16116 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i16117 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i16118 # CHECK-NEXT: linalg.yield %[[ADD]] : i16119 # CHECK-NEXT: -> tensor<4x8xi16>120 @func.FuncOp.from_py_func(121 RankedTensorType.get((4, 16), i32),122 RankedTensorType.get((16, 8), i32),123 RankedTensorType.get((4, 8), i16),124 )125 def test_i32i32i16_matmul(lhs, rhs, init_result):126 return matmul_poly(lhs, rhs, outs=[init_result])127 128 # CHECK-LABEL: @test_i8i8f32_matmul129 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32)130 # CHECK-NEXT: %[[A_CAST:.+]] = arith.sitofp %[[A_ARG]] : i8 to f32131 # CHECK-NEXT: %[[B_CAST:.+]] = arith.sitofp %[[B_ARG]] : i8 to f32132 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32133 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32134 # CHECK-NEXT: linalg.yield %[[ADD]] : f32135 # CHECK-NEXT: -> tensor<4x8xf32>136 @func.FuncOp.from_py_func(137 RankedTensorType.get((4, 16), i8),138 RankedTensorType.get((16, 8), i8),139 RankedTensorType.get((4, 8), f32),140 )141 def test_i8i8f32_matmul(lhs, rhs, init_result):142 return matmul_poly(lhs, rhs, outs=[init_result])143 144 # CHECK-LABEL: @test_i8i8f32_matmul_unsigned145 # CHECK: = arith.uitofp146 # CHECK: = arith.uitofp147 @func.FuncOp.from_py_func(148 RankedTensorType.get((4, 16), i8),149 RankedTensorType.get((16, 8), i8),150 RankedTensorType.get((4, 8), f32),151 )152 def test_i8i8f32_matmul_unsigned(lhs, rhs, init_result):153 return matmul_poly(lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned)154 155 # CHECK-LABEL: @test_f16f16f32_matmul156 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)157 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extf %[[A_ARG]] : f16 to f32158 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extf %[[B_ARG]] : f16 to f32159 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32160 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32161 # CHECK-NEXT: linalg.yield %[[ADD]] : f32162 # CHECK-NEXT: -> tensor<4x8xf32>163 @func.FuncOp.from_py_func(164 RankedTensorType.get((4, 16), f16),165 RankedTensorType.get((16, 8), f16),166 RankedTensorType.get((4, 8), f32),167 )168 def test_f16f16f32_matmul(lhs, rhs, init_result):169 return matmul_poly(lhs, rhs, outs=[init_result])170 171 # CHECK-LABEL: @test_f64f64f32_matmul172 # CHECK: ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)173 # CHECK-NEXT: %[[A_CAST:.+]] = arith.truncf %[[A_ARG]] : f64 to f32174 # CHECK-NEXT: %[[B_CAST:.+]] = arith.truncf %[[B_ARG]] : f64 to f32175 # CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32176 # CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32177 # CHECK-NEXT: linalg.yield %[[ADD]] : f32178 # CHECK-NEXT: -> tensor<4x8xf32>179 @func.FuncOp.from_py_func(180 RankedTensorType.get((4, 16), f64),181 RankedTensorType.get((16, 8), f64),182 RankedTensorType.get((4, 8), f32),183 )184 def test_f64f64f32_matmul(lhs, rhs, init_result):185 return matmul_poly(lhs, rhs, outs=[init_result])186 187 188print(module)189