Skip to content

Add ELU_FWD pointwise op support #311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rsuderman wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add ELU_FWD pointwise op support #311

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 2 additions & 1 deletion include/fusilli/attributes/pointwise_attributes.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ namespace fusilli {
OP(CMP_NEQ) \
OP(DIV) \
/* OP(ELU_BWD) */ \
/* OP(ELU_FWD) */ \
OP(ELU_FWD) \
/* OP(ERF) */ \
/* OP(EXP) */ \
/* OP(FLOOR) */ \
Expand Down Expand Up @@ -138,6 +138,7 @@ inline const std::unordered_map<PointwiseAttr::Mode, int>
{PointwiseAttr::Mode::CMP_GE, 2},
{PointwiseAttr::Mode::CMP_NEQ, 2},
{PointwiseAttr::Mode::DIV, 2},
{PointwiseAttr::Mode::ELU_FWD, 1},
{PointwiseAttr::Mode::MUL, 2},
{PointwiseAttr::Mode::RELU_FWD, 1},
{PointwiseAttr::Mode::SIGMOID_FWD, 1},
Expand Down
13 changes: 13 additions & 0 deletions include/fusilli/support/asm_emitter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1723,12 +1723,23 @@ inline ErrorOr<std::string> PointwiseNode::emitNodePreAsm() const {
{6}
)";

constexpr std::string_view kEluSchema = R"(
{0}
%elu_alpha_{7} = torch.constant.float 1.000000e+00
%elu_scale_{7} = torch.constant.float 1.000000e+00
%elu_input_scale_{7} = torch.constant.float 1.000000e+00
{1} = {6} {2}, %elu_alpha_{7}, %elu_scale_{7}, %elu_input_scale_{7} : {3}, !torch.float, !torch.float, !torch.float -> {4}
{5}
)";
Comment on lines +1726 to +1733
Copy link
Copy Markdown
Member

@sjain-stanford sjain-stanford Apr 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need alpha to be configurable and not hardcoded (OK to have a default when not set).


#define FUSILLI_DECLARE_UNARY_TORCH_EMITTER(PWOP, OPIR) \
FUSILLI_DECLARE_UNARY_POINTWISE_EMITTER(PWOP, kUnaryTorchSchema, OPIR)
#define FUSILLI_DECLARE_BINARY_TORCH_EMITTER(PWOP, OPIR) \
FUSILLI_DECLARE_BINARY_POINTWISE_EMITTER(PWOP, kBinaryTorchSchema, OPIR)
#define FUSILLI_DECLARE_SUB_ADD_TORCH_EMITTER(PWOP, OPIR) \
FUSILLI_DECLARE_BINARY_POINTWISE_EMITTER(PWOP, kSubAddSchema, OPIR)
#define FUSILLI_DECLARE_ELU_TORCH_EMITTER(PWOP, OPIR) \
FUSILLI_DECLARE_UNARY_POINTWISE_EMITTER(PWOP, kEluSchema, OPIR)

switch (pointwiseAttr.getMode()) {
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(ABS, torch.aten.abs)
Expand All @@ -1740,6 +1751,7 @@ inline ErrorOr<std::string> PointwiseNode::emitNodePreAsm() const {
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(CMP_GE, torch.aten.ge.Tensor)
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(CMP_NEQ, torch.aten.ne.Tensor)
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(DIV, torch.aten.div.Tensor)
FUSILLI_DECLARE_ELU_TORCH_EMITTER(ELU_FWD, torch.aten.elu)
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(MUL, torch.aten.mul.Tensor)
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(RELU_FWD, torch.aten.relu)
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(SIGMOID_FWD, torch.aten.sigmoid)
Expand All @@ -1757,6 +1769,7 @@ inline ErrorOr<std::string> PointwiseNode::emitNodePreAsm() const {
#undef FUSILLI_DECLARE_UNARY_TORCH_EMITTER
#undef FUSILLI_DECLARE_BINARY_TORCH_EMITTER
#undef FUSILLI_DECLARE_SUB_ADD_TORCH_EMITTER
#undef FUSILLI_DECLARE_ELU_TORCH_EMITTER

//===----------------------------------------------------------------------===//
//
Expand Down
6 changes: 6 additions & 0 deletions samples/pointwise/pointwise_unary_ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
const auto mode = GENERATE(
PointwiseAttr::Mode::ABS,
PointwiseAttr::Mode::CEIL,
PointwiseAttr::Mode::ELU_FWD,
PointwiseAttr::Mode::RELU_FWD,
PointwiseAttr::Mode::SIGMOID_FWD,
PointwiseAttr::Mode::TANH_FWD);
Expand Down Expand Up @@ -106,6 +107,11 @@ TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
y = std::abs(xD);
break;
}
case PointwiseAttr::Mode::ELU_FWD: {
double xD = static_cast<double>(x);
y = xD >= 0 ? xD : std::expm1(xD);
Copy link
Copy Markdown
Member

@sjain-stanford sjain-stanford Apr 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't matter but just to be precise with PyTorch semantics:

Suggested change
y = xD >= 0 ? xD : std::expm1(xD);
y = xD > 0 ? xD : std::expm1(xD);

https://docs.pytorch.org/docs/stable/generated/torch.nn.ELU.html#torch.nn.ELU

break;
}
case PointwiseAttr::Mode::RELU_FWD: {
y = std::max(x, T(0));
break;
Expand Down
1 change: 1 addition & 0 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,7 @@ add_fusilli_lit_tests(
lit/test_pointwise_asm_emitter_cmp_lt.cpp
lit/test_pointwise_asm_emitter_cmp_le.cpp
lit/test_pointwise_asm_emitter_div.cpp
lit/test_pointwise_asm_emitter_elu.cpp
lit/test_pointwise_asm_emitter_mul.cpp
lit/test_pointwise_asm_emitter_mul_scalar.cpp
lit/test_pointwise_asm_emitter_sigmoid.cpp
Expand Down
63 changes: 63 additions & 0 deletions tests/lit/test_pointwise_asm_emitter_elu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
// Copyright 2026 Advanced Micro Devices, Inc.
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// RUN: %{TEST_EXE} | iree-opt --verify-roundtrip
// RUN: %{TEST_EXE} | FileCheck %s --check-prefix=TORCH-CHECK
// RUN: %{TEST_EXE} stats | FileCheck %s --check-prefix=%{BACKEND}-STATS-CHECK

// clang-format off
//
// TORCH-CHECK: module @module {
// TORCH-CHECK: func.func @main(%result_: !torch.tensor<[16,256,64,32],f32>, %arg0: !torch.vtensor<[16,256,64,32],f32>) attributes {torch.assume_strict_symbolic_shapes} {
// TORCH-CHECK: %permute_IN_0_val_0_pointwise_elu = torch.constant.int 0
// TORCH-CHECK: %permute_IN_0_val_1_pointwise_elu = torch.constant.int 1
// TORCH-CHECK: %permute_IN_0_val_2_pointwise_elu = torch.constant.int 2
// TORCH-CHECK: %permute_IN_0_val_3_pointwise_elu = torch.constant.int 3
// TORCH-CHECK: %permute_IN_0_pointwise_elu = torch.prim.ListConstruct %permute_IN_0_val_0_pointwise_elu, %permute_IN_0_val_1_pointwise_elu, %permute_IN_0_val_2_pointwise_elu, %permute_IN_0_val_3_pointwise_elu : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
// TORCH-CHECK: %arg0_pointwise_elu_perm = torch.aten.permute %arg0, %permute_IN_0_pointwise_elu : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
// TORCH-CHECK: %elu_alpha_pointwise_elu = torch.constant.float 1.000000e+00
// TORCH-CHECK: %elu_scale_pointwise_elu = torch.constant.float 1.000000e+00
// TORCH-CHECK: %elu_input_scale_pointwise_elu = torch.constant.float 1.000000e+00
// TORCH-CHECK: %result_pointwise_elu_perm = torch.aten.elu %arg0_pointwise_elu_perm, %elu_alpha_pointwise_elu, %elu_scale_pointwise_elu, %elu_input_scale_pointwise_elu : !torch.vtensor<[16,256,64,32],f32>, !torch.float, !torch.float, !torch.float -> !torch.vtensor<[16,256,64,32],f32>
// TORCH-CHECK: %permute_OUT_0_val_0_pointwise_elu = torch.constant.int 0
// TORCH-CHECK: %permute_OUT_0_val_1_pointwise_elu = torch.constant.int 1
// TORCH-CHECK: %permute_OUT_0_val_2_pointwise_elu = torch.constant.int 2
// TORCH-CHECK: %permute_OUT_0_val_3_pointwise_elu = torch.constant.int 3
// TORCH-CHECK: %permute_OUT_0_pointwise_elu = torch.prim.ListConstruct %permute_OUT_0_val_0_pointwise_elu, %permute_OUT_0_val_1_pointwise_elu, %permute_OUT_0_val_2_pointwise_elu, %permute_OUT_0_val_3_pointwise_elu : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
// TORCH-CHECK: %result = torch.aten.permute %result_pointwise_elu_perm, %permute_OUT_0_pointwise_elu : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
// TORCH-CHECK: torch.overwrite.tensor.contents %result overwrites %result_ : !torch.vtensor<[16,256,64,32],f32>, !torch.tensor<[16,256,64,32],f32>
// TORCH-CHECK: return
// TORCH-CHECK: }
// TORCH-CHECK: }
//
// AMDGPU-STATS-CHECK: "transient-memory-size": 0
// AMDGPU-STATS-CHECK: "dispatch-count": 1
// CPU-STATS-CHECK: "transient-memory-size": 0
// CPU-STATS-CHECK: "dispatch-count": 1
//
// clang-format on

#include <fusilli.h>

#include "utils.h"

#include <iostream>
#include <string>

using namespace fusilli;

int main(int argc, char **argv) {
std::string mode = (argc > 1) ? argv[1] : "default";

auto status = testUnaryPointwiseAsmEmitter(
"pointwise_asm_emitter_elu", "pointwise_elu", mode,
PointwiseAttr::Mode::ELU_FWD, {16, 256, 64, 32});
if (isError(status)) {
std::cerr << "Test failed: " << status << std::endl;
return 1;
}
return 0;
}
Loading