Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions include/infinicore/ops.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,5 @@
#include "ops/ones.hpp"
#include "ops/rearrange.hpp"
#include "ops/rms_norm.hpp"
#include "ops/silu.hpp"
#include "ops/swiglu.hpp"
16 changes: 16 additions & 0 deletions include/infinicore/ops/silu.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
#pragma once

#include "../device.hpp"
#include "common/op.hpp"

namespace infinicore::op {
class Silu {
public:
using schema = void (*)(Tensor, Tensor);
static void execute(Tensor output, Tensor input);
static common::OpDispatcher<schema> &dispatcher();
};

Tensor silu(Tensor input);
void silu_(Tensor output, Tensor input);
} // namespace infinicore::op
2 changes: 2 additions & 0 deletions python/infinicore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from infinicore.ops.matmul import matmul
from infinicore.ops.rearrange import rearrange
from infinicore.ops.rms_norm import rms_norm
from infinicore.ops.silu import silu
from infinicore.ops.swiglu import swiglu
from infinicore.tensor import (
empty,
Expand Down Expand Up @@ -75,6 +76,7 @@
"matmul",
"rearrange",
"rms_norm",
"silu",
"swiglu",
"empty",
"from_blob",
Expand Down
9 changes: 9 additions & 0 deletions python/infinicore/ops/silu.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
from infinicore.lib import _infinicore
from infinicore.tensor import Tensor


def silu(input, *, out=None):
if out is None:
return Tensor(_infinicore.silu(input._underlying))

_infinicore.silu_(out._underlying, input._underlying)
32 changes: 32 additions & 0 deletions src/infinicore/ops/silu/silu.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
#include "infinicore/ops/silu.hpp"
#include <stdexcept>

namespace infinicore::op {

common::OpDispatcher<Silu::schema> &Silu::dispatcher() {
static common::OpDispatcher<Silu::schema> dispatcher_;
return dispatcher_;
};

void Silu::execute(Tensor output, Tensor input) {
auto device_type = context::getDevice().getType();
auto func = dispatcher().lookup(device_type);

if (func == nullptr) {
throw std::runtime_error("No Silu implementation found for device type: " + std::to_string(static_cast<int>(device_type)));
}

func(output, input);
}

Tensor silu(Tensor input) {
Shape shape = input->shape();
auto output = Tensor::empty(shape, input->dtype(), input->device());
silu_(output, input);
return output;
}

void silu_(Tensor output, Tensor input) {
Silu::execute(output, input);
}
} // namespace infinicore::op
52 changes: 52 additions & 0 deletions src/infinicore/ops/silu/silu_infiniop.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
#include "../../utils.hpp"
#include "infinicore/common/hash.hpp"
#include "infinicore/ops/common/cache.hpp"
#include "infinicore/ops/silu.hpp"
#include <infiniop.h>

namespace infinicore::op::silu_impl::infiniop {

thread_local common::OpCache<size_t, infiniopSiluDescriptor_t> caches(
100, // capacity
[](infiniopSiluDescriptor_t &desc) {
if (desc != nullptr) {
INFINICORE_CHECK_ERROR(infiniopDestroySiluDescriptor(desc));
desc = nullptr;
}
});

void calculate(Tensor output, Tensor input) {
size_t seed = hash_combine(output, input);

auto device_type = context::getDevice().getType();
auto device_index = context::getDevice().getIndex();

auto &cache = caches.getCache(device_type, device_index);

auto desc_opt = cache.get(seed);
infiniopSiluDescriptor_t desc = nullptr;

if (!desc_opt) {
INFINICORE_CHECK_ERROR(infiniopCreateSiluDescriptor(
context::getInfiniopHandle(), &desc,
output->desc(), input->desc()));
cache.put(seed, desc);
} else {
desc = *desc_opt;
}

size_t workspace_size = 0;
INFINICORE_CHECK_ERROR(infiniopGetSiluWorkspaceSize(desc, &workspace_size));
std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);

INFINICORE_CHECK_ERROR(infiniopSilu(
desc, workspace->data(), workspace_size,
output->data(), input->data(), context::getStream()));
}

static bool registered = []() {
Silu::dispatcher().registerAll(&calculate, false);
return true;
}();

} // namespace infinicore::op::silu_impl::infiniop
2 changes: 2 additions & 0 deletions src/infinicore/pybind11/ops.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#include "ops/matmul.hpp"
#include "ops/rearrange.hpp"
#include "ops/rms_norm.hpp"
#include "ops/silu.hpp"
#include "ops/swiglu.hpp"

namespace py = pybind11;
Expand All @@ -19,6 +20,7 @@ inline void bind(py::module &m) {
bind_matmul(m);
bind_rearrange(m);
bind_rms_norm(m);
bind_silu(m);
bind_swiglu(m);
}

Expand Down
24 changes: 24 additions & 0 deletions src/infinicore/pybind11/ops/silu.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#pragma once

#include <pybind11/pybind11.h>

#include "infinicore/ops/silu.hpp"

namespace py = pybind11;

namespace infinicore::ops {

inline void bind_silu(py::module &m) {
m.def("silu",
&op::silu,
py::arg("input"),
R"doc(SiLU (Swish) activation function.)doc");

m.def("silu_",
&op::silu_,
py::arg("output"),
py::arg("input"),
R"doc(In-place SiLU (Swish) activation function.)doc");
}

} // namespace infinicore::ops
112 changes: 112 additions & 0 deletions test/infinicore/ops/silu.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import torch
import infinicore
from framework.base import BaseOperatorTest, TensorSpec, TestCase
from framework.runner import GenericTestRunner

# ==============================================================================
# Operator-specific configuration
# ==============================================================================

# Test cases format: (operation_mode, shape, input_strides, output_strides)
# SiLU is a single-input activation function: output = input * sigmoid(input)
_TEST_CASES_DATA = [
# Basic 2D SiLU
(TestCase.BOTH, (2, 4), None, None),
(TestCase.BOTH, (128, 64), None, None),
# 3D SiLU
(TestCase.BOTH, (2, 4, 8), None, None),
(TestCase.BOTH, (4, 48, 6), None, None),
# Strided tensors
(TestCase.BOTH, (1, 2048), (4096, 1), (4096, 1)),
(TestCase.BOTH, (6, 2560), (2048, 1), (2560, 1)),
# Mixed cases
(TestCase.BOTH, (8, 16, 32), None, None),
# Large tensors
(TestCase.BOTH, (16, 5632), None, None),
(TestCase.BOTH, (4, 4, 5632), None, None),
]


def parse_test_cases(data):
"""
Parse silu test case data according to format:
(operation_mode, shape, input_strides, output_strides)
"""
operation_mode = data[0]
shape = data[1]
input_strides = data[2] if len(data) > 2 else None
output_strides = data[3] if len(data) > 3 else None

# Create input specifications
inputs = []

# Tensor input
if input_strides is not None:
inputs.append(TensorSpec.from_strided_tensor(shape, input_strides))
else:
inputs.append(TensorSpec.from_tensor(shape))

# Output tensor
if output_strides is not None:
output = TensorSpec.from_strided_tensor(shape, output_strides)
else:
output = TensorSpec.from_tensor(shape)

return TestCase(operation_mode, inputs, output)


# Parse test cases
_TEST_CASES = [parse_test_cases(data) for data in _TEST_CASES_DATA]

# Data types
_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32]

# Tolerance
_TOLERANCE_MAP = {
infinicore.float16: {"atol": 1e-3, "rtol": 1e-3},
infinicore.float32: {"atol": 1e-5, "rtol": 1e-5},
infinicore.bfloat16: {"atol": 5e-3, "rtol": 1e-2},
}


class OpTest(BaseOperatorTest):
"""SiLU test with simplified test case parsing"""

def __init__(self):
super().__init__("SiLU")

def get_test_cases(self):
return _TEST_CASES

def get_tensor_dtypes(self):
return _TENSOR_DTYPES

def get_tolerance_map(self):
return _TOLERANCE_MAP

def torch_operator(self, input, out=None, **kwargs):
# SiLU implementation: input * sigmoid(input)
sigmoid_input = torch.sigmoid(input)
result = input * sigmoid_input
if out is not None:
out.copy_(result)
return out
return result

def infinicore_operator(self, input, out=None, **kwargs):
return infinicore.silu(input, out=out)


def main():
"""Main entry point"""
runner = GenericTestRunner(OpTest)
runner.run_and_exit()


if __name__ == "__main__":
main()