See the IREE docs/using_colab.md document for instructions.

This notebook shows off some concepts of the low level IREE python bindings.


In [0]:
import numpy as np
from pyiree.tf import compiler as ireec
from pyiree import rt as ireert

In [0]:
# Compile a module.
ctx = ireec.Context()
input_module = ctx.parse_asm("""
  module @arithmetic {
    func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
          attributes { iree.module.export } {
        %0 = "mhlo.multiply"(%arg0, %arg1) {name = "mul.1"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
        return %0 : tensor<4xf32>
    } 
  }
  """)
blob = input_module.compile()
m = ireert.VmModule.from_flatbuffer(blob)

In [36]:
# Register the module with a runtime context.
# Use the CPU interpreter (which has the most implementation done):
driver_name = "interpreter"
# Live on the edge and give the vulkan driver a try:
# driver_name = "vulkan"
config = ireert.Config(driver_name)
ctx = ireert.SystemContext(config=config)
ctx.add_module(m)

# Invoke the function and print the result.
print("INVOKE simple_mul")
arg0 = np.array([1., 2., 3., 4.], dtype=np.float32)
arg1 = np.array([4., 5., 6., 7.], dtype=np.float32)
f = ctx.modules.arithmetic["simple_mul"]
results = f(arg0, arg1)
print("Results:", results)


INVOKE simple_mul
Results: [ 4. 10. 18. 28.]
Created IREE driver interpreter: <pyiree.rt.binding.HalDriver object at 0x7fae493257f0>
SystemContext driver=<pyiree.rt.binding.HalDriver object at 0x7fae493257f0>