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 *
7*36550692SRiver Riddleimport mlir.dialects.func as func
825d17349SUday Bondhugulaimport mlir.dialects.math as mlir_math
925d17349SUday Bondhugula
1025d17349SUday Bondhuguladef run(f):
1125d17349SUday Bondhugula  print("\nTEST:", f.__name__)
1225d17349SUday Bondhugula  f()
1325d17349SUday Bondhugula
1425d17349SUday Bondhugula# CHECK-LABEL: TEST: testMathOps
1525d17349SUday Bondhugula@run
1625d17349SUday Bondhuguladef testMathOps():
1725d17349SUday Bondhugula  with Context() as ctx, Location.unknown():
1825d17349SUday Bondhugula    module = Module.create()
1925d17349SUday Bondhugula    with InsertionPoint(module.body):
20*36550692SRiver Riddle      @func.FuncOp.from_py_func(F32Type.get())
2125d17349SUday Bondhugula      def emit_sqrt(arg):
2225d17349SUday Bondhugula        return mlir_math.SqrtOp(arg)
2325d17349SUday Bondhugula
2425d17349SUday Bondhugula    # CHECK-LABEL: func @emit_sqrt(
2525d17349SUday Bondhugula    # CHECK-SAME:                  %[[ARG:.*]]: f32) -> f32 {
2625d17349SUday Bondhugula    # CHECK:         math.sqrt %[[ARG]] : f32
2725d17349SUday Bondhugula    # CHECK:         return
2825d17349SUday Bondhugula    # CHECK:       }
2925d17349SUday Bondhugula    print(module)
30