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