This project is a functional compiler and code-gen tool that will accelerate writing AI algorithms as well as making them easier. The core is a lisp-like IR that can be translated from high-level languages, and can be linked to a variety of backends to generate code.
Currently implemented frontends are
- TorchScript: ts2k
- Julia: j2ks
- F#: f2k
- ONNX: onnx2ks
- KS-Lisp: The IR itself is exchanged in a lisp-like text format (see below).
Current backends:
- CPU/C++: In src/ksc and src/runtime
- GPU/Futhark: Also in src/ksc
- MLIR: In mlir
Current transformers:
- KSC: Various Autodiff and optimization transforms, in Haskell
- RLO: A rewriting-based optimizer using reinforcement learning
The text form of the IR is "low sugar", focused more on simplicity than user-friendliness. However it is human writable, and lispers in particular may like to play with it. There's a VS Code syntax highlighter extension in etc/ks-vscode.
The AST has just a few concepts: Lambda, Call, Let binding, Conditionals, Constants, Assert, Top-level function definitions, and Rewrite rules.
It's also strongly typed, using a bare set of types: Tensor, Tuple, and Values (Float, Integer, String). However, unlike say MLIR or TorchScript, these are composable, so e.g. we can have a (Tensor 1 (Tuple Float (Tensor 2 Float)))
. Tensors are annotated with their "rank" or number of dimensions.
The IR is pure functional, so functions may be called more than once or not at all depending on the use of their outputs. Values are passed and returned "as if" by value, although the optimizer reserves the right implement this using refcounting/copy/move, as long as the semantics are preserved.
The lisp-like IR is extremely simple -- all the language builtins are in this code:
;; Externally defined function "sqrt" returns a Float, takes two Float
(edef atan2 Float (Float Float))
#| Block comments
-- User-defined function myfun
takes an Integer and Tensor of (Float Float) pairs
and returns a pair of String and Float
|#
(def myfun ; function name
(Tuple String Float) ; return type
((i : Integer) ; argument 1: int
(v : Tensor 1 (Tuple Float Float))) ; argument 2: tensor of tuple
(assert (gt i 0) ; (assert TEST BODY)
(if (eq i 0) ; (if TEST TEXPR FEXPR)
; "then" br
(tuple "fred" ; tuple constructor
(let (tmp (index 0 v)) ; (let (VAR VAL) BODY)
(mul (get$1$2 tmp) 2.0)) ) ; no builtins -- even mul is a function
; "else" br
(let ((t1 (index 0 v)) ; (let ((VAR1 VAL1) ... (VARn VALn)) BODY)
(t2 (myfun (sub i 1) v))) ; Recursive call to myfun
t2))))
;; Rewrite rule
(rule "mul.commute" ; rule name
((a : Float) (b : Float)) ; "template_vars": free variables in the template
(mul a b) ; "template": a pattern to match
(mul b a)) ; replacement
; Rules are *not* applied exhaustively, that's very intentionally left up to the compiler.
; For example the rule above, if matched exhaustively would not terminate.
;; And compilation produces f and various derivatives, as if
(edef rev$myfun ; classic reverse mode. If f :: S -> T, then rev$f :: (S, dT) -> dS
(Tuple (Tuple) (Tensor 1 (Tuple Float Float))) ; dS is tangent-type of inputs (dInteger = void)
((#|s |# (Tuple Integer (Tensor 1 (Tuple Float Float)))) ; inputs in a single tuple
(#|dt |# (Tuple (Tuple) Float)))) ; dT is tangent type of returns
(edef D$myfun ; linear map as per Elliot. If f :: S->T, D$f :: S -> (LM dS dT) where LM is linear map
(LM
(Tuple (Tuple) (Tensor 1 (Tuple Float Float))) ; dS
(Tuple (Tuple) Float) ; dT
)
(Integer (Tensor 1 (Tuple Float Float)))) ; inputs as normal (not single-tupled)
See the ksc syntax primer for a longer
introduction to the ks
language. The ksc test
directory provides more examples of the constructs
available when writing .ks
files.
Knossos will only be a successful project if the onboarding experience is straightforward. We consider any difficulty getting started whilst following these instructions to be a critical issue that we should fix urgently. Therefore if you experience any difficulty getting started please follow these steps:
-
File an issue with the title "Difficulty onboarding" that explains as much as possible about the difficulty you are having.
-
Email knossos@service.microsoft.com with the subject "Urgent: difficulty onboarding to knossos-ksc" with a link to the new issue you just filed.
We will respond to you as a matter of top priority.
Please report your experience of onboarding, regardless of whether it was good or bad. It is hard to test onboarding automatically and so we rely on new users to tell us about their experience. After following this guide, whether you were successful or not, please
- File an issue with the title "Experience report: new user onboarding" describing how you found your onboarding experience.
Many thanks, the Knossos team.
Run the ksc
executable as follows to differentiate, compile and run
a .ks
program. This example runs hello-world.ks
.
./build/bin/ksc --compile-and-run \
--ks-source-file src/runtime/prelude.ks \
--ks-source-file test/ksc/hello-world.ks \
--ks-output-file obj/test/ksc/hello-world.kso \
--cpp-output-file obj/test/ksc/hello-world.cpp \
--c++ g++ \
--exe-output-file obj/test/ksc/hello-world.exe
or with PowerShell syntax:
./build/bin/ksc --compile-and-run `
--ks-source-file src/runtime/prelude.ks `
--ks-source-file test/ksc/hello-world.ks `
--ks-output-file obj/test/ksc/hello-world.kso `
--cpp-output-file obj/test/ksc/hello-world.cpp `
--c++ g++ `
--exe-output-file obj/test/ksc/hello-world.exe
Which should produce output like this:
read decls
Linted Typechecked defs
Linted Grad
Linted Grad tupled
Linted Optgrad
Linted Optgrad tupled
Linted Diffs
Linted OptDiffs
Linted CSE
Writing to obj/test/ksc/hello-world.kso
Writing to obj/test/ksc/hello-world.cpp
Compiling: g++ -fmax-errors=5 -fdiagnostics-color=always -Wall -Wno-unused -Wno-maybe-uninitialized -Isrc/runtime -O3 -g -std=c++17 -o obj/test/ksc/hello-world.exe obj/test/ksc/hello-world.cpp
Running
Hello world!
If you are seeing this output then knossos-ksc has successfully compiled and run the hello-world.ks program!
To run the ksc self-tests use the command line
./build/bin/ksc --test --fs-test out.fs
(Don't worry if the final test, of out.fs
, fails. It is a test for
F#-to-ks, which most users will not have set up.)
To generate a .kso
file from a .ks
file without differentiating,
i.e. to type check and apply ksc's heuristic optimisations, use the
command line
./build/bin/ksc --generate-cpp-without-diffs \
--ks-source-file src/runtime/prelude.ks \
--ks-source-file input.ks \
--ks-output-file output.ks \
--cpp-output-file output.cpp
In the compiler, the IR is defined in Lang.hs
.
The syntax is defined by the parser in
Parse.hs
and the pretty-printer in
Lang.hs
. testRoundTrip
in
Main.hs
checks that they agree.
knossos-ksc has a continuous integration build set up on Azure DevOps. To manually run a CI build click "Run pipeline", type the name of your branch under "Branch/tag", and click "Run".
Collaboration on this project is subject to the Microsoft Open Source Code of Conduct.