View Source EXLA.MLIR.Module (EXLA v0.7.3)

Representation of an MLIR module.

Summary

Functions

Adds a new function to the MLIR module.

Creates a new MLIR module.

Functions

Link to this function

add_function(module, name, arg_shapes, return_shapes)

View Source

Adds a new function to the MLIR module.

Link to this function

compile(module, client, argument_shapes, return_shape, options \\ [])

View Source

Compiles a module into an executable.

Options

  • :device_id - the device id to compile to and run the executable on. Defaults to the :default_device_id on the client. If :num_replicas or :num_partitions are given, this option is ignored and the device id is set to -1.

  • :num_replicas - the number of replicas this computation will run on. It defaults to 1 but you can set it if you want to enable single-program multiple data.

  • :use_spmd - enables Single-Program Multiple-Data partioning. This is set to true if :num_partitions is more than one, otherwise is false.

Currently those options do not have an effect as they related to running the same compiled executable on multiple replicas.

Some options apply to TPU only:

  • :num_partitions - the number of partitions this computation will run on.
Link to this function

new(arg_shapes, return_shape, fun)

View Source

Creates a new MLIR module.