xref: /llvm-project/mlir/test/python/dialects/math_dialect.py (revision f9008e6366c2496b1ca1785b891d5578174ad63e)
125d17349SUday Bondhugula# RUN: %PYTHON %s | FileCheck %s
225d17349SUday Bondhugula
325d17349SUday Bondhugula# Naming this file with a `_dialect` suffix to avoid a naming conflict with
425d17349SUday Bondhugula# python package's math module (coming in from random.py).
525d17349SUday Bondhugula
625d17349SUday Bondhugulafrom mlir.ir import *
736550692SRiver Riddleimport mlir.dialects.func as func
825d17349SUday Bondhugulaimport mlir.dialects.math as mlir_math
925d17349SUday Bondhugula
10*f9008e63STobias Hieta
1125d17349SUday Bondhuguladef run(f):
1225d17349SUday Bondhugula    print("\nTEST:", f.__name__)
1325d17349SUday Bondhugula    f()
1425d17349SUday Bondhugula
15*f9008e63STobias Hieta
1625d17349SUday Bondhugula# CHECK-LABEL: TEST: testMathOps
1725d17349SUday Bondhugula@run
1825d17349SUday Bondhuguladef testMathOps():
1925d17349SUday Bondhugula    with Context() as ctx, Location.unknown():
2025d17349SUday Bondhugula        module = Module.create()
2125d17349SUday Bondhugula        with InsertionPoint(module.body):
22*f9008e63STobias Hieta
2336550692SRiver Riddle            @func.FuncOp.from_py_func(F32Type.get())
2425d17349SUday Bondhugula            def emit_sqrt(arg):
2525d17349SUday Bondhugula                return mlir_math.SqrtOp(arg)
2625d17349SUday Bondhugula
2725d17349SUday Bondhugula        # CHECK-LABEL: func @emit_sqrt(
2825d17349SUday Bondhugula        # CHECK-SAME:                  %[[ARG:.*]]: f32) -> f32 {
2925d17349SUday Bondhugula        # CHECK:         math.sqrt %[[ARG]] : f32
3025d17349SUday Bondhugula        # CHECK:         return
3125d17349SUday Bondhugula        # CHECK:       }
3225d17349SUday Bondhugula        print(module)
33