[CK_TILE] Support for elementwise kernel #2246

amd-yashagar · 2025-05-27T09:53:10Z

Proposed changes

Implementation for the elementwise kernel that supports arbitrary number of input tensors for arbitrary dimensions.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

I have added tests relevant to the introduced functionality, and the unit tests are passing locally
I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
I have added inline documentation which enables the maintainers with understanding the motivation
I have removed the stale documentation which is no longer relevant after this pull request
(If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
I have run clang-format on all changed files
Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

AviralGoelAMD · 2025-05-27T18:41:15Z

Hey team, I made a few suggestions in another PR (feel free to keep/discard).
#2248

aosewski

Good start! Please add unit-tests (functional) here: https://github.com/ROCm/composable_kernel/tree/develop/test/ck_tile

aosewski · 2025-05-28T14:28:58Z

example/ck_tile/37_elementwise/CMakeLists.txt

Can you please change the example number? To 19 ? (At this moment last is 18) I don't understand why sb started with thirty sth...

Fixed in commit 4145fbf.

aosewski · 2025-05-28T14:31:41Z

example/ck_tile/37_elementwise/elementwise_example.cpp

Please add copyright header here and in other files.

Fixed in commit bcccb72.

aosewski · 2025-05-29T07:19:03Z

example/ck_tile/37_elementwise/CMakeLists.txt

+# Elementwise example targets 2D inputs
+set(TARGET_NAME_2D_INPUT tile_example_elementwise)
+add_executable(${TARGET_NAME_2D_INPUT} elementwise_example.cpp)
+rocm_install(TARGETS ${TARGET_NAME_2D_INPUT} COMPONENT examples)


I'm not sure we need this at the moment. @illsilin Could you please comment on this?

rocm_install lines removed from our CMakeLists.txt in commit 96d1a6b. I went ahead with this since none of the other examples have a corresponding line.

aosewski · 2025-05-29T07:43:59Z

example/ck_tile/37_elementwise/elementwise_example.cpp

+    int warmup            = arg_parser.get_int("warmup");
+    int repeat            = arg_parser.get_int("repeat");
+
+    assert(stride >= N);


The assert function is a no-op (disabled) in a release build mode. If you want to validate it, throw an error if has incorrect value.

Fixed in commit 05bbba7.

aosewski · 2025-05-29T07:48:49Z

example/ck_tile/37_elementwise/elementwise_example.cpp

+    // Copy data from host tensors to device buffers.
+    x_buf_a.ToDevice(x_host_a.data());
+    x_buf_b.ToDevice(x_host_b.data());
+    y_buf.ToDevice(y_host.data());


You'd rather want to set this to zero (if it's needed, or just for sanity). Right now you copy from uninitialized (or in best scenario default initialized) memory region.

Addressed in commit 64a7d3f where we no longer copy y_buf to the device. Please advise if this is not what you meant.

aosewski · 2025-05-29T11:49:49Z

include/ck_tile/ops/elementwise/kernel/elementwise.hpp

+
+            auto y = y_tile(tile_idx);
+
+            apply_operation(ElementWiseOperation{}, y, x_values);


Although this approach seem to be nice generic impl for associative operations - we do not always run such simple operations. Therefore we'd rather pass all inputs to the elementwise functor as arguments.

This is now addressed in commit 8c5d714.

aosewski · 2025-05-29T11:53:11Z

include/ck_tile/ops/elementwise/kernel/elementwise.hpp

+        // Move windows to next block
+        move_tile_windows(x_windows);
+        move_tile_window(y_window, {S::kBlockM});


This is not needed if each of your workgroups is processing just single tile.

Removed in commit 876b057.

aosewski · 2025-05-29T11:56:31Z

include/ck_tile/ops/elementwise/pipeline/elementwise_operators.hpp

This is the place for unary elementwise ops: https://github.com/ROCm/composable_kernel/blob/develop/include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp
You may want to add eg binary_elementwise_operation.hpp

Added in commit fc1ff7f.

aosewski · 2025-05-29T11:57:53Z

include/ck_tile/ops/elementwise/pipeline/elementwise_shape.hpp

You might want to reuse this one: https://github.com/ROCm/composable_kernel/blob/develop/include/ck_tile/ops/common/generic_2d_block_shape.hpp

As per our discussion on Teams, we will keep using the existing ElementWiseShape because it is slightly more appropriate for our implementation.

aosewski · 2025-05-29T11:59:46Z

include/ck_tile/ops/elementwise/pipeline/elementwise_pipeline_default_policy.hpp

+#include "ck_tile/ops/norm_reduce/block/block_norm_reduce_problem.hpp"
+#include "ck_tile/ops/norm_reduce/block/block_norm_reduce.hpp"


These are not used here.

Removed in commit abd0d35.

Co-authored-by: Sami Aario <[email protected]> Co-authored-by: Mohsen Saffari <[email protected]> Co-authored-by: yashagar <[email protected]>

…eration with the exact number of arguments in xs

…y_tile

…es a single tile

…rid size also for the unary example

… the other examples

…nce in a call to make_naive_tensor_view

…entwise_operation.hpp

…ing the transpose reference

- Refer to AMD terminology except when suggesting NVIDIA alternatives in parentheses - ElementWiseTraits was renamed to ElementWiseShape - Adopt suggestions made by Copilot when prompted to check for factual or typographical errors

aosewski

In general looks good. I got mostly cosmetic changes.

aosewski · 2025-06-04T10:44:10Z

include/ck_tile/host/reference/reference_transpose.hpp

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.


Suggested change

// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.

Fixed in 99eac8e.

aosewski · 2025-06-04T10:48:44Z

include/ck_tile/ops/elementwise/kernel/elementwise.hpp

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.


I don't see this fixed.

aosewski · 2025-06-04T10:50:16Z

include/ck_tile/ops/elementwise/kernel/elementwise.hpp

Please rename to elementwise_kernel.hpp

Fixed in 99eac8e.

aosewski · 2025-06-04T10:53:45Z

include/ck_tile/ops/elementwise/kernel/elementwise.hpp

+        auto make_tile_windows = [&](const auto& tensors, const auto& iM_) {
+            return generate_tuple(
+                [&](auto idx) {
+                    auto tensor_view = make_naive_tensor_view<address_space_enum::global>(


I think we could configure here cache coherence as well. The data won't be reused so we could set up this.
https://github.com/ROCm/composable_kernel/blob/develop/include/ck_tile/core/tensor/tensor_view.hpp#L433

Now I am not sure anymore what you meant by this: #2246 (comment)

We switched this call to make_naive_tensor_view to use the default template parameters based on that note. In any case, now we have two calls to make_naive_tensor_view in our kernel with different template parameters, because I neglected to change the other one on line 67. The two calls should still use the same parameters, am I right?

Please confirm if you want the call on line 42 to be switched to:

auto tensor_view = make_naive_tensor_view<address_space_enum::global, memory_operation_enum::set, amd_buffer_coherence_enum::slc>(

aosewski · 2025-06-04T10:58:18Z

include/ck_tile/ops/elementwise/pipeline/elementwise_pipeline_default_policy.hpp

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.


Please fix the copyright year.

Fixed in 99eac8e.

aosewski · 2025-06-04T11:00:52Z

include/ck_tile/ops/elementwise/pipeline/elementwise_shape.hpp

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.


Copyright year.

Fixed in 99eac8e.

aosewski · 2025-06-04T11:02:34Z

include/ck_tile/ops/elementwise/pipeline/elementwise_pipeline_problem.hpp

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.


Copyright year

Fixed in 99eac8e.

aosewski · 2025-06-04T11:09:19Z

include/ck_tile/core/container/tuple.hpp

+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get() const &          { TP_COM_(); return impl::getv<I>(*this); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get(number<I>) const & { TP_COM_(); return get<I>(); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get() &                { TP_COM_(); return impl::getv<I>(*this); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get(number<I>) &       { TP_COM_(); return get<I>(); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get() &&               { TP_COM_(); return impl::getv<I>(std::move(*this)); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get(number<I>) &&      { TP_COM_(); return std::move(*this).template get<I>(); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get() const &&         { TP_COM_(); return impl::getv<I>(std::move(*this)); }
+    template<index_t I> CK_TILE_HOST_DEVICE constexpr decltype(auto) get(number<I>) const &&{ TP_COM_(); return std::move(*this).template get<I>(); }


Would that impact performance or register usage?

amd-yashagar force-pushed the LWPCK-3170-elementwise-kernel-general branch from e4271dd to 02bfa62 Compare May 28, 2025 08:27

aosewski requested changes May 29, 2025

View reviewed changes

Damien Lejeune and others added 12 commits May 30, 2025 11:49

Elementwise kernel implementation

44dcdad

Co-authored-by: Sami Aario <[email protected]> Co-authored-by: Mohsen Saffari <[email protected]> Co-authored-by: yashagar <[email protected]>

Elementwise with generalized nDims

4d41cc2

Adding the n-ary input tensor feature

b2de0b1

Generalize dimensions on top of inputs

685f915

Add TFLOPS + remove std usage for tuples

8645970

1D basecase optimization

984198c

Cleanup code + refactoring to a common interface

e149526

Generalize to unary and add an example

f0f1c5a

Cleanup, refactoring and commenting

55fe333

Suggestions for LWPCK-3170: elementwise kernel improvements

5b225a1

Clang-format: remod.py

bc8e9ae

Replace InputTensorType with XDataType as the type of input_tensors

b885fb2

SamiAario-AMD force-pushed the LWPCK-3170-elementwise-kernel-general branch from 559af12 to 5f16d7e Compare May 30, 2025 11:50

Add Tuple::apply and use it in ElementWiseKernel::operator to call op…

8c5d714

…eration with the exact number of arguments in xs

SamiAario-AMD force-pushed the LWPCK-3170-elementwise-kernel-general branch from 5f16d7e to 8c5d714 Compare May 30, 2025 12:26

SamiAario-AMD added 3 commits May 30, 2025 12:38

Move examples to folder 19_elementwise

4145fbf

Add missing copyright headers and fix some existing ones

bcccb72

Replace an assert with throw std::runtime_error in elementwise example

05bbba7

SamiAario-AMD force-pushed the LWPCK-3170-elementwise-kernel-general branch from 2964594 to 05bbba7 Compare May 30, 2025 13:19

SamiAario-AMD and others added 6 commits June 2, 2025 08:47

Avoid reading the output by using make_static_distributed_tensor for …

2c46907

…y_tile

Removed two unused includes

abd0d35

No need to move windows to the next block when each workgroup process…

876b057

…es a single tile

Only copy input tensors to the device

64a7d3f

Use get_warp_size to obtain warp size, and use ceiling division for g…

a33a41e

…rid size also for the unary example

Adding output strides to the kernel, transposition example and update…

b52400c

… the other examples

SamiAario-AMD force-pushed the LWPCK-3170-elementwise-kernel-general branch from c1f4b26 to b52400c Compare June 2, 2025 12:01

Changes made by remod.py

891f701

SamiAario-AMD added 2 commits June 3, 2025 06:46

Use default template parameter values for memory operation and cohere…

86519ac

…nce in a call to make_naive_tensor_view

Move binary operations to include/ck_tile/ops/elementwise/binary_elem…

fc1ff7f

…entwise_operation.hpp

SamiAario-AMD force-pushed the LWPCK-3170-elementwise-kernel-general branch from eb11196 to fc1ff7f Compare June 3, 2025 06:55

Damien Lejeune and others added 3 commits June 3, 2025 14:33

Reuse generic reference binary/unary operation in examples + refactor…

1046ff9

…ing the transpose reference

Fix comments in elementwise_example.cpp

821019d

- Refer to AMD terminology except when suggesting NVIDIA alternatives in parentheses - ElementWiseTraits was renamed to ElementWiseShape - Adopt suggestions made by Copilot when prompted to check for factual or typographical errors

Simplify CMakeLists.txt and remove the unused variables this uncovers

96d1a6b

aosewski requested changes Jun 4, 2025

View reviewed changes

aosewski marked this pull request as ready for review June 4, 2025 11:10

aosewski requested review from illsilin, carlushuang, qianfengz, poyenc, geyyer, bartekxk, andriy-ca, afagaj, asleepzzz, ThomasNing, coderfeli, a team and tenpercent as code owners June 4, 2025 11:10

SamiAario-AMD and others added 3 commits June 4, 2025 11:34

Rename a file and fix some copyright statements

99eac8e

Changes made by script/clang-format-overwrite.sh

6cbf189

Add basic unit test for ElementWiseKernel

2a30295


		auto y = y_tile(tile_idx);

		apply_operation(ElementWiseOperation{}, y, x_values);

		#include "ck_tile/ops/norm_reduce/block/block_norm_reduce_problem.hpp"
		#include "ck_tile/ops/norm_reduce/block/block_norm_reduce.hpp"

		@@ -0,0 +1,33 @@
		// SPDX-License-Identifier: MIT
		// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

	// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
	// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,142 @@
		// SPDX-License-Identifier: MIT
		// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,29 @@
		// SPDX-License-Identifier: MIT
		// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,31 @@
		// SPDX-License-Identifier: MIT
		// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

		@@ -0,0 +1,24 @@
		// SPDX-License-Identifier: MIT
		// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.

[CK_TILE] Support for elementwise kernel #2246

Are you sure you want to change the base?

[CK_TILE] Support for elementwise kernel #2246

Conversation

amd-yashagar commented May 27, 2025

Proposed changes

Checklist

Discussion

Uh oh!

AviralGoelAMD commented May 27, 2025

Uh oh!

aosewski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SamiAario-AMD May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SamiAario-AMD Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aosewski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

SamiAario-AMD May 30, 2025 •

edited

Loading

SamiAario-AMD Jun 4, 2025 •

edited

Loading