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.func as func
   8 import mlir.dialects.math as mlir_math
   9
  10
  11 def run(f):
  12     print("\nTEST:", f.__name__)
  13     f()
  14
  15
  16 # CHECK-LABEL: TEST: testMathOps
  17 @run
  18 def 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]] : f32
  30         # CHECK:         return
  31         # CHECK:       }
  32         print(module)