Параметризация: метапрограммирование на этапе компиляции

Во многих языках есть средства для метапрограммирования, то есть для написания кода, который генерирует или изменяет код на этапе компиляции. В Python есть средства для динамического метапрограммирования: такие функции, как декораторы, метаклассы и многое другое. Эти функции делают Python очень гибким и производительным, но, поскольку они динамичны, они сопряжены с большими затратами времени выполнения. Другие языки имеют статические функции метапрограммирования или функции метапрограммирования на этапе компиляции, такие как макросы препроцессора C и шаблоны C++. Они могут быть ограниченными и сложными в использовании.

Чтобы поддержать работу Modular в области искусственного интеллекта, Mojo стремится обеспечить мощное, простое в использовании метапрограммирование с нулевыми затратами на выполнение. Это метапрограммирование на этапе компиляции использует тот же язык, что и программы во время выполнения, поэтому вам не нужно изучать новый язык - только несколько новых функций.

Основная новая функция - это параметры. Вы можете представить параметр как переменную во время компиляции, которая становится константой во время выполнения. Такое использование термина "параметр", вероятно, отличается от того, к которому вы привыкли в других языках, где "параметр" и "аргумент" часто используются взаимозаменяемо. В Mojo "параметр" и "выражение параметра" относятся к значениям времени компиляции, а "аргумент" и "выражение" относятся к значениям времени выполнения.

В Mojo вы можете добавлять параметры к структуре или функции. Вы также можете определить выражения именованных параметров — comptime значения, которые можно использовать в качестве констант времени выполнения.

Параметризованные функции

Чтобы определить параметризованную функцию, добавьте параметры в квадратных скобках перед списком аргументов. Каждый параметр оформляется так же, как и аргумент: имя параметра, за которым следует двоеточие и тип (который является обязательным). В следующем примере функция имеет единственный параметр count типа Int.

fn repeat[count: Int](msg: String):
    @parameter
    for i in range(count):
        print(msg)

Показанный здесь декоратор @parameter вызывает вычисление цикла for во время компиляции. Декоратор работает только в том случае, если пределы цикла являются константами времени компиляции. Поскольку count является параметром, range(count) может быть вычислен во время компиляции.

Вызывая параметризованную функцию, вы вводите значения для параметров точно так же, как и аргументы функции:

repeat[3]("Hello")

Hello
Hello
Hello

Компилятор определяет значения параметров во время компиляции и создает конкретную версию функции repeat[]() для каждого уникального значения параметра. После определения значений параметров и завершения цикла функция repeat[3]() будет примерно эквивалентна этой:

fn repeat_3(msg: String):
    print(msg)
    print(msg)
    print(msg)

Это не соответствует фактическому коду, сгенерированному компилятором. К моменту определения параметров код Mojo уже был преобразован в промежуточное представление в MLIR.

Если компилятор не может преобразовать все значения параметров в постоянные, компиляция завершается неудачей.

Структура списка параметров

Параметры функции или структуры заключаются в квадратные скобки после названия функции или структуры. Для параметров всегда требуется указание типа.

Когда вы просматриваете определение функции или структуры, вы можете увидеть некоторые специальные символы, такие как / и *, в списке параметров. Вот пример:

def my_sort[
    # infer-only parameters
    dtype: DType,
    width: Int,
    //,
    # positional-only parameter
    values: SIMD[dtype, width],
    /,
    # positional-or-keyword parameter
    compare: fn (Scalar[dtype], Scalar[dtype]) -> Int,
    *,
    # keyword-only parameter
    reverse: Bool = False,
]() -> SIMD[dtype, width]:

Вот краткий обзор специальных символов в списке параметров:

• Двойная косая черта (//): параметры, объявленные перед двойной косой чертой, являются параметрами только для вывода.
• Косая черта (/): параметры, объявленные перед косой чертой, являются параметрами только для определения положения. Параметры, относящиеся только к позиции и ключевому слову, подчиняются тем же правилам, что и аргументы, относящиеся только к позиции и ключевому слову.
• Имя параметра, перед которым ставится звездочка, например *Types, идентифицирует переменный параметр (не показанный в примере выше). Все параметры, следующие за вариативным параметром, относятся только к ключевым словам.
• Звездочка (*): в списке параметров, не содержащих вариативного параметра, звездочка сама по себе указывает, что следующие параметры относятся только к ключевым словам.
• Знак равенства (=) указывает на значение по умолчанию для необязательного параметра.

Параметры и обобщённые типы(дженерики)

"Дженерики" относятся к функциям, которые могут работать с несколькими типами значений, или контейнерам, которые могут содержать несколько типов значений. Например, List может содержать значения разных типов, поэтому у вас может быть список значений типа Int или список значений String.

В Mojo дженерики используют параметры для указания типов. Например, List принимает параметр типа, поэтому вектор целых чисел записывается как List[Int]. Таким образом, все дженерики используют параметры, но не все, что использует параметры, является дженериком.

Например, функция repeat[](), описанная в предыдущем разделе, содержит параметр типа Int и аргумент типа String. Она параметризована, но не универсальна. Универсальная функция или структура параметризуется в зависимости от типа. Например, мы могли бы переписать repeat[](), чтобы он принимал любой тип аргумента, соответствующий признаку Stringable:

fn repeat[MsgType: Stringable, //, count: Int](msg: MsgType):
    @parameter
    for i in range(count):
        print(String(msg))


def main():
    # MsgType is always inferred, so first positional keyword `2` is
    # passed to `count`
    repeat[2](42)

42
42

Эта обновленная функция принимает любой Stringable тип, поэтому вы можете передать ей значение Int, String или Bool.

Обратите внимание, что в списке параметров после MsgType есть двойная косая черта (//), указывающая, что это параметр только для вывода, поэтому вам не нужно указывать его явно. Вместо этого компилятор видит, что аргументом msg является Int, и определяет тип по значению.

Поддержка универсальных функций Mojo еще только начинается. Вы можете писать универсальные функции, подобные этой, используя признаки и параметры. Вы также можете создавать универсальные коллекции, такие как List и Dict. Если вам интересно узнать, как работают эти типы, вы можете найти исходный код для стандартных библиотечных коллекций на GitHub.

Параметризованные структуры

Вы также можете добавлять параметры к структурам. Вы можете использовать параметризованные структуры для создания универсальных коллекций. Например, тип универсального массива может содержать код, подобный этому:

struct GenericArray[ElementType: Copyable & ImplicitlyDestructible]:
    var data: UnsafePointer[Self.ElementType, MutExternalOrigin]
    var size: Int

    fn __init__(out self, var *elements: Self.ElementType):
        self.size = len(elements)
        self.data = alloc[Self.ElementType](self.size)
        for i in range(self.size):
            (self.data + i).init_pointee_move(elements[i].copy())

    fn __del__(deinit self):
        for i in range(self.size):
            (self.data + i).destroy_pointee()
        self.data.free()

    fn __getitem__(self, i: Int) raises -> ref [self] Self.ElementType:
        if i < self.size:
            return self.data[i]
        else:
            raise Error("Out of bounds")

Эта структура имеет единственный параметр ElementType, который является заполнителем для типа данных, который вы хотите сохранить в массиве, иногда называемый параметром типа. ElementType соответствует трейту Copyable и, следовательно, трейту Moveable.

Как и в случае с параметризованными функциями, при использовании параметризованной структуры вам необходимо передавать значения параметров. В этом случае, когда вы создаете экземпляр GenericArray, вам нужно указать тип, который вы хотите сохранить, например, Int или Float64. (Это немного сбивает с толку, потому что значение параметра, которое вы передаете в данном случае, является типом. Все в порядке: тип Mojo является допустимым значением во время компиляции.)

Вы увидите это Self.ElementType используется во всей структуре, где вы обычно видите имя типа. Например, в качестве формального типа для element в конструкторе и возвращаемого типа метода __getitem__().

Вот пример использования GenericArray:

var array = GenericArray(1, 2, 3, 4)
for i in range(array.size):
    end = ", " if i < array.size - 1 else "\n"
    print(array[i], end=end)

1, 2, 3, 4

Параметризованная структура может использовать тип Self для представления конкретного экземпляра структуры (то есть со всеми заданными параметрами). Например, вы могли бы добавить статический фабричный метод в GenericArray со следующей сигнатурой:

struct GenericArray[ElementType: Copyable & ImplicitlyDestructible]:
    ...

    @staticmethod
    fn splat(count: Int, value: Self.ElementType) -> Self:
        # Create a new array with count instances of the given value

Здесь Self эквивалентен записи GenericArray[Self.ElementType]. То есть вы можете вызвать метод splat() следующим образом:

GenericArray[Float64].splat(8, 0)

Метод возвращает экземпляр GenericArray[Float64].

Ссылочные параметры структуры

Как показано в предыдущем разделе, вы ссылаетесь на параметр структуры, используя точечный синтаксис, точно так же, как на метод структуры или поле (например, Self.ElementType).

Этот доступ к параметру структуры работает где угодно, а не только внутри методов структуры. Вы можете получить доступ к параметрам в качестве атрибутов самого типа:

fn on_type():
    print(SIMD[DType.float32, 2].size)  # prints 2

Или в качестве атрибутов экземпляра типа:

fn on_instance():
    var x = SIMD[DType.int32, 2](4, 8)
    print(x.dtype)  # prints int32

Работает и для элементов comptime тоже Хотя вы, возможно, еще не видели элементы comptime, они представляют собой другой тип значений во время компиляции, связанных со структурой, и к ним можно получить доступ таким же образом.

Условное соответствие

При создании универсальной структуры может потребоваться определить некоторые методы, для которых требуются дополнительные функции. Например, рассмотрим коллекцию, подобную GenericArray, которая содержит экземпляры типа, соответствующего трейту Copyable. Это накладывает множество ограничений: вы не можете реализовать метод sort(), потому что вы не можете гарантировать, что сохраненный тип поддерживает операторы сравнения; вы не можете написать полезный метод __str__() или `__repr__(), потому что вы не можете гарантировать, что сохраненный тип поддерживает операторы сравнения. сохраненный тип поддерживает преобразование в строку.

Ответом на эти вопросы является условное соответствие, которое позволяет определить метод, требующий дополнительных функций. Для этого вы определяете значение self, которое имеет более конкретную привязку к одному или нескольким параметрам.

Например, следующий код определяет тип контейнера, который содержит экземпляр типа, соответствующего Movable. Он также определяет метод __str__(), который может быть вызван только в том случае, если сохраненный ElementType соответствует Writable, так и Movable:

@fieldwise_init
struct Container[ElementType: Movable & ImplicitlyDestructible]:
    var element: Self.ElementType

    def __str__[
        StrElementType: Writable & Movable, //
    ](self: Container[StrElementType]) -> String:
        return String(self.element)


def main():
    float_container = Container(5.0)
    string_container = Container("Hello")
    print(float_container.__str__())
    print(string_container.__str__())

5.0
Hello

Обратите внимание на сигнатуру метода __str__(), который объявляет аргумент self с более конкретным типом. В частности, он объявляет, что принимает контейнер с типом ElementType, который соответствует параметрам, доступным для записи и перемещения.

def __str__[StrElementType: Writable & Movable, //](
        self: Container[StrElementType]) -> String:

Признак, привязанный к StrElementType, должен включать исходный признак ElementType (в данном случае Movable) либо по составу, либо по наследованию. Композиция трейтов Writable & Movable, включает исходный признак. Вы также можете использовать признак, который наследуется от Movable.

Обратите внимание, что функция main() вызывает метод __str__() напрямую, а не вызывает String(float_container). Одним из текущих ограничений условного соответствия является то, что Mojo не может распознать структуру Container[Int] как соответствующий Stringable, даже если метод __str__() реализован для любого ElementType, который также является Stringable.

Конкретный пример: тип SIMD

В качестве реального примера параметризованного типа давайте рассмотрим тип SIMD из стандартной библиотеки Mojo.

SIMD - это технология параллельной обработки данных, встроенная во многие современные процессоры, GPU и пользовательские ускорители. SIMD позволяет выполнять одну операцию над несколькими фрагментами данных одновременно. Например, если вы хотите извлечь квадратный корень из каждого элемента массива, вы можете использовать SIMD для распараллеливания работы.

Процессоры реализуют SIMD, используя низкоуровневые векторные регистры в аппаратном обеспечении, которые содержат несколько экземпляров скалярного типа данных. Чтобы использовать инструкции SIMD на этих процессорах, данные должны быть преобразованы в соответствующие SIMD ширину (тип данных) и длину (размер вектора). Процессоры могут поддерживать 512-разрядные или более длинные SIMD-векторы и поддерживать множество типов данных от 8-разрядных целых чисел до 64-разрядных чисел с плавающей запятой, поэтому нецелесообразно определять все возможные варианты SIMD.

SIMD-тип Mojo(определенный как struct) предоставляет общие SIMD-операции с помощью своих методов и делает значения типа данных SIMD и размера параметрическими. Это позволяет напрямую сопоставлять ваши данные с SIMD-векторами на любом оборудовании.

Вот сокращенная (нефункциональная) версия определения SIMD-типа Mojo:

struct SIMD[dtype: DType, size: Int]:
    var value: … # Здесь есть кое-что из низкоуровневого MLIR

    # Создайте новый SIMD из нескольких скаляров
    fn __init__(out self, *elems: SIMD[Self.dtype, 1]):  ...

    # Заполняем SIMD дублированным скалярным значением.
    @staticmethod
    fn splat(x: SIMD[Self.dtype, 1]) -> SIMD[Self.dtype, Self.size]: ...

    # Приводим элементы SIMD к другому типу elt.
    fn cast[target: DType](self) -> SIMD[target, Self.size]: ...

    # Поддерживается множество стандартных операторов.
    fn __add__(self, rhs: Self) -> Self: ...

Таким образом, вы можете создать и использовать SIMD-вектор, подобный этому:

var vector = SIMD[DType.int16, 4](1, 2, 3, 4)
vector = vector * vector
for i in range(4):
    print(vector[i], end=" ")

1 4 9 16

Как вы можете видеть, простой арифметический оператор типа *, применяемый к паре SIMD-векторов, воздействует на соответствующие элементы в каждом векторе.

Определение каждого варианта SIMD с помощью параметров отлично подходит для повторного использования кода, поскольку тип SIMD может выражать все различные векторные варианты статически, вместо того чтобы требовать от языка предварительного определения каждого варианта.

Поскольку SIMD является параметризованным типом, аргумент self в его функциях содержит эти параметры — полное имя типа SIMD[type, size]. Хотя это допустимо записать (как показано в типе возвращаемого значения splat()), это может быть многословным, поэтому мы рекомендуем использовать тип Self (из PEP673), как в примере __add__.

Перегрузка параметров

Функции и методы могут быть перегружены по своим сигнатурам параметров.

Использование параметризованных типов и функций

Вы можете использовать параметрические типы и функции, передавая значения параметров в квадратных скобках. Например, для приведенного выше SIMD тип type указывает на тип данных, а size - длину SIMD-вектора (она должна быть в степени 2).:

def main():
    # Создаем вектор из 4 поплавков.
    var small_vec = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0)

    # Создаем большой вектор, содержащий 1.0 в формате float16.
    var big_vec = SIMD[DType.float16, 32](1.0)

    # Произведите некоторую математику и преобразуйте элементы в float32.
    var bigger_vec = (big_vec + big_vec).cast[DType.float32]()

    # Конечно, можно явно указать типы, если хотите.
    var bigger_vec2: SIMD[DType.float32, 32] = bigger_vec

    print("small_vec DType:", small_vec.dtype, "size:", small_vec.size)
    print(
        "bigger_vec2 DType:",
        bigger_vec2.dtype,
        "size:",
        bigger_vec2.size,
    )

small_vec type: float32 length: 4
bigger_vec2 type: float32 length: 32

Обратите внимание, что методу cast() также требуется параметр, чтобы указать тип, который вы хотите получить из приведенного выше определения метода (в приведенном выше определении метода ожидается целевое параметрическое значение). Таким образом, так же, как структура SIMD является определением универсального типа, метод cast() является определением универсального метода. Во время компиляции компилятор создает конкретную версию метода cast() с параметром, привязанным к DType.float32.

В приведенном выше коде показано использование конкретных типов(то есть все параметры привязаны к известным значениям). Но главная сила параметров заключается в возможности определять параметрические алгоритмы и типы(код, который использует значения параметров). Например, вот как определить параметрический алгоритм с помощью Scalar, который не зависит от типа данных:

from math import sqrt

fn rsqrt[dt: DType](x: Scalar[dt]) -> Scalar[dt]:
    return 1 / sqrt(x)

def main():
    var v = Scalar[DType.float16](42)
    print(rsqrt(v))

0.154296875

Обратите внимание, что аргумент x является скалярным типом, который параметризуется на основе параметра функции dt. И обратите внимание, что фактический вызов rsqrt() не указывает параметр dt, поскольку он может быть выведен из типа аргумента (здесь - Scalar[DType.float16]).

Вывод параметров

Компилятор Mojo часто может выводить значения параметров, поэтому вам не всегда нужно их указывать. Например, вы можете вызвать функцию rsqrt(), определенную выше, без указания параметра dt:

from math import sqrt

fn rsqrt[dt: DType](x: Scalar[dt]) -> Scalar[dt]:
    return 1 / sqrt(x)

def main():
    var v = Scalar[DType.float16](33)
    print(rsqrt(v))

0.174072265625

Компилятор выводит параметр dt на основе переданного ему параметрического значения v, как если бы вы явно написали rsqrt[DType.float16](v). На рисунке 1 показана ментальная модель того, как работает вывод параметров.

Рисунок. 1. Вывод параметров

Вывод параметров может показаться немного запутанным: может показаться, что компилятор выводит значения параметров во время компиляции из значений аргументов во время выполнения. Но на самом деле он выводит параметры из статически известных типов аргументов.

Ошибки логического вывода Если при выводе параметра происходит сбой, компилятор сообщает об ошибке, обычно "не удалось вывести параметр 'param_name'". К сожалению, компилятор также иногда сообщает об этой ошибке неправильно, например, когда фактической ошибкой является несоответствие типа. В таких случаях явное указание недостающих параметров часто позволяет Mojo сообщить о правильной ошибке.

Mojo также может выводить значения параметров структуры из аргументов, передаваемых конструктору или статическому методу.

Например, рассмотрим следующую структуру:

struct One[Type: Writable & Copyable]:
    var value: Self.Type

    fn __init__(out self, value: Self.Type):
        self.value = value.copy()


def use_one():
    s1 = One(123)  # equivalent to One[Int](123)
    s2 = One("Hello")  # equivalent to One[String]("Hello")

Обратите внимание, что вы можете создать экземпляр одного из них без указания параметра Type — Mojo может вывести его из аргумента value.

Вы также можете вывести параметры из параметризованного типа, передаваемого конструктору или статическому методу:

struct Two[Type: Writable & Copyable]:
    var val1: Self.Type
    var val2: Self.Type

    fn __init__(out self, one: One[Self.Type], another: One[Self.Type]):
        self.val1 = one.value.copy()
        self.val2 = another.value.copy()
        print(String(self.val1), String(self.val2))

    @staticmethod
    fn fire(thing1: One[Self.Type], thing2: One[Self.Type]):
        print("🔥", String(thing1.value), String(thing2.value))

def use_two():
    s3 = Two(One("infer"), One("me"))
    Two.fire(One(1), One(2))
    # Two.fire(One("mixed"), One(0)) # Error: parameter inferred to two
                                     # different values

use_two()

infer me
🔥 1 2

Two принимает параметр Type, а его конструктор принимает значения типа One[Type]. При создании экземпляра Two вам не нужно указывать параметр Type, поскольку он может быть выведен из аргументов.

Аналогично, статический метод fire() принимает значения типа One[Type], поэтому Mojo может определить значение типа во время компиляции. Обратите внимание, что передача двух экземпляров One с разными типами не работает.

Если вы знакомы с C++, вы можете понять, что это похоже на вывод аргументов по шаблону класса (Class Template Argument Deduction (CTAD)).

Необязательные и именованные параметры

Точно так же, как вы можете указать необязательные аргументы в сигнатурах функций, вы также можете определить необязательный параметр, присвоив ему значение по умолчанию.

Вы также можете передавать параметры по ключевому слову, точно так же, как вы можете использовать аргументы по ключевому слову. Для функции или структуры с несколькими необязательными параметрами использование ключевых слов позволяет передавать только те параметры, которые вы хотите указать, независимо от их позиции в сигнатуре функции.

Например, вот функция с двумя параметрами, каждый из которых имеет значение по умолчанию:

fn speak[a: Int = 3, msg: String = "woof"]():
    print(msg, a)


fn use_defaults():
    speak()  # prints 'woof 3'
    speak[5]()  # prints 'woof 5'
    speak[7, "meow"]()  # prints 'meow 7'
    speak[msg="baaa"]()  # prints 'baaa 3'

Напомним, что при вызове параметрической функции Mojo может выводить значения параметров. То есть он может определять значения своих параметров из параметров, присоединенных к аргументу. Если параметрическая функция также имеет значение по умолчанию, то выводимое значение параметра имеет приоритет.

Например, в следующем коде мы обновляем параметрическую функцию speak[](), чтобы она принимала аргумент с параметрическим типом. Хотя функция имеет значение параметра по умолчанию для a, Mojo вместо этого использует значение параметра, полученное из аргумента bar (как уже говорилось, значение a по умолчанию никогда не может быть использовано, но это только в демонстрационных целях).:

@fieldwise_init
struct Bar[v: Int]:
    pass


fn speak[a: Int = 3, msg: String = "woof"](bar: Bar[a]):
    print(msg, a)


fn use_inferred():
    speak(Bar[9]())  # prints 'woof 9'

Как упоминалось выше, вы также можете использовать необязательные параметры и параметры ключевых слов в структуре:

struct KwParamStruct[greeting: String = "Hello", name: String = "🔥mojo🔥"]:
    fn __init__(out self):
        print(Self.greeting, Self.name)

fn use_kw_params():
    var a = KwParamStruct[]()                 # prints 'Hello 🔥mojo🔥'
    var b = KwParamStruct[name="World"]()     # prints 'Hello World'
    var c = KwParamStruct[greeting="Hola"]()  # prints 'Hola 🔥mojo🔥'

Mojo поддерживает параметры только для позиций и ключевых слов, следуя тем же правилам, что и аргументы только для позиций и ключевых слов.

Параметры, доступные только для вывода

Иногда вам нужно объявлять функции, параметры которых зависят от других параметров. Поскольку сигнатура обрабатывается слева направо, параметр может зависеть только от параметра, указанного ранее в списке параметров. Например:

fn dependent_type[dtype: DType, value: Scalar[dtype]]():
    print("Value: ", value)
    print("Value is floating-point: ", dtype.is_floating_point())

dependent_type[DType.float64, Float64(2.2)]()

Value:  2.2000000000000002
Value is floating-point:  True

Вы не можете изменить положение параметров dtype и value на противоположное, поскольку value зависит от dtype. Однако, поскольку dtype является обязательным параметром, вы не можете исключить его из списка параметров и позволить Mojo выводить его из value:

dependent_type[Float64(2.2)]() # Error!

Параметры только для вывода - это особый класс параметров, которые всегда либо выводятся из контекста, либо задаются ключевым словом. Параметры только для вывода располагаются в начале списка параметров и отделяются от других параметров символом //:

fn example[type: Copyable, //, list: List[type]]()

Преобразование dtype в параметр, доступный только для вывода, решает эту проблему:

fn dependent_type[dtype: DType, //, value: Scalar[dtype]]():
    print("Value: ", value)
    print("Value is floating-point: ", dtype.is_floating_point())

dependent_type[Float64(2.2)]()

Value:  2.2000000000000002
Value is floating-point:  True

Поскольку параметры, предназначенные только для вывода, объявляются в начале списка параметров, от них могут зависеть другие параметры, и компилятор всегда будет пытаться вывести значения, предназначенные только для вывода, из связанных параметров или аргументов.

Иногда бывает полезно указать параметр только для вывода с помощью ключевого слова. Например, тип Span является параметрическим origin:

struct Span[mut: Bool, //, T: Copyable, origin: Origin[mut], ...]:
    ...

Здесь параметр mut доступен только для вывода. Значение обычно выводится при создании экземпляра Span. Привязка параметра mut к ключевому слову позволяет определить диапазон, для которого требуется изменяемый источник.

def mutate_span(span: Span[mut=True, Byte]):
    for i in range(0, len(span), 2):
        if i + 1 < len(span):
            span.swap_elements(i, i + 1)

Если компилятор не может вывести значение параметра, предназначенного только для вывода, и оно не задано ключевым словом, компиляция завершается неудачей.

Параметры с переменным числом аргументов

Mojo также поддерживает параметры с переменным числом аргументов, аналогичные аргументам с переменным числом элементов:

struct MyTensor[*dimensions: Int]:
    pass

Параметры с переменным числом аргументов в настоящее время имеют некоторые ограничения, которых нет у аргументов:

• Параметры с переменным числом аргументов должны быть однородными, то есть все значения должны быть одного типа.

• Тип параметра должен быть доступен для регистрации.

• Значения параметров автоматически не проецируются в VariadicList, поэтому вам необходимо создать этот список явно:

fn sum_params[*values: Int]() -> Int:
    comptime list = VariadicList(values)
    var sum = 0
    for v in list:
        sum += v
    return sum

Переменные параметры ключевых слов (например, **kwparams) пока не поддерживаются.

Выражения параметров - это просто Mojo код

Выражение параметра - это любое кодовое выражение (например, a+b), которое встречается там, где ожидается параметр. Выражения параметров поддерживают операторы и вызовы функций, как и в коде во время выполнения, и все типы параметров используют ту же систему типов, что и программа во время выполнения (например, Int и DType).

Поскольку выражения параметров используют ту же грамматику и типы, что и код Mojo во время выполнения, вы можете использовать множество функций "зависимого типа". Например, вы можете захотеть определить вспомогательную функцию для объединения двух SIMD-векторов:

fn concat[
    dtype: DType, ls_size: Int, rh_size: Int, //
](lhs: SIMD[dtype, ls_size], rhs: SIMD[dtype, rh_size]) -> SIMD[
    dtype, ls_size + rh_size
]:
    var result = SIMD[dtype, ls_size + rh_size]()

    @parameter
    for i in range(ls_size):
        result[i] = lhs[i]

    @parameter
    for j in range(rh_size):
        result[ls_size + j] = rhs[j]
    return result

result type: float32 length: 4

Обратите внимание, что результирующая длина является суммой длин входных векторов, и это выражается простой операцией +.

Мощное программирование во время компиляции

Хотя простые выражения полезны, иногда требуется написать императивную логику во время компиляции с использованием потока управления. Вы даже можете выполнить рекурсию во время компиляции. Например, вот пример алгоритма "сокращения дерева", который рекурсивно суммирует все элементы вектора в скаляр:

fn slice[
    dtype: DType, size: Int, //
](x: SIMD[dtype, size], offset: Int) -> SIMD[dtype, size // 2]:
    comptime new_size = size // 2
    var result = SIMD[dtype, new_size]()
    for i in range(new_size):
        result[i] = SIMD[dtype, 1](x[i + offset])
    return result


fn reduce_add(x: SIMD) -> Int:
    @parameter
    if x.size == 1:
        return Int(x[0])
    elif x.size == 2:
        return Int(x[0]) + Int(x[1])

    # Extract the top/bottom halves, add them, sum the elements.
    comptime half_size = x.size // 2
    var lhs = slice(x, 0)
    var rhs = slice(x, half_size)
    return reduce_add(lhs + rhs)


def main():
    var x = SIMD[DType.int, 4](1, 2, 3, 4)
    print(x)
    print("Elements sum:", reduce_add(x))

[1, 2, 3, 4]
Elements sum: 10

При этом используется декоратор @parameter для создания параметрического условия if, которое представляет собой оператор if, выполняющийся во время компиляции. Требуется, чтобы его условие было допустимым выражением параметра, и гарантирует, что в программу будет скомпилирована только текущая ветвь оператора if. (Это похоже на использование декоратора @parameter с циклом for, показанным ранее).

comptime: выражения именованных параметров

Очень часто возникает желание присвоить имена значениям времени компиляции. В то время как var определяет значение во время выполнения, нам нужен способ определить временное значение во время компиляции. Для этого Mojo использует объявление comptime. В самом простом случае comptime можно использовать для определения постоянного значения:

comptime rows = 512

Значение comptime всегда вычисляется во время компиляции, поэтому вы можете использовать comptime для принудительного запуска функции во время компиляции. Вы можете использовать это для вычисления постоянных значений на основе информации, доступной во время компиляции, такой как параметры оборудования.

comptime block_size = _calculate_block_size()

Типы - это еще одно распространенное использование значений comptime. Поскольку типы являются выражениями времени компиляции, вы можете использовать comptime в качестве сокращения (псевдоним типа или "typedef") для параметризованного типа:

comptime Float16 = SIMD[DType.float16, 1]
comptime UInt8 = SIMD[DType.uint8, 1]

var x: Float16 = 0  # Float16 works like a "typedef"

(Эти и другие псевдонимы на самом деле определены в модуле simd.)

Как и переменные var, значения comptime подчиняются области видимости, и вы можете использовать локальные значения comptime внутри функций, как и следовало ожидать.

Параметрические comptime значения

Параметрическое comptime значение - это выражение времени компиляции, которое принимает список параметров и возвращает значение константы времени компиляции:

comptime AddOne[a: Int] : Int = a + 1

comptime nine = AddOne[8]

Как вы можете видеть из предыдущего примера, параметрическое comptime значение немного похоже на функцию, используемую только во время компиляции. Обычная функция или метод также могут быть вызваны во время компиляции:

fn add_one(a: Int) -> Int:
    return a + 1

comptime ten = add_one(9)

Основное различие между функцией и параметрическим значением comptime заключается в том, что значение выражения comptime может быть типом, в то время как функция не может возвращать тип в качестве значения.

# Does not work—-dynamic type values not permitted
fn int_type() -> AnyType:
    return Int

# Works
comptime IntType = Int

Поскольку значение comptime может быть типом, вы можете использовать параметрические значения comptime для выражения новых типов:

comptime TwoOfAKind[dt: DType] = SIMD[dt, 2]
twoFloats = TwoOfAKind[DType.float32](1.0, 2.0)

comptime StringKeyDict[ValueType: Copyable & ImplicitlyDestructible] = Dict[String, ValueType]
var b: StringKeyDict[UInt8] = {"answer": 42}

Параметрические объявления comptime поддерживают те же функции, что и параметризованные структуры или функции: параметры только для вывода, параметры только для ключевых слов и необязательные параметры, автоматическая параметризация и так далее.

    comptime Floats[size: Int, half_width: Bool = False] = SIMD[
        (DType.float16 if half_width else DType.float32), size
    ]
    var floats = Floats[2](6.0, 8.0)
    var half_floats = Floats[2, True](10.0, 12.0)

comptime члены

Вы также можете определить comptime значения как элементы struct или trait:

struct Circle[radius: Float64]:
    comptime pi = 3.14159265359
    comptime circumference = 2 * Self.pi * Self.radius

Эти элементы comptime могут использоваться по-разному:

• Постоянные значения, специфичные для данного типа.
• Постоянные значения вычисляются на основе параметров структуры.
• Связанные типы, основанные на параметрах структуры.

Разница между параметрами и элементами comptime заключается в том, что значения параметров задаются пользователем, но элементы comptime представляют собой либо постоянные значения, либо значения, полученные из входных параметров.

Признак может объявлять элемент comptime, который должен быть определен всеми соответствующими структурами.

Ссылки на элементы comptime Ссылки на элементы comptime работают так же, как ссылки на параметры — вы можете ссылаться на элемент, используя точечный синтаксис (например, Self.IteratorType).

comptime члены как перечисления

Некоторые типы Mojo используют элементы comptime для выражения перечислений. Например, в следующем коде определяется тип настройки, который определяет константы comptime для различных значений настройки:

@fieldwise_init
struct Sentiment(Equatable, ImplicitlyCopyable):
    var _value: Int

    comptime NEGATIVE = Sentiment(0)
    comptime NEUTRAL = Sentiment(1)
    comptime POSITIVE = Sentiment(2)

    fn __eq__(self, other: Self) -> Bool:
        return self._value == other._value

    fn __ne__(self, other: Self) -> Bool:
        return not (self == other)

fn is_happy(s: Sentiment):
    if s == Sentiment.POSITIVE:
        print("Yes. 😀")
    else:
        print("No. ☹️")

Этот шаблон обеспечивает типобезопасное перечисление.

Структура DType реализует простое перечисление, используя элементы comptime, подобные этому. Это позволяет клиентам использовать значения, подобные DType.float32, в выражениях параметров или во время выполнения.

Элементы comptime как связанные типы

Связанные типы обычно используются для элементов comptime. Например, структура List[T] содержит значения типа T. Метод list __iter__() возвращает итератор списка, который возвращает значения типа T. List использует элемент comptime, IteratorType, для определения типа возвращаемого итератора.

В следующем фрагменте кода показана упрощенная версия некоторого кода списка, показывающая список и связанный с ним тип IteratorType:

@fieldwise_init
struct _ListIter[
    mut: Bool,
    //,
    T: Copyable,
    origin: Origin[mut],
](ImplicitlyCopyable, Iterable, Iterator):

    comptime Element = Self.T  # Required by the Iterator trait

    var index: Int
    var src: Pointer[List[Self.Element], Self.origin]

    ...

struct List[T: Copyable](
    Boolable, Copyable, Defaultable, Iterable, Sized
):
    comptime IteratorType[
        iterable_mut: Bool, //, iterable_origin: Origin[iterable_mut]
    ]: Iterator = _ListIter[Self.T, iterable_origin]

    ...

    fn __iter__(ref self) -> Self.IteratorType[origin_of(self)]:
        return {0, Pointer(to=self)}
    ...

Элемент IteratorType параметризуется в начале координат, поэтому он может представлять как изменяемые, так и неизменяемые итераторы.

Полностью связанные, частично связанные и несвязанные типы

Параметрический тип с заданными параметрами называется полностью привязанным. То есть все его параметры привязаны к значениям. Как упоминалось ранее, вы можете создать экземпляр только полностью привязанного типа (иногда называемого конкретным типом).

Однако в некоторых контекстах параметрические типы могут быть несвязанными или частично связанными. Например, вы можете использовать comptime для создания псевдонима частично связанного типа, чтобы создать новый тип, требующий меньшего количества параметров:

comptime StringKeyDict = Dict[String, _]
var b: StringKeyDict[UInt8] = {"answer": 42}

Частично привязанные типы в сравнении с параметрическими значениями времени вычисления Вы можете заметить, что этот пример очень похож на пример из раздела, посвященного параметрическим значениям времени вычисления. Для псевдонимов простых типов, подобных этому, вы можете использовать либо частично привязанный тип, либо параметрическое значение времени вычисления. Параметрические значения comptime предоставляют более гибкий способ определения псевдонимов типов, поскольку вы можете определить порядок следования параметров, добавить значения по умолчанию и так далее. Частично связанные и несвязанные типы могут обеспечить удобный способ сокращения времени при определении параметрических функций и значений времени выполнения, называемый автоматической параметризацией.

Вы также можете использовать частично привязанные типы в качестве привязки типа для аргумента или параметра.

Например, учитывая следующий тип:

struct MyType[s: String, i: Int, i2: Int, b: Bool = True]:
    pass

Он может отображаться в коде в следующих формах:

• Полностью привязанный, со всеми заданными параметрами:

  def my_fn1(m1: MyType["Hello", 3, 4, True]):
      pass
  

• Частично привязанный, с заданными некоторыми, но не всеми параметрами:

  def my_fn2(m2: MyType["Hola", _, _, True]):
      pass
  

• Несвязанный, без заданных параметров:

  def my_fn3(m3: MyType[_, _, _, _]):
      pass
  

Вы также можете использовать выражение трех точек ... чтобы отменить привязку произвольного количества параметров в конце списка параметров:

# Эти два типа эквивалентны
MyType["Hello", ...]
MyType["Hello", _, _, _]

Если параметр явно не привязан к выражениям _, или ... Чтобы использовать этот тип, необходимо указать значение для этого параметра. Значения явно несвязанных параметров по умолчанию игнорируются.

Частично связанные и несвязанные параметрические типы могут использоваться в некоторых контекстах, где отсутствующие (несвязанные) параметры будут предоставлены позже - например, в значениях comptime и автоматически параметризованных функциях.

Пропущенные параметры

Mojo также поддерживает альтернативный формат для несвязанных параметров, когда параметры просто опускаются в выражении:

@fieldwise_init
struct MyComplicatedType[a: Int = 7, /, b: Int = 8, *, c: Int, d: Int = 9]:
    pass

# Несвязанный
fn my_func(t: MyComplicatedType):
    pass

Это эквивалентно fn my_func(t: MyComplicatedType[...]): pass. То есть все параметры (только позиционные, позиционные или ключевые, только для ключевых слов) не привязаны, а их значения по умолчанию (если таковые имеются) игнорируются.

Обратите внимание, что при частичной привязке типа аргумента значения по умолчанию будут привязаны:

# Partially bound
MyComplicatedType[1]
# Equivalent to
MyComplicatedType[1, 8, c=_, d=9]  # Uses default values for `b` and `d`.

Такое поведение с пропущенными параметрами в настоящее время поддерживается для обеспечения обратной совместимости. В будущем мы намерены согласовать поведение пропущенных параметров и явно несвязанных параметров.

Автоматическая параметризация

Mojo поддерживает "автоматическую" параметризацию функций и параметрических значений времени выполнения. Если тип аргумента функции или параметра частично привязан или несвязан, несвязанные параметры автоматически добавляются в качестве параметров функции. Это проще понять на примере:

fn print_params(vec: SIMD):
    print(vec.dtype)
    print(vec.size)

var v = SIMD[DType.float64, 4](1.0, 2.0, 3.0, 4.0)
print_params(v)

float64
4

В приведенном выше примере функция print_params() автоматически параметризуется. Аргумент vec принимает аргумент типа SIMD[...]. Это несвязанный параметризованный тип, то есть в нем не указаны значения параметров для этого типа. Mojo рассматривает несвязанные параметры в vec как параметры функции, доступные только для вывода. Это примерно эквивалентно следующему коду:

fn print_params2[t: DType, s: Int, //](vec: SIMD[t, s]):
    print(vec.dtype)
    print(vec.size)

Когда вы вызываете print_params(), вы должны передать ему конкретный экземпляр типа SIMD, то есть экземпляр со всеми указанными параметрами, например SIMD[DType.float64, 4]. Компилятор Mojo выводит значения параметров из входного аргумента.

При использовании функции с параметризацией вручную вы можете обращаться к параметрам по имени (например, t и s в предыдущем примере), что не является опцией в автоматически параметризованной функции.

Однако вы всегда можете получить доступ к параметрам типа и элементам comptime, используя точечный синтаксис (например, vec.dtype), как описано в разделе Ссылки на параметры структуры. Эта возможность доступа к параметрам типа и элементам comptime не является специфичной для автоматически параметризованных функций, вы можете использовать ее где угодно.

Вы даже можете использовать этот синтаксис в сигнатуре функции, чтобы определить аргументы функции и тип возвращаемого значения на основе параметров аргумента или элементов comptime.

Например, если вы хотите, чтобы ваша функция принимала два SIMD-вектора одинакового типа и размера, вы можете написать код следующим образом:

fn interleave(v1: SIMD, v2: type_of(v1)) -> SIMD[v1.dtype, v1.size*2]:
    var result = SIMD[v1.dtype, v1.size*2]()
    for i in range(v1.size):
        result[i*2] = SIMD[v1.dtype, 1](v1[i])
        result[i*2+1] = SIMD[v1.dtype, 1](v2[i])
    return result

var a = SIMD[DType.int16, 4](1, 2, 3, 4)
var b = SIMD[DType.int16, 4](0, 0, 0, 0)
var c = interleave(a, b)
print(c)

[1, 0, 2, 0, 3, 0, 4, 0]

Как показано в примере, вы можете использовать волшебный вызов type_of(x), если вы просто хотите сопоставить тип аргумента. В этом случае это более удобно и компактно, чем написание эквивалентного SIMD[v1.dtype, v1.size].

Автоматическая параметризация параметров

Вы также можете воспользоваться преимуществами автоматической параметризации в списке параметров функции или параметрического значения comptime. Например:

fn foo[value: SIMD]():
    pass

# Equivalent to:
fn foo[dtype: DType, size: Int, //, value: SIMD[dtype, size]]():
    pass

Вот еще один пример использования параметрического значения comptime:

comptime Foo[S: SIMD] = Bar[S]

# Equivalent to:
comptime Foo[dtype: DType, size: Int, //, S: SIMD[dtype, size]] = Bar[S]

Автоматическая параметризация с частично привязанными типами

Mojo также поддерживает автоматическую параметризацию: с частично привязанными параметризованными типами (то есть типами с некоторыми, но не со всеми указанными параметрами).

Например, предположим, что у нас есть структура Fudge с тремя параметрами:

@fieldwise_init
struct Fudge[sugar: Int, cream: Int, chocolate: Int = 7](Stringable):
    fn __str__(self) -> String:
        return String.write(
            "Fudge (", Self.sugar, ",", Self.cream, ",", Self.chocolate, ")"
        )

Мы можем написать функцию, которая принимает аргумент Fudge только с одним связанным параметром(он частично связан):

fn eat(f: Fudge[5, ...]):
    print("Ate " + String(f))

Функция eat() принимает структуру Fudge, первый параметр которой (sugar) привязан к значению 5. Второй и третий параметры, cream и chocolate, не привязаны.

Несвязанные параметры cream и chocolate становятся неявными параметрами функции eat. На практике это примерно эквивалентно написанию:

fn eat[cr: Int, ch: Int](f: Fudge[5, cr, ch]):
    print("Ate", String(f))

В обоих случаях мы можем вызвать функцию, передав экземпляр с привязанными параметрами cream и chocolate:

eat(Fudge[5, 5, 7]())
eat(Fudge[5, 8, 9]())

Ate Fudge (5,5,7)
Ate Fudge (5,8,9)

Если вы попытаетесь передать аргумент со значением sugar, отличным от 5, компиляция завершится ошибкой, поскольку оно не соответствует типу аргумента:

eat(Fudge[12, 5, 7]())
# ERROR: invalid call to 'eat': argument #0 cannot be converted from 'Fudge[12, 5, 7]' to 'Fudge[5, 5, 7]'

Вы также можете явно отменить привязку отдельных параметров. Это дает вам больше свободы при указании несвязанных параметров.

Например, вы можете разрешить пользователю указывать значения для sugar и chocolate, а cream оставить неизменными. Чтобы сделать это, замените каждое значение несвязанного параметра одним символом подчеркивания (_).:

fn devour(f: Fudge[_, 6, _]):
    print("Devoured",  String(f))

Опять же, несвязанные параметры (sugar и chocolate) добавляются в функцию в качестве неявных параметров. Эта версия примерно эквивалентна следующему коду, где эти два значения явно привязаны к входным параметрам su и ch:

fn devour[su: Int, ch: Int](f: Fudge[su, 6, ch]):
    print("Devoured", String(f))

Вы также можете указать параметры по ключевому слову или комбинировать позиционные параметры и параметры ключевых слов, поэтому следующая функция примерно эквивалентна предыдущей: первый параметр, sugar, явно не привязан с помощью символа подчеркивания. Параметр chocolate не привязан с помощью синтаксиса ключевого слова, chocolate=_. И cream явно привязан к значению 6:

fn devour(f: Fudge[_, chocolate=_, cream=6]):
    print("Devoured", String(f))

Все три версии функции devour() работают со следующими вызовами:

devour(Fudge[3, 6, 9]())
devour(Fudge[4, 6, 8]())

Devoured Fudge (3,6,9)
Devoured Fudge (4,6,8)

Устаревший синтаксис (пропущенные параметры)

Вы также можете указать несвязанный или частично связанный тип, опустив параметры: например:

fn nibble(f: Fudge[5]):
    print("Ate", String(f))

nibble(Fudge[5, 4, 7]())

Ate Fudge (5,4,7)

Здесь Fudge[5] работает как Fudge[5, ...], за исключением обработки параметров со значениями по умолчанию. Вместо того чтобы отбрасывать значение chocolate по умолчанию, Fudge[5] немедленно привязывает значение по умолчанию, что делает его эквивалентным: Fudge[5, _, 7].

Это означает, что следующий код не будет скомпилирован с предыдущим определением функции nibble(), поскольку в нем не используется значение по умолчанию для chocolate:

nibble(Fudge[5, 5, 9]())
# ERROR: invalid call to 'nibble': argument #0 cannot be converted from 'Fudge[5, 5, 9]' to 'Fudge[5, 5, 7]'

TODO В будущем мы намерены согласовать поведение пропущенных параметров и явно несвязанных параметров.

Встроенная функция rebind()

Одним из последствий того, что Mojo не выполняет создание экземпляров функций в синтаксическом анализаторе, подобном C++, является то, что Mojo не всегда может определить, равны ли некоторые параметрические типы, и жалуется на недопустимое преобразование. Обычно это происходит в статических шаблонах диспетчеризации. Например, следующий код не будет скомпилирован:

fn take_simd8(x: SIMD[DType.float32, 8]):
    pass

fn generic_simd[nelts: Int](x: SIMD[DType.float32, nelts]):
    @parameter
    if nelts == 8:
        take_simd8(x)

Синтаксический анализатор будет жаловаться:

error: invalid call to 'take_simd8': argument #0 cannot be converted from
'SIMD[f32, nelts]' to 'SIMD[f32, 8]'
        take_simd8(x)
        ~~~~~~~~~~^~~

Это происходит потому, что синтаксический анализатор полностью проверяет тип функции без создания экземпляра, и тип x по-прежнему SIMD[f32, nelts], а не SIMD[f32, 8], несмотря на статическое условие. Решение состоит в том, чтобы вручную "rebind" тип x, используя встроенную функцию rebind, которая вставляет утверждение во время компиляции о том, что типы ввода и результата преобразуются в один и тот же тип после создания экземпляра функции:

fn take_simd8(x: SIMD[DType.float32, 8]):
    pass

fn generic_simd[nelts: Int](x: SIMD[DType.float32, nelts]):
    @parameter
    if nelts == 8:
        take_simd8(rebind[SIMD[DType.float32, 8]](x))