Skip to content

Commit 25d64ad

Browse files
rsudermanclaude
rsuderman
and
claude
authored
Add RECIPROCAL pointwise op support (#303)
Signed-off-by: Rob Suderman <rob.suderman@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 4f94f51 commit 25d64ad

5 files changed

Lines changed: 103 additions & 3 deletions

File tree

‎include/fusilli/attributes/pointwise_attributes.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ namespace fusilli {
5656
OP(MIN_OP) \
5757
OP(MUL) \
5858
OP(NEG) \
59-
/* OP(RECIPROCAL) */ \
59+
OP(RECIPROCAL) \
6060
/* OP(RELU_BWD) */ \
6161
OP(RELU_FWD) \
6262
/* OP(RSQRT) */ \
@@ -146,6 +146,7 @@ inline const std::unordered_map<PointwiseAttr::Mode, int>
146146
{PointwiseAttr::Mode::MIN_OP, 2},
147147
{PointwiseAttr::Mode::MUL, 2},
148148
{PointwiseAttr::Mode::NEG, 1},
149+
{PointwiseAttr::Mode::RECIPROCAL, 1},
149150
{PointwiseAttr::Mode::RELU_FWD, 1},
150151
{PointwiseAttr::Mode::SIGMOID_FWD, 1},
151152
{PointwiseAttr::Mode::SUB, 2},

‎include/fusilli/support/asm_emitter.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1780,6 +1780,7 @@ inline ErrorOr<std::string> PointwiseNode::emitNodePreAsm() const {
17801780
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(MIN_OP, torch.aten.minimum)
17811781
FUSILLI_DECLARE_BINARY_TORCH_EMITTER(MUL, torch.aten.mul.Tensor)
17821782
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(NEG, torch.aten.neg)
1783+
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(RECIPROCAL, torch.aten.reciprocal)
17831784
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(RELU_FWD, torch.aten.relu)
17841785
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(SIGMOID_FWD, torch.aten.sigmoid)
17851786
FUSILLI_DECLARE_UNARY_TORCH_EMITTER(TANH_FWD, torch.aten.tanh)

‎samples/pointwise/pointwise_unary_ops.cpp‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,35 @@ static std::string generateName(PointwiseAttr::Mode mode, DataType type,
3939
TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
4040
const auto dim = std::vector<int64_t>{2, 16, 64, 64};
4141

42+
auto supportsInteger = [](PointwiseAttr::Mode m) {
43+
switch (m) {
44+
case PointwiseAttr::Mode::ABS:
45+
case PointwiseAttr::Mode::NEG:
46+
case PointwiseAttr::Mode::RELU_FWD:
47+
return true;
48+
default:
49+
return false;
50+
}
51+
};
52+
53+
auto supportsFloat = [](PointwiseAttr::Mode m) {
54+
switch (m) {
55+
case PointwiseAttr::Mode::ABS:
56+
case PointwiseAttr::Mode::CEIL:
57+
case PointwiseAttr::Mode::ERF:
58+
case PointwiseAttr::Mode::EXP:
59+
case PointwiseAttr::Mode::FLOOR:
60+
case PointwiseAttr::Mode::NEG:
61+
case PointwiseAttr::Mode::RECIPROCAL:
62+
case PointwiseAttr::Mode::RELU_FWD:
63+
case PointwiseAttr::Mode::SIGMOID_FWD:
64+
case PointwiseAttr::Mode::TANH_FWD:
65+
return true;
66+
default:
67+
return false;
68+
}
69+
};
70+
4271
// clang-format off
4372
const auto mode = GENERATE(
4473
PointwiseAttr::Mode::ABS,
@@ -47,6 +76,7 @@ TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
4776
PointwiseAttr::Mode::EXP,
4877
PointwiseAttr::Mode::FLOOR,
4978
PointwiseAttr::Mode::NEG,
79+
PointwiseAttr::Mode::RECIPROCAL,
5080
PointwiseAttr::Mode::RELU_FWD,
5181
PointwiseAttr::Mode::SIGMOID_FWD,
5282
PointwiseAttr::Mode::TANH_FWD);
@@ -148,6 +178,11 @@ TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
148178
y = -x;
149179
break;
150180
}
181+
case PointwiseAttr::Mode::RECIPROCAL: {
182+
double xD = static_cast<double>(x);
183+
y = 1.0 / xD;
184+
break;
185+
}
151186
default:
152187
FAIL(
153188
"Unsupported pointwise mode: " << PointwiseAttr::kModeToStr.at(mode));
@@ -185,7 +220,9 @@ TEST_CASE("Pointwise unary ops", "[pointwise][graph]") {
185220
FUSILLI_REQUIRE_ASSIGN(Handle handle, Handle::create(kDefaultBackend));
186221

187222
// int32
188-
execute(handle, DataType::Int32, int(-128));
223+
if (supportsInteger(mode))
224+
execute(handle, DataType::Int32, int(-128));
189225
// fp16
190-
execute(handle, DataType::Half, half(3.14));
226+
if (supportsFloat(mode))
227+
execute(handle, DataType::Half, half(3.14));
191228
}

‎tests/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ add_fusilli_lit_tests(
162162
lit/test_pointwise_asm_emitter_mul.cpp
163163
lit/test_pointwise_asm_emitter_mul_scalar.cpp
164164
lit/test_pointwise_asm_emitter_neg.cpp
165+
lit/test_pointwise_asm_emitter_reciprocal.cpp
165166
lit/test_pointwise_asm_emitter_sigmoid.cpp
166167
lit/test_pointwise_asm_emitter_tanh.cpp
167168
lit/test_pointwise_asm_emitter_sub.cpp

‎tests/lit/test_pointwise_asm_emitter_reciprocal.cpp‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2026 Advanced Micro Devices, Inc.
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
// RUN: %{TEST_EXE} | iree-opt --verify-roundtrip
8+
// RUN: %{TEST_EXE} | FileCheck %s --check-prefix=TORCH-CHECK
9+
// RUN: %{TEST_EXE} stats | FileCheck %s --check-prefix=%{BACKEND}-STATS-CHECK
10+
11+
// clang-format off
12+
//
13+
// TORCH-CHECK: module @module {
14+
// TORCH-CHECK: func.func @main(%[[RESULT0:.+]]: !torch.tensor<[16,256,64,32],f32>, %[[ARG0:.+]]: !torch.vtensor<[16,256,64,32],f32>) attributes {torch.assume_strict_symbolic_shapes} {
15+
// TORCH-CHECK: %[[PERM0_0:.+]] = torch.constant.int 0
16+
// TORCH-CHECK: %[[PERM0_1:.+]] = torch.constant.int 1
17+
// TORCH-CHECK: %[[PERM0_2:.+]] = torch.constant.int 2
18+
// TORCH-CHECK: %[[PERM0_3:.+]] = torch.constant.int 3
19+
// TORCH-CHECK: %[[PERM0_LIST:.+]] = torch.prim.ListConstruct %[[PERM0_0]], %[[PERM0_1]], %[[PERM0_2]], %[[PERM0_3]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
20+
// TORCH-CHECK: %[[PERMUTE0:.+]] = torch.aten.permute %[[ARG0]], %[[PERM0_LIST]] : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
21+
// TORCH-CHECK: %[[RECIPROCAL:.+]] = torch.aten.reciprocal %[[PERMUTE0]] : !torch.vtensor<[16,256,64,32],f32> -> !torch.vtensor<[16,256,64,32],f32>
22+
// TORCH-CHECK: %[[PERM_OUT_0:.+]] = torch.constant.int 0
23+
// TORCH-CHECK: %[[PERM_OUT_1:.+]] = torch.constant.int 1
24+
// TORCH-CHECK: %[[PERM_OUT_2:.+]] = torch.constant.int 2
25+
// TORCH-CHECK: %[[PERM_OUT_3:.+]] = torch.constant.int 3
26+
// TORCH-CHECK: %[[PERM_OUT_LIST:.+]] = torch.prim.ListConstruct %[[PERM_OUT_0]], %[[PERM_OUT_1]], %[[PERM_OUT_2]], %[[PERM_OUT_3]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
27+
// TORCH-CHECK: %[[PERM_OUT:.+]] = torch.aten.permute %[[RECIPROCAL]], %[[PERM_OUT_LIST]] : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
28+
// TORCH-CHECK: torch.overwrite.tensor.contents %[[PERM_OUT]] overwrites %[[RESULT0]] : !torch.vtensor<[16,256,64,32],f32>, !torch.tensor<[16,256,64,32],f32>
29+
// TORCH-CHECK: return
30+
// TORCH-CHECK: }
31+
// TORCH-CHECK: }
32+
//
33+
// AMDGPU-STATS-CHECK: "transient-memory-size": 0
34+
// AMDGPU-STATS-CHECK: "dispatch-count": 1
35+
// CPU-STATS-CHECK: "transient-memory-size": 0
36+
// CPU-STATS-CHECK: "dispatch-count": 1
37+
//
38+
// clang-format on
39+
40+
#include <fusilli.h>
41+
42+
#include "utils.h"
43+
44+
#include <iostream>
45+
#include <string>
46+
47+
using namespace fusilli;
48+
49+
int main(int argc, char **argv) {
50+
std::string mode = (argc > 1) ? argv[1] : "default";
51+
52+
auto status = testUnaryPointwiseAsmEmitter(
53+
"pointwise_asm_emitter_reciprocal", "reciprocal", mode,
54+
PointwiseAttr::Mode::RECIPROCAL, {16, 256, 64, 32});
55+
if (isError(status)) {
56+
std::cerr << "Test failed: " << status << std::endl;
57+
return 1;
58+
}
59+
return 0;
60+
}

0 commit comments

Comments
 (0)