feat: add timestep embedding operator

include/infinicore/ops.hpp

@@ -90,6 +90,7 @@
 #include "ops/swap.hpp"
 #include "ops/swiglu.hpp"
 #include "ops/tanh.hpp"
+#include "ops/timestep_embedding.hpp"
 #include "ops/topksoftmax.hpp"
 #include "ops/vocab_parallel_embedding.hpp"
 #include "ops/zeros.hpp"

include/infinicore/ops/timestep_embedding.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(TimestepEmbedding, Tensor, const Tensor &, float);
+
+Tensor timestep_embedding(const Tensor &timestep,
+                          size_t embedding_dim = 256,
+                          float max_period = 10000.0f);
+
+void timestep_embedding_(Tensor output,
+                         const Tensor &timestep,
+                         float max_period = 10000.0f);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -155,6 +155,7 @@
 #include "infiniop/ops/tan.h"
 #include "infiniop/ops/tanh.h"
 #include "infiniop/ops/tanhshrink.h"
+#include "infiniop/ops/timestep_embedding.h"
 #include "infiniop/ops/topk.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"

include/infiniop/ops/timestep_embedding.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_TIMESTEP_EMBEDDING_API_H__
+#define __INFINIOP_TIMESTEP_EMBEDDING_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTimestepEmbeddingDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateTimestepEmbeddingDescriptor(
+    infiniopHandle_t handle,
+    infiniopTimestepEmbeddingDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output_desc,
+    infiniopTensorDescriptor_t timestep_desc);
+
+__INFINI_C __export infiniStatus_t infiniopTimestepEmbedding(
+    infiniopTimestepEmbeddingDescriptor_t desc,
+    void *output,
+    const void *timestep,
+    float max_period,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyTimestepEmbeddingDescriptor(
+    infiniopTimestepEmbeddingDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -165,6 +165,7 @@
 from infinicore.ops.scatter import scatter
 from infinicore.ops.sinh import sinh
 from infinicore.ops.situ_and_mul import situ_and_mul
+from infinicore.ops.timestep_embedding import timestep_embedding
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.sum import sum
 from infinicore.ops.swap import swap
@@ -280,6 +281,7 @@
     "mul",
     "mul_scalar",
     "situ_and_mul",
+    "timestep_embedding",
     "diff",
     "digamma",
     "dist",

python/infinicore/ops/timestep_embedding.py

@@ -0,0 +1,21 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def timestep_embedding(timestep, embedding_dim=256, max_period=10000.0, *, out=None):
+    max_period = float(max_period)
+    if out is None:
+        return Tensor(
+            _infinicore.timestep_embedding(
+                timestep._underlying,
+                int(embedding_dim),
+                max_period,
+            )
+        )
+
+    _infinicore.timestep_embedding_(
+        out._underlying,
+        timestep._underlying,
+        max_period,
+    )
+    return out

src/infinicore/ops/timestep_embedding/timestep_embedding.cc

@@ -0,0 +1,57 @@
+#include "infinicore/ops/timestep_embedding.hpp"
+
+#include "../../utils.hpp"
+
+#include <stdexcept>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(TimestepEmbedding);
+
+TimestepEmbedding::TimestepEmbedding(Tensor output,
+                                     const Tensor &timestep,
+                                     float max_period) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, timestep);
+    INFINICORE_GRAPH_OP_DISPATCH(
+        output->device().getType(), output, timestep, max_period);
+}
+
+void TimestepEmbedding::execute(Tensor output,
+                                const Tensor &timestep,
+                                float max_period) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(
+        TimestepEmbedding, output, timestep, max_period);
+}
+
+Tensor timestep_embedding(const Tensor &timestep,
+                          size_t embedding_dim,
+                          float max_period) {
+    auto output = Tensor::empty(
+        {timestep->numel(), embedding_dim},
+        DataType::F32,
+        timestep->device());
+    timestep_embedding_(output, timestep, max_period);
+    return output;
+}
+
+void timestep_embedding_(Tensor output,
+                         const Tensor &timestep,
+                         float max_period) {
+    if (timestep->ndim() != 1) {
+        throw std::runtime_error("timestep_embedding expects timestep shape [N]");
+    }
+    if (output->ndim() != 2 || output->size(0) != timestep->size(0)
+        || output->size(1) == 0 || output->size(1) % 2 != 0) {
+        throw std::runtime_error(
+            "timestep_embedding expects output shape [N, even embedding_dim]");
+    }
+    if (output->dtype() != DataType::F32) {
+        throw std::runtime_error("timestep_embedding output must be float32");
+    }
+    if (max_period <= 0.0f) {
+        throw std::runtime_error("timestep_embedding max_period must be positive");
+    }
+    TimestepEmbedding::execute(output, timestep, max_period);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/timestep_embedding/timestep_embedding_infiniop.cc

@@ -0,0 +1,46 @@
+#include "../infiniop_impl.hpp"
+#include "infinicore/ops/timestep_embedding.hpp"
+
+namespace infinicore::op::timestep_embedding_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, TimestepEmbedding, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor output, timestep;
+    float max_period;
+};
+
+void *plan(Tensor output, const Tensor &timestep, float max_period) {
+    size_t seed = hash_combine(output, timestep);
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, TimestepEmbedding,
+        seed, output->desc(), timestep->desc());
+
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(output),
+        graph::GraphTensor(timestep),
+        max_period};
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+    INFINICORE_CHECK_ERROR(infiniopTimestepEmbedding(
+        planned->descriptor->desc,
+        planned->output->data(),
+        planned->timestep->data(),
+        planned->max_period,
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(
+    TimestepEmbedding, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::timestep_embedding_impl::infiniop

src/infinicore/pybind11/ops.hpp

@@ -135,6 +135,7 @@
 #include "ops/tan.hpp"
 #include "ops/tanh.hpp"
 #include "ops/tanhshrink.hpp"
+#include "ops/timestep_embedding.hpp"
 #include "ops/topk.hpp"
 #include "ops/topksoftmax.hpp"
 #include "ops/triplet_margin_loss.hpp"
@@ -318,6 +319,7 @@ inline void bind(py::module &m) {
     bind_linear_w8a8i8(m);
     bind_silu_and_mul(m);
     bind_situ_and_mul(m);
+    bind_timestep_embedding(m);
     bind_sum(m);
     bind_var_mean(m);
     bind_var(m);

src/infinicore/pybind11/ops/timestep_embedding.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/timestep_embedding.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_timestep_embedding(py::module &m) {
+    m.def("timestep_embedding",
+          &op::timestep_embedding,
+          py::arg("timestep"),
+          py::arg("embedding_dim") = 256,
+          py::arg("max_period") = 10000.0f,
+          R"doc(Build sinusoidal timestep embeddings on the active device.)doc");
+
+    m.def("timestep_embedding_",
+          &op::timestep_embedding_,
+          py::arg("output"),
+          py::arg("timestep"),
+          py::arg("max_period") = 10000.0f,
+          R"doc(Build sinusoidal timestep embeddings into output.)doc");
+}
+
+} // namespace infinicore::ops

src/infiniop/ops/timestep_embedding/nvidia/timestep_embedding_nvidia.cu

@@ -0,0 +1,130 @@
+#include "timestep_embedding_nvidia.cuh"
+
+#include "../../../devices/nvidia/nvidia_common.cuh"
+#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
+#include "../../../tensor.h"
+
+#include <cmath>
+#include <cuda_runtime.h>
+
+namespace {
+
+template <typename T>
+INFINIOP_CUDA_KERNEL timestepEmbeddingKernel(
+    float *__restrict__ output,
+    const T *__restrict__ timestep,
+    size_t num_timesteps,
+    size_t embedding_dim,
+    float log_max_period) {
+    const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+    const size_t numel = num_timesteps * embedding_dim;
+    if (index >= numel) {
+        return;
+    }
+
+    const size_t half_dim = embedding_dim / 2;
+    const size_t timestep_index = index / embedding_dim;
+    const size_t output_dim = index % embedding_dim;
+    const size_t frequency_index = output_dim % half_dim;
+    const float frequency = expf(
+        -log_max_period * static_cast<float>(frequency_index)
+        / static_cast<float>(half_dim));
+    const float angle = static_cast<float>(timestep[timestep_index]) * frequency;
+    output[index] = output_dim < half_dim ? cosf(angle) : sinf(angle);
+}
+
+} // namespace
+
+namespace op::timestep_embedding::nvidia {
+
+struct Descriptor::Opaque {
+    std::shared_ptr<device::nvidia::Handle::Internal> internal;
+};
+
+Descriptor::~Descriptor() {
+    delete _opaque;
+}
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t output_desc,
+    infiniopTensorDescriptor_t timestep_desc) {
+    CHECK_OR_RETURN(timestep_desc->shape().size() == 1,
+                    INFINI_STATUS_BAD_TENSOR_SHAPE);
+    CHECK_OR_RETURN(output_desc->shape().size() == 2,
+                    INFINI_STATUS_BAD_TENSOR_SHAPE);
+    CHECK_OR_RETURN(output_desc->shape()[0] == timestep_desc->shape()[0]
+                        && output_desc->shape()[1] > 0
+                        && output_desc->shape()[1] % 2 == 0,
+                    INFINI_STATUS_BAD_TENSOR_SHAPE);
+    CHECK_OR_RETURN(output_desc->dtype() == INFINI_DTYPE_F32,
+                    INFINI_STATUS_BAD_TENSOR_DTYPE);
+    CHECK_DTYPE(timestep_desc->dtype(),
+                INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
+    CHECK_OR_RETURN(output_desc->isContiguous() && timestep_desc->isContiguous(),
+                    INFINI_STATUS_BAD_TENSOR_STRIDES);
+
+    *desc_ptr = new Descriptor(
+        timestep_desc->shape()[0],
+        output_desc->shape()[1],
+        timestep_desc->dtype(),
+        new Opaque{reinterpret_cast<device::nvidia::Handle *>(handle)->internal()},
+        handle->device,
+        handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *output,
+    const void *timestep,
+    float max_period,
+    void *stream) const {
+    if (max_period <= 0.0f) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+    const size_t numel = _num_timesteps * _embedding_dim;
+    if (numel == 0) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+    constexpr size_t block_size = 256;
+    const size_t grid_size = (numel + block_size - 1) / block_size;
+    const float log_max_period = std::log(max_period);
+    auto cuda_stream = reinterpret_cast<cudaStream_t>(stream);
+
+    switch (_input_dtype) {
+    case INFINI_DTYPE_F16:
+        timestepEmbeddingKernel<half><<<grid_size, block_size, 0, cuda_stream>>>(
+            reinterpret_cast<float *>(output),
+            reinterpret_cast<const half *>(timestep),
+            _num_timesteps,
+            _embedding_dim,
+            log_max_period);
+        break;
+    case INFINI_DTYPE_BF16:
+        timestepEmbeddingKernel<cuda_bfloat16><<<grid_size, block_size, 0, cuda_stream>>>(
+            reinterpret_cast<float *>(output),
+            reinterpret_cast<const cuda_bfloat16 *>(timestep),
+            _num_timesteps,
+            _embedding_dim,
+            log_max_period);
+        break;
+    case INFINI_DTYPE_F32:
+        timestepEmbeddingKernel<float><<<grid_size, block_size, 0, cuda_stream>>>(
+            reinterpret_cast<float *>(output),
+            reinterpret_cast<const float *>(timestep),
+            _num_timesteps,
+            _embedding_dim,
+            log_max_period);
+        break;
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return cudaGetLastError() == cudaSuccess
+             ? INFINI_STATUS_SUCCESS
+             : INFINI_STATUS_INTERNAL_ERROR;
+}
+
+} // namespace op::timestep_embedding::nvidia

src/infiniop/ops/timestep_embedding/nvidia/timestep_embedding_nvidia.cuh

@@ -0,0 +1,8 @@
+#ifndef __TIMESTEP_EMBEDDING_NVIDIA_CUH__
+#define __TIMESTEP_EMBEDDING_NVIDIA_CUH__
+
+#include "../timestep_embedding.h"
+
+TIMESTEP_EMBEDDING_DESCRIPTOR(nvidia)
+
+#endif