Skip to content

Commit f745fc5

Browse files
rsudermanclaude
rsuderman
and
claude
authored
Add ADD_SQUARE pointwise op (#358)
Emits y = x0 + x1^2 as torch.aten.mul.Tensor followed by torch.aten.add.Tensor. Adds sample and lit tests. --------- Signed-off-by: Rob Suderman <rob.suderman@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 2a79d33 commit f745fc5

6 files changed

Lines changed: 127 additions & 19 deletions

File tree

‎include/fusilli/attributes/pointwise_attributes.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ namespace fusilli {
2727
#define FUSILLI_POINTWISE_OPS(OP) \
2828
OP(ABS) \
2929
OP(ADD) \
30-
/* OP(ADD_SQUARE) */ \
30+
OP(ADD_SQUARE) \
3131
/* OP(BINARY_SELECT) */ \
3232
OP(CEIL) \
3333
OP(CMP_EQ) \
@@ -151,6 +151,7 @@ inline const std::unordered_map<PointwiseAttr::Mode, int>
151151
PointwiseAttr::kModeToRequiredInputCount = {
152152
{PointwiseAttr::Mode::ABS, 1},
153153
{PointwiseAttr::Mode::ADD, 2},
154+
{PointwiseAttr::Mode::ADD_SQUARE, 2},
154155
{PointwiseAttr::Mode::CEIL, 1},
155156
{PointwiseAttr::Mode::CMP_EQ, 2},
156157
{PointwiseAttr::Mode::CMP_LT, 2},

‎include/fusilli/node/pointwise_node.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ class PointwiseNode : public NodeCRTP<PointwiseNode> {
5252
std::string emitNodePreAsm() const override final;
5353
std::string getPermuteInputOpsAsm(int inputIndex) const;
5454
std::string getPermuteOut0OpsAsm() const;
55+
std::string getInputNameAsm(size_t inputIndex) const;
56+
std::string getInputTypeAsm(size_t inputIndex) const;
5557
std::string getOperandNamesAsm() const;
5658
std::string getOperandTypesAsm() const;
5759
std::string getResultNamesAsm() const;

‎include/fusilli/support/asm_emitter.h‎

Lines changed: 49 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1660,35 +1660,43 @@ inline std::string MatmulNode::emitNodePreAsm() const {
16601660
//
16611661
//===----------------------------------------------------------------------===//
16621662

1663-
// Emits PointwiseNode's operand names in MLIR assembly format.
1663+
// Emits the SSA name for a single PointwiseNode input in MLIR assembly
1664+
// format.
16641665
//
16651666
// The unique suffix is included to ensure SSA uniqueness when the same
16661667
// tensor is used by multiple operations in a graph.
1668+
inline std::string PointwiseNode::getInputNameAsm(size_t inputIndex) const {
1669+
const auto &in = pointwiseAttr.inputs.at(
1670+
static_cast<PointwiseAttr::InputNames>(inputIndex));
1671+
return in->getValueNameAsm() + "_" + pointwiseAttr.getName() + "_perm";
1672+
}
1673+
1674+
// Emits the type for a single PointwiseNode input in MLIR assembly format.
1675+
inline std::string PointwiseNode::getInputTypeAsm(size_t inputIndex) const {
1676+
const auto &in = pointwiseAttr.inputs.at(
1677+
static_cast<PointwiseAttr::InputNames>(inputIndex));
1678+
return in->getTensorTypeAsm(/*isValueTensor=*/true, /*useLogicalDims=*/true);
1679+
}
1680+
1681+
// Emits PointwiseNode's operand names in MLIR assembly format.
16671682
inline std::string PointwiseNode::getOperandNamesAsm() const {
16681683
std::ostringstream oss;
1669-
std::string suffix = pointwiseAttr.getName();
1670-
const auto &in0 = pointwiseAttr.getIN_0();
1671-
oss << in0->getValueNameAsm() << "_" << suffix << "_perm";
1672-
if (const auto &in1 = pointwiseAttr.getIN_1())
1673-
oss << ", " << in1->getValueNameAsm() << "_" << suffix << "_perm";
1674-
if (const auto &in2 = pointwiseAttr.getIN_2())
1675-
oss << ", " << in2->getValueNameAsm() << "_" << suffix << "_perm";
1684+
oss << getInputNameAsm(0);
1685+
if (pointwiseAttr.getIN_1())
1686+
oss << ", " << getInputNameAsm(1);
1687+
if (pointwiseAttr.getIN_2())
1688+
oss << ", " << getInputNameAsm(2);
16761689
return oss.str();
16771690
}
16781691

16791692
// Emits PointwiseNode's operand types in MLIR assembly format.
16801693
inline std::string PointwiseNode::getOperandTypesAsm() const {
16811694
std::ostringstream oss;
1682-
const auto &in0 = pointwiseAttr.getIN_0();
1683-
oss << in0->getTensorTypeAsm(/*isValueTensor=*/true, /*useLogicalDims=*/true);
1684-
if (const auto &in1 = pointwiseAttr.getIN_1())
1685-
oss << ", "
1686-
<< in1->getTensorTypeAsm(/*isValueTensor=*/true,
1687-
/*useLogicalDims=*/true);
1688-
if (const auto &in2 = pointwiseAttr.getIN_2())
1689-
oss << ", "
1690-
<< in2->getTensorTypeAsm(/*isValueTensor=*/true,
1691-
/*useLogicalDims=*/true);
1695+
oss << getInputTypeAsm(0);
1696+
if (pointwiseAttr.getIN_1())
1697+
oss << ", " << getInputTypeAsm(1);
1698+
if (pointwiseAttr.getIN_2())
1699+
oss << ", " << getInputTypeAsm(2);
16921700
return oss.str();
16931701
}
16941702

@@ -1778,6 +1786,15 @@ inline std::string PointwiseNode::emitNodePreAsm() const {
17781786
{6}
17791787
)";
17801788

1789+
constexpr std::string_view kAddSquareSchema = R"(
1790+
{0}
1791+
{1}
1792+
%add_square_sq_{8} = torch.aten.mul.Tensor {3}, {3} : {4}, {4} -> {4}
1793+
%alpha_{8} = torch.constant.int 1
1794+
{5} = torch.aten.add.Tensor {2}, %add_square_sq_{8}, %alpha_{8} : {9}, !torch.int -> {6}
1795+
{7}
1796+
)";
1797+
17811798
constexpr std::string_view kIdentitySchema = R"(
17821799
{0}
17831800
%none_{7} = torch.constant.none
@@ -1894,6 +1911,20 @@ inline std::string PointwiseNode::emitNodePreAsm() const {
18941911
FUSILLI_DECLARE_SUB_ADD_TORCH_EMITTER(ADD, torch.aten.add.Tensor)
18951912
FUSILLI_DECLARE_SUB_ADD_TORCH_EMITTER(SUB, torch.aten.sub.Tensor)
18961913

1914+
case PointwiseAttr::Mode::ADD_SQUARE: {
1915+
return std::format(kAddSquareSchema, permuteIN0, /* {0} */
1916+
permuteIN1, /* {1} */
1917+
getInputNameAsm(0), /* {2} */
1918+
getInputNameAsm(1), /* {3} */
1919+
getInputTypeAsm(1), /* {4} */
1920+
getResultNamesAsm(), /* {5} */
1921+
getResultTypesAsm(), /* {6} */
1922+
permuteOUT0, /* {7} */
1923+
getName(), /* {8} */
1924+
getOperandTypesAsm() /* {9} */
1925+
);
1926+
}
1927+
18971928
default:
18981929
assert(false && "Unsupported pointwise mode");
18991930
return "";

‎samples/pointwise/pointwise_binary_ops.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ TEST_CASE("Pointwise binary ops", "[pointwise][graph]") {
4646
// clang-format off
4747
const auto mode = GENERATE(
4848
PointwiseAttr::Mode::ADD,
49+
PointwiseAttr::Mode::ADD_SQUARE,
4950
PointwiseAttr::Mode::DIV,
5051
PointwiseAttr::Mode::MAX_OP,
5152
PointwiseAttr::Mode::MIN_OP,
@@ -120,6 +121,10 @@ TEST_CASE("Pointwise binary ops", "[pointwise][graph]") {
120121
y = x0 + x1;
121122
break;
122123
}
124+
case PointwiseAttr::Mode::ADD_SQUARE: {
125+
y = x0 + x1 * x1;
126+
break;
127+
}
123128
case PointwiseAttr::Mode::DIV: {
124129
y = x0 / x1;
125130
break;

‎tests/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ add_fusilli_lit_tests(
147147
lit/test_pointwise_asm_emitter_abs.cpp
148148
lit/test_pointwise_asm_emitter_relu.cpp
149149
lit/test_pointwise_asm_emitter_add.cpp
150+
lit/test_pointwise_asm_emitter_add_square.cpp
150151
lit/test_pointwise_asm_emitter_add_transposed.cpp
151152
lit/test_pointwise_asm_emitter_ceil.cpp
152153
lit/test_pointwise_asm_emitter_cmp_eq.cpp

‎tests/lit/test_pointwise_asm_emitter_add_square.cpp‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
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(%result_: !torch.tensor<[16,256,64,32],f32>, %arg0: !torch.vtensor<[16,256,64,32],f32>, %arg1: !torch.vtensor<[1,256,1,1],f32>) attributes {torch.assume_strict_symbolic_shapes} {
15+
// TORCH-CHECK: %permute_IN_0_val_0_pointwise_add_square = torch.constant.int 0
16+
// TORCH-CHECK: %permute_IN_0_val_1_pointwise_add_square = torch.constant.int 1
17+
// TORCH-CHECK: %permute_IN_0_val_2_pointwise_add_square = torch.constant.int 2
18+
// TORCH-CHECK: %permute_IN_0_val_3_pointwise_add_square = torch.constant.int 3
19+
// TORCH-CHECK: %permute_IN_0_pointwise_add_square = torch.prim.ListConstruct %permute_IN_0_val_0_pointwise_add_square, %permute_IN_0_val_1_pointwise_add_square, %permute_IN_0_val_2_pointwise_add_square, %permute_IN_0_val_3_pointwise_add_square : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
20+
// TORCH-CHECK: %arg0_pointwise_add_square_perm = torch.aten.permute %arg0, %permute_IN_0_pointwise_add_square : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
21+
// TORCH-CHECK: %permute_IN_1_val_0_pointwise_add_square = torch.constant.int 0
22+
// TORCH-CHECK: %permute_IN_1_val_1_pointwise_add_square = torch.constant.int 1
23+
// TORCH-CHECK: %permute_IN_1_val_2_pointwise_add_square = torch.constant.int 2
24+
// TORCH-CHECK: %permute_IN_1_val_3_pointwise_add_square = torch.constant.int 3
25+
// TORCH-CHECK: %permute_IN_1_pointwise_add_square = torch.prim.ListConstruct %permute_IN_1_val_0_pointwise_add_square, %permute_IN_1_val_1_pointwise_add_square, %permute_IN_1_val_2_pointwise_add_square, %permute_IN_1_val_3_pointwise_add_square : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
26+
// TORCH-CHECK: %arg1_pointwise_add_square_perm = torch.aten.permute %arg1, %permute_IN_1_pointwise_add_square : !torch.vtensor<[1,256,1,1],f32>, !torch.list<int> -> !torch.vtensor<[1,256,1,1],f32>
27+
// TORCH-CHECK: %add_square_sq_pointwise_add_square = torch.aten.mul.Tensor %arg1_pointwise_add_square_perm, %arg1_pointwise_add_square_perm : !torch.vtensor<[1,256,1,1],f32>, !torch.vtensor<[1,256,1,1],f32> -> !torch.vtensor<[1,256,1,1],f32>
28+
// TORCH-CHECK: %alpha_pointwise_add_square = torch.constant.int 1
29+
// TORCH-CHECK: %result_pointwise_add_square_perm = torch.aten.add.Tensor %arg0_pointwise_add_square_perm, %add_square_sq_pointwise_add_square, %alpha_pointwise_add_square : !torch.vtensor<[16,256,64,32],f32>, !torch.vtensor<[1,256,1,1],f32>, !torch.int -> !torch.vtensor<[16,256,64,32],f32>
30+
// TORCH-CHECK: %permute_OUT_0_val_0_pointwise_add_square = torch.constant.int 0
31+
// TORCH-CHECK: %permute_OUT_0_val_1_pointwise_add_square = torch.constant.int 1
32+
// TORCH-CHECK: %permute_OUT_0_val_2_pointwise_add_square = torch.constant.int 2
33+
// TORCH-CHECK: %permute_OUT_0_val_3_pointwise_add_square = torch.constant.int 3
34+
// TORCH-CHECK: %permute_OUT_0_pointwise_add_square = torch.prim.ListConstruct %permute_OUT_0_val_0_pointwise_add_square, %permute_OUT_0_val_1_pointwise_add_square, %permute_OUT_0_val_2_pointwise_add_square, %permute_OUT_0_val_3_pointwise_add_square : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
35+
// TORCH-CHECK: %result = torch.aten.permute %result_pointwise_add_square_perm, %permute_OUT_0_pointwise_add_square : !torch.vtensor<[16,256,64,32],f32>, !torch.list<int> -> !torch.vtensor<[16,256,64,32],f32>
36+
// TORCH-CHECK: torch.overwrite.tensor.contents %result overwrites %result_ : !torch.vtensor<[16,256,64,32],f32>, !torch.tensor<[16,256,64,32],f32>
37+
// TORCH-CHECK: return
38+
// TORCH-CHECK: }
39+
// TORCH-CHECK: }
40+
//
41+
// AMDGPU-STATS-CHECK: "transient-memory-size": 0
42+
// AMDGPU-STATS-CHECK: "dispatch-count": 1
43+
// CPU-STATS-CHECK: "transient-memory-size": 0
44+
// CPU-STATS-CHECK: "dispatch-count": 1
45+
//
46+
// clang-format on
47+
48+
#include <fusilli.h>
49+
50+
#include "pointwise_utils.h"
51+
52+
#include <iostream>
53+
#include <string>
54+
55+
using namespace fusilli;
56+
57+
int main(int argc, char **argv) {
58+
std::string mode = (argc > 1) ? argv[1] : "default";
59+
60+
auto status = testBinaryPointwiseAsmEmitter(
61+
"pointwise_asm_emitter_add_square", "pointwise_add_square", mode,
62+
PointwiseAttr::Mode::ADD_SQUARE, {16, 256, 64, 32}, {1, 256, 1, 1});
63+
if (isError(status)) {
64+
std::cerr << "Test failed: " << status << std::endl;
65+
return 1;
66+
}
67+
return 0;
68+
}

0 commit comments

Comments
 (0)