1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import builtin
5from mlir.dialects import linalg
6from mlir.dialects import std
7
8from mlir.dialects.linalg.opdsl.lang import *
9
10# This tests miscellaneous features of the emitter that are not tested by the
11# matmul, convolution, or, pooling tests. The features include:
12# - constant defined in the body
13# - fix/predefined types
14# - exponential functions
15# - custom op names.
16
17
18@linalg_structured_op
19def fill_rng_poly(
20    min=ScalarDef(F64),
21    max=ScalarDef(F64),
22    seed=ScalarDef(I32),
23    O=TensorDef(T, S.M, S.N, output=True)):
24  multiplier = TypeFn.cast(I32, const(1103515245))
25  increment = TypeFn.cast(I32, const(12345))
26  rand1 = (TypeFn.cast(I32, index(D.m)) + seed) * multiplier + increment
27  rand2 = (TypeFn.cast(I32, index(D.n)) + rand1) * multiplier + increment
28  inv_range = TypeFn.cast(F64, const(2.3283064e-10))
29  offset = TypeFn.cast(F64, const(2147483647))
30  scaling = (max - min) * inv_range
31  O[D.m, D.n] = TypeFn.cast(T,
32                            (offset + TypeFn.cast(F64, rand2)) * scaling + min)
33
34
35@linalg_structured_op
36def soft_plus_poly(
37    I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
38  O[D.m, D.n] = ArithFn.log(
39      TypeFn.cast(U, const(1.0)) + TypeFn.cast(U, ArithFn.exp(I[D.m, D.n])))
40
41
42@linalg_structured_op(op_name="custom_op_name")
43def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)):
44  O[D.n] = I[D.n]
45
46
47with Context() as ctx, Location.unknown():
48  module = Module.create()
49  f32 = F32Type.get()
50  f64 = F64Type.get()
51  i32 = IntegerType.get_signless(32)
52  with InsertionPoint(module.body):
53
54    # CHECK-LABEL: @test_i32_fill_rng
55    # CHECK:      ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %{{.*}}
56    # CHECK-DAG:    %[[IDX0:.+]] = linalg.index 0 : index
57    # CHECK-DAG:    %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32
58    # CHECK-DAG:    %[[RND0:.+]] = arith.addi %[[IDX0_CAST]], %[[SEED]] : i32
59    # CHECK-DAG:    %[[CST0:.+]] = arith.constant 1103515245 : i64
60    # CHECK-DAG:    %[[CST0_CAST:.+]] = arith.trunci %[[CST0]] : i64 to i32
61    # Skip the remaining random number computation and match the scaling logic.
62    # CHECK-DAG:    %[[DIFF:.+]] = arith.subf %[[MAX]], %[[MIN]] : f64
63    # CHECK-DAG:    %[[CST3:.+]] = arith.constant 2.3283063999999999E-10 : f64
64    # CHECK-DAG:    %[[FACT:.+]] = arith.mulf %[[DIFF]], %[[CST3]] : f64
65    # CHECK-DAG:    %[[RND4:.+]] = arith.mulf %{{.+}}, %[[FACT]] : f64
66    # CHECK-DAG:    %[[RND5:.+]] = arith.addf %[[RND4]], %[[MIN]] : f64
67    # CHECK-DAG:    %{{.*}} = arith.fptosi %[[RND5]] : f64 to i32
68    @builtin.FuncOp.from_py_func(f64, f64, i32,
69                                 RankedTensorType.get((4, 16), i32))
70    def test_i32_fill_rng(min, max, seed, init_result):
71      return fill_rng_poly(min, max, seed, outs=[init_result])
72
73    # CHECK-LABEL: @test_f32_soft_plus
74    # CHECK:      ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32)
75    # CHECK-NEXT:   %[[C1:.+]] = arith.constant 1.000000e+00 : f64
76    # CHECK-NEXT:   %[[C1_CAST:.+]] = arith.truncf %[[C1]] : f64 to f32
77    # CHECK-NEXT:   %[[EXP:.+]] = math.exp %[[IN]] : f32
78    # CHECK-NEXT:   %[[SUM:.+]] = arith.addf %[[C1_CAST]], %[[EXP]] : f32
79    # CHECK-NEXT:   %[[LOG:.+]] = math.log %[[SUM]] : f32
80    # CHECK-NEXT:   linalg.yield %[[LOG]] : f32
81    # CHECK-NEXT: -> tensor<4x16xf32>
82    @builtin.FuncOp.from_py_func(
83        RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
84    def test_f32_soft_plus(input, init_result):
85      return soft_plus_poly(input, outs=[init_result])
86
87    # Just check that we don't assert out on name mismatch.
88    # CHECK-LABEL: @test_non_default_op_name
89    @builtin.FuncOp.from_py_func(
90        RankedTensorType.get((42,), f32), RankedTensorType.get((42,), f32))
91    def test_non_default_op_name(input, init_result):
92      return non_default_op_name(input, outs=[init_result])
93
94
95print(module)
96