Intrinsics

intrinsics.py

# RUN: %PYTHON %s | iree-dialects-opt -split-input-file | FileCheck --enable-var-scope --dump-input-filter=all %s

from typing import List
from mlir.dialects.iree_pydm.importer import *
from mlir.dialects.iree_pydm.importer.test_util import *

from mlir.dialects import iree_pydm as d
from mlir import ir


################################################################################
# Pyfunc intrinsics
################################################################################

@def_pyfunc_intrinsic(symbol="__return_one")
def intrinsic_return_one() -> int:
  return 1


@def_pyfunc_intrinsic(symbol="__return_first_true")
def intrinsic_return_first_true(a: int, b: int) -> int:
  return a or b


# CHECK-LABEL: @test_intrinsic_function_no_args
# CHECK: iree_pydm.dynamic_call @__return_one() : () -> (!iree_pydm.exception_result, !iree_pydm.object)
# CHECK: iree_pydm.func private @__return_one()
@test_import_global
def test_intrinsic_function_no_args():
  value = intrinsic_return_one()
  return value


# CHECK-LABEL: @test_intrinsic_function_args
# CHECK: %[[ZERO:.*]] = iree_pydm.constant 0 : si64 -> !iree_pydm.integer
# CHECK: %[[ONE:.*]] = iree_pydm.constant 1 : si64 -> !iree_pydm.integer
# CHECK: iree_pydm.dynamic_call @__return_first_true(%[[ZERO]], %[[ONE]]) : (!iree_pydm.integer, !iree_pydm.integer) -> (!iree_pydm.exception_result, !iree_pydm.object)
# CHECK: iree_pydm.func private @__return_first_true
@test_import_global
def test_intrinsic_function_args():
  value = intrinsic_return_first_true(0, 1)
  return value

################################################################################
# IR macro intrinsics
################################################################################

@def_ir_macro_intrinsic
def macro_return_none(stage: ImportStage) -> ir.Value:
  return d.NoneOp(d.NoneType.get()).result


# Boxing isn't load bearing here: It is just something we can do/test.
@def_ir_macro_intrinsic
def macro_box_arg(stage: ImportStage, arg: ir.Value) -> ir.Value:
  return stage.ic.box(arg)


# CHECK-LABEL: @test_intrinsic_macro_no_args
# CHECK: %[[ONE:.*]] = iree_pydm.constant 1
# CHECK: iree_pydm.box %[[ONE]] : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
@test_import_global
def test_intrinsic_macro_no_args() -> int:
  return macro_box_arg(1)


################################################################################
# Test multi func intrinsic.
# There is nothing special about a logical not. It is just something we can
# test.
################################################################################
@def_pyfunc_intrinsic(symbol="__logical_not_bool")
def logical_not_bool(x: bool) -> bool:
  return not x


@def_pyfunc_intrinsic(symbol="__logical_not_generic")
def logical_not_generic(x):
  return not x


logical_not = def_pattern_call_intrinsic(match_generic=[logical_not_generic],
                                         match_specific=[logical_not_bool])


# CHECK-LABEL: @test_pattern_call
# CHECK: %[[TRUE:.*]] = iree_pydm.constant true
# CHECK: iree_pydm.pattern_match_call(%[[TRUE]]) : (!iree_pydm.bool) -> (!iree_pydm.exception_result, !iree_pydm.object)
# CHECK-SAME:   matching generic [@__logical_not_generic] specific [@__logical_not_bool]
# CHECK-DAG: iree_pydm.func private @__logical_not_generic
# CHECK-DAG: iree_pydm.func private @__logical_not_bool
@test_import_global
def test_pattern_call():
  return logical_not(True)

output.mlir

// -----
builtin.module  {
  iree_pydm.func @test_intrinsic_function_no_args() -> (!iree_pydm.exception_result, !iree_pydm.object) attributes {arg_names = [], cell_vars = [], free_vars = ["value"]} {
    %exc_result, %result = iree_pydm.dynamic_call @__return_one() : () -> (!iree_pydm.exception_result, !iree_pydm.object)
    iree_pydm.raise_on_failure %exc_result : !iree_pydm.exception_result
    iree_pydm.store_free_var "value", %result : !iree_pydm.object
    %0 = iree_pydm.load_free_var "value" -> !iree_pydm.object
    iree_pydm.return %0 : !iree_pydm.object
  }
  iree_pydm.func private @__return_one() -> (!iree_pydm.exception_result, !iree_pydm.integer) attributes {arg_names = [], cell_vars = [], free_vars = []} {
    %0 = iree_pydm.constant 1 : si64 -> !iree_pydm.integer
    iree_pydm.return %0 : !iree_pydm.integer
  }
}

// -----
builtin.module  {
  iree_pydm.func @test_intrinsic_function_args() -> (!iree_pydm.exception_result, !iree_pydm.object) attributes {arg_names = [], cell_vars = [], free_vars = ["value"]} {
    %0 = iree_pydm.constant 0 : si64 -> !iree_pydm.integer
    %1 = iree_pydm.constant 1 : si64 -> !iree_pydm.integer
    %exc_result, %result = iree_pydm.dynamic_call @__return_first_true(%0, %1) : (!iree_pydm.integer, !iree_pydm.integer) -> (!iree_pydm.exception_result, !iree_pydm.object)
    iree_pydm.raise_on_failure %exc_result : !iree_pydm.exception_result
    iree_pydm.store_free_var "value", %result : !iree_pydm.object
    %2 = iree_pydm.load_free_var "value" -> !iree_pydm.object
    iree_pydm.return %2 : !iree_pydm.object
  }
  iree_pydm.func private @__return_first_true(%arg0: !iree_pydm.integer, %arg1: !iree_pydm.integer) -> (!iree_pydm.exception_result, !iree_pydm.integer) attributes {arg_names = ["a", "b"], cell_vars = [], free_vars = ["a", "b"]} {
    %0 = iree_pydm.box %arg0 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    iree_pydm.store_free_var "a", %0 : !iree_pydm.object<!iree_pydm.integer>
    %1 = iree_pydm.box %arg1 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    iree_pydm.store_free_var "b", %1 : !iree_pydm.object<!iree_pydm.integer>
    %2 = iree_pydm.load_free_var "a" -> !iree_pydm.object
    %3 = iree_pydm.as_bool %2 : !iree_pydm.object -> !iree_pydm.bool
    %4 = iree_pydm.bool_to_pred %3 : !iree_pydm.bool -> i1
    %5 = scf.if %4 -> (!iree_pydm.object) {
      scf.yield %2 : !iree_pydm.object
    } else {
      %6 = iree_pydm.load_free_var "b" -> !iree_pydm.object
      scf.yield %6 : !iree_pydm.object
    }
    %status, %primitive = iree_pydm.unbox %5 : !iree_pydm.object -> !iree_pydm.integer
    iree_pydm.raise_on_failure %status : !iree_pydm.exception_result
    iree_pydm.return %primitive : !iree_pydm.integer
  }
}

// -----
builtin.module  {
  iree_pydm.func @test_intrinsic_macro_no_args() -> (!iree_pydm.exception_result, !iree_pydm.integer) attributes {arg_names = [], cell_vars = [], free_vars = []} {
    %0 = iree_pydm.constant 1 : si64 -> !iree_pydm.integer
    %1 = iree_pydm.box %0 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    %status, %primitive = iree_pydm.unbox %1 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    iree_pydm.raise_on_failure %status : !iree_pydm.exception_result
    iree_pydm.return %primitive : !iree_pydm.integer
  }
}

// -----
builtin.module  {
  iree_pydm.func @test_pattern_call() -> (!iree_pydm.exception_result, !iree_pydm.object) attributes {arg_names = [], cell_vars = [], free_vars = []} {
    %0 = iree_pydm.constant true -> !iree_pydm.bool
    %exc_result, %result = iree_pydm.pattern_match_call(%0) : (!iree_pydm.bool) -> (!iree_pydm.exception_result, !iree_pydm.object) matching generic [@__logical_not_generic] specific [@__logical_not_bool]
    iree_pydm.raise_on_failure %exc_result : !iree_pydm.exception_result
    iree_pydm.return %result : !iree_pydm.object
  }
  iree_pydm.func private @__logical_not_generic(%arg0: !iree_pydm.object) -> (!iree_pydm.exception_result, !iree_pydm.object) attributes {arg_names = ["x"], cell_vars = [], free_vars = ["x"]} {
    iree_pydm.store_free_var "x", %arg0 : !iree_pydm.object
    %0 = iree_pydm.load_free_var "x" -> !iree_pydm.object
    %1 = iree_pydm.as_bool %0 : !iree_pydm.object -> !iree_pydm.bool
    %2 = iree_pydm.constant true -> !iree_pydm.bool
    %3 = iree_pydm.constant false -> !iree_pydm.bool
    %4 = iree_pydm.select %1, %3, %2 : !iree_pydm.bool
    %5 = iree_pydm.box %4 : !iree_pydm.bool -> !iree_pydm.object<!iree_pydm.bool>
    iree_pydm.return %5 : !iree_pydm.object<!iree_pydm.bool>
  }
  iree_pydm.func private @__logical_not_bool(%arg0: !iree_pydm.bool) -> (!iree_pydm.exception_result, !iree_pydm.bool) attributes {arg_names = ["x"], cell_vars = [], free_vars = ["x"]} {
    %0 = iree_pydm.box %arg0 : !iree_pydm.bool -> !iree_pydm.object<!iree_pydm.bool>
    iree_pydm.store_free_var "x", %0 : !iree_pydm.object<!iree_pydm.bool>
    %1 = iree_pydm.load_free_var "x" -> !iree_pydm.object
    %2 = iree_pydm.as_bool %1 : !iree_pydm.object -> !iree_pydm.bool
    %3 = iree_pydm.constant true -> !iree_pydm.bool
    %4 = iree_pydm.constant false -> !iree_pydm.bool
    %5 = iree_pydm.select %2, %4, %3 : !iree_pydm.bool
    iree_pydm.return %5 : !iree_pydm.bool
  }
}