33 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3# Naming this file with a `_dialect` suffix to avoid a naming conflict with4# python package's math module (coming in from random.py).5 6from mlir.ir import *7import mlir.dialects.func as func8import mlir.dialects.math as mlir_math9 10 11def run(f):12 print("\nTEST:", f.__name__)13 f()14 15 16# CHECK-LABEL: TEST: testMathOps17@run18def testMathOps():19 with Context() as ctx, Location.unknown():20 module = Module.create()21 with InsertionPoint(module.body):22 23 @func.FuncOp.from_py_func(F32Type.get())24 def emit_sqrt(arg):25 return mlir_math.SqrtOp(arg)26 27 # CHECK-LABEL: func @emit_sqrt(28 # CHECK-SAME: %[[ARG:.*]]: f32) -> f32 {29 # CHECK: math.sqrt %[[ARG]] : f3230 # CHECK: return31 # CHECK: }32 print(module)33