nvfuser-cu121-torch25 0.2.25.dev20250201__cp310-cp310-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/expr_simplifier.h

@@ -0,0 +1,218 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <exceptions.h>
+#include <ir/all_nodes.h>
+#include <visibility.h>
+
+#include <vector>
+
+// Note: [The Mathematics of Integer Arithmetic]
+//
+// We learnt arithmetic from as early as elementary school, and have been used
+// to simplify expressions using rules like (a+b)/c = a/c+b/c. However, when we
+// are dealing with integer arithmetic, which is the case for index and
+// predicate simplification, lots of rules we learnt in elementary school no
+// longer hold. For example, (1+1)/2 != 1/2+1/2 because the left hand side is 1
+// and the right hand side is 0 + 0 = 0. So when considering adding a new
+// simplification rule, we need to be very careful to make sure the rule is
+// mathematically correct.
+//
+// Suggested reading matherials:
+// - doc/math/abstract-algebra.md reviews abstract algebra, a theory that tells
+//   us which rule we are used to is still valid, and which is not.
+// - doc/math/integer-division.md reviews the definitions and properties of div
+//   and mod in textbooks, it also describes some theorems that we proved
+//   ourselves that is useful for simplifying integer expressions.
+// - doc/math/monotonic-function.md reviews the definition and properties of
+//   monotonic function.
+//
+// We can use the following rules to simplify integer expressions:
+//
+// A) Associativity of +: a + (b + c) = (a + b) + c
+// B) Associativity of *: a * (b * c) = (a * b) * c
+// C) Commutativity of +: a + b = b + a
+// D) Commutativity of *: a * b = b * a
+// E) Distributivity of * over +: a * (b + c) = (a * b) + (a * c)
+// F) Distributivity of * over +: (a + b) * c = (a * c) + (b * c)
+// G) (-a) / b = -(a / b) = a / (-b)
+// H) (-a) % b = -(a % b) = a % (-b)
+// I) If -|a| < r < |a|, then r % a = r, r / a = 0
+// J) Distributivity of % over +:
+//    If compatible_sign(a, b), then (a + b) % c = (a % c + b % c) % c
+// J.1) If compatible_sign(a, b) and a % c = 0, then (a + b) % c = b % c
+// J.2) Let g = gcd(a, c). If compatible_sign(a, b), and -|g| < b < |g|
+//      then (a + b) % c = a % c + b
+// K) Distributivity of % over *:
+//    If compatible_sign(a, b), then (a * b) % c = (a % c * b % c) % c
+// L) If a is a multiple of b, then a % b = 0
+// M) If b is a multiple of c, then we have: a*(b/c) = (a*b)/c
+// N) a / (b * c) = (a / b) / c
+// O) If d divides a and b, then a % b = ((a / d) % (b / d)) * d
+// P) If b is a multiple of c, then a/(b/c) = (a*c)/b
+// Q) If compatible_sign(a, b) and -|c| < a % c + b % c < |c|, then
+//    (a+b)/c = a/c + b/c
+// Q.1) If compatible_sign(a, b) and a % c = 0, then (a+b)/c = a/c + b/c
+// Q.2) Let g = gcd(a, c). If compatible_sign(a, b), and -|g| < b < |g|
+//      then (a + b) / c = a/c
+//
+// See doc/math/integer-division.md for proofs of these rules.
+//
+// Some examples on applying the above rules to simplify expressions:
+//
+// Example 7.1: Given that a >= 0 and b >= 0, simplify (a*4 + b) % 4
+// Answer: (a*4 + b) % 4 = ((a*4)%4 + b%4) % 4 (Rule J)
+// = (0 + b%4) % 4 (Rule L)
+// = b % 4 % 4 (Basic math)
+// = b % 4 (Rule I)
+//
+// Example 7.2: Given that 0 <= a < 3, simplify a % 4
+// Answer: a % 4 = a (Rule I)
+//
+// Example 7.3: Simplify (a * 256) / 4
+// Answer: (a * 256) / 4 = a * (256 / 4) (Rule M)
+// = a * 64 (Basic math)
+//
+// Example 7.4: Simplify (a / 4) / 64
+// Answer: (a / 4) / 64 = a / (4 * 64) (Rule N)
+// = a / 256 (Basic math)
+//
+// Example 7.5: Simplify (a * 64) % 256 / 4
+// Answer: (a * 64) % 256 / 4 = ((a % 4) * 64) / 4 (Rule O)
+// = (a % 4) * (64 / 4) (Rule M)
+// = (a % 4) * 16 (Basic math)
+//
+// Example 7.6: Simplify (a * 4) / 256
+// Answer: (a * 4) / 256 = a / (256 / 4) (Rule P)
+// = a / 64 (Basic math)
+//
+// Example 7.7: Given that a >= 0 and b >= 0, simplify (a * 256 + b) / 4
+// Answer: because (a * 256) % 4 = 0, we have
+// (a * 256 + b) / 4 = a * 256 / 4 + b / 4 (Rule Q)
+// = a * (256 / 4) + b / 4 (Rule M)
+// = a * 64 + b / 4 (Basic math)
+//
+// Example 7.8: Given that a >= 0 and 0 <= b < 4, simplify (a * 4 + b) / 4
+// Answer: Similar to above, we have
+// (a * 4 + b) / 4 = a + b / 4
+// = a + 0 (Rule I)
+// = a
+
+namespace nvfuser {
+
+// Information for a single variable. Possible values that this variable can
+// take is: start, start + step, start + 2 * step, ... (< stop), which is
+// similar to the loop variable of for loop:
+//   for variable in range(start, stop, step)
+struct VarInfo {
+  Val* variable = nullptr;
+  // If this variable is an unrolled loop index. It is important to know this
+  // because unrolled loop index is compile constant to nvRTC. Note that a
+  // constant to nvRTC might not be a constant to nvFuser. For example, if I
+  // have loop
+  //   #pragma unroll
+  //   FOR i1 in ...:
+  //     ...
+  // Then `i1` is a compile constant to nvRTC, but not a compile time constant
+  // to nvFuser.
+  bool is_unrolled_loop_index = false;
+};
+
+// Analyze expression register usage
+enum class RegisterType { GeneralPurpose, Uniform, Immediate, Unknown };
+RegisterType getRegisterType(Val* value);
+
+// Simplify expressions with the given information of variables.
+//
+// The argument `variables` specifies which scalar are considered variable and
+// some information about these variables. Any scalar not contained in
+// `variables` are considered constants. Tensors are always considered as
+// variables, regardless of if it is specified in `variables`.
+//
+// Note that in `variables`, the order matters. This order specifies how we
+// should organize associative and commutative expressions. For example, if the
+// `variables` is {a, b, c, d}, then we will simplify (a + d) + (c + b) as
+// ((a + b) + c) + d. Tensors are always considered as at the right of all
+// scalars, regardless of if it is inside `variables` or not.
+// See note [Reordering associative and commutative operators] for detailed
+// information about this reordering.
+//
+// Some simplifications like a*b/b -> a is always correct in valid case, but
+// when there is an error (e.g. division-by-zero), these simplifications could
+// potentially hide the error. The argument `preserve_error` specifies whether
+// we should disable these optimization, unless we can prove there won't be an
+// error.
+NVF_API Val* simplifyExpr(
+    Val* value,
+    const std::list<VarInfo>& variables = {},
+    std::vector<Val*> assumptions = {},
+    bool preserve_error = false);
+
+class Context;
+namespace assoc_comm {
+// The expression type that represents the flattened ops. For example, if I have
+// out = a + b + 3 + c + 5, then I will have:
+//   FlattenedAssocCommOp {
+//     inputs: [a, b, 3, c, 5]
+//     outputs: [out]
+//   }
+class FlattenedAssocCommOp : public Expr {
+ public:
+  using Expr::Expr;
+
+  FlattenedAssocCommOp(
+      IrBuilderPasskey passkey,
+      BinaryOpType op,
+      Val* out,
+      std::vector<Val*> terms);
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  const char* getOpString() const override;
+
+  // FlattenedAssocCommOp is unordered, so we should have
+  // FlattenedAdd(a, b)->sameAs(FlattenedAdd(b, a))
+  bool sameAs(const Statement* other) const override;
+
+  std::string toString(int indent_size = 0) const override;
+
+  std::string toInlineString(int indent_size = 0) const override;
+
+  DataType dtype() const {
+    return *output(0)->getDataType();
+  }
+
+  BinaryOpType getOpType() const {
+    return attribute<BinaryOpType>(0);
+  }
+
+  // Get a vector of inputs, sorted as the order given by `variables`. Note that
+  // the sorting key is the rightmost variable that an input depends on. For
+  // example, if I have two inputs.
+  // v1 = a * c
+  // v2 = b
+  // and variables is [a, b, c], then v2 < v1 because the rightmost depending
+  // variable of v2 is b, and the rightmost depending variable of v1 is c,
+  // and b < c. So in this example, this function will return [v2, v1].
+  // Tensors are always considered as variables and they are always considered
+  // as the rightmost.
+  std::vector<Val*> sortedInputs(const Context& context);
+
+  bool isTrivial() const {
+    return inputs().size() == 1;
+  }
+
+  std::vector<PolymorphicValue> evaluate(
+      const ExpressionEvaluator& ee,
+      const std::vector<PolymorphicValue>& inputs) const override;
+};
+
+} // namespace assoc_comm
+
+} // namespace nvfuser

nvfuser/include/nvfuser/flatbuffers/allocator.h

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2021 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FLATBUFFERS_ALLOCATOR_H_
+#define FLATBUFFERS_ALLOCATOR_H_
+
+#include "flatbuffers/base.h"
+
+namespace flatbuffers {
+
+// Allocator interface. This is flatbuffers-specific and meant only for
+// `vector_downward` usage.
+class Allocator {
+ public:
+  virtual ~Allocator() {}
+
+  // Allocate `size` bytes of memory.
+  virtual uint8_t *allocate(size_t size) = 0;
+
+  // Deallocate `size` bytes of memory at `p` allocated by this allocator.
+  virtual void deallocate(uint8_t *p, size_t size) = 0;
+
+  // Reallocate `new_size` bytes of memory, replacing the old region of size
+  // `old_size` at `p`. In contrast to a normal realloc, this grows downwards,
+  // and is intended specifcally for `vector_downward` use.
+  // `in_use_back` and `in_use_front` indicate how much of `old_size` is
+  // actually in use at each end, and needs to be copied.
+  virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size,
+                                       size_t new_size, size_t in_use_back,
+                                       size_t in_use_front) {
+    FLATBUFFERS_ASSERT(new_size > old_size);  // vector_downward only grows
+    uint8_t *new_p = allocate(new_size);
+    memcpy_downward(old_p, old_size, new_p, new_size, in_use_back,
+                    in_use_front);
+    deallocate(old_p, old_size);
+    return new_p;
+  }
+
+ protected:
+  // Called by `reallocate_downward` to copy memory from `old_p` of `old_size`
+  // to `new_p` of `new_size`. Only memory of size `in_use_front` and
+  // `in_use_back` will be copied from the front and back of the old memory
+  // allocation.
+  void memcpy_downward(uint8_t *old_p, size_t old_size, uint8_t *new_p,
+                       size_t new_size, size_t in_use_back,
+                       size_t in_use_front) {
+    memcpy(new_p + new_size - in_use_back, old_p + old_size - in_use_back,
+           in_use_back);
+    memcpy(new_p, old_p, in_use_front);
+  }
+};
+
+}  // namespace flatbuffers
+
+#endif  // FLATBUFFERS_ALLOCATOR_H_

nvfuser/include/nvfuser/flatbuffers/array.h

@@ -0,0 +1,253 @@
+/*
+ * Copyright 2021 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FLATBUFFERS_ARRAY_H_
+#define FLATBUFFERS_ARRAY_H_
+
+#include <memory>
+
+#include "flatbuffers/base.h"
+#include "flatbuffers/stl_emulation.h"
+#include "flatbuffers/vector.h"
+
+namespace flatbuffers {
+
+// This is used as a helper type for accessing arrays.
+template<typename T, uint16_t length> class Array {
+  // Array<T> can carry only POD data types (scalars or structs).
+  typedef typename flatbuffers::bool_constant<flatbuffers::is_scalar<T>::value>
+      scalar_tag;
+  typedef
+      typename flatbuffers::conditional<scalar_tag::value, T, const T *>::type
+          IndirectHelperType;
+
+ public:
+  typedef uint16_t size_type;
+  typedef typename IndirectHelper<IndirectHelperType>::return_type return_type;
+  typedef VectorConstIterator<T, return_type> const_iterator;
+  typedef VectorReverseIterator<const_iterator> const_reverse_iterator;
+
+  // If T is a LE-scalar or a struct (!scalar_tag::value).
+  static FLATBUFFERS_CONSTEXPR bool is_span_observable =
+      (scalar_tag::value && (FLATBUFFERS_LITTLEENDIAN || sizeof(T) == 1)) ||
+      !scalar_tag::value;
+
+  FLATBUFFERS_CONSTEXPR uint16_t size() const { return length; }
+
+  return_type Get(uoffset_t i) const {
+    FLATBUFFERS_ASSERT(i < size());
+    return IndirectHelper<IndirectHelperType>::Read(Data(), i);
+  }
+
+  return_type operator[](uoffset_t i) const { return Get(i); }
+
+  // If this is a Vector of enums, T will be its storage type, not the enum
+  // type. This function makes it convenient to retrieve value with enum
+  // type E.
+  template<typename E> E GetEnum(uoffset_t i) const {
+    return static_cast<E>(Get(i));
+  }
+
+  const_iterator begin() const { return const_iterator(Data(), 0); }
+  const_iterator end() const { return const_iterator(Data(), size()); }
+
+  const_reverse_iterator rbegin() const {
+    return const_reverse_iterator(end());
+  }
+  const_reverse_iterator rend() const {
+    return const_reverse_iterator(begin());
+  }
+
+  const_iterator cbegin() const { return begin(); }
+  const_iterator cend() const { return end(); }
+
+  const_reverse_iterator crbegin() const { return rbegin(); }
+  const_reverse_iterator crend() const { return rend(); }
+
+  // Get a mutable pointer to elements inside this array.
+  // This method used to mutate arrays of structs followed by a @p Mutate
+  // operation. For primitive types use @p Mutate directly.
+  // @warning Assignments and reads to/from the dereferenced pointer are not
+  //  automatically converted to the correct endianness.
+  typename flatbuffers::conditional<scalar_tag::value, void, T *>::type
+  GetMutablePointer(uoffset_t i) const {
+    FLATBUFFERS_ASSERT(i < size());
+    return const_cast<T *>(&data()[i]);
+  }
+
+  // Change elements if you have a non-const pointer to this object.
+  void Mutate(uoffset_t i, const T &val) { MutateImpl(scalar_tag(), i, val); }
+
+  // The raw data in little endian format. Use with care.
+  const uint8_t *Data() const { return data_; }
+
+  uint8_t *Data() { return data_; }
+
+  // Similarly, but typed, much like std::vector::data
+  const T *data() const { return reinterpret_cast<const T *>(Data()); }
+  T *data() { return reinterpret_cast<T *>(Data()); }
+
+  // Copy data from a span with endian conversion.
+  // If this Array and the span overlap, the behavior is undefined.
+  void CopyFromSpan(flatbuffers::span<const T, length> src) {
+    const auto p1 = reinterpret_cast<const uint8_t *>(src.data());
+    const auto p2 = Data();
+    FLATBUFFERS_ASSERT(!(p1 >= p2 && p1 < (p2 + length)) &&
+                       !(p2 >= p1 && p2 < (p1 + length)));
+    (void)p1;
+    (void)p2;
+    CopyFromSpanImpl(flatbuffers::bool_constant<is_span_observable>(), src);
+  }
+
+ protected:
+  void MutateImpl(flatbuffers::true_type, uoffset_t i, const T &val) {
+    FLATBUFFERS_ASSERT(i < size());
+    WriteScalar(data() + i, val);
+  }
+
+  void MutateImpl(flatbuffers::false_type, uoffset_t i, const T &val) {
+    *(GetMutablePointer(i)) = val;
+  }
+
+  void CopyFromSpanImpl(flatbuffers::true_type,
+                        flatbuffers::span<const T, length> src) {
+    // Use std::memcpy() instead of std::copy() to avoid performance degradation
+    // due to aliasing if T is char or unsigned char.
+    // The size is known at compile time, so memcpy would be inlined.
+    std::memcpy(data(), src.data(), length * sizeof(T));
+  }
+
+  // Copy data from flatbuffers::span with endian conversion.
+  void CopyFromSpanImpl(flatbuffers::false_type,
+                        flatbuffers::span<const T, length> src) {
+    for (size_type k = 0; k < length; k++) { Mutate(k, src[k]); }
+  }
+
+  // This class is only used to access pre-existing data. Don't ever
+  // try to construct these manually.
+  // 'constexpr' allows us to use 'size()' at compile time.
+  // @note Must not use 'FLATBUFFERS_CONSTEXPR' here, as const is not allowed on
+  //  a constructor.
+#if defined(__cpp_constexpr)
+  constexpr Array();
+#else
+  Array();
+#endif
+
+  uint8_t data_[length * sizeof(T)];
+
+ private:
+  // This class is a pointer. Copying will therefore create an invalid object.
+  // Private and unimplemented copy constructor.
+  Array(const Array &);
+  Array &operator=(const Array &);
+};
+
+// Specialization for Array[struct] with access using Offset<void> pointer.
+// This specialization used by idl_gen_text.cpp.
+template<typename T, uint16_t length> class Array<Offset<T>, length> {
+  static_assert(flatbuffers::is_same<T, void>::value, "unexpected type T");
+
+ public:
+  typedef const void *return_type;
+
+  const uint8_t *Data() const { return data_; }
+
+  // Make idl_gen_text.cpp::PrintContainer happy.
+  return_type operator[](uoffset_t) const {
+    FLATBUFFERS_ASSERT(false);
+    return nullptr;
+  }
+
+ private:
+  // This class is only used to access pre-existing data.
+  Array();
+  Array(const Array &);
+  Array &operator=(const Array &);
+
+  uint8_t data_[1];
+};
+
+template<class U, uint16_t N>
+FLATBUFFERS_CONSTEXPR_CPP11 flatbuffers::span<U, N> make_span(Array<U, N> &arr)
+    FLATBUFFERS_NOEXCEPT {
+  static_assert(
+      Array<U, N>::is_span_observable,
+      "wrong type U, only plain struct, LE-scalar, or byte types are allowed");
+  return span<U, N>(arr.data(), N);
+}
+
+template<class U, uint16_t N>
+FLATBUFFERS_CONSTEXPR_CPP11 flatbuffers::span<const U, N> make_span(
+    const Array<U, N> &arr) FLATBUFFERS_NOEXCEPT {
+  static_assert(
+      Array<U, N>::is_span_observable,
+      "wrong type U, only plain struct, LE-scalar, or byte types are allowed");
+  return span<const U, N>(arr.data(), N);
+}
+
+template<class U, uint16_t N>
+FLATBUFFERS_CONSTEXPR_CPP11 flatbuffers::span<uint8_t, sizeof(U) * N>
+make_bytes_span(Array<U, N> &arr) FLATBUFFERS_NOEXCEPT {
+  static_assert(Array<U, N>::is_span_observable,
+                "internal error, Array<T> might hold only scalars or structs");
+  return span<uint8_t, sizeof(U) * N>(arr.Data(), sizeof(U) * N);
+}
+
+template<class U, uint16_t N>
+FLATBUFFERS_CONSTEXPR_CPP11 flatbuffers::span<const uint8_t, sizeof(U) * N>
+make_bytes_span(const Array<U, N> &arr) FLATBUFFERS_NOEXCEPT {
+  static_assert(Array<U, N>::is_span_observable,
+                "internal error, Array<T> might hold only scalars or structs");
+  return span<const uint8_t, sizeof(U) * N>(arr.Data(), sizeof(U) * N);
+}
+
+// Cast a raw T[length] to a raw flatbuffers::Array<T, length>
+// without endian conversion. Use with care.
+// TODO: move these Cast-methods to `internal` namespace.
+template<typename T, uint16_t length>
+Array<T, length> &CastToArray(T (&arr)[length]) {
+  return *reinterpret_cast<Array<T, length> *>(arr);
+}
+
+template<typename T, uint16_t length>
+const Array<T, length> &CastToArray(const T (&arr)[length]) {
+  return *reinterpret_cast<const Array<T, length> *>(arr);
+}
+
+template<typename E, typename T, uint16_t length>
+Array<E, length> &CastToArrayOfEnum(T (&arr)[length]) {
+  static_assert(sizeof(E) == sizeof(T), "invalid enum type E");
+  return *reinterpret_cast<Array<E, length> *>(arr);
+}
+
+template<typename E, typename T, uint16_t length>
+const Array<E, length> &CastToArrayOfEnum(const T (&arr)[length]) {
+  static_assert(sizeof(E) == sizeof(T), "invalid enum type E");
+  return *reinterpret_cast<const Array<E, length> *>(arr);
+}
+
+template<typename T, uint16_t length>
+bool operator==(const Array<T, length> &lhs,
+                const Array<T, length> &rhs) noexcept {
+  return std::addressof(lhs) == std::addressof(rhs) ||
+         (lhs.size() == rhs.size() &&
+          std::memcmp(lhs.Data(), rhs.Data(), rhs.size() * sizeof(T)) == 0);
+}
+
+}  // namespace flatbuffers
+
+#endif  // FLATBUFFERS_ARRAY_H_