mlir/test/python/dialects/math_dialect.py

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