feat!: align non-InfiniLM operator interfaces

scripts/torch_ops.yaml

@@ -5,13 +5,11 @@
 # generator emits one InfiniOps wrapper per overload, so this
 # file lists 500+ base names and produces 500+ wrappers.
 #
-# To exclude an op, comment out its line.  Ops whose hand-written
-# `src/base/<op>.h` signature does not match the ATen-derived one
-# (currently `add`, `linear`, `matmul`, `mul` — they pre-date this
-# codegen and use a different parameter shape) must stay excluded:
-# the generator skips emitting their base, but would still emit a
-# torch backend declaring `operator()` with the ATen signature, and
-# that override would not compile against the hand-written base.
+# To exclude an op, comment out its line. Existing base operators may remain
+# on this list when their public overloads match the ATen schema; the generator
+# then emits their Torch backend. `cat`, `linear`, and `matmul` use this path.
+# Operators with handwritten Torch backends, such as `add`, stay excluded to
+# avoid generating a duplicate backend.
 
 - abs
 - absolute
@@ -68,6 +66,7 @@
 - bitwise_xor
 - bmm
 - bucketize
+- cat
 - ceil
 - chain_matmul
 - cholesky
@@ -235,6 +234,7 @@
 - linalg_tensorsolve
 - linalg_vecdot
 - linalg_vector_norm
+- linear
 - linspace
 - log
 - log10
@@ -259,6 +259,7 @@
 - lu_solve
 - lu_unpack
 - masked_select
+- matmul
 - matrix_power
 - max
 - max_pool2d_with_indices

src/CMakeLists.txt

@@ -419,7 +419,8 @@ set(INFINI_OPS_OPS "" CACHE STRING
 set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL
     "Build only the smoke-test operator subset")
 set(_infini_ops_smoke_ops
-    add mul cast cat gemm matmul linear rms_norm swiglu causal_softmax abs clamp exp)
+    add mul cast cat gemm matmul linear rms_norm swiglu silu_and_mul
+    causal_softmax causal_softmax_infinilm abs clamp exp)
 set(_infini_ops_smoke_torch_ops abs clamp exp)
 
 if(INFINI_OPS_SMOKE_BUILD)

src/base/add.h

@@ -1,28 +1,33 @@
 #ifndef INFINI_OPS_BASE_ADD_H_
 #define INFINI_OPS_BASE_ADD_H_
 
-#include <optional>
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <vector>
 
 #include "operator.h"
 
 namespace infini::ops {
 
 class Add : public Operator<Add> {
  public:
-  Add(const Tensor input, const Tensor other, Tensor out)
+  Add(const Tensor input, const Tensor other, const double alpha, Tensor out)
       : ndim_{out.ndim()},
         output_size_{out.numel()},
         input_type_{input.dtype()},
         other_type_{other.dtype()},
         out_type_{out.dtype()},
-        input_shape_{input.shape()},
-        other_shape_{other.shape()},
+        input_shape_{out.shape()},
+        other_shape_{out.shape()},
         out_shape_{out.shape()},
-        input_strides_{input.strides()},
-        other_strides_{other.strides()},
+        input_strides_{BroadcastStrides(input, out)},
+        other_strides_{BroadcastStrides(other, out)},
         out_strides_{out.strides()},
-        is_input_contiguous_{input.IsContiguous()},
-        is_other_contiguous_{other.IsContiguous()},
+        is_input_contiguous_{input.shape() == out.shape() &&
+                             input.IsContiguous()},
+        is_other_contiguous_{other.shape() == out.shape() &&
+                             other.IsContiguous()},
         is_out_contiguous_{out.IsContiguous()} {
     assert(!out.HasBroadcastDim() &&
            "the output of `Add` should NOT have broadcasted dim!");
@@ -31,18 +36,89 @@ class Add : public Operator<Add> {
     assert(input_type_ == other_type_ && other_type_ == out_type_ &&
            "operator `Add` requires all input and output Tensors to have the "
            "same dtype");
+    if (input_type_ != DataType::kFloat16 &&
+        input_type_ != DataType::kBFloat16 &&
+        input_type_ != DataType::kFloat32 &&
+        input_type_ != DataType::kFloat64) {
+      assert(alpha == static_cast<int64_t>(alpha) &&
+             "operator `Add` requires integral `alpha` for integer tensors");
+    }
+    ValidateBroadcast(input, other, out);
   }
 
+  Add(const Tensor input, const Tensor other, Tensor out)
+      : Add{input, other, 1.0, out} {}
+
   virtual void operator()(const Tensor input, const Tensor other,
-                          Tensor out) const = 0;
+                          const double alpha, Tensor out) const = 0;
+
+  void operator()(const Tensor input, const Tensor other, Tensor out) const {
+    (*this)(input, other, 1.0, out);
+  }
 
   template <typename TensorLike>
   static auto MakeReturnValue(const TensorLike& input,
                               const TensorLike& other) {
-    return TensorLike::Empty(input.shape(), input.dtype(), input.device());
+    return MakeReturnValue(input, other, 1.0);
+  }
+
+  template <typename TensorLike>
+  static auto MakeReturnValue(const TensorLike& input, const TensorLike& other,
+                              const double /*alpha*/) {
+    auto input_shape = input.shape();
+    auto other_shape = other.shape();
+    auto ndim = std::max(input_shape.size(), other_shape.size());
+    std::vector<std::size_t> out_shape(ndim, 1);
+
+    for (std::size_t i = 0; i < ndim; ++i) {
+      auto input_dim = i < ndim - input_shape.size()
+                           ? 1
+                           : input_shape[i + input_shape.size() - ndim];
+      auto other_dim = i < ndim - other_shape.size()
+                           ? 1
+                           : other_shape[i + other_shape.size() - ndim];
+      assert((input_dim == other_dim || input_dim == 1 || other_dim == 1) &&
+             "operator `Add` requires broadcast-compatible input shapes");
+      out_shape[i] = std::max(input_dim, other_dim);
+    }
+
+    return TensorLike::Empty(out_shape, input.dtype(), input.device());
   }
 
  protected:
+  static Tensor::Strides BroadcastStrides(const Tensor input,
+                                          const Tensor out) {
+    assert(input.ndim() <= out.ndim() &&
+           "operator `Add` input rank must not exceed output rank");
+    Tensor::Strides strides(out.ndim(), 0);
+    auto offset = out.ndim() - input.ndim();
+
+    for (Tensor::Size i = 0; i < input.ndim(); ++i) {
+      auto out_dim = i + offset;
+      assert((input.size(i) == 1 || input.size(i) == out.size(out_dim)) &&
+             "operator `Add` input shape is not broadcast-compatible with "
+             "output shape");
+      strides[out_dim] = input.size(i) == 1 ? 0 : input.stride(i);
+    }
+
+    return strides;
+  }
+
+  static void ValidateBroadcast(const Tensor input, const Tensor other,
+                                const Tensor out) {
+    for (Tensor::Size i = 0; i < out.ndim(); ++i) {
+      auto input_dim = i < out.ndim() - input.ndim()
+                           ? 1
+                           : input.size(i + input.ndim() - out.ndim());
+      auto other_dim = i < out.ndim() - other.ndim()
+                           ? 1
+                           : other.size(i + other.ndim() - out.ndim());
+      assert(out.size(i) == std::max(input_dim, other_dim) &&
+             "operator `Add` output shape must equal the broadcasted input "
+             "shape");
+    }
+  }
+
   Tensor::Size ndim_{0};
 
   Tensor::Size output_size_{0};

src/base/add_rms_norm.h

@@ -9,7 +9,9 @@
 
 namespace infini::ops {
 
-// Fused add + RMSNorm aligned with vLLM `fused_add_rms_norm`.
+// Legacy out-of-place fused add + RMSNorm interface.
+/// \deprecated Use `FusedAddRmsNorm`. This interface will be removed in a
+/// future release.
 class AddRmsNorm : public Operator<AddRmsNorm> {
  public:
   AddRmsNorm(const Tensor input, const Tensor residual, const Tensor weight,

src/base/cat.h

@@ -9,22 +9,46 @@ namespace infini::ops {
 
 class Cat : public Operator<Cat> {
  public:
-  Cat(const Tensor first_input, std::vector<Tensor> rest_inputs, int64_t dim,
-      Tensor out)
-      : input_count_{1 + rest_inputs.size()} {
-    assert(input_count_ >= 2 && "`Cat` requires at least 2 input tensors");
+  Cat(const std::vector<Tensor> tensors, const int64_t dim, Tensor out)
+      : out_shape_{out.shape()},
+        out_strides_{out.strides()},
+        out_type_{out.dtype()},
+        input_count_{tensors.size()} {
+    assert(!tensors.empty() && "`Cat` requires a non-empty tensor list");
 
     auto ndim = static_cast<int64_t>(out.ndim());
-    // Normalize negative dim (e.g. -1 means last dimension).
     dim_ = dim < 0 ? dim + ndim : dim;
     assert(dim_ >= 0 && dim_ < ndim && "`Cat` dim out of range");
+
+    Tensor::Size cat_size = 0;
+    for (const auto& tensor : tensors) {
+      assert(tensor.ndim() == out.ndim() &&
+             "`Cat` requires all tensors to have the output rank");
+      assert(tensor.dtype() == out.dtype() &&
+             "`Cat` requires all tensors to have the output dtype");
+
+      for (Tensor::Size axis = 0; axis < out.ndim(); ++axis) {
+        if (axis != static_cast<Tensor::Size>(dim_)) {
+          assert(tensor.size(axis) == out.size(axis) &&
+                 "`Cat` input dimensions must match outside `dim`");
+        }
+      }
+      cat_size += tensor.size(dim_);
+    }
+    assert(cat_size == out.size(dim_) &&
+           "`Cat` output size along `dim` must equal the input sum");
   }
 
-  virtual void operator()(const Tensor first_input,
-                          std::vector<Tensor> rest_inputs, int64_t dim,
+  virtual void operator()(const std::vector<Tensor> tensors, const int64_t dim,
                           Tensor out) const = 0;
 
  protected:
+  Tensor::Shape out_shape_;
+
+  Tensor::Strides out_strides_;
+
+  DataType out_type_;
+
   int64_t dim_;
 
   size_t input_count_;

src/base/causal_softmax.h

@@ -9,6 +9,8 @@
 
 namespace infini::ops {
 
+/// \deprecated Use `CausalSoftmaxInfinilm`. This interface will be removed in
+/// a future release.
 class CausalSoftmax : public Operator<CausalSoftmax> {
  public:
   CausalSoftmax(const Tensor input, Tensor out)

src/base/causal_softmax_infinilm.h

@@ -0,0 +1,52 @@
+#ifndef INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_
+#define INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_
+
+#include <cassert>
+#include <cstddef>
+
+#include "operator.h"
+#include "tensor.h"
+
+namespace infini::ops {
+
+class CausalSoftmaxInfinilm : public Operator<CausalSoftmaxInfinilm> {
+ public:
+  CausalSoftmaxInfinilm(const Tensor input, Tensor out)
+      : dtype_{input.dtype()},
+        ndim_{out.ndim()},
+        batch_size_{ndim_ == 2 ? 1 : out.size(-3)},
+        seq_len_{out.size(-2)},
+        total_seq_len_{out.size(-1)},
+        input_strides_{input.strides()},
+        out_strides_{out.strides()} {
+    assert(input.shape() == out.shape() &&
+           "`CausalSoftmaxInfinilm` requires `input` and `out` same shape");
+    assert(input.dtype() == out.dtype() &&
+           "`CausalSoftmaxInfinilm` requires `input` and `out` same dtype");
+    assert((ndim_ == 2 || ndim_ == 3) &&
+           "`CausalSoftmaxInfinilm` requires 2D or 3D tensor");
+    assert(seq_len_ <= total_seq_len_ &&
+           "`CausalSoftmaxInfinilm` requires shape[-2] <= shape[-1]");
+  }
+
+  virtual void operator()(const Tensor input, Tensor out) const = 0;
+
+ protected:
+  const DataType dtype_;
+
+  Tensor::Size ndim_{0};
+
+  Tensor::Size batch_size_{0};
+
+  Tensor::Size seq_len_{0};
+
+  Tensor::Size total_seq_len_{0};
+
+  Tensor::Strides input_strides_;
+
+  Tensor::Strides out_strides_;
+};
+
+}  // namespace infini::ops
+
+#endif