joelberkeley / spidr

Accelerated machine learning with dependent types
Apache License 2.0
73 stars 4 forks source link

Improve type safety of Idris API to XLA #273

Open joelberkeley opened 2 years ago

joelberkeley commented 2 years ago

The Tensor backend is unscalable. Most apparently, it completely lack type safety. Rewrite the backend to make it type safe where possible.

Areas that come to mind:

  1. create named wrappers for [GC]AnyPtrs. These can be as simple as
    data XlaOp : Type where
        MkXlaOp : GCAnyPtr -> XlaOp

    though we may find it useful to embellish these with other features. See 3.

  2. wrap all foreign functions in Idris functions that capture the requirements of the functions. For example, functions that require a list and the length of that list can be wrapped to just require the list. Similarly, positive integers can be exposed as Nat then converted to Int when passed to FFI. For example, functions that require a list and the length of that list can be wrapped to just require the list. Similarly, positive integers can be exposed as Nat then converted to Int when passed to FFI
  3. capture XLA rules in types. For example,
    • each XlaOp can only be registered to one XlaBuilder
    • a computation cannot be re-used once it has been built (we may already have this due to michael messer's suggestion)
joelberkeley commented 2 years ago

This was partially completed in #274