nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/ops/composite.h

@@ -0,0 +1,130 @@
+// 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 <visibility.h>
+
+#include <ir/interface_nodes.h>
+#include <type.h>
+
+//
+// The operations defined in this header is intended as user facing functions.
+// The user will provide the necessary input TensorViews and the function will
+// create the correct intermediate nodes and return the output TensorViews.
+//
+
+namespace nvfuser {
+
+struct ForwardDropoutResult {
+  TensorView* output = nullptr;
+  TensorView* mask = nullptr;
+};
+
+NVF_API ForwardDropoutResult dropout(TensorView* x, Val* prob);
+
+NVF_API ForwardDropoutResult dropout(TensorView* x, Val* prob, Val* scale);
+
+NVF_API TensorView* dropout_backward(
+    TensorView* dy,
+    TensorView* mask,
+    Val* scale);
+
+NVF_API TensorView* triu(TensorView* tv, Val* offset);
+
+struct LstmResult {
+  TensorView* cell = nullptr;
+  TensorView* hidden = nullptr;
+};
+
+NVF_API LstmResult lstm(
+    TensorView* prev_cell,
+    TensorView* in_x,
+    TensorView* forget_x,
+    TensorView* cell_x,
+    TensorView* out_x);
+
+// Linear functions which takes in two tensors of shapes input[* , in_features],
+// weight[out_features, in_features] / [in_features] and an optional bias of
+// shape [out_features] or 0D scalar. Bias can only be given if weight is a 2-D
+// tensor.
+TensorView* linear(TensorView* input, TensorView* weight, TensorView* bias);
+// This is an implementation detail to reflect when linear is called
+// without a bias. This calls the above function. We use this function
+// since it simplifies creating a Python API which takes optional arguments.
+// Other options include using lambdas or creating a new RecordFunctor for
+// Linear.
+TensorView* linear(TensorView* input, TensorView* weight);
+
+NVF_API TensorView* sign(TensorView* x);
+NVF_API Val* sign(Val* x);
+TensorView* softplus(TensorView* x, Val* beta, Val* threshold);
+NVF_API TensorView* gelu(TensorView* x);
+NVF_API TensorView* gelu_backward(TensorView* dy, TensorView* x);
+TensorView* tanh_gelu(TensorView* x);
+TensorView* tanh_gelu_backward(TensorView* dy, TensorView* x);
+TensorView* tanh_backward(TensorView* dy, TensorView* tanh_x);
+TensorView* leaky_relu(TensorView* x, Val* negative_slope);
+
+NVF_API TensorView* view_as_real(TensorView* x);
+
+// Matmul function which takes in tensors with the shapes
+// A[*, M, K] / A[K] and B[*, K, N] / B[K], but the tensors may have different
+// layouts via strides. This has the same functionality as torch.matmul
+TensorView* matmul(TensorView* tv_a, TensorView* tv_b);
+
+// Scaled Dot Product Flash Attention Forward Result
+struct SdpfaFwdResult {
+  TensorView* output = nullptr;
+  TensorView* log_sumexp = nullptr;
+  TensorView* philox_seed = nullptr;
+  TensorView* philox_offset = nullptr;
+};
+
+// Scaled Dot Product Flash Attention Forward API.
+// Returns the same output as at::_scaled_dot_product_flash_attention
+SdpfaFwdResult sdpfa_fwd(
+    TensorView* query,
+    TensorView* key,
+    TensorView* value,
+    Val* dropout_p,
+    Val* is_causal,
+    Val* scale);
+
+// Scaled Dot Product Flash Attention Backward Result
+struct SdpfaBwdResult {
+  TensorView* grad_query = nullptr;
+  TensorView* grad_key = nullptr;
+  TensorView* grad_value = nullptr;
+};
+
+// Scaled Dot Product Flash Attention Backward API.
+// Returns the same output as at::_scaled_dot_product_flash_attention_backward
+SdpfaBwdResult sdpfa_bwd(
+    TensorView* grad_output,
+    TensorView* query,
+    TensorView* key,
+    TensorView* value,
+    TensorView* output,
+    TensorView* log_sumexp,
+    Val* dropout_p,
+    Val* is_causal,
+    TensorView* philox_seed,
+    TensorView* philox_offset,
+    Val* scale);
+
+TensorView* embedding_fwd(
+    TensorView* input,
+    TensorView* weight,
+    Val* padding_idx,
+    Val* max_norm,
+    Val* norm_type,
+    Val* scale_grad_by_freq,
+    Val* sparse);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ops/indexing.h

@@ -0,0 +1,55 @@
+// 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 <visibility.h>
+
+#include <ir/interface_nodes.h>
+#include <type.h>
+
+namespace nvfuser {
+
+NVF_API TensorView* select(TensorView* tv, int64_t dim, Val* index);
+
+// torch.index_select
+NVF_API TensorView* indexSelect(
+    TensorView* input,
+    int64_t dim,
+    TensorView* index);
+
+// torch.gather
+NVF_API TensorView* torchGather(
+    TensorView* input,
+    int64_t dim,
+    TensorView* index);
+
+// torch.scatter
+TensorView* scatterOp(
+    ScatterOpType type,
+    TensorView* self,
+    int64_t dim,
+    TensorView* index,
+    TensorView* src);
+
+NVF_API TensorView* scatter(
+    TensorView* self,
+    int64_t dim,
+    TensorView* index,
+    TensorView* src);
+
+//! numpy.take_along_axis
+//! (https://numpy.org/doc/stable/reference/generated/numpy.take_along_axis.html)
+//! Note the order of the parameters follows the numpy order, which is
+//! different from torchGather.
+NVF_API TensorView* takeAlongAxis(
+    TensorView* input,
+    TensorView* index,
+    int64_t dim);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ops/normalization.h

@@ -0,0 +1,263 @@
+// 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 <visibility.h>
+
+#include <ir/interface_nodes.h>
+#include <type.h>
+
+#include <tuple>
+#include <vector>
+
+//
+// The operations defined in this header is intended as user facing functions.
+// The user will provide the necessary input TensorViews and the function will
+// create the correct intermediate nodes and return the output TensorViews.
+//
+
+namespace nvfuser {
+
+struct ForwardNormResult {
+  TensorView* output = nullptr;
+  TensorView* mean = nullptr;
+  TensorView* invstd = nullptr;
+};
+
+struct BackwardNormResult {
+  TensorView* grad_input = nullptr;
+  TensorView* grad_weight = nullptr;
+  TensorView* grad_bias = nullptr;
+};
+
+struct ForwardRMSNormResult {
+  TensorView* output = nullptr;
+  TensorView* invstd = nullptr;
+};
+
+struct BackwardRMSNormResult {
+  TensorView* grad_input = nullptr;
+  TensorView* grad_weight = nullptr;
+};
+
+struct VarMeanResult {
+  TensorView* var = nullptr;
+  TensorView* mean = nullptr;
+};
+
+} // namespace nvfuser
+
+namespace std {
+
+// Make these results behave like a std::tuple
+using nvfuser::BackwardNormResult;
+using nvfuser::BackwardRMSNormResult;
+using nvfuser::ForwardNormResult;
+using nvfuser::ForwardRMSNormResult;
+using nvfuser::TensorView;
+using nvfuser::VarMeanResult;
+
+template <int i>
+constexpr TensorView* get(const ForwardNormResult& results) {
+  if (i == 0) {
+    return results.output;
+  }
+  if (i == 1) {
+    return results.mean;
+  }
+  if (i == 2) {
+    return results.invstd;
+  }
+  return nullptr;
+}
+
+template <int i>
+constexpr TensorView* get(const BackwardNormResult& results) {
+  if (i == 0) {
+    return results.grad_input;
+  }
+  if (i == 1) {
+    return results.grad_weight;
+  }
+  if (i == 2) {
+    return results.grad_bias;
+  }
+  return nullptr;
+}
+
+template <int i>
+constexpr TensorView* get(const ForwardRMSNormResult& results) {
+  if (i == 0) {
+    return results.output;
+  }
+  if (i == 1) {
+    return results.invstd;
+  }
+  return nullptr;
+}
+
+template <int i>
+constexpr TensorView* get(const BackwardRMSNormResult& results) {
+  if (i == 0) {
+    return results.grad_input;
+  }
+  if (i == 1) {
+    return results.grad_weight;
+  }
+  return nullptr;
+}
+
+template <int i>
+constexpr TensorView* get(const VarMeanResult& results) {
+  if (i == 0) {
+    return results.var;
+  }
+  if (i == 1) {
+    return results.mean;
+  }
+  return nullptr;
+}
+
+} // namespace std
+
+namespace nvfuser {
+
+TensorView* mean(TensorView* x, const std::vector<int64_t>& dims, bool keepdim);
+
+NVF_API TensorView* variance(
+    TensorView* x,
+    const std::vector<int64_t>& dims,
+    bool unbiased,
+    bool keepdim);
+
+NVF_API TensorView* variance(
+    TensorView* x,
+    const std::vector<int64_t>& dims,
+    int64_t correction,
+    bool keepdim);
+
+NVF_API VarMeanResult variance_mean(
+    TensorView* x,
+    const std::vector<int64_t>& dims,
+    int64_t correction,
+    bool keepdim);
+
+NVF_API TensorView* standard_deviation(
+    TensorView* x,
+    const std::vector<int64_t>& dims,
+    bool unbiased,
+    bool keepdim);
+
+NVF_API TensorView* softmax(TensorView* x, int64_t dim);
+
+NVF_API TensorView* softmax_backward(
+    TensorView* dy,
+    TensorView* y,
+    const int64_t dim);
+
+NVF_API TensorView* log_softmax(TensorView* x, int64_t dim);
+
+NVF_API TensorView* log_softmax_backward(
+    TensorView* dy,
+    TensorView* y,
+    const int64_t dim);
+
+NVF_API ForwardNormResult layer_norm(
+    TensorView* x,
+    const std::vector<int64_t>& norm_shape,
+    TensorView* weight,
+    TensorView* bias,
+    Val* eps);
+
+NVF_API ForwardNormResult layer_norm(
+    TensorView* x,
+    const int64_t kNormShapeNumDims,
+    TensorView* weight,
+    TensorView* bias,
+    Val* eps);
+
+NVF_API ForwardRMSNormResult rms_norm(
+    TensorView* x,
+    const std::vector<int64_t>& norm_shape,
+    TensorView* weight,
+    Val* eps);
+
+NVF_API ForwardRMSNormResult rms_norm(
+    TensorView* x,
+    const int64_t kNormShapeNumDims,
+    TensorView* weight,
+    Val* eps);
+
+NVF_API BackwardNormResult layer_norm_backward(
+    TensorView* dy,
+    TensorView* x,
+    const std::vector<int64_t>& norm_shape,
+    TensorView* mean,
+    TensorView* rstd,
+    TensorView* weight,
+    TensorView* bias,
+    const std::vector<bool>& output_mask);
+
+NVF_API BackwardRMSNormResult rms_norm_backward(
+    TensorView* dy,
+    TensorView* x,
+    const std::vector<int64_t>& norm_shape,
+    TensorView* rstd,
+    TensorView* weight,
+    const std::vector<bool>& output_mask);
+
+NVF_API ForwardNormResult batch_norm(
+    TensorView* x,
+    TensorView* weight,
+    TensorView* bias,
+    TensorView* running_mean,
+    TensorView* running_var,
+    const bool kTraining,
+    Val* momentum,
+    Val* eps,
+    bool channels_last = false);
+
+NVF_API BackwardNormResult batch_norm_backward(
+    TensorView* x,
+    TensorView* dy,
+    TensorView* weight,
+    TensorView* running_mean,
+    TensorView* running_var,
+    TensorView* save_mean,
+    TensorView* save_invstd,
+    const bool kTraining,
+    Val* eps,
+    const std::vector<bool>& output_mask,
+    bool channels_last = false);
+
+NVF_API ForwardNormResult instance_norm(
+    TensorView* x,
+    TensorView* weight,
+    TensorView* bias,
+    TensorView* running_mean,
+    TensorView* running_var,
+    const bool kUseInputStats, // kTraining?
+    Val* momentum,
+    Val* eps,
+    bool channels_last = false);
+
+NVF_API BackwardNormResult instance_norm_backward(
+    TensorView* x,
+    TensorView* dy,
+    TensorView* weight,
+    TensorView* running_mean,
+    TensorView* running_var,
+    TensorView* save_mean,
+    TensorView* save_invstd,
+    const bool kTraining,
+    Val* eps,
+    const std::vector<bool>& output_mask,
+    bool channels_last = false);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ops/utils.h

@@ -0,0 +1,127 @@
+// 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/base_nodes.h>
+#include <ir/interface_nodes.h>
+#include <scheduler/matmul_utils.h>
+#include <type.h>
+#include <visibility.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+enum class AttnRole { Q = 0, K, V, Mask };
+
+namespace ops {
+
+TensorView* maybe_broadcast_inner_to_rank(TensorView* t, size_t rank);
+
+// A utility function that broadcasts index TensorView to the rank of the other
+// TensorView.
+TensorView* maybeBroadcastIndexTv(TensorView* t, size_t dim, size_t rank);
+
+// A utility function that checks if index tv is already broadcasted to correct
+// shape for index_select
+bool isIndexAlreadyBroadcast(
+    const std::vector<IterDomain*>& index_domain,
+    size_t dim,
+    size_t rank);
+
+Val* simplifiedInt(Val* val);
+
+// If one size is nullptr, return the other. If both symbolic just return v1. If
+// one's concrete, prefer that one (simplified). If both concrete make sure
+// they're the same size.
+Val* promoteSize(Val* v1, Val* v2);
+
+// Will return a new value of type val with the DataType dtype.
+Val* newScalar(ValType vtype, DataType dtype);
+
+IterType promoteIterType(IterType type1, IterType type2);
+
+// For MatmulOp, the input iterdomains at a given index do not necessarily map
+// to the output iterdomain at that index This function aligns the input
+// iterdomain to the output and returns a vector where each element is the input
+// iterdomain corresponding to the output iterdomain at that index. If the
+// element is nullptr, there is no mapping between input-output at that index.
+// Based on the input dimensions following cases are possible:
+// 1. A/B is 1D: [M, K] x [K] -> [M] (Mapping A: {id_M}, Mapping B: {nullptr})
+// or [K] x [N, K] -> [N] (Mapping A: {nullptr}, Mapping B: {id_N})
+// 2. A and B are 2D: [M, K] x [K, N] -> [M, N] (Mapping A: {id_M, nullptr},
+// Mapping B: {nullptr, id_N})
+// 3. A/B are atleast 1D and one of them is > 2D: [B, M, K] x [K, N] -> [B, M,
+// N] (Mapping A: {id_B, id_M, nullptr}, Mapping B: {nullptr, nullptr, id_N})
+// Args:
+// 1. input_domain: root/logical domain without reductions for any input to
+// MatmulOp
+// 2. input_position: Specifies if the input is A / B (0 or 1)
+// 3: out_size: MatmulOp output dimension (input and output may not be the same
+// size).
+std::vector<IterDomain*> mapMatmulOpIterDomains(
+    const std::vector<IterDomain*>& input_domain,
+    int64_t input_position,
+    size_t out_size);
+
+// For LinearOp, the output is the same as the first input (A[*,
+// in_features])for all but the last dimension. If the second input is 2D
+// (B[out_features, in_features]), the last dimension of output is out_features.
+// If bias is 1D (bias[out_features]) it maps to the last dimension of the
+// output. Args:
+// 1. input_domain: root/logical domain without reductions for any input to
+// LinearOp
+// 2. input_position: Specifies if the input is A / B / Bias (0, 1, or 2)
+// (MatmulTensorRole::Input_A/Input_B/Input_C) 3: out_size: LinearOp output
+// dimension (input and output may not be the same size).
+std::vector<IterDomain*> mapLinearOpIterDomains(
+    const std::vector<IterDomain*>& input_domain,
+    int64_t input_position,
+    size_t out_size,
+    bool k_bcast);
+
+// Takes a vector of aligned input iterdomains to create the output iterdomain.
+// This is used if the input iterdomains are not trivially mapped to the output
+// iterdomains. For eg: MatmulOp. If given, the forced_iter_type argument will
+// be the output IterType regardless of the inputs; otherwise the output
+// IterType is inferred from ids.
+IterDomain* newOutputIterDomain(
+    const std::vector<IterDomain*>& ids,
+    const std::optional<IterType> force_iter_type = std::nullopt);
+
+// Takes a vector of `Val*`s and assumes they are all aligned to create the
+// output tensorview, e.g., for BinaryOp. `vals` can contain scalars, e.g, when
+// creating the output TensorView for `tv0+scalar`. This is for convenience and
+// scalars will be ignored.
+std::vector<IterDomain*> newOutputDomain(const std::vector<Val*>& vals);
+
+TensorView* newOutputTV(const std::vector<Val*>& vals, DataType dtype);
+
+std::vector<Val*> maybeBroadcast(const std::vector<Val*>& vals);
+
+NVF_API Val* newValLike(Val* val, DataType dtype);
+
+// returns the minimum init value for reduction:
+//   -inf for floating type;
+//   lowest value for integer type;
+//   false for bool.
+Val* getMinimumValue(DataType v);
+
+// returns the maximum init value for reduction:
+//   inf for floating type;
+//   highest value for integer type;
+//   true for bool.
+Val* getMaximumValue(DataType v);
+
+std::vector<unsigned int> canonicalizeAxes(
+    const std::vector<int64_t>& axes,
+    int64_t ndims);
+
+} // namespace ops
+} // namespace nvfuser