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 conv_poly( 16 I=TensorDef(T1, S.N, S.IH, S.IW, S.C), 17 K=TensorDef(T2, S.KH, S.KW, S.C), 18 O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), 19 strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), 20 dilations=IndexAttrDef(S.DH, S.DW, default=[1, 2]), 21): 22 domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c) 23 O[D.n, D.oh, D.ow, D.c] += TypeFn.cast_signed( 24 U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c] 25 ) * TypeFn.cast_signed(U, K[D.kh, D.kw, D.c]) 26 27 28with Context() as ctx, Location.unknown(): 29 module = Module.create() 30 f32 = F32Type.get() 31 i32 = IntegerType.get_signless(32) 32 with InsertionPoint(module.body): 33 34 # Convolution indexing maps. 35 # CHECK: #[[$CONV_MAP_I:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d3, d2 * 4 + d4 * 2, d5)> 36 # CHECK: #[[$CONV_MAP_K:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)> 37 # CHECK: #[[$CONV_MAP_O:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d5)> 38 39 # CHECK-LABEL: @test_f32i32_conv 40 # CHECK: linalg.generic 41 # CHECK-SAME: indexing_maps = [#[[$CONV_MAP_I]], #[[$CONV_MAP_K]], #[[$CONV_MAP_O]]] 42 # CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "parallel"] 43 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[FILTER:.+]]: f32, %[[OUT:.+]]: i32) 44 # CHECK-NEXT: %[[IN_CAST:.+]] = arith.fptosi %[[IN:.+]] : f32 to i32 45 # CHECK-NEXT: %[[FILTER_CAST:.+]] = arith.fptosi %[[FILTER:.+]] : f32 to i32 46 # CHECK-NEXT: %[[PROD:.+]] = arith.muli %[[IN_CAST]], %[[FILTER_CAST]] : i32 47 # CHECK-NEXT: %[[SUM:.+]] = arith.addi %[[OUT]], %[[PROD]] : i32 48 # CHECK-NEXT: linalg.yield %[[SUM]] : i32 49 # CHECK-NEXT: -> tensor<1x2x4x1xi32> 50 @func.FuncOp.from_py_func( 51 RankedTensorType.get((1, 4, 16, 1), f32), 52 RankedTensorType.get((2, 2, 1), f32), 53 RankedTensorType.get((1, 2, 4, 1), i32), 54 ) 55 def test_f32i32_conv(input, filter, init_result): 56 # Use default dilations and set non-default strides. 57 return conv_poly(input, filter, outs=[init_result], strides=[2, 4]) 58 59 60print(module) 61