An introduction to Mojo 🔥

28th May 2023

A new programming language

"for AI/....... developers"

  • Ease of Python
  • Performance of C++ and Rust
  • Ablility to leverage the Entire Python Ecosystem

Why ?

Some (Marketing) Highlights

  • compiler technologies with integrated caching, multithreading, and cloud distribution technologies
  • autotuning and compile-time metaprogramming features allow you to write code that is portable to even the most exotic hardware

PSA

  • new language, I probably am going to make some mistakes here
  • my analogies might be very poor, so throw them away as soon as it breaks

Mojo🔥 what ?

  • allows you to leverage the entire Python ecosystem
  • designed to become a superset of Python over time
    • preserves Python’s dynamic features
    • adds new primitives for systems programming
  • bringing the best of dynamic lang. and systems lang.
    • unified programming model
    • scripting, applications programming, systems programming, accelerators etc.

"It's just not there today"

  • At the moment, Mojo is still a work in progress and the documentation is targeted to developers with systems programming experience.
  • As the language grows and becomes more broadly available, we intend for it to be friendly and accessible to everyone, including beginner programmers. It’s just not there today.

as quoted from the Mojo manual - 28th May '23

Fine fine ! What does this all mean ?

Let's dig in

  • We will be following the Mojo manual
  • btw, will not be covering all language features today
  • will stop right before Parameterization : compile-time metaprogramming, the rest for the next round

Mojo lets you switch modes

  • normally it acts (semantically) like Python

 

BUT, when you invoke the spells

  • it switches into a whole different language

 

At least, that's my experience/understanding

Let's dig into these spells, and see the building blocks.

Where do we begin ?

  • Mojo is based on MLIR and LLVM
    • cutting-edge compiler and code generation system
    • better control over data organization, direct access to data fields, and other ways to improve performance. (Think pointers in C)
  • However, in modern prog. langs. we want to build high-level and safe abstractions on top of the lower level operations without loss of performance

let & var, walked into a bar

  • let is for immutable, and var for mutable
    • declare that a value is immutable for type-safety and performance.
    • support late initialization, type specifiers
    • get an error if they mistype a variable name in an assignment
      • especially late assignment
      • unused var (not the second time though, esp. in def declaration)
    • completely opt-in when in def declarations

we start with a struct type

  • high level abstraction for data access
  • A Mojo struct is very similar to Python class
  • Similarities
    • methods
    • fields
    • operator overloading
    • decorators for metaprogramming
    • etc.
  • Not the fun part, lets look at what's different

Mojo struct vs Python class

  • Python classes are dynamic
    • allows for dynamic dispatch (selecting which implementation of method/function to call at runtime)
    • monkey patching (dynamically update behaviour of code at runtime)
    • dynamically binding properties at runtime
  • Mojo structs are static
    • bound at compile time (cannot add methods at runtime)
    • trading flexibility for performance

struct-aaast-ic

  • Because of its static nature
    • Mojo will run the code FASTER
    • the program knows where to find the structs information
    • it knows how to use it without extra steps
    • no need for inspection at runtime, all the information is already baked in

...Mojo struct things

  • All instance properties on struct MUST be explicitly declared with let or var
    • contents and structure set(locked) at compile time
    • can’t be modified when running (unlike Python)
    • cannot use del to remove method or change its value
  • struct works well with operator oveloading
    • overload symbols for your own data

...more Mojo struct things

  • All “standard types” (Int, Bool, String, Tuple) are made up of struct
    • not hardwired into the language
    • nothing special about them
    • included in the "standard set of tools"

Int vs. int ! FIGHT !!!

  • Int vs int
    • Python int type can handle really big numbers and has some extra features, like checking if two numbers are the same object.
      • comes with some extra baggage that can slow things down
    • Mojo’s Int is designed to be simple, fast, and tuned for your computer’s hardware to handle quickly.

Strong type checking

  • mix Python’s flexible types & Mojo's strict type checks
  • primary method is to use struct
    • A struct defines a compile-time-bound name
    • references to that name in a type context are treated as a strong specification for the value being defined
  • Not only for Type checking
    • once we know the types are accurate
    • we can optimize the code based on those types
    • the foundation of the safety and predictability provided by Mojo

Overloaded functions and methods

  • leave your argument names without type definitions
    • then the function behaves just like Python
  • when type safety is needed, full support for overloaded functions and methods
    • define multiple functions with the same name but with different arguments
  • As soon as you define a single argument type
    • Mojo will look for overloaded candidates and resolve function calls

let's talk about def first

  • defined to be very dynamic, flexible and generally compatible with Python
  • arguments are mutable
  • local variables are implicitly declared on first use
  • and scoping isn’t enforced.
  • great for high level programming and scripting
    • but not always great for systems programming

strict but fn

  • fn is “strict mode” for def, the callee doesn't care
  • fn has the following limitations
    • arg. values default to being immutable (like let)
    • arg. values require a type specification
      • missing return type is interpreted as returning None (instead of unknown return type)
    • implicit declaration of local variable is disabled
      • all locals must be declared (let or var)
    • raising exceptions must be explicitly declared with a raises keyword on the fn

It's all about the semantics

  • Value semantics vs Reference semantics
  • Let's also chat a bit about lifetimes and "move"s, somewhat Rusty before we head further

__copyinit__, __moveinit__, innit ?

  • __copyinit__ and __moveinit__ special methods
  • Mojo supports full “value semantics” as present in langs. like C++, Swift etc.
  • Custom copy and move constructors can be supplied
    • __copyinit__ ( __init__ but "how to copy")
    • __moveinit__ ( __init__ but "how to move")
    • useful when doing low level systems programming
    • control the lifetime of a value, including making types copyable, move-only, and not-movable.
    • more control than languages like Swift & Rust offer, which require values to at least be movable

Argument passing control and memory ownership

  • the whole language is function calls (passing arguments to functions)
  • lots of behaviour implemented with dunder(magic) calls
  • in this dunder calls a lot of memory ownership is determined with argument passing
  • let's compare Python and Mojo semantics further

def (initely) fn defaults

  • All values passed into a Python def function use reference semantics.
  • All values passed into a Mojo def function use value semantics by default
    • the function receives a copy of the arguments, can modify arguments but will not be seen modified outside, unlike Python’s mutable arguments
  • All values passed into a Mojo fn function are immutable references by default.
    • also known as “borrowing”
  • How do you change the defaults, though ?

Argument conventions : Python

  • all fundamental values are references to objects
  • a Python function can modify the original object
  • Python devs. are “used to thinking” of everything in reference semantics
  • at the CPython level, you can see that the references themselves are actually passed by-copy . Python copies a pointer and adjusts reference counts.
  • comfortable programming paradigm, but requires everything to be allocated on the heap
  • Should we talk about heap vs. stack vs. registers

Argument conventions : Mojo

  • we want the values to live on the stack OR even on hardware registers (FAAAAST)
  • structs are always inlined to their container
    • whether they’re a field of another type
    • or into the stack frame of the containing function
  • classes(when they’re ready) will do the same as Python (reference semantics)
    • but isn’t practical for simple well defined types like integers

Argument conventions : Mojo 🤔

  • how do you implement methods that needs to mutate self (as __iadd__)
  • how does let work, prevent mutation & how are lifetimes of these values controlled for Mojo to be memory safe ?
  • Mojo compiler uses dataflow analysis and type annotations to provide full control over value copies, aliasing of references, and mutation control
  • Similar to Rust lifetimes, but work a bit differently to make Mojo easier to learn and integrates better in Python ecosystem (without huge annotation burden)

Memory ownership : fn arguments

  • memory ownership for objects passed into Mojo fn as arguments
  • not just the defaults
  • can be specified for each argument object as needed
    • borrowed
    • inout
    • owned

borrowed : Immutable arguments

  • borrowed object is an immutable reference to an object that a function receives
  • callee function has full read-and-execute access to the object, but cannot modify it
  • the caller still has exclusive “ownership” of the object
  • DEFAULT behaviour for fn arguments, made explicit with borrowed keyword
  • applies to all arguments uniformly, including the self argument of methods
  • SEE, you cannot just update self values

still borrowed : Immutable arguments

  • very efficient when passing large values or expensive values (like reference counted pointers)
    • because the copy constructor and destructor don’t have to be invoked when passing the argument
  • implements a borrow checker that prevents multiple mutable references to the same value
    • Multiple borrows okay, but mutable reference and a borrow at the same time not okay
  • Small values like Int, Float, and SIMD are passed directly in machine registers (FAAST)
    • decorated with @register_passable

a quick borrow : Rust vs. Mojo

  • In both cases, borrow checker enforces the exclusivity of invariants.
  • Mojo does not require a sigil(sign) on the caller side to pass by borrow
  • Mojo is more efficient when passing small values (machine registers)
  • Rust defaults to moving values while Mojo passes them around by borrow
  • policy and syntax decisions allow Mojo to provide an easier-to-use programming model, probably !!!

 

inout : Mutable arguments

  • changes to the argument inside the function are visible outside the function
  • declare  as mutable with the inout keyword
  • multiple inouts possible
  • also for non-trival things, like array subscript-ing
    • compiler will reject setting array subscript if __getitem__ is implemented but not __setitem__
  • not called “pass by reference” even if the inout convention is the same, since the internal impl. may actually pass values using pointers

owned & ^ : Transfer arguments

  • take exclusive ownership over a value via argument
  • often used with the postfix ^ operator at the call site on the value that we want to transfer
  • very useful for things like unique pointers, or avoid any sort of copies
  • owned convention on destructors and on consuming move constructor
  • Specifying owned in the del function is important because you must own a value to destroy it.

def vs. fn argument passing : Mojo

  • Mojos def is just sugaring for the fn function
  • def argument without an explicit type annotation defaults to Object
  • def argument without a convention keyword (such as inout or owned) is passed by implicit copy into a shadow mutable var with the same name as the argument
  • requires that the type have a __copyinit__ method
  • shadow copies typically add no overhead, because references for small types like Object are cheap to copy

more values, more lifetimes

  • Mojo provides very good lifetimes support
  • comparable to other mainstream languages with lifetimes, hinted to be perhaps even more powerful in the future
  • customizable as needed
  • less syntax burden that promises ease of usage

if you 🔥 can't beat them 🐍

  • import any other Python module
    • Python.import_module("module_name")
    • let mod_name = Python.import_module("module_name") is equivalent to import module_name as mod_name
  • Currently, CANNOT import members of the module
    • must import the whole Python module and then access members through the module name
  • Local Python modules
    • Add the directory to Python path
      • Python.add_to_path("path/to/module")

mo' interop, no mo' problems 🤞

  • no need to worry about memory management when using Python in Mojo.
  • BUT also works the other way around
  • Mojo primitive types implicitly convert to Python objects
    • lists, tuples, integers, floats, booleans, and strings
  • Mojo doesn’t have a standard Dictionary yet
    • but you can work with Python dictionaries in Mojo though!
    • use the dict method

only the foundations

  • there's lot more to Mojo
  • only scratched the surface today
  • if enough interest, we can continue with the rest of language features in the coming weeks
  • let's try these features out in the playground now

...continued

Compile Time Metaprogramming

  • "program the program" - metaprogramming
  • built into the compiler as a separate stage of compilation
  • after parsing, semantic analysis, and IR generation
    • ---- HAPPENING HERE ----
  • but before lowering to target-specific code.
  • uses the same host language for runtime programs as it does for metaprograms
  • leverages MLIR to represent and evaluate these programs predictably

 

Some naming thing before we proceed

  • In Python
    • “arguments” and “parameters” used fairly interchangeably
    • for “things that are passed into functions.”
  • In Mojo
    • “parameter” and “parameter expression” to represent compile-time values
    • “argument” and “expression” to refer to runtime values
  • use the words “parameterized” and “parametric” for compile-time metaprogramming

Parameterized types and functions

  • Mojo structs and functions both can be parameterized
  • Let's motivate this with an example
    • SIMD type
      • a low-level vector register in hardware
      • holds multiple instances of a scalar data-type
    • Hardware accelerators these days are getting exotic
      • some CPUs have 512-bit or longer SIMD vectors
      • diversity in hardware(eg. SSE, AVX-512, NEON, SVE ...)
      • but many operations are common and used by numerics and ML kernel developers
      • the SIMD type exposes them all with same api to Mojo programmers

Parameters in Mojo

  • declared in square brackets
  • are named and have types like the rest of Mojo program
  • BUT, they are evaluated at compile time
  • the runtime program may use the value of parameters
    • because the parameters are resolved at compile-time (before needed by the runtime)
  • but the compile-time parameter expressions may not use runtime values
    • which makes sense if you think about the compile time not being able to depend on run time behaviour
  • Let's look at an example next

Parameterization : Example

Text