torch.fx.passes.runtime_assert.insert_deferred_runtime_asserts

torch.fx.passes.runtime_assert.insert_deferred_runtime_asserts(gm, shape_env, name, export=False) [исходный код]

Во время трассировки мы могли обнаружить, что для некоторых значений, зависящих от данных, выполняется проверка во время выполнения; например, torch.empty(x.item()) подразумевает проверку во время выполнения, что x.item() >= 0. Эти проверки могут возникать непредсказуемо при распространении фиктивных тензоров, поэтому мы не можем удобно вставлять их в граф FX в момент возникновения. Вместо этого мы накапливаем их в ShapeEnv и на этом этапе вставляем в граф как полноценные проверки.

Этот этап также устраняет дублирование вычислений, связанных с размерностями, выполняя CSE для операций, которые создают символьные значения и/или участвуют в проверках во время выполнения. Кроме того, вызовы для получения формы (size/stride/storage_offset) по возможности преобразуются в вычисления на основе размеров входных данных, что позволяет раньше освобождать промежуточные тензоры. Например, здесь dynamo удалит вызовы cat и repeat:

z = torch.cat([x, x], dim=0) # 2*s0 w = z.repeat(y.shape[0]) # 2*s0*s1 _w = w.shape[0] # something with _w, but not w …

# turns into -> _w0 = 2 * s0 _w = _w0 * s1

# where s0, s1 are either SymInt graph inputs, or the result of added size calls

Также устраняются дублирующиеся проверки, сгенерированные компилятором и проверяющие одно и то же выражение, а также дублирующиеся вызовы constrain_range. Кроме того, поскольку проверки границ для отдельных символов (например, u0 >= 0, u0 <= 5) накапливают информацию в ShapeEnv, ShapeEnv содержит минимальные и максимальные границы для каждого символа, и в конце этого этапа мы добавляем проверки границ, если в графе ещё нет эквивалентной явной проверки.

Примечание

Обратная совместимость этого API гарантируется.