PyDM with CPA Type Inference Fibonacci

input_pydm.mlir

module @module  {
  iree_pydm.func @main(%arg0: !iree_pydm.integer) -> (!iree_pydm.exception_result, !iree_pydm.integer) attributes {arg_names = ["n"], cell_vars = [], free_vars = ["n", "a", "b", "i", "c"]} {
    %i = alloc_free_var "i" -> !iree_pydm.free_var_ref
    %b = alloc_free_var "b" -> !iree_pydm.free_var_ref
    %a = alloc_free_var "a" -> !iree_pydm.free_var_ref
    %c = alloc_free_var "c" -> !iree_pydm.free_var_ref
    %n = alloc_free_var "n" -> !iree_pydm.free_var_ref
    %0 = box %arg0 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    store_var %n = %0 : !iree_pydm.free_var_ref, !iree_pydm.object<!iree_pydm.integer>
    %1 = constant 0 : i64 -> !iree_pydm.integer
    %2 = box %1 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    store_var %a = %2 : !iree_pydm.free_var_ref, !iree_pydm.object<!iree_pydm.integer>
    %3 = constant 1 : i64 -> !iree_pydm.integer
    %4 = box %3 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    store_var %b = %4 : !iree_pydm.free_var_ref, !iree_pydm.object<!iree_pydm.integer>
    %5 = constant false -> !iree_pydm.bool
    %6 = load_var %n : !iree_pydm.free_var_ref -> !iree_pydm.object
    %7 = constant 0 : i64 -> !iree_pydm.integer
    %left_prime, %right_prime = dynamic_binary_promote %6, %7 : !iree_pydm.object, !iree_pydm.integer
    %8 = apply_compare "eq", %left_prime, %right_prime : !iree_pydm.object, !iree_pydm.object
    %9 = as_bool %8 : !iree_pydm.bool -> !iree_pydm.bool
    %10 = bool_to_pred %9
    std.cond_br %10, ^bb1, ^bb2
  ^bb1:  // pred: ^bb0
    %11 = load_var %a : !iree_pydm.free_var_ref -> !iree_pydm.object
    %exc_result, %primitive = unbox %11 : !iree_pydm.object -> !iree_pydm.integer
    raise_on_failure %exc_result : !iree_pydm.exception_result
    return %primitive : !iree_pydm.integer
  ^bb2:  // pred: ^bb0
    %12 = constant 2 : i64 -> !iree_pydm.integer
    %13 = box %12 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    store_var %i = %13 : !iree_pydm.free_var_ref, !iree_pydm.object<!iree_pydm.integer>
    std.br ^bb3
  ^bb3:  // 2 preds: ^bb2, ^bb4
    %14 = constant false -> !iree_pydm.bool
    %15 = load_var %i : !iree_pydm.free_var_ref -> !iree_pydm.object
    %16 = load_var %n : !iree_pydm.free_var_ref -> !iree_pydm.object
    %left_prime_0, %right_prime_1 = dynamic_binary_promote %15, %16 : !iree_pydm.object, !iree_pydm.object
    %17 = apply_compare "le", %left_prime_0, %right_prime_1 : !iree_pydm.object, !iree_pydm.object
    %18 = as_bool %17 : !iree_pydm.bool -> !iree_pydm.bool
    %19 = bool_to_pred %18
    std.cond_br %19, ^bb4, ^bb5
  ^bb4:  // pred: ^bb3
    %20 = load_var %a : !iree_pydm.free_var_ref -> !iree_pydm.object
    %21 = load_var %b : !iree_pydm.free_var_ref -> !iree_pydm.object
    %left_prime_2, %right_prime_3 = dynamic_binary_promote %20, %21 : !iree_pydm.object, !iree_pydm.object
    %22 = apply_binary "add", %left_prime_2, %right_prime_3 : !iree_pydm.object, !iree_pydm.object -> !iree_pydm.object
    store_var %c = %22 : !iree_pydm.free_var_ref, !iree_pydm.object
    %23 = load_var %b : !iree_pydm.free_var_ref -> !iree_pydm.object
    store_var %a = %23 : !iree_pydm.free_var_ref, !iree_pydm.object
    %24 = load_var %c : !iree_pydm.free_var_ref -> !iree_pydm.object
    store_var %b = %24 : !iree_pydm.free_var_ref, !iree_pydm.object
    %25 = load_var %i : !iree_pydm.free_var_ref -> !iree_pydm.object
    %26 = constant 1 : i64 -> !iree_pydm.integer
    %left_prime_4, %right_prime_5 = dynamic_binary_promote %25, %26 : !iree_pydm.object, !iree_pydm.integer
    %27 = apply_binary "add", %left_prime_4, %right_prime_5 : !iree_pydm.object, !iree_pydm.object -> !iree_pydm.object
    store_var %i = %27 : !iree_pydm.free_var_ref, !iree_pydm.object
    std.br ^bb3
  ^bb5:  // pred: ^bb3
    %28 = load_var %b : !iree_pydm.free_var_ref -> !iree_pydm.object
    %exc_result_6, %primitive_7 = unbox %28 : !iree_pydm.object -> !iree_pydm.integer
    raise_on_failure %exc_result_6 : !iree_pydm.exception_result
    return %primitive_7 : !iree_pydm.integer
  }
}

input_var_type_refined.mlir

module @module  {
  iree_pydm.func @main(%arg0: !iree_pydm.integer) -> (!iree_pydm.exception_result, !iree_pydm.integer) attributes {arg_names = ["n"], cell_vars = [], free_vars = ["n", "a", "b", "i", "c"]} {
    %0 = constant 2 : i64 -> !iree_pydm.integer
    %1 = constant 1 : i64 -> !iree_pydm.integer
    %2 = constant 0 : i64 -> !iree_pydm.integer
    %3 = box %arg0 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    %4 = box %2 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    %5 = box %1 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    %6 = apply_compare "eq", %arg0, %2 : !iree_pydm.integer, !iree_pydm.integer
    %7 = bool_to_pred %6
    std.cond_br %7, ^bb1(%4 : !iree_pydm.object<!iree_pydm.integer>), ^bb2(%3, %4, %5 : !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>)
  ^bb1(%8: !iree_pydm.object<!iree_pydm.integer>):  // 2 preds: ^bb0, ^bb3
    %exc_result, %primitive = unbox %8 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    raise_on_failure %exc_result : !iree_pydm.exception_result
    return %primitive : !iree_pydm.integer
  ^bb2(%9: !iree_pydm.object<!iree_pydm.integer>, %10: !iree_pydm.object<!iree_pydm.integer>, %11: !iree_pydm.object<!iree_pydm.integer>):  // pred: ^bb0
    %12 = box %0 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    std.br ^bb3(%9, %12, %11, %10 : !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>)
  ^bb3(%13: !iree_pydm.object<!iree_pydm.integer>, %14: !iree_pydm.object<!iree_pydm.integer>, %15: !iree_pydm.object<!iree_pydm.integer>, %16: !iree_pydm.object<!iree_pydm.integer>):  // 2 preds: ^bb2, ^bb4
    %exc_result_0, %primitive_1 = unbox %14 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    %exc_result_2, %primitive_3 = unbox %13 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    %17 = apply_compare "le", %primitive_1, %primitive_3 : !iree_pydm.integer, !iree_pydm.integer
    %18 = bool_to_pred %17
    std.cond_br %18, ^bb4(%16, %14, %15, %13 : !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>), ^bb1(%15 : !iree_pydm.object<!iree_pydm.integer>)
  ^bb4(%19: !iree_pydm.object<!iree_pydm.integer>, %20: !iree_pydm.object<!iree_pydm.integer>, %21: !iree_pydm.object<!iree_pydm.integer>, %22: !iree_pydm.object<!iree_pydm.integer>):  // pred: ^bb3
    %exc_result_4, %primitive_5 = unbox %19 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    %exc_result_6, %primitive_7 = unbox %21 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    %23 = apply_binary "add", %primitive_5, %primitive_7 : !iree_pydm.integer, !iree_pydm.integer -> !iree_pydm.integer
    %24 = box %23 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    %exc_result_8, %primitive_9 = unbox %20 : !iree_pydm.object<!iree_pydm.integer> -> !iree_pydm.integer
    %25 = apply_binary "add", %primitive_9, %1 : !iree_pydm.integer, !iree_pydm.integer -> !iree_pydm.integer
    %26 = box %25 : !iree_pydm.integer -> !iree_pydm.object<!iree_pydm.integer>
    std.br ^bb3(%22, %26, %24, %21 : !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>, !iree_pydm.object<!iree_pydm.integer>)
  }
}

iree.mlir

module @module  {
  func @main(%arg0: i32) -> (i32, i32) {
    %c0_i32 = arith.constant 0 : i32
    %c1_i32 = arith.constant 1 : i32
    %c2_i32 = arith.constant 2 : i32
    %0 = iree.list.create : !iree.list<!iree.variant>
    %c2 = arith.constant 2 : index
    iree.list.resize %0, %c2 : !iree.list<!iree.variant>
    %c0 = arith.constant 0 : index
    %c78_i32 = arith.constant 78 : i32
    iree.list.set %0[%c0], %c78_i32 : !iree.list<!iree.variant>, i32
    %c1 = arith.constant 1 : index
    iree.list.set %0[%c1], %arg0 : !iree.list<!iree.variant>, i32
    %1 = iree.list.create : !iree.list<!iree.variant>
    %c2_0 = arith.constant 2 : index
    iree.list.resize %1, %c2_0 : !iree.list<!iree.variant>
    %c0_1 = arith.constant 0 : index
    %c78_i32_2 = arith.constant 78 : i32
    iree.list.set %1[%c0_1], %c78_i32_2 : !iree.list<!iree.variant>, i32
    %c1_3 = arith.constant 1 : index
    iree.list.set %1[%c1_3], %c0_i32 : !iree.list<!iree.variant>, i32
    %2 = iree.list.create : !iree.list<!iree.variant>
    %c2_4 = arith.constant 2 : index
    iree.list.resize %2, %c2_4 : !iree.list<!iree.variant>
    %c0_5 = arith.constant 0 : index
    %c78_i32_6 = arith.constant 78 : i32
    iree.list.set %2[%c0_5], %c78_i32_6 : !iree.list<!iree.variant>, i32
    %c1_7 = arith.constant 1 : index
    iree.list.set %2[%c1_7], %c1_i32 : !iree.list<!iree.variant>, i32
    %3 = arith.cmpi eq, %arg0, %c0_i32 : i32
    cond_br %3, ^bb1(%1 : !iree.list<!iree.variant>), ^bb5(%0, %1, %2 : !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>)
  ^bb1(%4: !iree.list<!iree.variant>):  // 2 preds: ^bb0, ^bb10
    %c0_8 = arith.constant 0 : index
    %c78_i32_9 = arith.constant 78 : i32
    %5 = iree.list.get %4[%c0_8] : !iree.list<!iree.variant> -> i32
    %6 = arith.cmpi eq, %c78_i32_9, %5 : i32
    cond_br %6, ^bb2, ^bb23
  ^bb2:  // pred: ^bb1
    %c1_10 = arith.constant 1 : index
    %c0_i32_11 = arith.constant 0 : i32
    %7 = iree.list.get %4[%c1_10] : !iree.list<!iree.variant> -> i32
    br ^bb3(%c0_i32_11, %7 : i32, i32)
  ^bb3(%8: i32, %9: i32):  // 2 preds: ^bb2, ^bb23
    %c0_i32_12 = arith.constant 0 : i32
    %10 = arith.cmpi eq, %c0_i32_12, %8 : i32
    cond_br %10, ^bb4, ^bb24
  ^bb4:  // pred: ^bb3
    %c0_i32_13 = arith.constant 0 : i32
    return %c0_i32_13, %9 : i32, i32
  ^bb5(%11: !iree.list<!iree.variant>, %12: !iree.list<!iree.variant>, %13: !iree.list<!iree.variant>):  // pred: ^bb0
    %14 = iree.list.create : !iree.list<!iree.variant>
    %c2_14 = arith.constant 2 : index
    iree.list.resize %14, %c2_14 : !iree.list<!iree.variant>
    %c0_15 = arith.constant 0 : index
    %c78_i32_16 = arith.constant 78 : i32
    iree.list.set %14[%c0_15], %c78_i32_16 : !iree.list<!iree.variant>, i32
    %c1_17 = arith.constant 1 : index
    iree.list.set %14[%c1_17], %c2_i32 : !iree.list<!iree.variant>, i32
    br ^bb6(%11, %14, %13, %12 : !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>)
  ^bb6(%15: !iree.list<!iree.variant>, %16: !iree.list<!iree.variant>, %17: !iree.list<!iree.variant>, %18: !iree.list<!iree.variant>):  // 2 preds: ^bb5, ^bb17
    %c0_18 = arith.constant 0 : index
    %c78_i32_19 = arith.constant 78 : i32
    %19 = iree.list.get %16[%c0_18] : !iree.list<!iree.variant> -> i32
    %20 = arith.cmpi eq, %c78_i32_19, %19 : i32
    cond_br %20, ^bb7, ^bb18
  ^bb7:  // pred: ^bb6
    %c1_20 = arith.constant 1 : index
    %c0_i32_21 = arith.constant 0 : i32
    %21 = iree.list.get %16[%c1_20] : !iree.list<!iree.variant> -> i32
    br ^bb8(%c0_i32_21, %21 : i32, i32)
  ^bb8(%22: i32, %23: i32):  // 2 preds: ^bb7, ^bb18
    %c0_22 = arith.constant 0 : index
    %c78_i32_23 = arith.constant 78 : i32
    %24 = iree.list.get %15[%c0_22] : !iree.list<!iree.variant> -> i32
    %25 = arith.cmpi eq, %c78_i32_23, %24 : i32
    cond_br %25, ^bb9, ^bb19
  ^bb9:  // pred: ^bb8
    %c1_24 = arith.constant 1 : index
    %c0_i32_25 = arith.constant 0 : i32
    %26 = iree.list.get %15[%c1_24] : !iree.list<!iree.variant> -> i32
    br ^bb10(%c0_i32_25, %26 : i32, i32)
  ^bb10(%27: i32, %28: i32):  // 2 preds: ^bb9, ^bb19
    %29 = arith.cmpi sle, %23, %28 : i32
    cond_br %29, ^bb11(%18, %16, %17, %15 : !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>), ^bb1(%17 : !iree.list<!iree.variant>)
  ^bb11(%30: !iree.list<!iree.variant>, %31: !iree.list<!iree.variant>, %32: !iree.list<!iree.variant>, %33: !iree.list<!iree.variant>):  // pred: ^bb10
    %c0_26 = arith.constant 0 : index
    %c78_i32_27 = arith.constant 78 : i32
    %34 = iree.list.get %30[%c0_26] : !iree.list<!iree.variant> -> i32
    %35 = arith.cmpi eq, %c78_i32_27, %34 : i32
    cond_br %35, ^bb12, ^bb20
  ^bb12:  // pred: ^bb11
    %c1_28 = arith.constant 1 : index
    %c0_i32_29 = arith.constant 0 : i32
    %36 = iree.list.get %30[%c1_28] : !iree.list<!iree.variant> -> i32
    br ^bb13(%c0_i32_29, %36 : i32, i32)
  ^bb13(%37: i32, %38: i32):  // 2 preds: ^bb12, ^bb20
    %c0_30 = arith.constant 0 : index
    %c78_i32_31 = arith.constant 78 : i32
    %39 = iree.list.get %32[%c0_30] : !iree.list<!iree.variant> -> i32
    %40 = arith.cmpi eq, %c78_i32_31, %39 : i32
    cond_br %40, ^bb14, ^bb21
  ^bb14:  // pred: ^bb13
    %c1_32 = arith.constant 1 : index
    %c0_i32_33 = arith.constant 0 : i32
    %41 = iree.list.get %32[%c1_32] : !iree.list<!iree.variant> -> i32
    br ^bb15(%c0_i32_33, %41 : i32, i32)
  ^bb15(%42: i32, %43: i32):  // 2 preds: ^bb14, ^bb21
    %44 = arith.addi %38, %43 : i32
    %45 = iree.list.create : !iree.list<!iree.variant>
    %c2_34 = arith.constant 2 : index
    iree.list.resize %45, %c2_34 : !iree.list<!iree.variant>
    %c0_35 = arith.constant 0 : index
    %c78_i32_36 = arith.constant 78 : i32
    iree.list.set %45[%c0_35], %c78_i32_36 : !iree.list<!iree.variant>, i32
    %c1_37 = arith.constant 1 : index
    iree.list.set %45[%c1_37], %44 : !iree.list<!iree.variant>, i32
    %c0_38 = arith.constant 0 : index
    %c78_i32_39 = arith.constant 78 : i32
    %46 = iree.list.get %31[%c0_38] : !iree.list<!iree.variant> -> i32
    %47 = arith.cmpi eq, %c78_i32_39, %46 : i32
    cond_br %47, ^bb16, ^bb22
  ^bb16:  // pred: ^bb15
    %c1_40 = arith.constant 1 : index
    %c0_i32_41 = arith.constant 0 : i32
    %48 = iree.list.get %31[%c1_40] : !iree.list<!iree.variant> -> i32
    br ^bb17(%c0_i32_41, %48 : i32, i32)
  ^bb17(%49: i32, %50: i32):  // 2 preds: ^bb16, ^bb22
    %51 = arith.addi %50, %c1_i32 : i32
    %52 = iree.list.create : !iree.list<!iree.variant>
    %c2_42 = arith.constant 2 : index
    iree.list.resize %52, %c2_42 : !iree.list<!iree.variant>
    %c0_43 = arith.constant 0 : index
    %c78_i32_44 = arith.constant 78 : i32
    iree.list.set %52[%c0_43], %c78_i32_44 : !iree.list<!iree.variant>, i32
    %c1_45 = arith.constant 1 : index
    iree.list.set %52[%c1_45], %51 : !iree.list<!iree.variant>, i32
    br ^bb6(%33, %52, %45, %32 : !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>, !iree.list<!iree.variant>)
  ^bb18:  // pred: ^bb6
    %c-4_i32 = arith.constant -4 : i32
    %c0_i32_46 = arith.constant 0 : i32
    br ^bb8(%c-4_i32, %c0_i32_46 : i32, i32)
  ^bb19:  // pred: ^bb8
    %c-4_i32_47 = arith.constant -4 : i32
    %c0_i32_48 = arith.constant 0 : i32
    br ^bb10(%c-4_i32_47, %c0_i32_48 : i32, i32)
  ^bb20:  // pred: ^bb11
    %c-4_i32_49 = arith.constant -4 : i32
    %c0_i32_50 = arith.constant 0 : i32
    br ^bb13(%c-4_i32_49, %c0_i32_50 : i32, i32)
  ^bb21:  // pred: ^bb13
    %c-4_i32_51 = arith.constant -4 : i32
    %c0_i32_52 = arith.constant 0 : i32
    br ^bb15(%c-4_i32_51, %c0_i32_52 : i32, i32)
  ^bb22:  // pred: ^bb15
    %c-4_i32_53 = arith.constant -4 : i32
    %c0_i32_54 = arith.constant 0 : i32
    br ^bb17(%c-4_i32_53, %c0_i32_54 : i32, i32)
  ^bb23:  // pred: ^bb1
    %c-4_i32_55 = arith.constant -4 : i32
    %c0_i32_56 = arith.constant 0 : i32
    br ^bb3(%c-4_i32_55, %c0_i32_56 : i32, i32)
  ^bb24:  // pred: ^bb3
    %c0_i32_57 = arith.constant 0 : i32
    return %8, %c0_i32_57 : i32, i32
  }
}

sample.py

  def test_spaceopt(self):
    @jit(debug=2)
    def compute(n: int) -> int:
      a = 0
      b = 1
      if n == 0:
        return a
      i = 2
      while i <= n:  # TODO: Upgrade to for...range
        c = a + b
        a = b
        b = c
        i = i + 1  # TODO: Support AugAssign
      return b
    print(compute(20))